Merge pull request #45 from AeneasVerif/son_merge_types

Big cleanup

1 files changed, 149 insertions, 130 deletions

diff --git a/compiler/Interpreter.ml b/compiler/Interpreter.ml
index 24ff4808..c2e47da9 100644
--- a/compiler/Interpreter.ml
+++ b/compiler/Interpreter.ml
@@ -4,19 +4,22 @@ open InterpreterProjectors
 open InterpreterBorrows
 open InterpreterStatements
 open LlbcAstUtils
-module L = Logging
-module T = Types
-module A = LlbcAst
+open Types
+open TypesUtils
+open Values
+open LlbcAst
+open Contexts
+open SynthesizeSymbolic
 module SA = SymbolicAst
 
 (** The local logger *)
-let log = L.interpreter_log
+let log = Logging.interpreter_log
 
-let compute_contexts (m : A.crate) : C.decls_ctx =
+let compute_contexts (m : crate) : decls_ctx =
   let type_decls_list, _, _, _, _ = split_declarations m.declarations in
-  let type_decls = m.types in
-  let fun_decls = m.functions in
-  let global_decls = m.globals in
+  let type_decls = m.type_decls in
+  let fun_decls = m.fun_decls in
+  let global_decls = m.global_decls in
   let trait_decls = m.trait_decls in
   let trait_impls = m.trait_impls in
   let type_decls_groups, _, _, _, _ =
@@ -25,15 +28,19 @@ let compute_contexts (m : A.crate) : C.decls_ctx =
   let type_infos =
     TypesAnalysis.analyze_type_declarations type_decls type_decls_list
   in
-  let type_ctx = { C.type_decls_groups; type_decls; type_infos } in
+  let type_ctx = { type_decls_groups; type_decls; type_infos } 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
-  { C.type_ctx; fun_ctx; global_ctx; trait_decls_ctx; trait_impls_ctx }
+  let regions_hierarchies =
+    RegionsHierarchy.compute_regions_hierarchies type_decls fun_decls
+      global_decls trait_decls trait_impls
+  in
+  let fun_ctx = { fun_decls; fun_infos; regions_hierarchies } in
+  let global_ctx = { global_decls } in
+  let trait_decls_ctx = { trait_decls } in
+  let trait_impls_ctx = { trait_impls } in
+  { type_ctx; fun_ctx; global_ctx; trait_decls_ctx; trait_impls_ctx }
 
 (** Small helper.
 
@@ -41,15 +48,14 @@ let compute_contexts (m : A.crate) : C.decls_ctx =
     to compute a normalization map (for the associated types) and that we added
     it in the context.
  *)
-let normalize_inst_fun_sig (ctx : C.eval_ctx) (sg : A.inst_fun_sig) :
-    A.inst_fun_sig =
-  let { A.regions_hierarchy = _; trait_type_constraints = _; inputs; output } =
+let normalize_inst_fun_sig (ctx : eval_ctx) (sg : inst_fun_sig) : inst_fun_sig =
+  let { regions_hierarchy = _; trait_type_constraints = _; inputs; output } =
     sg
   in
-  let norm = AssociatedTypes.ctx_normalize_rty ctx in
+  let norm = AssociatedTypes.ctx_normalize_ty ctx in
   let inputs = List.map norm inputs in
   let output = norm output in
-  { sg with A.inputs; output }
+  { sg with inputs; output }
 
 (** Instantiate a function signature for a symbolic execution.
 
@@ -61,27 +67,28 @@ let normalize_inst_fun_sig (ctx : C.eval_ctx) (sg : A.inst_fun_sig) :
     clauses (we are not considering a function call, so we don't need to
     normalize because a trait clause was instantiated with a specific trait ref).
  *)
-let symbolic_instantiate_fun_sig (ctx : C.eval_ctx) (sg : A.fun_sig)
-    (kind : A.fun_kind) : C.eval_ctx * A.inst_fun_sig =
+let symbolic_instantiate_fun_sig (ctx : eval_ctx) (sg : fun_sig)
+    (regions_hierarchy : region_groups) (kind : fun_kind) :
+    eval_ctx * inst_fun_sig =
   let tr_self =
     match kind with
-    | RegularKind | TraitMethodImpl _ -> T.UnknownTrait __FUNCTION__
-    | TraitMethodDecl _ | TraitMethodProvided _ -> T.Self
+    | RegularKind | TraitMethodImpl _ -> UnknownTrait __FUNCTION__
+    | TraitMethodDecl _ | TraitMethodProvided _ -> Self
   in
   let generics =
-    let { T.regions; types; const_generics; trait_clauses } = sg.generics in
-    let regions = List.map (fun _ -> T.Erased) regions in
-    let types = List.map (fun (v : T.type_var) -> T.TypeVar v.T.index) types in
+    let { regions; types; const_generics; trait_clauses } = sg.generics in
+    let regions = List.map (fun _ -> RErased) regions in
+    let types = List.map (fun (v : type_var) -> TVar v.index) types in
     let const_generics =
-      List.map
-        (fun (v : T.const_generic_var) -> T.ConstGenericVar v.T.index)
-        const_generics
+      List.map (fun (v : const_generic_var) -> CgVar v.index) const_generics
     in
     (* Annoying that we have to generate this substitution here *)
     let r_subst _ = raise (Failure "Unexpected region") in
-    let ty_subst = Subst.make_type_subst_from_vars sg.generics.types types in
+    let ty_subst =
+      Substitute.make_type_subst_from_vars sg.generics.types types
+    in
     let cg_subst =
-      Subst.make_const_generic_subst_from_vars sg.generics.const_generics
+      Substitute.make_const_generic_subst_from_vars sg.generics.const_generics
         const_generics
     in
     (* TODO: some clauses may use the types of other clauses, so we may have to
@@ -110,42 +117,39 @@ let symbolic_instantiate_fun_sig (ctx : C.eval_ctx) (sg : A.fun_sig)
        ]}
     *)
     (* We will need to update the trait refs map while we perform the instantiations *)
-    let mk_tr_subst
-        (tr_map : T.erased_region T.trait_instance_id T.TraitClauseId.Map.t)
-        clause_id : T.erased_region T.trait_instance_id =
-      match T.TraitClauseId.Map.find_opt clause_id tr_map with
+    let mk_tr_subst (tr_map : trait_instance_id TraitClauseId.Map.t) clause_id :
+        trait_instance_id =
+      match TraitClauseId.Map.find_opt clause_id tr_map with
       | Some tr -> tr
       | None -> raise (Failure "Local trait clause not found")
     in
     let mk_subst tr_map =
       let tr_subst = mk_tr_subst tr_map in
-      { Subst.r_subst; ty_subst; cg_subst; tr_subst; tr_self }
+      { Substitute.r_subst; ty_subst; cg_subst; tr_subst; tr_self }
     in
     let _, trait_refs =
       List.fold_left_map
-        (fun tr_map (c : T.trait_clause) ->
+        (fun tr_map (c : trait_clause) ->
           let subst = mk_subst tr_map in
-          let { T.trait_id = trait_decl_id; generics; _ } = c in
-          let generics = Subst.generic_args_substitute subst generics in
-          let trait_decl_ref = { T.trait_decl_id; decl_generics = generics } in
+          let { trait_id = trait_decl_id; clause_generics; _ } = c in
+          let generics =
+            Substitute.generic_args_substitute subst clause_generics
+          in
+          let trait_decl_ref = { trait_decl_id; decl_generics = generics } in
           (* Note that because we directly refer to the clause, we give it
              empty generics *)
-          let trait_id = T.Clause c.clause_id in
+          let trait_id = Clause c.clause_id in
           let trait_ref =
-            {
-              T.trait_id;
-              generics = TypesUtils.mk_empty_generic_args;
-              trait_decl_ref;
-            }
+            { trait_id; generics = empty_generic_args; trait_decl_ref }
           in
           (* Update the traits map *)
-          let tr_map = T.TraitClauseId.Map.add c.T.clause_id trait_id tr_map in
+          let tr_map = TraitClauseId.Map.add c.clause_id trait_id tr_map in
           (tr_map, trait_ref))
-        T.TraitClauseId.Map.empty trait_clauses
+        TraitClauseId.Map.empty trait_clauses
     in
-    { T.regions; types; const_generics; trait_refs }
+    { regions; types; const_generics; trait_refs }
   in
-  let inst_sg = instantiate_fun_sig ctx generics tr_self sg in
+  let inst_sg = instantiate_fun_sig ctx generics tr_self sg regions_hierarchy in
   (* Compute the normalization maps *)
   let ctx =
     AssociatedTypes.ctx_add_norm_trait_types_from_preds ctx
@@ -169,8 +173,8 @@ let symbolic_instantiate_fun_sig (ctx : C.eval_ctx) (sg : A.fun_sig)
     - the list of symbolic values introduced for the input values
     - the instantiated function signature
  *)
-let initialize_symbolic_context_for_fun (ctx : C.decls_ctx) (fdef : A.fun_decl)
-    : C.eval_ctx * V.symbolic_value list * A.inst_fun_sig =
+let initialize_symbolic_context_for_fun (ctx : decls_ctx) (fdef : fun_decl) :
+    eval_ctx * symbolic_value list * inst_fun_sig =
   (* The abstractions are not initialized the same way as for function
    * calls: they contain *loan* projectors, because they "provide" us
    * with the input values (which behave as if they had been returned
@@ -184,8 +188,11 @@ let initialize_symbolic_context_for_fun (ctx : C.decls_ctx) (fdef : A.fun_decl)
    * *)
   let sg = fdef.signature in
   (* Create the context *)
+  let regions_hierarchy =
+    FunIdMap.find (FRegular fdef.def_id) ctx.fun_ctx.regions_hierarchies
+  in
   let region_groups =
-    List.map (fun (g : T.region_var_group) -> g.id) sg.regions_hierarchy
+    List.map (fun (g : region_group) -> g.id) regions_hierarchy
   in
   let ctx =
     initialize_eval_context ctx region_groups sg.generics.types
@@ -195,17 +202,17 @@ let initialize_symbolic_context_for_fun (ctx : C.decls_ctx) (fdef : A.fun_decl)
      at the same time the normalization map for the associated types.
   *)
   let ctx, inst_sg =
-    symbolic_instantiate_fun_sig ctx fdef.signature fdef.kind
+    symbolic_instantiate_fun_sig ctx fdef.signature regions_hierarchy fdef.kind
   in
   (* Create fresh symbolic values for the inputs *)
   let input_svs =
-    List.map (fun ty -> mk_fresh_symbolic_value V.SynthInput ty) inst_sg.inputs
+    List.map (fun ty -> mk_fresh_symbolic_value SynthInput ty) inst_sg.inputs
   in
   (* Initialize the abstractions as empty (i.e., with no avalues) abstractions *)
-  let call_id = C.fresh_fun_call_id () in
-  assert (call_id = V.FunCallId.zero);
-  let compute_abs_avalues (abs : V.abs) (ctx : C.eval_ctx) :
-      C.eval_ctx * V.typed_avalue list =
+  let call_id = fresh_fun_call_id () in
+  assert (call_id = FunCallId.zero);
+  let compute_abs_avalues (abs : abs) (ctx : eval_ctx) :
+      eval_ctx * typed_avalue list =
     (* Project over the values - we use *loan* projectors, as explained above *)
     let avalues =
       List.map (mk_aproj_loans_value_from_symbolic_value abs.regions) input_svs
@@ -215,8 +222,8 @@ let initialize_symbolic_context_for_fun (ctx : C.decls_ctx) (fdef : A.fun_decl)
   let region_can_end _ = false in
   let ctx =
     create_push_abstractions_from_abs_region_groups
-      (fun rg_id -> V.SynthInput rg_id)
-      inst_sg.A.regions_hierarchy region_can_end compute_abs_avalues ctx
+      (fun rg_id -> SynthInput rg_id)
+      inst_sg.regions_hierarchy region_can_end compute_abs_avalues ctx
   in
   (* Split the variables between return var, inputs and remaining locals *)
   let body = Option.get fdef.body in
@@ -225,12 +232,12 @@ let initialize_symbolic_context_for_fun (ctx : C.decls_ctx) (fdef : A.fun_decl)
     Collections.List.split_at (List.tl body.locals) body.arg_count
   in
   (* Push the return variable (initialized with ⊥) *)
-  let ctx = C.ctx_push_uninitialized_var ctx ret_var in
+  let ctx = ctx_push_uninitialized_var ctx ret_var in
   (* Push the input variables (initialized with symbolic values) *)
   let input_values = List.map mk_typed_value_from_symbolic_value input_svs in
-  let ctx = C.ctx_push_vars ctx (List.combine input_vars input_values) in
+  let ctx = ctx_push_vars ctx (List.combine input_vars input_values) in
   (* Push the remaining local variables (initialized with ⊥) *)
-  let ctx = C.ctx_push_uninitialized_vars ctx local_vars in
+  let ctx = ctx_push_uninitialized_vars ctx local_vars in
   (* Return *)
   (ctx, input_svs, inst_sg)
 
@@ -246,20 +253,19 @@ let initialize_symbolic_context_for_fun (ctx : C.decls_ctx) (fdef : A.fun_decl)
 
     [inside_loop]: [true] if we are *inside* a loop (result [EndContinue]).
 *)
-let evaluate_function_symbolic_synthesize_backward_from_return
-    (config : C.config) (fdef : A.fun_decl) (inst_sg : A.inst_fun_sig)
-    (back_id : T.RegionGroupId.id) (loop_id : V.LoopId.id option)
-    (is_regular_return : bool) (inside_loop : bool) (ctx : C.eval_ctx) :
-    SA.expression option =
+let evaluate_function_symbolic_synthesize_backward_from_return (config : config)
+    (fdef : fun_decl) (inst_sg : inst_fun_sig) (back_id : RegionGroupId.id)
+    (loop_id : LoopId.id option) (is_regular_return : bool) (inside_loop : bool)
+    (ctx : eval_ctx) : SA.expression option =
   log#ldebug
     (lazy
       ("evaluate_function_symbolic_synthesize_backward_from_return:"
      ^ "\n- fname: "
-      ^ Print.fun_name_to_string fdef.name
+      ^ Print.EvalCtx.name_to_string ctx fdef.name
       ^ "\n- back_id: "
-      ^ T.RegionGroupId.to_string back_id
+      ^ RegionGroupId.to_string back_id
       ^ "\n- loop_id: "
-      ^ Print.option_to_string V.LoopId.to_string loop_id
+      ^ Print.option_to_string LoopId.to_string loop_id
       ^ "\n- is_regular_return: "
       ^ Print.bool_to_string is_regular_return
       ^ "\n- inside_loop: "
@@ -270,9 +276,11 @@ let evaluate_function_symbolic_synthesize_backward_from_return
    * the return type. Note that it is important to re-generate
    * an instantiation of the signature, so that we use fresh
    * region ids for the return abstractions. *)
-  let sg = fdef.signature in
+  let regions_hierarchy =
+    FunIdMap.find (FRegular fdef.def_id) ctx.fun_context.regions_hierarchies
+  in
   let _, ret_inst_sg =
-    symbolic_instantiate_fun_sig ctx fdef.signature fdef.kind
+    symbolic_instantiate_fun_sig ctx fdef.signature regions_hierarchy fdef.kind
   in
   let ret_rty = ret_inst_sg.output in
   (* Move the return value out of the return variable *)
@@ -283,11 +291,11 @@ let evaluate_function_symbolic_synthesize_backward_from_return
    * will end - this will allow us to, first, mark the other return
    * regions as non-endable, and, second, end those parent regions in
    * proper order. *)
-  let parent_rgs = list_ancestor_region_groups sg back_id in
+  let parent_rgs = list_ancestor_region_groups regions_hierarchy back_id in
   let parent_input_abs_ids =
-    T.RegionGroupId.mapi
+    RegionGroupId.mapi
       (fun rg_id rg ->
-        if T.RegionGroupId.Set.mem rg_id parent_rgs then Some rg.T.id else None)
+        if RegionGroupId.Set.mem rg_id parent_rgs then Some rg.id else None)
       inst_sg.regions_hierarchy
   in
   let parent_input_abs_ids =
@@ -296,12 +304,12 @@ let evaluate_function_symbolic_synthesize_backward_from_return
 
   (* Insert the return value in the return abstractions (by applying
    * borrow projections) *)
-  let cf_consume_ret (ret_value : V.typed_value option) ctx =
+  let cf_consume_ret (ret_value : typed_value option) ctx =
     let ctx =
       if is_regular_return then (
         let ret_value = Option.get ret_value in
-        let compute_abs_avalues (abs : V.abs) (ctx : C.eval_ctx) :
-            C.eval_ctx * V.typed_avalue list =
+        let compute_abs_avalues (abs : abs) (ctx : eval_ctx) :
+            eval_ctx * typed_avalue list =
           let ctx, avalue =
             apply_proj_borrows_on_input_value config ctx abs.regions
               abs.ancestors_regions ret_value ret_rty
@@ -315,18 +323,15 @@ let evaluate_function_symbolic_synthesize_backward_from_return
          * that this is important for soundness: this is part of the borrow checking).
          * Also see the documentation of the [can_end] field of [abs] for more
          * information. *)
-        let parent_and_current_rgs =
-          T.RegionGroupId.Set.add back_id parent_rgs
-        in
+        let parent_and_current_rgs = RegionGroupId.Set.add back_id parent_rgs in
         let region_can_end rid =
-          T.RegionGroupId.Set.mem rid parent_and_current_rgs
+          RegionGroupId.Set.mem rid parent_and_current_rgs
         in
         assert (region_can_end back_id);
         let ctx =
           create_push_abstractions_from_abs_region_groups
-            (fun rg_id -> V.SynthRet rg_id)
-            ret_inst_sg.A.regions_hierarchy region_can_end compute_abs_avalues
-            ctx
+            (fun rg_id -> SynthRet rg_id)
+            ret_inst_sg.regions_hierarchy region_can_end compute_abs_avalues ctx
         in
         ctx)
       else ctx
@@ -347,16 +352,16 @@ let evaluate_function_symbolic_synthesize_backward_from_return
      *)
     let current_abs_id, end_fun_synth_input =
       let fun_abs_id =
-        (T.RegionGroupId.nth inst_sg.regions_hierarchy back_id).id
+        (RegionGroupId.nth inst_sg.regions_hierarchy back_id).id
       in
       if not inside_loop then (fun_abs_id, true)
       else
-        let pred (abs : V.abs) =
+        let pred (abs : abs) =
           match abs.kind with
-          | V.Loop (_, rg_id', kind) ->
+          | Loop (_, rg_id', kind) ->
               let rg_id' = Option.get rg_id' in
               let is_ret =
-                match kind with V.LoopSynthInput -> true | V.LoopCall -> false
+                match kind with LoopSynthInput -> true | LoopCall -> false
               in
               rg_id' = back_id && is_ret
           | _ -> false
@@ -378,24 +383,24 @@ let evaluate_function_symbolic_synthesize_backward_from_return
              }
            ]}
         *)
-        let abs = Option.get (C.ctx_find_abs ctx pred) in
+        let abs = Option.get (ctx_find_abs ctx pred) in
         (abs.abs_id, false)
     in
     log#ldebug
       (lazy
         ("evaluate_function_symbolic_synthesize_backward_from_return: ending \
           input abstraction: "
-        ^ V.AbstractionId.to_string current_abs_id));
+        ^ AbstractionId.to_string current_abs_id));
 
     (* Set the proper abstractions as endable *)
     let ctx =
       let visit_loop_abs =
         object
-          inherit [_] C.map_eval_ctx
+          inherit [_] map_eval_ctx
 
           method! visit_abs _ abs =
             match abs.kind with
-            | V.Loop (loop_id', rg_id', V.LoopSynthInput) ->
+            | Loop (loop_id', rg_id', LoopSynthInput) ->
                 (* We only allow to end the loop synth input abs for the region
                    group [rg_id] *)
                 assert (
@@ -406,11 +411,11 @@ let evaluate_function_symbolic_synthesize_backward_from_return
                 if rg_id' = back_id && inside_loop then
                   { abs with can_end = true }
                 else abs
-            | V.Loop (loop_id', _, V.LoopCall) ->
+            | Loop (loop_id', _, LoopCall) ->
                 (* We can end all the loop call abstractions *)
                 assert (loop_id = Some loop_id');
                 { abs with can_end = true }
-            | V.SynthInput rg_id' ->
+            | SynthInput rg_id' ->
                 if rg_id' = back_id && end_fun_synth_input then
                   { abs with can_end = true }
                 else abs
@@ -447,18 +452,26 @@ let evaluate_function_symbolic_synthesize_backward_from_return
       for the synthesis)
     - the symbolic AST generated by the symbolic execution
  *)
-let evaluate_function_symbolic (synthesize : bool) (ctx : C.decls_ctx)
-    (fdef : A.fun_decl) : V.symbolic_value list * SA.expression option =
+let evaluate_function_symbolic (synthesize : bool) (ctx : decls_ctx)
+    (fdef : fun_decl) : symbolic_value list * SA.expression option =
   (* Debug *)
-  let name_to_string () = Print.fun_name_to_string fdef.A.name in
+  let name_to_string () =
+    Print.Types.name_to_string
+      (Print.Contexts.decls_ctx_to_fmt_env ctx)
+      fdef.name
+  in
   log#ldebug (lazy ("evaluate_function_symbolic: " ^ name_to_string ()));
 
   (* Create the evaluation context *)
   let ctx, input_svs, inst_sg = initialize_symbolic_context_for_fun ctx fdef in
 
+  let regions_hierarchy =
+    FunIdMap.find (FRegular fdef.def_id) ctx.fun_context.regions_hierarchies
+  in
+
   (* Create the continuation to finish the evaluation *)
-  let config = C.mk_config C.SymbolicMode in
-  let cf_finish res ctx =
+  let config = mk_config SymbolicMode in
+  let cf_finish (res : statement_eval_res) (ctx : eval_ctx) =
     let ctx0 = ctx in
     log#ldebug
       (lazy
@@ -510,13 +523,13 @@ let evaluate_function_symbolic (synthesize : bool) (ctx : C.decls_ctx)
                  fdef inst_sg back_id loop_id is_regular_return inside_loop ctx)
           in
           let back_el =
-            T.RegionGroupId.mapi
+            RegionGroupId.mapi
               (fun gid _ -> (gid, finish_back_eval gid))
-              fdef.signature.regions_hierarchy
+              regions_hierarchy
           in
-          let back_el = T.RegionGroupId.Map.of_list back_el in
+          let back_el = RegionGroupId.Map.of_list back_el in
           (* Put everything together *)
-          S.synthesize_forward_end ctx0 None fwd_e back_el
+          synthesize_forward_end ctx0 None fwd_e back_el
         else None
     | EndEnterLoop (loop_id, loop_input_values)
     | EndContinue (loop_id, loop_input_values) ->
@@ -554,13 +567,13 @@ let evaluate_function_symbolic (synthesize : bool) (ctx : C.decls_ctx)
                  inside_loop ctx)
           in
           let back_el =
-            T.RegionGroupId.mapi
+            RegionGroupId.mapi
               (fun gid _ -> (gid, finish_back_eval gid))
-              fdef.signature.regions_hierarchy
+              regions_hierarchy
           in
-          let back_el = T.RegionGroupId.Map.of_list back_el in
+          let back_el = RegionGroupId.Map.of_list back_el in
           (* Put everything together *)
-          S.synthesize_forward_end ctx0 (Some loop_input_values) fwd_e back_el
+          synthesize_forward_end ctx0 (Some loop_input_values) fwd_e back_el
         else None
     | Panic ->
         (* Note that as we explore all the execution branches, one of
@@ -573,7 +586,7 @@ let evaluate_function_symbolic (synthesize : bool) (ctx : C.decls_ctx)
 
   (* Evaluate the function *)
   let symbolic =
-    eval_function_body config (Option.get fdef.A.body).body cf_finish ctx
+    eval_function_body config (Option.get fdef.body).body cf_finish ctx
   in
 
   (* Return *)
@@ -583,29 +596,33 @@ module Test = struct
   (** Test a unit function (taking no arguments) by evaluating it in an empty
       environment.
    *)
-  let test_unit_function (crate : A.crate) (decls_ctx : C.decls_ctx)
-      (fid : A.FunDeclId.id) : unit =
+  let test_unit_function (crate : crate) (decls_ctx : decls_ctx)
+      (fid : FunDeclId.id) : unit =
     (* Retrieve the function declaration *)
-    let fdef = A.FunDeclId.Map.find fid crate.functions in
+    let fdef = FunDeclId.Map.find fid crate.fun_decls in
     let body = Option.get fdef.body in
 
     (* Debug *)
     log#ldebug
-      (lazy ("test_unit_function: " ^ Print.fun_name_to_string fdef.A.name));
+      (lazy
+        ("test_unit_function: "
+        ^ Print.Types.name_to_string
+            (Print.Contexts.decls_ctx_to_fmt_env decls_ctx)
+            fdef.name));
 
     (* Sanity check - *)
-    assert (fdef.A.signature.generics = TypesUtils.mk_empty_generic_params);
-    assert (body.A.arg_count = 0);
+    assert (fdef.signature.generics = empty_generic_params);
+    assert (body.arg_count = 0);
 
     (* Create the evaluation context *)
     let ctx = initialize_eval_context decls_ctx [] [] [] in
 
     (* Insert the (uninitialized) local variables *)
-    let ctx = C.ctx_push_uninitialized_vars ctx body.A.locals in
+    let ctx = ctx_push_uninitialized_vars ctx body.locals in
 
     (* Create the continuation to check the function's result *)
-    let config = C.mk_config C.ConcreteMode in
-    let cf_check res ctx =
+    let config = mk_config ConcreteMode in
+    let cf_check (res : statement_eval_res) (ctx : eval_ctx) =
       match res with
       | Return ->
           (* Ok: drop the local variables and finish *)
@@ -615,7 +632,9 @@ module Test = struct
           raise
             (Failure
                ("Unit test failed (concrete execution) on: "
-               ^ Print.fun_name_to_string fdef.A.name))
+               ^ Print.Types.name_to_string
+                   (Print.Contexts.decls_ctx_to_fmt_env decls_ctx)
+                   fdef.name))
     in
 
     (* Evaluate the function *)
@@ -624,21 +643,21 @@ module Test = struct
 
   (** Small helper: return true if the function is a *transparent* unit function
       (no parameters, no arguments) - TODO: move *)
-  let fun_decl_is_transparent_unit (def : A.fun_decl) : bool =
+  let fun_decl_is_transparent_unit (def : fun_decl) : bool =
     Option.is_some def.body
-    && def.A.signature.generics = TypesUtils.mk_empty_generic_params
-    && def.A.signature.inputs = []
+    && def.signature.generics = empty_generic_params
+    && def.signature.inputs = []
 
   (** Test all the unit functions in a list of function definitions *)
-  let test_unit_functions (crate : A.crate) : unit =
+  let test_unit_functions (crate : crate) : unit =
     let unit_funs =
-      A.FunDeclId.Map.filter
+      FunDeclId.Map.filter
         (fun _ -> fun_decl_is_transparent_unit)
-        crate.functions
+        crate.fun_decls
     in
     let decls_ctx = compute_contexts crate in
-    let test_unit_fun _ (def : A.fun_decl) : unit =
-      test_unit_function crate decls_ctx def.A.def_id
+    let test_unit_fun _ (def : fun_decl) : unit =
+      test_unit_function crate decls_ctx def.def_id
     in
-    A.FunDeclId.Map.iter test_unit_fun unit_funs
+    FunDeclId.Map.iter test_unit_fun unit_funs
 end