Formatting, generalization of readonly lids, and a peephole optimizat…

…ion in Structs
FStarLang · Jul 31, 2024 · 5a2983a · 5a2983a
1 parent d5759a8
commit 5a2983a
Show file tree

Hide file tree

Showing 6 changed files with 93 additions and 75 deletions.
diff --git a/lib/AstToCStar.ml b/lib/AstToCStar.ml
@@ -248,7 +248,7 @@ let rec unit_to_void env e es extra =
     trivially as `(uint8_t)200 * (uint8_t)225 * (uint8_t)200 * (uint8_t)250`.
     (The effect of the casts is immediately undone by the promotion to int.)
   - (255uy `FStar.UInt8.add_mod` 1uy) / FStar.UInt8.div 2uy gives 0uy in F*, but
-    ((uint8_t) 255 * (uint8_t) 1) / (uint8_t) 2 gives 128 in C
+    ((uint8_t) 255 + (uint8_t) 1) / (uint8_t) 2 gives 128 in C
   - same kinds of issues with shifts, bitwise complement (which materializes
     "more bytes" in the intermediary expression), and so on.
 

diff --git a/lib/Helpers.ml b/lib/Helpers.ml
@@ -289,14 +289,20 @@ let pure_builtin_lids = [
   [ "FStar"; "UInt128" ], "";
 ]
 
-let is_readonly_builtin_lid_ = ref (fun lid ->
+let is_readonly_builtin_lid_ = ref (fun lid _t ->
   List.exists (fun lid' ->
     let lid = Idents.string_of_lident lid in
     let lid' = Idents.string_of_lident lid' in
     KString.starts_with lid lid'
   ) pure_builtin_lids)
 
-let is_readonly_builtin_lid lid = !is_readonly_builtin_lid_ lid
+let is_readonly_builtin_lid e =
+  let lid = match e.node with
+    | EQualified lid
+    | ETApp ({ node = EQualified lid; _ }, _, _, _) -> lid
+    | _ -> failwith "not an lid"
+  in
+  !is_readonly_builtin_lid_ lid e.typ
 
 class ['self] closed_term_visitor = object (_: 'self)
   inherit [_] reduce
@@ -338,9 +344,11 @@ class ['self] readonly_visitor = object (self: 'self)
     | EPolyComp _
     | EOp _ ->
         List.for_all (self#visit_expr_w ()) es
-    | EQualified lid when is_readonly_builtin_lid lid ->
+    | EQualified _ 
+    when is_readonly_builtin_lid e ->
         List.for_all (self#visit_expr_w ()) es
-    | ETApp ({ node = EQualified lid; _ }, _, _, _) when is_readonly_builtin_lid lid ->
+    | ETApp ({ node = EQualified _; _ } as e, _, _, _)
+    when is_readonly_builtin_lid e ->
         List.for_all (self#visit_expr_w ()) es
     | _ ->
         false

diff --git a/lib/Simplify.ml b/lib/Simplify.ml
@@ -75,8 +75,8 @@ class ['self] safe_use = object (self: 'self)
   method! visit_EApp env e es =
     match e.node with
     | EOp _ -> super#visit_EApp env e es
-    | EQualified lid when Helpers.is_readonly_builtin_lid lid -> super#visit_EApp env e es
-    | ETApp ({ node = EQualified lid; _ }, _, _, _) when Helpers.is_readonly_builtin_lid lid -> super#visit_EApp env e es
+    | EQualified _ when Helpers.is_readonly_builtin_lid e -> super#visit_EApp env e es
+    | ETApp ({ node = EQualified _; _ } as e, _, _, _) when Helpers.is_readonly_builtin_lid e -> super#visit_EApp env e es
     | _ -> self#unordered env (e :: es)
 
   method! visit_ELet env _ e1 e2 = self#sequential env e1 (Some e2)

diff --git a/lib/Structs.ml b/lib/Structs.ml
@@ -229,45 +229,54 @@ let pass_by_ref (should_rewrite: _ -> policy) = object (self)
         None
     ) (List.combine binders (args_are_structs @ (if ret_is_struct then [ false ] else []))) in
 
-    let body = self#visit_expr_w to_be_starred body in
-
-    (* Step 4: if the function now takes an extra argument for the output struct. *)
     let body =
-      if ret_is_struct then
-        let assign_into_ret e =
-          match e.node with
-          | EAbort _ ->
-              e
-          | _ ->
-              if ret_is_array then
-                with_type TUnit (EAssign (Option.get ret_atom, e))
-              else
-                with_type TUnit (EBufWrite (Option.get ret_atom, Helpers.zerou32, e))
-        in
-        (* Step 4.1: early-returns `return e` become `dst := e; return` *)
-        let body = (object
-          inherit [_] map
-          method! visit_EReturn _ e =
-            ESequence [
-              assign_into_ret e;
-              with_type e.typ (EReturn Helpers.eunit)
-            ]
-        end)#visit_expr_w () body in
-        (* Step 4.2: the overall value computed by the function is also assigned into the
-           destination. This relies on the invariant that functions are either in the style of 4.1
-           where returns in terminal position have been removed earlier (eurydice), or in
-           expression-style. *)
-        Helpers.nest_in_return_pos TUnit (fun _ e ->
-          match e.node with
-          | EWhile _ -> e
-            (* ret is a struct type, not unit -- therefore, it must be the case that this is an
-               unreachable loop -- bail *)
-          | EReturn _ -> e
-            (* handled above *)
-          | _ -> assign_into_ret e
-        ) body
-      else
-        body
+      match body.node with
+      | EApp (e, es) when ret_is_struct ->
+          (* Fast-path: the body is a single application node, meaning that the we can optimize and
+             directly pass our own destination parameter to the callee, rather than generating a
+             temporary that ends up memcpy'd into our own destination parameter. *)
+          self#rewrite_app to_be_starred e es ret_atom
+
+      | _ ->
+          (* General case: recursively visit *)
+          let body = self#visit_expr_w to_be_starred body in
+
+          (* Step 4: if the function now takes an extra argument for the output struct. *)
+          if ret_is_struct then
+            let assign_into_ret e =
+              match e.node with
+              | EAbort _ ->
+                  e
+              | _ ->
+                  if ret_is_array then
+                    with_type TUnit (EAssign (Option.get ret_atom, e))
+                  else
+                    with_type TUnit (EBufWrite (Option.get ret_atom, Helpers.zerou32, e))
+            in
+            (* Step 4.1: early-returns `return e` become `dst := e; return` *)
+            let body = (object
+              inherit [_] map
+              method! visit_EReturn _ e =
+                ESequence [
+                  assign_into_ret e;
+                  with_type e.typ (EReturn Helpers.eunit)
+                ]
+            end)#visit_expr_w () body in
+            (* Step 4.2: the overall value computed by the function is also assigned into the
+               destination. This relies on the invariant that functions are either in the style of 4.1
+               where returns in terminal position have been removed earlier (eurydice), or in
+               expression-style. *)
+            Helpers.nest_in_return_pos TUnit (fun _ e ->
+              match e.node with
+              | EWhile _ -> e
+                (* ret is a struct type, not unit -- therefore, it must be the case that this is an
+                   unreachable loop -- bail *)
+              | EReturn _ -> e
+                (* handled above *)
+              | _ -> assign_into_ret e
+            ) body
+          else
+            body
     in
     let body = DeBruijn.close_binders binders body in
     DFunction (cc, flags, n_cg, n, ret, lid, binders, body)

diff --git a/lib/dune b/lib/dune
@@ -1,28 +1,30 @@
 (include_subdirs unqualified)
 
 (library
-  (name krml)
-  (libraries
-    ppx_deriving.std
-    ppx_deriving_yojson
-    zarith
-    pprint
-    unix
-    menhirLib
-    sedlex
-    process
-    fix
-    wasm
-    visitors.ppx
-    visitors.runtime
-    sedlex.ppx
-    uucp
-  )
-  (preprocess
-    (pps ppx_deriving.std ppx_deriving_yojson sedlex.ppx visitors.ppx)
-  )
-; Reinstate @39 once https://github.com/whitequark/ppx_deriving_yojson/issues/21
-; is fixed
-; VD: Fix @40 (only really needed for GenCtypes.ml)
-  (flags (:standard -warn-error -A -w @1-2@3-7@8..12@14..21@23..29-30@31..38-39-40-41@43@57))
-)
+ (name krml)
+ (libraries
+  ppx_deriving.std
+  ppx_deriving_yojson
+  zarith
+  pprint
+  unix
+  menhirLib
+  sedlex
+  process
+  fix
+  wasm
+  visitors.ppx
+  visitors.runtime
+  sedlex.ppx
+  uucp)
+ (preprocess
+  (pps ppx_deriving.std ppx_deriving_yojson sedlex.ppx visitors.ppx))
+ ; Reinstate @39 once https://github.com/whitequark/ppx_deriving_yojson/issues/21
+ ; is fixed
+ ; VD: Fix @40 (only really needed for GenCtypes.ml)
+ (flags
+  (:standard
+   -warn-error
+   -A
+   -w
+   @1-2@3-7@8..12@14..21@23..29-30@31..38-39-40-41@43@57)))
diff --git a/lib/parser/dune b/lib/parser/dune
@@ -1,5 +1,4 @@
-  (menhir
-    (modules KParser)
-    (flags --explain --dump)
-    (infer true)
-  )
+(menhir
+ (modules KParser)
+ (flags --explain --dump)
+ (infer true))