Merge trans_ctx and decls_ctx

7 files changed, 53 insertions, 120 deletions

diff --git a/compiler/Contexts.ml b/compiler/Contexts.ml
index 65760d94..a5bc7dc0 100644
--- a/compiler/Contexts.ml
+++ b/compiler/Contexts.ml
@@ -255,7 +255,11 @@ type type_context = {
 }
 [@@deriving show]
 
-type fun_context = { fun_decls : fun_decl FunDeclId.Map.t } [@@deriving show]
+type fun_context = {
+  fun_decls : fun_decl FunDeclId.Map.t;
+  fun_infos : FunsAnalysis.fun_info FunDeclId.Map.t;
+}
+[@@deriving show]
 
 type global_context = { global_decls : global_decl GlobalDeclId.Map.t }
 [@@deriving show]
diff --git a/compiler/Extract.ml b/compiler/Extract.ml
index e841082b..596fa013 100644
--- a/compiler/Extract.ml
+++ b/compiler/Extract.ml
@@ -856,9 +856,7 @@ let mk_formatter (ctx : trans_ctx) (crate_name : string)
             | Assumed Range -> "r"
             | Assumed State -> ConstStrings.state_basename
             | AdtId adt_id ->
-                let def =
-                  TypeDeclId.Map.find adt_id ctx.type_context.type_decls
-                in
+                let def = TypeDeclId.Map.find adt_id ctx.type_ctx.type_decls in
                 (* Derive the var name from the last ident of the type name
                  * Ex.: ["hashmap"; "HashMap"] ~~> "HashMap" -> "hash_map" -> "hm"
                  *)
@@ -3115,9 +3113,7 @@ and extract_StructUpdate (ctx : extraction_ctx) (fmt : F.formatter)
   let extract_as_unit =
     match (!backend, supd.struct_id) with
     | HOL4, AdtId adt_id ->
-        let d =
-          TypeDeclId.Map.find adt_id ctx.trans_ctx.type_context.type_decls
-        in
+        let d = TypeDeclId.Map.find adt_id ctx.trans_ctx.type_ctx.type_decls in
         d.kind = Struct []
     | _ -> false
   in
diff --git a/compiler/ExtractBase.ml b/compiler/ExtractBase.ml
index a81ec052..1586e6ed 100644
--- a/compiler/ExtractBase.ml
+++ b/compiler/ExtractBase.ml
@@ -709,10 +709,10 @@ type extraction_ctx = {
     instance).
  *)
 let id_to_string (id : id) (ctx : extraction_ctx) : string =
-  let global_decls = ctx.trans_ctx.global_context.global_decls in
-  let fun_decls = ctx.trans_ctx.fun_context.fun_decls in
-  let type_decls = ctx.trans_ctx.type_context.type_decls in
-  let trait_decls = ctx.trans_ctx.trait_decls_context.trait_decls in
+  let global_decls = ctx.trans_ctx.global_ctx.global_decls in
+  let fun_decls = ctx.trans_ctx.fun_ctx.fun_decls in
+  let type_decls = ctx.trans_ctx.type_ctx.type_decls in
+  let trait_decls = ctx.trans_ctx.trait_decls_ctx.trait_decls in
   let trait_decl_id_to_string (id : A.TraitDeclId.id) : string =
     let trait_name =
       Print.fun_name_to_string (A.TraitDeclId.Map.find id trait_decls).name
@@ -1293,9 +1293,7 @@ let ctx_compute_fun_name (trans_group : pure_fun_translation) (def : fun_decl)
     (ctx : extraction_ctx) : string =
   (* Lookup the LLBC def to compute the region group information *)
   let def_id = def.def_id in
-  let llbc_def =
-    A.FunDeclId.Map.find def_id ctx.trans_ctx.fun_context.fun_decls
-  in
+  let llbc_def = A.FunDeclId.Map.find def_id ctx.trans_ctx.fun_ctx.fun_decls in
   let sg = llbc_def.signature in
   let num_rgs = List.length sg.regions_hierarchy in
   let { keep_fwd; fwd = _; backs } = trans_group in
diff --git a/compiler/Interpreter.ml b/compiler/Interpreter.ml
index 4ce6dae8..752d6f2f 100644
--- a/compiler/Interpreter.ml
+++ b/compiler/Interpreter.ml
@@ -26,7 +26,10 @@ let compute_contexts (m : A.crate) : C.decls_ctx =
     TypesAnalysis.analyze_type_declarations type_decls type_decls_list
   in
   let type_ctx = { C.type_decls_groups; type_decls; type_infos } in
-  let fun_ctx = { C.fun_decls } in
+  let fun_infos =
+    FunsAnalysis.analyze_module m fun_decls global_decls !Config.use_state
+  in
+  let fun_ctx = { C.fun_decls; fun_infos } in
   let global_ctx = { C.global_decls } in
   let trait_decls_ctx = { C.trait_decls } in
   let trait_impls_ctx = { C.trait_impls } in
@@ -567,7 +570,8 @@ module Test = struct
   (** Test a unit function (taking no arguments) by evaluating it in an empty
       environment.
    *)
-  let test_unit_function (crate : A.crate) (fid : A.FunDeclId.id) : unit =
+  let test_unit_function (crate : A.crate) (decls_ctx : C.decls_ctx)
+      (fid : A.FunDeclId.id) : unit =
     (* Retrieve the function declaration *)
     let fdef = A.FunDeclId.Map.find fid crate.functions in
     let body = Option.get fdef.body in
@@ -581,7 +585,6 @@ module Test = struct
     assert (body.A.arg_count = 0);
 
     (* Create the evaluation context *)
-    let decls_ctx = compute_contexts crate in
     let ctx = initialize_eval_context decls_ctx [] [] [] in
 
     (* Insert the (uninitialized) local variables *)
@@ -620,8 +623,9 @@ module Test = struct
         (fun _ -> fun_decl_is_transparent_unit)
         crate.functions
     in
+    let decls_ctx = compute_contexts crate in
     let test_unit_fun _ (def : A.fun_decl) : unit =
-      test_unit_function crate def.A.def_id
+      test_unit_function crate decls_ctx def.A.def_id
     in
     A.FunDeclId.Map.iter test_unit_fun unit_funs
 end
diff --git a/compiler/PureMicroPasses.ml b/compiler/PureMicroPasses.ml
index 53148dbb..2130d5c2 100644
--- a/compiler/PureMicroPasses.ml
+++ b/compiler/PureMicroPasses.ml
@@ -586,9 +586,7 @@ let intro_struct_updates (ctx : trans_ctx) (def : fun_decl) : fun_decl =
                   generics = _;
                 } ->
                 (* Lookup the def *)
-                let decl =
-                  TypeDeclId.Map.find adt_id ctx.type_context.type_decls
-                in
+                let decl = TypeDeclId.Map.find adt_id ctx.type_ctx.type_decls in
                 (* Check that there are as many arguments as there are fields - note
                    that the def should have a body (otherwise we couldn't use the
                    constructor) *)
@@ -597,8 +595,7 @@ let intro_struct_updates (ctx : trans_ctx) (def : fun_decl) : fun_decl =
                   (* Check if the definition is recursive *)
                   let is_rec =
                     match
-                      TypeDeclId.Map.find adt_id
-                        ctx.type_context.type_decls_groups
+                      TypeDeclId.Map.find adt_id ctx.type_ctx.type_decls_groups
                     with
                     | NonRec _ -> false
                     | Rec _ -> true
@@ -796,7 +793,7 @@ let expression_contains_child_call_in_all_paths (ctx : trans_ctx)
                       | FunId fun_id -> fun_id
                       | TraitMethod (_, _, fun_decl_id) -> A.Regular fun_decl_id
                     in
-                    LlbcAstUtils.lookup_fun_sig id0 ctx.fun_context.fun_decls
+                    LlbcAstUtils.lookup_fun_sig id0 ctx.fun_ctx.fun_decls
                   in
                   (* Compute the set of ancestors of the function in call1 *)
                   let call1_ancestors =
@@ -1094,7 +1091,7 @@ let simplify_aggregates (ctx : trans_ctx) (def : fun_decl) : fun_decl =
                 (* This is a struct *)
                 (* Retrieve the definiton, to find how many fields there are *)
                 let adt_decl =
-                  TypeDeclId.Map.find adt_id ctx.type_context.type_decls
+                  TypeDeclId.Map.find adt_id ctx.type_ctx.type_decls
                 in
                 let fields =
                   match adt_decl.kind with
diff --git a/compiler/Translate.ml b/compiler/Translate.ml
index 4a4affea..13e339ea 100644
--- a/compiler/Translate.ml
+++ b/compiler/Translate.ml
@@ -29,34 +29,12 @@ let translate_function_to_symbolics (trans_ctx : trans_ctx) (fdef : A.fun_decl)
       ("translate_function_to_symbolics: "
       ^ Print.fun_name_to_string fdef.A.name));
 
-  let {
-    type_context;
-    fun_context;
-    global_context;
-    trait_decls_context;
-    trait_impls_context;
-  } =
-    trans_ctx
-  in
-  let fun_context = { C.fun_decls = fun_context.fun_decls } in
-
-  (* TODO: we should merge trans_ctx and decls_ctx *)
-  let decls_ctx : C.decls_ctx =
-    {
-      C.type_ctx = type_context;
-      fun_ctx = fun_context;
-      global_ctx = global_context;
-      trait_decls_ctx = trait_decls_context;
-      trait_impls_ctx = trait_impls_context;
-    }
-  in
-
   match fdef.body with
   | None -> None
   | Some _ ->
       (* Evaluate *)
       let synthesize = true in
-      let inputs, symb = evaluate_function_symbolic synthesize decls_ctx fdef in
+      let inputs, symb = evaluate_function_symbolic synthesize trans_ctx fdef in
       Some (inputs, Option.get symb)
 
 (** Translate a function, by generating its forward and backward translations.
@@ -74,15 +52,6 @@ let translate_function_to_pure (trans_ctx : trans_ctx)
     (lazy
       ("translate_function_to_pure: " ^ Print.fun_name_to_string fdef.A.name));
 
-  let {
-    type_context;
-    fun_context;
-    global_context;
-    trait_decls_context;
-    trait_impls_context;
-  } =
-    trans_ctx
-  in
   let def_id = fdef.def_id in
 
   (* Compute the symbolic ASTs, if the function is transparent *)
@@ -107,25 +76,25 @@ let translate_function_to_pure (trans_ctx : trans_ctx)
       (List.filter_map
          (fun (tid, g) ->
            match g with Charon.GAst.NonRec _ -> None | Rec _ -> Some tid)
-         (T.TypeDeclId.Map.bindings trans_ctx.type_context.type_decls_groups))
+         (T.TypeDeclId.Map.bindings trans_ctx.type_ctx.type_decls_groups))
   in
   let type_context =
     {
-      SymbolicToPure.type_infos = type_context.type_infos;
-      llbc_type_decls = type_context.type_decls;
+      SymbolicToPure.type_infos = trans_ctx.type_ctx.type_infos;
+      llbc_type_decls = trans_ctx.type_ctx.type_decls;
       type_decls = pure_type_decls;
       recursive_decls = recursive_type_decls;
     }
   in
   let fun_context =
     {
-      SymbolicToPure.llbc_fun_decls = fun_context.fun_decls;
+      SymbolicToPure.llbc_fun_decls = trans_ctx.fun_ctx.fun_decls;
       fun_sigs;
-      fun_infos = fun_context.fun_infos;
+      fun_infos = trans_ctx.fun_ctx.fun_infos;
     }
   in
   let global_context =
-    { SymbolicToPure.llbc_global_decls = global_context.global_decls }
+    { SymbolicToPure.llbc_global_decls = trans_ctx.global_ctx.global_decls }
   in
 
   (* Compute the set of loops, and find better ids for them (starting at 0).
@@ -173,8 +142,8 @@ let translate_function_to_pure (trans_ctx : trans_ctx)
       type_context;
       fun_context;
       global_context;
-      trait_decls_ctx = trait_decls_context.trait_decls;
-      trait_impls_ctx = trait_impls_context.trait_impls;
+      trait_decls_ctx = trans_ctx.trait_decls_ctx.trait_decls;
+      trait_impls_ctx = trans_ctx.trait_impls_ctx.trait_impls;
       fun_decl = fdef;
       forward_inputs = [];
       (* Empty for now *)
@@ -311,22 +280,8 @@ let translate_crate_to_pure (crate : A.crate) :
   (* Debug *)
   log#ldebug (lazy "translate_crate_to_pure");
 
-  (* Compute the type and function contexts *)
-  let decls_ctx = compute_contexts crate in
-  let fun_infos =
-    FA.analyze_module crate decls_ctx.fun_ctx.C.fun_decls
-      decls_ctx.global_ctx.C.global_decls !Config.use_state
-  in
-  let fun_context = { fun_decls = decls_ctx.fun_ctx.fun_decls; fun_infos } in
-  let trans_ctx =
-    {
-      type_context = decls_ctx.type_ctx;
-      fun_context;
-      global_context = decls_ctx.global_ctx;
-      trait_decls_context = decls_ctx.trait_decls_ctx;
-      trait_impls_context = decls_ctx.trait_impls_ctx;
-    }
-  in
+  (* Compute the translation context *)
+  let trans_ctx = compute_contexts crate in
 
   (* Translate all the type definitions *)
   let type_decls =
@@ -362,8 +317,8 @@ let translate_crate_to_pure (crate : A.crate) :
   in
   let sigs = List.append assumed_sigs local_sigs in
   let fun_sigs =
-    SymbolicToPure.translate_fun_signatures fun_context.fun_infos
-      decls_ctx.type_ctx.type_infos sigs
+    SymbolicToPure.translate_fun_signatures trans_ctx.fun_ctx.fun_infos
+      trans_ctx.type_ctx.type_infos sigs
   in
 
   (* Translate all the *transparent* functions *)
@@ -374,18 +329,18 @@ let translate_crate_to_pure (crate : A.crate) :
   in
 
   (* Translate the trait declarations *)
-  let type_infos = trans_ctx.type_context.type_infos in
+  let type_infos = trans_ctx.type_ctx.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)
+      (T.TraitDeclId.Map.values trans_ctx.trait_decls_ctx.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)
+      (T.TraitImplId.Map.values trans_ctx.trait_impls_ctx.trait_impls)
   in
 
   (* Apply the micro-passes *)
@@ -554,7 +509,7 @@ let export_types_group (fmt : Format.formatter) (config : gen_config)
  *)
 let export_global (fmt : Format.formatter) (config : gen_config) (ctx : gen_ctx)
     (id : A.GlobalDeclId.id) : unit =
-  let global_decls = ctx.trans_ctx.global_context.global_decls in
+  let global_decls = ctx.trans_ctx.global_ctx.global_decls in
   let global = A.GlobalDeclId.Map.find id global_decls in
   let trans = A.FunDeclId.Map.find global.body_id ctx.trans_funs in
   assert (trans.fwd.loops = []);
diff --git a/compiler/TranslateCore.ml b/compiler/TranslateCore.ml
index f31dc458..3427fd43 100644
--- a/compiler/TranslateCore.ml
+++ b/compiler/TranslateCore.ml
@@ -10,27 +10,7 @@ module FA = FunsAnalysis
 (** The local logger *)
 let log = L.translate_log
 
-type type_context = C.type_context [@@deriving show]
-
-type fun_context = {
-  fun_decls : A.fun_decl A.FunDeclId.Map.t;
-  fun_infos : FA.fun_info A.FunDeclId.Map.t;
-}
-[@@deriving show]
-
-type trait_decls_context = C.trait_decls_context [@@deriving show]
-type trait_impls_context = C.trait_impls_context [@@deriving show]
-type global_context = C.global_context [@@deriving show]
-
-(* TODO: we should use Contexts.decls_ctx *)
-type trans_ctx = {
-  type_context : type_context;
-  fun_context : fun_context;
-  global_context : global_context;
-  trait_decls_context : trait_decls_context;
-  trait_impls_context : trait_impls_context;
-}
-
+type trans_ctx = C.decls_ctx [@@deriving show]
 type fun_and_loops = { f : Pure.fun_decl; loops : Pure.fun_decl list }
 type pure_fun_translation_no_loops = Pure.fun_decl * Pure.fun_decl list
 
@@ -50,10 +30,10 @@ let trans_ctx_to_type_formatter (ctx : trans_ctx)
     (type_params : Pure.type_var list)
     (const_generic_params : Pure.const_generic_var list) :
     PrintPure.type_formatter =
-  let type_decls = ctx.type_context.type_decls in
-  let global_decls = ctx.global_context.global_decls in
-  let trait_decls = ctx.trait_decls_context.trait_decls in
-  let trait_impls = ctx.trait_impls_context.trait_impls in
+  let type_decls = ctx.type_ctx.type_decls in
+  let global_decls = ctx.global_ctx.global_decls in
+  let trait_decls = ctx.trait_decls_ctx.trait_decls in
+  let trait_impls = ctx.trait_impls_ctx.trait_impls in
   PrintPure.mk_type_formatter type_decls global_decls trait_decls trait_impls
     type_params const_generic_params
 
@@ -66,17 +46,17 @@ let type_decl_to_string (ctx : trans_ctx) (def : Pure.type_decl) : string =
 
 let type_id_to_string (ctx : trans_ctx) (id : Pure.TypeDeclId.id) : string =
   Print.fun_name_to_string
-    (Pure.TypeDeclId.Map.find id ctx.type_context.type_decls).name
+    (Pure.TypeDeclId.Map.find id ctx.type_ctx.type_decls).name
 
 let trans_ctx_to_ast_formatter (ctx : trans_ctx)
     (type_params : Pure.type_var list)
     (const_generic_params : Pure.const_generic_var list) :
     PrintPure.ast_formatter =
-  let type_decls = ctx.type_context.type_decls in
-  let fun_decls = ctx.fun_context.fun_decls in
-  let global_decls = ctx.global_context.global_decls in
-  let trait_decls = ctx.trait_decls_context.trait_decls in
-  let trait_impls = ctx.trait_impls_context.trait_impls in
+  let type_decls = ctx.type_ctx.type_decls in
+  let fun_decls = ctx.fun_ctx.fun_decls in
+  let global_decls = ctx.global_ctx.global_decls in
+  let trait_decls = ctx.trait_decls_ctx.trait_decls in
+  let trait_impls = ctx.trait_impls_ctx.trait_impls in
   PrintPure.mk_ast_formatter type_decls fun_decls global_decls trait_decls
     trait_impls type_params const_generic_params
 
@@ -95,5 +75,4 @@ let fun_decl_to_string (ctx : trans_ctx) (def : Pure.fun_decl) : string =
   PrintPure.fun_decl_to_string fmt def
 
 let fun_decl_id_to_string (ctx : trans_ctx) (id : A.FunDeclId.id) : string =
-  Print.fun_name_to_string
-    (A.FunDeclId.Map.find id ctx.fun_context.fun_decls).name
+  Print.fun_name_to_string (A.FunDeclId.Map.find id ctx.fun_ctx.fun_decls).name