Filter some type arguments for the builtin types/functions

4 files changed, 87 insertions, 10 deletions

diff --git a/compiler/Extract.ml b/compiler/Extract.ml
index 275cb3b9..bf90a411 100644
--- a/compiler/Extract.ml
+++ b/compiler/Extract.ml
@@ -1175,6 +1175,26 @@ let rec extract_ty (ctx : extraction_ctx) (fmt : F.formatter)
               (* TODO: for now, only the opaque *functions* are extracted in the
                  opaque module. The opaque *types* are assumed. *)
               F.pp_print_string fmt (ctx_get_type type_id ctx);
+              (* We might need to filter the type arguments, if the type
+                 is builtin (for instance, we filter the global allocator type
+                 argument for `Vec`). *)
+              let generics =
+                match type_id with
+                | AdtId id -> (
+                    match
+                      TypeDeclId.Map.find_opt id ctx.types_filter_type_args_map
+                    with
+                    | None -> generics
+                    | Some filter ->
+                        let types = List.combine filter generics.types in
+                        let types =
+                          List.filter_map
+                            (fun (b, ty) -> if b then Some ty else None)
+                            types
+                        in
+                        { generics with types })
+                | _ -> generics
+              in
               extract_generic_args ctx fmt no_params_tys generics;
               if print_paren then F.pp_print_string fmt ")"
           | HOL4 ->
@@ -1335,6 +1355,17 @@ let extract_type_decl_register_names (ctx : extraction_ctx) (def : type_decl) :
   let open ExtractBuiltin in
   let sname = name_to_simple_name def.name in
   let info = SimpleNameMap.find_opt sname (builtin_types_map ()) in
+  (* Register the filtering information, if there is *)
+  let ctx =
+    match info with
+    | Some { keep_params = Some keep; _ } ->
+        {
+          ctx with
+          types_filter_type_args_map =
+            TypeDeclId.Map.add def.def_id keep ctx.types_filter_type_args_map;
+        }
+    | _ -> ctx
+  in
   (* Compute and register the type def name *)
   let def_name =
     match info with
@@ -2361,7 +2392,19 @@ let extract_fun_decl_register_names (ctx : extraction_ctx)
       in
       (* Use the builtin names if necessary *)
       match builtin with
-      | Some (_filter, info) ->
+      | Some (filter_info, info) ->
+          (* Register the filtering information, if there is *)
+          let ctx =
+            match filter_info with
+            | Some keep ->
+                {
+                  ctx with
+                  funs_filter_type_args_map =
+                    FunDeclId.Map.add def.fwd.f.def_id keep
+                      ctx.funs_filter_type_args_map;
+                }
+            | _ -> ctx
+          in
           let backs = List.map (fun f -> f.f) def.backs in
           let funs = if def.keep_fwd then def.fwd.f :: backs else backs in
           List.fold_left
@@ -2734,7 +2777,27 @@ and extract_function_call (ctx : extraction_ctx) (fmt : F.formatter)
 
       (* Sanity check: HOL4 doesn't support const generics *)
       assert (generics.const_generics = [] || !backend <> HOL4);
-      (* Print the generics *)
+      (* Print the generics.
+
+         We might need to filter some of the type arguments, if the type
+         is builtin (for instance, we filter the global allocator type
+         argument for `Vec::new`).
+      *)
+      let generics =
+        match fun_id with
+        | FromLlbc (FunId (Regular id), _, _) -> (
+            match FunDeclId.Map.find_opt id ctx.funs_filter_type_args_map with
+            | None -> generics
+            | Some filter ->
+                let types = List.combine filter generics.types in
+                let types =
+                  List.filter_map
+                    (fun (b, ty) -> if b then Some ty else None)
+                    types
+                in
+                { generics with types })
+        | _ -> generics
+      in
       extract_generic_args ctx fmt TypeDeclId.Set.empty generics;
       (* Print the arguments *)
       List.iter
diff --git a/compiler/ExtractBase.ml b/compiler/ExtractBase.ml
index f26beeb6..e004aba8 100644
--- a/compiler/ExtractBase.ml
+++ b/compiler/ExtractBase.ml
@@ -640,6 +640,19 @@ type extraction_ctx = {
   functions_with_decreases_clause : PureUtils.FunLoopIdSet.t;
   trans_trait_decls : Pure.trait_decl Pure.TraitDeclId.Map.t;
   trans_trait_impls : Pure.trait_impl Pure.TraitImplId.Map.t;
+  types_filter_type_args_map : bool list TypeDeclId.Map.t;
+      (** The map to filter the type arguments for the builtin type
+          definitions.
+
+          We need this for type `Vec`, for instance, which takes a useless
+          (in the context of the type translation) type argument for the
+          allocator which is used, and which we want to remove.
+
+          TODO: it would be cleaner to filter those types in a micro-pass,
+          rather than at code generation time.
+        *)
+  funs_filter_type_args_map : bool list FunDeclId.Map.t;
+      (** Same as {!types_filter_type_args_map}, but for functions *)
 }
 
 (** Debugging function, used when communicating name collisions to the user,
diff --git a/compiler/ExtractBuiltin.ml b/compiler/ExtractBuiltin.ml
index d3cea54e..65c18efd 100644
--- a/compiler/ExtractBuiltin.ml
+++ b/compiler/ExtractBuiltin.ml
@@ -203,9 +203,8 @@ type builtin_fun_info = {
 let builtin_funs () :
     (string list * bool list option * builtin_fun_info list) list =
   let rg0 = Some Types.RegionGroupId.zero in
-  (* TODO: fix the names below *)
   [
-    ( [ "core::mem::replace" ],
+    ( [ "core"; "mem"; "replace" ],
       None,
       [
         {
@@ -223,7 +222,7 @@ let builtin_funs () :
             | Lean -> "mem.replace_back");
         };
       ] );
-    ( [ "alloc::vec::Vec::new" ],
+    ( [ "alloc"; "vec"; "Vec"; "new" ],
       Some [ true; false ],
       [
         {
@@ -241,7 +240,7 @@ let builtin_funs () :
             | Lean -> "Vec.new_back");
         };
       ] );
-    ( [ "alloc::vec::Vec::push" ],
+    ( [ "alloc"; "vec"; "Vec"; "push" ],
       Some [ true; false ],
       [
         (* The forward function shouldn't be used *)
@@ -260,7 +259,7 @@ let builtin_funs () :
             | Lean -> "Vec.push");
         };
       ] );
-    ( [ "alloc::vec::Vec::insert" ],
+    ( [ "alloc"; "vec"; "Vec"; "insert" ],
       Some [ true; false ],
       [
         (* The forward function shouldn't be used *)
@@ -279,7 +278,7 @@ let builtin_funs () :
             | Lean -> "Vec.insert");
         };
       ] );
-    ( [ "alloc::vec::Vec::len" ],
+    ( [ "alloc"; "vec"; "Vec"; "len" ],
       Some [ true; false ],
       [
         {
@@ -290,7 +289,7 @@ let builtin_funs () :
             | Lean -> "Vec.len");
         };
       ] );
-    ( [ "alloc::vec::Vec::index" ],
+    ( [ "alloc"; "vec"; "Vec"; "index" ],
       Some [ true; false ],
       [
         {
@@ -309,7 +308,7 @@ let builtin_funs () :
             | Lean -> "Vec.index_shared_back");
         };
       ] );
-    ( [ "alloc::vec::Vec::index_mut" ],
+    ( [ "alloc"; "vec"; "Vec"; "index_mut" ],
       Some [ true; false ],
       [
         {
diff --git a/compiler/Translate.ml b/compiler/Translate.ml
index b3269aa2..35dff9e6 100644
--- a/compiler/Translate.ml
+++ b/compiler/Translate.ml
@@ -1074,6 +1074,8 @@ let translate_crate (filename : string) (dest_dir : string) (crate : A.crate) :
       trans_types;
       trans_funs;
       functions_with_decreases_clause = rec_functions;
+      types_filter_type_args_map = Pure.TypeDeclId.Map.empty;
+      funs_filter_type_args_map = Pure.FunDeclId.Map.empty;
     }
   in