Make more progress

3 files changed, 67 insertions, 24 deletions

diff --git a/compiler/Extract.ml b/compiler/Extract.ml
index e140ea1c..17f850a3 100644
--- a/compiler/Extract.ml
+++ b/compiler/Extract.ml
@@ -1650,7 +1650,7 @@ let insert_req_space (fmt : F.formatter) (space : bool ref) : unit =
     We add the trait self clause for provided methods (see {!TraitSelfClauseId}).
  *)
 let extract_trait_self_clause (insert_req_space : unit -> unit)
-    (ctx : extraction_ctx) (fmt : F.formatter) (trait_decl : A.trait_decl)
+    (ctx : extraction_ctx) (fmt : F.formatter) (trait_decl : trait_decl)
     (params : string list) : unit =
   insert_req_space ();
   F.pp_print_string fmt "(";
@@ -1673,7 +1673,7 @@ let extract_trait_self_clause (insert_req_space : unit -> unit)
  *)
 let extract_generic_params (ctx : extraction_ctx) (fmt : F.formatter)
     (no_params_tys : TypeDeclId.Set.t) (use_forall : bool) (as_implicits : bool)
-    (space : bool ref option) (trait_decl : A.trait_decl option)
+    (space : bool ref option) (trait_decl : trait_decl option)
     (generics : generic_params) (type_params : string list)
     (cg_params : string list) (trait_clauses : string list) : unit =
   let all_params = List.concat [ type_params; cg_params; trait_clauses ] in
@@ -3111,10 +3111,7 @@ let extract_fun_parameters (space : bool ref) (ctx : extraction_ctx)
   let ctx, trait_decl =
     match def.kind with
     | TraitMethodProvided (decl_id, _) ->
-        let trait_decl =
-          T.TraitDeclId.Map.find decl_id
-            ctx.trans_ctx.trait_decls_context.trait_decls
-        in
+        let trait_decl = T.TraitDeclId.Map.find decl_id ctx.trans_trait_decls in
         let ctx, _ = ctx_add_trait_self_clause ctx in
         let ctx = { ctx with is_provided_method = true } in
         (ctx, Some trait_decl)
@@ -3887,6 +3884,16 @@ let extract_global_decl (ctx : extraction_ctx) (fmt : F.formatter)
       (* Add a break to insert lines between declarations *)
       F.pp_print_break fmt 0 0
 
+(** Similar to {!extract_type_decl_register_names} *)
+let extract_trait_decl_register_names (ctx : extraction_ctx) (d : trait_decl) :
+    extraction_ctx =
+  raise (Failure "TODO")
+
+(** Similar to {!extract_type_decl_register_names} *)
+let extract_trait_impl_register_names (ctx : extraction_ctx) (d : trait_impl) :
+    extraction_ctx =
+  raise (Failure "TODO")
+
 (** Extract a trait declaration *)
 let extract_trait_decl (ctx : extraction_ctx) (fmt : F.formatter)
     (trait_decl : trait_decl) : unit =
diff --git a/compiler/ExtractBase.ml b/compiler/ExtractBase.ml
index 251d8b36..2855b3b9 100644
--- a/compiler/ExtractBase.ml
+++ b/compiler/ExtractBase.ml
@@ -654,6 +654,8 @@ type extraction_ctx = {
   trait_decl_id : trait_decl_id option;
       (** If we are extracting a trait declaration, identifies it *)
   is_provided_method : bool;
+  trans_trait_decls : Pure.trait_decl Pure.TraitDeclId.Map.t;
+  trans_trait_impls : Pure.trait_impl Pure.TraitImplId.Map.t;
 }
 
 (** Debugging function, used when communicating name collisions to the user,
diff --git a/compiler/Translate.ml b/compiler/Translate.ml
index 790dbe14..8df69961 100644
--- a/compiler/Translate.ml
+++ b/compiler/Translate.ml
@@ -301,8 +301,13 @@ let translate_function_to_pure (trans_ctx : trans_ctx)
   (* Return *)
   (pure_forward, pure_backwards)
 
+(* TODO: factor out the return type *)
 let translate_crate_to_pure (crate : A.crate) :
-    trans_ctx * Pure.type_decl list * (bool * pure_fun_translation) list =
+    trans_ctx
+    * Pure.type_decl list
+    * (bool * pure_fun_translation) list
+    * Pure.trait_decl list
+    * Pure.trait_impl list =
   (* Debug *)
   log#ldebug (lazy "translate_crate_to_pure");
 
@@ -368,13 +373,28 @@ let translate_crate_to_pure (crate : A.crate) :
       (A.FunDeclId.Map.values crate.functions)
   in
 
+  (* Translate the trait declarations *)
+  let type_infos = trans_ctx.type_context.type_infos in
+  let trait_decls =
+    List.map
+      (SymbolicToPure.translate_trait_decl type_infos)
+      (T.TraitDeclId.Map.values trans_ctx.trait_decls_context.trait_decls)
+  in
+
+  (* Translate the trait implementations *)
+  let trait_impls =
+    List.map
+      (SymbolicToPure.translate_trait_impl type_infos)
+      (T.TraitImplId.Map.values trans_ctx.trait_impls_context.trait_impls)
+  in
+
   (* Apply the micro-passes *)
   let pure_translations =
     Micro.apply_passes_to_pure_fun_translations trans_ctx pure_translations
   in
 
   (* Return *)
-  (trans_ctx, type_decls, pure_translations)
+  (trans_ctx, type_decls, pure_translations, trait_decls, trait_impls)
 
 (** Extraction context *)
 type gen_ctx = {
@@ -735,14 +755,8 @@ let export_functions_group (fmt : Format.formatter) (config : gen_config)
 let export_trait_decl (fmt : Format.formatter) (_config : gen_config)
     (ctx : gen_ctx) (trait_decl_id : Pure.trait_decl_id) : unit =
   let trait_decl =
-    T.TraitDeclId.Map.find trait_decl_id
-      ctx.extract_ctx.trans_ctx.trait_decls_context.trait_decls
+    T.TraitDeclId.Map.find trait_decl_id ctx.extract_ctx.trans_trait_decls
   in
-  (* We translate the trait declaration on the fly (note that
-     trait declarations do not directly contain functions, constants,
-     etc.: they simply refer to them). *)
-  let type_infos = ctx.extract_ctx.trans_ctx.type_context.type_infos in
-  let trait_decl = SymbolicToPure.translate_trait_decl type_infos trait_decl in
   let ctx = ctx.extract_ctx in
   let ctx = { ctx with trait_decl_id = Some trait_decl.def_id } in
   Extract.extract_trait_decl ctx fmt trait_decl
@@ -751,14 +765,8 @@ let export_trait_decl (fmt : Format.formatter) (_config : gen_config)
 let export_trait_impl (fmt : Format.formatter) (_config : gen_config)
     (ctx : gen_ctx) (trait_impl_id : Pure.trait_impl_id) : unit =
   let trait_impl =
-    T.TraitImplId.Map.find trait_impl_id
-      ctx.extract_ctx.trans_ctx.trait_impls_context.trait_impls
+    T.TraitImplId.Map.find trait_impl_id ctx.extract_ctx.trans_trait_impls
   in
-  (* We translate the trait implementation on the fly (note that
-     trait implementations do not directly contain functions, constants,
-     etc.: they simply refer to them). *)
-  let type_infos = ctx.extract_ctx.trans_ctx.type_context.type_infos in
-  let trait_impl = SymbolicToPure.translate_trait_impl type_infos trait_impl in
   Extract.extract_trait_impl ctx.extract_ctx fmt trait_impl
 
 (** A generic utility to generate the extracted definitions: as we may want to
@@ -978,7 +986,9 @@ let extract_file (config : gen_config) (ctx : gen_ctx) (fi : extract_file_info)
 let translate_crate (filename : string) (dest_dir : string) (crate : A.crate) :
     unit =
   (* Translate the module to the pure AST *)
-  let trans_ctx, trans_types, trans_funs = translate_crate_to_pure crate in
+  let trans_ctx, trans_types, trans_funs, trans_trait_decls, trans_trait_impls =
+    translate_crate_to_pure crate
+  in
 
   (* Initialize the extraction context - for now we extract only to F*.
    * We initialize the names map by registering the keywords used in the
@@ -997,6 +1007,18 @@ let translate_crate (filename : string) (dest_dir : string) (crate : A.crate) :
   in
   (* Put everything in the context *)
   let ctx =
+    let trans_trait_decls =
+      T.TraitDeclId.Map.of_list
+        (List.map
+           (fun (d : Pure.trait_decl) -> (d.def_id, d))
+           trans_trait_decls)
+    in
+    let trans_trait_impls =
+      T.TraitImplId.Map.of_list
+        (List.map
+           (fun (d : Pure.trait_impl) -> (d.def_id, d))
+           trans_trait_impls)
+    in
     {
       ExtractBase.trans_ctx;
       names_map;
@@ -1008,6 +1030,8 @@ let translate_crate (filename : string) (dest_dir : string) (crate : A.crate) :
       fun_name_info = PureUtils.RegularFunIdMap.empty;
       trait_decl_id = None (* None by default *);
       is_provided_method = false (* false by default *);
+      trans_trait_decls;
+      trans_trait_impls;
     }
   in
 
@@ -1034,7 +1058,7 @@ let translate_crate (filename : string) (dest_dir : string) (crate : A.crate) :
   in
   let rec_functions = PureUtils.FunLoopIdSet.of_list rec_functions in
 
-  (* Register unique names for all the top-level types, globals and functions.
+  (* Register unique names for all the top-level types, globals, functions...
    * Note that the order in which we generate the names doesn't matter:
    * we just need to generate a mapping from identifier to name, and make
    * sure there are no name clashes. *)
@@ -1071,6 +1095,16 @@ let translate_crate (filename : string) (dest_dir : string) (crate : A.crate) :
       (A.GlobalDeclId.Map.values crate.globals)
   in
 
+  let ctx =
+    List.fold_left Extract.extract_trait_decl_register_names ctx
+      trans_trait_decls
+  in
+
+  let ctx =
+    List.fold_left Extract.extract_trait_impl_register_names ctx
+      trans_trait_impls
+  in
+
   (* Open the output file *)
   (* First compute the filename by replacing the extension and converting the
    * case (rust module names are snake case) *)