Local references in Steel.ST

src/extraction/FStar.Extraction.Krml.fst

@@ -563,7 +563,10 @@ and translate_expr env e: expr =
       EBufCreateNoInit (Stack, translate_expr env elen)
 
   | MLE_App ({ expr = MLE_TApp({ expr = MLE_Name p }, _) } , [ init ])
-    when (string_of_mlpath p = "FStar.HyperStack.ST.salloc") ->
+    when (
+      string_of_mlpath p = "FStar.HyperStack.ST.salloc" ||
+      string_of_mlpath p = "Steel.ST.Reference._alloca"
+    ) ->
       EBufCreate (Stack, translate_expr env init, EConstant (UInt32, "1"))
 
   | MLE_App ({ expr = MLE_TApp({ expr = MLE_Name p }, _) }, [ e2 ])
@@ -690,10 +693,15 @@ and translate_expr env e: expr =
     when string_of_mlpath p = "Steel.ST.Reference.write" ->
       EBufWrite (translate_expr env e1, EConstant (UInt32, "0"), translate_expr env e2)
 
-  | MLE_App ({ expr = MLE_Name p }, [ _ ]) when (string_of_mlpath p = "FStar.HyperStack.ST.push_frame") ->
+  | MLE_App ({ expr = MLE_Name p }, [ _ ]) when (
+        string_of_mlpath p = "FStar.HyperStack.ST.push_frame" ||
+        string_of_mlpath p = "Steel.ST.Reference._push_frame"
+      ) ->
       EPushFrame
   | MLE_App ({ expr = MLE_Name p }, [ _ ]) when (string_of_mlpath p = "FStar.HyperStack.ST.pop_frame") ->
       EPopFrame
+  | MLE_App ({ expr = MLE_TApp({ expr = MLE_Name p }, _) }, [ _; _ ]) when (string_of_mlpath p = "Steel.ST.Reference._free_and_pop_frame") ->
+      EPopFrame
   | MLE_App ({ expr = MLE_TApp({ expr = MLE_Name p }, _) }, [ e1; e2; e3; e4; e5 ]) when (
       string_of_mlpath p = "FStar.Buffer.blit" ||
       string_of_mlpath p = "LowStar.Monotonic.Buffer.blit" ||

ulib/experimental/Steel.ST.GhostReference.fsti

@@ -119,3 +119,32 @@ val free (#a:Type0)
   : STGhostT unit u
       (pts_to r full_perm v)
       (fun _ -> emp)
+
+/// Executes a code block with a ghost reference temporarily
+/// allocated. This function is declared in the `STF` effect so
+/// that the pre- and post-resources can be properly inferred by the
+/// Steel tactic from the caller's context.
+///
+/// By virtue of `alloc` and `free` being STGhost
+/// functions, `with_local init body` simply extracts as `body _`.
+///
+/// This combinator is defined only to mirror its Steel.ST.Reference
+/// counterpart.
+inline_for_extraction
+let with_local
+  (#t: Type)
+  (init: Ghost.erased t)
+  (#pre: vprop)
+  (#ret_t: Type)
+  (#post: ret_t -> vprop)
+  (body: (r: ref t) ->
+    STT ret_t
+    (pts_to r full_perm init `star` pre)
+    (fun v -> exists_ (pts_to r full_perm) `star` post v)
+  )
+: STF ret_t pre post True (fun _ -> True)
+= let r = alloc init in
+  let v = body r in
+  let _ = elim_exists () in
+  free r;
+  return v

ulib/experimental/Steel.ST.Loops.Util.fst

@@ -55,7 +55,7 @@ let repeat_until_body
     intro_exists b (repeat_until_inv p r)
 
 let repeat_until p $body
-  = let r = R.alloc true in
+= R.with_local true (fun r ->
     rewrite (R.pts_to r full_perm true `star` p true)
             (repeat_until_inv p r true);
     intro_exists true (repeat_until_inv p r);
@@ -65,4 +65,5 @@ let repeat_until p $body
       (repeat_until_body p r body);
     rewrite (repeat_until_inv p r false)
             (R.pts_to r full_perm false `star` p false);
-    R.free r
+    noop () // to handle intro_exists
+  )

ulib/experimental/Steel.ST.Reference.fst

@@ -97,6 +97,73 @@ let free (#a:Type0)
   = coerce_steel(fun _ -> R.free_pt r);
     return ()
 
+/// Local primitive, to be extracted to Low* EPushFrame.  To remember
+/// that we need to call some pop_frame later, we insert some dummy
+/// vprop into the context.
+let _stack_frame : vprop = pure True
+let _push_frame () : STT unit emp (fun _ -> _stack_frame) =
+  rewrite (pure True) _stack_frame
+
+/// Local primitive, to be extracted to Low* EBufCreate
+let _alloca (#a:Type) (x:a)
+  : ST (ref a)
+      emp
+      (fun r -> pts_to r full_perm x)
+      (requires True)
+      (ensures fun r -> not (is_null r))
+= alloc x
+
+/// Local primitive, to be extracted to Low* EPopFrame
+let _free_and_pop_frame
+  (#a:Type0)
+  (#v:erased a)
+  (r:ref a)
+: STT unit
+    (pts_to r full_perm v `star` _stack_frame)
+    (fun _ -> emp)
+= free r;
+  rewrite _stack_frame (pure True);
+  elim_pure _
+
+let with_local
+  (#t: Type)
+  (init: t)
+  (#pre: vprop)
+  (#ret_t: Type)
+  (#post: ret_t -> vprop)
+  (body: (r: ref t) ->
+    STT ret_t
+    (pts_to r full_perm init `star` pre)
+    (fun v -> exists_ (pts_to r full_perm) `star` post v)
+  )
+: STF ret_t pre post True (fun _ -> True)
+= _push_frame ();
+  let r = _alloca init in
+  let v = body r in
+  let _ = elim_exists () in
+  _free_and_pop_frame r;
+  return v
+
+let with_named_local
+  (#t: Type)
+  (init: t)
+  (#pre: vprop)
+  (#ret_t: Type)
+  (#post: ret_t -> vprop)
+  (name: string)
+  (body: (r: ref t) ->
+    STT ret_t
+    (pts_to r full_perm init `star` pre)
+    (fun v -> exists_ (pts_to r full_perm) `star` post v)
+  )
+: STF ret_t pre post True (fun _ -> True)
+= _push_frame ();
+  [@(rename_let name)]
+  let r = _alloca init in
+  let v = body r in
+  let _ = elim_exists () in
+  _free_and_pop_frame r;
+  return v
 
 let share (#a:Type0)
           (#uses:_)

ulib/experimental/Steel.ST.Reference.fsti

@@ -120,6 +120,66 @@ val free (#a:Type0)
   : STT unit
     (pts_to r full_perm v) (fun _ -> emp)
 
+/// Executes a code block with a local variable temporarily allocated
+/// on the stack. This function is declared in the `STF` effect so
+/// that the pre- and post-resources can be properly inferred by the
+/// Steel tactic from the caller's context.
+///
+/// From the extraction point of view, `with_local init body` is to behave
+/// similarly as the following Low* code:
+///
+/// <<<
+/// push_frame ();
+/// let r = alloca 1ul init in
+/// let res = body r in
+/// pop_frame ();
+/// r
+/// >>>
+///
+/// and thus, is to be extracted to C as:
+/// <<<
+/// ret_t res;
+/// {
+///   t r = init;
+///   res = <body r>;
+/// }
+/// >>>
+///
+/// To this end, we mimic the Low* behavior by defining local
+/// primitives with primitive extraction in the `.fst`, and have them
+/// called by `with_local`. This is why we mark `with_local` as
+/// `inline_for_extraction`.
+inline_for_extraction
+val with_local
+  (#t: Type)
+  (init: t)
+  (#pre: vprop)
+  (#ret_t: Type)
+  (#post: ret_t -> vprop)
+  (body: (r: ref t) ->
+    STT ret_t
+    (pts_to r full_perm init `star` pre)
+    (fun v -> exists_ (pts_to r full_perm) `star` post v)
+  )
+: STF ret_t pre post True (fun _ -> True)
+
+/// Same as with_local, with an additional string argument to set the
+/// name of the local variable in the extracted C code.
+inline_for_extraction
+val with_named_local
+  (#t: Type)
+  (init: t)
+  (#pre: vprop)
+  (#ret_t: Type)
+  (#post: ret_t -> vprop)
+  (name: string)
+  (body: (r: ref t) ->
+    STT ret_t
+    (pts_to r full_perm init `star` pre)
+    (fun v -> exists_ (pts_to r full_perm) `star` post v)
+  )
+: STF ret_t pre post True (fun _ -> True)
+
 /// Splits the permission on reference [r] into two. This function is
 /// computationally irrelevant (it has effect SteelGhost)
 val share (#a:Type0)