Handle properly the builtin, non fallible functions

3 files changed, 54 insertions, 18 deletions

diff --git a/compiler/ExtractBase.ml b/compiler/ExtractBase.ml
index e004aba8..8f32ba44 100644
--- a/compiler/ExtractBase.ml
+++ b/compiler/ExtractBase.ml
@@ -22,8 +22,8 @@ type region_group_info = {
        *)
 }
 
-module StringSet = Collections.MakeSet (Collections.OrderedString)
-module StringMap = Collections.MakeMap (Collections.OrderedString)
+module StringSet = Collections.StringSet
+module StringMap = Collections.StringMap
 
 type name = Names.name
 type type_name = Names.type_name
diff --git a/compiler/ExtractBuiltin.ml b/compiler/ExtractBuiltin.ml
index 9cc7c226..4c6fe014 100644
--- a/compiler/ExtractBuiltin.ml
+++ b/compiler/ExtractBuiltin.ml
@@ -376,7 +376,39 @@ let mk_builtin_funs_map () =
 let builtin_funs_map = mk_memoized mk_builtin_funs_map
 
 let builtin_non_fallible_funs =
-  [ "alloc::boxed::Box::deref"; "alloc::boxed::Box::deref_mut" ]
+  let int_names =
+    [
+      "usize";
+      "u8";
+      "u16";
+      "u32";
+      "u64";
+      "u128";
+      "isize";
+      "i8";
+      "i16";
+      "i32";
+      "i64";
+      "i128";
+    ]
+  in
+  let int_ops =
+    [ "wrapping_add"; "wrapping_sub"; "rotate_left"; "rotate_right" ]
+  in
+  let int_funs =
+    List.map
+      (fun int_name ->
+        List.map (fun op -> "core::num::" ^ int_name ^ "::" ^ op) int_ops)
+      int_names
+  in
+  let int_funs = List.concat int_funs in
+  [
+    "alloc::boxed::Box::deref";
+    "alloc::boxed::Box::deref_mut";
+    "core::mem::replace";
+    "core::mem::take";
+  ]
+  @ int_funs
 
 let builtin_non_fallible_funs_set =
   SimpleNameSet.of_list
diff --git a/compiler/FunsAnalysis.ml b/compiler/FunsAnalysis.ml
index 1273f57d..3ba5d35d 100644
--- a/compiler/FunsAnalysis.ml
+++ b/compiler/FunsAnalysis.ml
@@ -57,21 +57,16 @@ let analyze_module (m : crate) (funs_map : fun_decl FunDeclId.Map.t)
     let stateful = ref false in
     let can_diverge = ref false in
     let is_rec = ref false in
+    let is_builtin_non_fallible_group = ref false in
 
     (* We have some specialized knowledge of some library functions; we don't
        have any more custom treatment than this, and these functions can be modeled
        suitably in Primitives.fst, rather than special-casing for them all the
        way. *)
-    let module M = struct type opaque_info = { fallible: bool; stateful: bool } end in
-    let open M in
-    let opaque_info (f: fun_decl) =
-      match f.name with
-      | [ Ident "core"; Ident "num"; Ident "u32"; _; Ident "wrapping_add" ]
-      | [ Ident "core"; Ident "num"; Ident "u32"; _; Ident "rotate_left" ] ->
-          { fallible = false; stateful = false }
-      | _ ->
-          (* Opaque function: we consider they fail by default *)
-          { fallible = true; stateful = true }
+    let is_builtin_non_fallible (f : fun_decl) : bool =
+      let open ExtractBuiltin in
+      let name = name_to_simple_name f.name in
+      SimpleNameSet.mem name builtin_non_fallible_funs_set
     in
 
     (* JP: Why not use a reduce visitor here with a tuple of the values to be
@@ -124,11 +119,16 @@ let analyze_module (m : crate) (funs_map : fun_decl FunDeclId.Map.t)
       in
       (* Sanity check: global bodies don't contain stateful calls *)
       assert ((not f.is_global_decl_body) || not !stateful);
+      let is_builtin_non_fallible = is_builtin_non_fallible f in
+      is_builtin_non_fallible_group :=
+        !is_builtin_non_fallible_group || is_builtin_non_fallible;
       match f.body with
       | None ->
-          let info = opaque_info f in
-          obj#may_fail info.fallible;
-          stateful := (not f.is_global_decl_body) && use_state && info.stateful
+          obj#may_fail (not is_builtin_non_fallible);
+          stateful :=
+            (not f.is_global_decl_body)
+            && use_state
+            && not is_builtin_non_fallible
       | Some body -> obj#visit_statement () body.body
     in
     List.iter visit_fun d;
@@ -136,12 +136,16 @@ let analyze_module (m : crate) (funs_map : fun_decl FunDeclId.Map.t)
      * groups containing globals contain exactly one declaration *)
     let is_global_decl_body = List.exists (fun f -> f.is_global_decl_body) d in
     assert ((not is_global_decl_body) || List.length d = 1);
+    assert ((not !is_builtin_non_fallible_group) || List.length d = 1);
     (* We ignore on purpose functions that cannot fail and consider they *can*
      * fail: the result of the analysis is not used yet to adjust the translation
      * so that the functions which syntactically can't fail don't use an error monad.
-     * However, we do keep the result of the analysis for global bodies.
+     * However, we do keep the result of the analysis for global bodies and for
+     * builtin functions which are marked as non-fallible.
      * *)
-    can_fail := (not is_global_decl_body) || !can_fail;
+    can_fail :=
+      ((not is_global_decl_body) && not !is_builtin_non_fallible_group)
+      || !can_fail;
     {
       can_fail = !can_fail;
       stateful = !stateful;