Update type_context to have to have a type info field, use maps instead

of lists to store the types/functions definitions

6 files changed, 62 insertions, 36 deletions

diff --git a/src/Contexts.ml b/src/Contexts.ml
index f7a24d27..cd433890 100644
--- a/src/Contexts.ml
+++ b/src/Contexts.ml
@@ -138,13 +138,14 @@ let config_of_partial (mode : interpreter_mode) (config : partial_config) :
 
 type type_context = {
   type_defs_groups : M.type_declaration_group TypeDefId.Map.t;
-  type_defs : type_def list;
+  type_defs : type_def TypeDefId.Map.t;
+  type_infos : TypesAnalysis.type_infos;
 }
 [@@deriving show]
 
 type eval_ctx = {
   type_context : type_context;
-  fun_context : fun_def list;
+  fun_context : fun_def FunDefId.Map.t;
   type_vars : type_var list;
   env : env;
   ended_regions : RegionId.Set.t;
@@ -172,11 +173,11 @@ let ctx_lookup_binder (ctx : eval_ctx) (vid : VarId.id) : binder =
 
 (** TODO: make this more efficient with maps *)
 let ctx_lookup_type_def (ctx : eval_ctx) (tid : TypeDefId.id) : type_def =
-  TypeDefId.nth ctx.type_context.type_defs tid
+  TypeDefId.Map.find tid ctx.type_context.type_defs
 
 (** TODO: make this more efficient with maps *)
 let ctx_lookup_fun_def (ctx : eval_ctx) (fid : FunDefId.id) : fun_def =
-  FunDefId.nth ctx.fun_context fid
+  FunDefId.Map.find fid ctx.fun_context
 
 (** Retrieve a variable's value in an environment *)
 let env_lookup_var_value (env : env) (vid : VarId.id) : typed_value =
diff --git a/src/Interpreter.ml b/src/Interpreter.ml
index 3c467fbe..ec1e6373 100644
--- a/src/Interpreter.ml
+++ b/src/Interpreter.ml
@@ -23,14 +23,19 @@ let log = L.interpreter_log
 module Test = struct
   let initialize_context (m : M.cfim_module) (type_vars : T.type_var list) :
       C.eval_ctx =
+    let type_decls, _ = M.split_declarations m.declarations in
+    let type_defs, fun_defs = M.compute_defs_maps m in
     let type_defs_groups, _funs_defs_groups =
       M.split_declarations_to_group_maps m.declarations
     in
-    let type_context = { C.type_defs_groups; type_defs = m.types } in
+    let type_infos =
+      TypesAnalysis.analyze_type_declarations type_defs type_decls
+    in
+    let type_context = { C.type_defs_groups; type_defs; type_infos } in
     C.reset_global_counters ();
     {
       C.type_context;
-      C.fun_context = m.functions;
+      C.fun_context = fun_defs;
       C.type_vars;
       C.env = [];
       C.ended_regions = T.RegionId.Set.empty;
diff --git a/src/InterpreterPaths.ml b/src/InterpreterPaths.ml
index 4482a507..cf02fc23 100644
--- a/src/InterpreterPaths.ml
+++ b/src/InterpreterPaths.ml
@@ -346,14 +346,14 @@ let write_place_unwrap (config : C.config) (access : access_kind) (p : E.place)
   | Ok ctx -> ctx
 
 (** Compute an expanded ADT bottom value *)
-let compute_expanded_bottom_adt_value (tyctx : T.type_def list)
+let compute_expanded_bottom_adt_value (tyctx : T.type_def T.TypeDefId.Map.t)
     (def_id : T.TypeDefId.id) (opt_variant_id : T.VariantId.id option)
     (regions : T.erased_region list) (types : T.ety list) : V.typed_value =
   (* Lookup the definition and check if it is an enumeration - it
      should be an enumeration if and only if the projection element
      is a field projection with *some* variant id. Retrieve the list
      of fields at the same time. *)
-  let def = T.TypeDefId.nth tyctx def_id in
+  let def = T.TypeDefId.Map.find def_id tyctx in
   assert (List.length regions = List.length def.T.region_params);
   (* Compute the field types *)
   let field_types =
diff --git a/src/Modules.ml b/src/Modules.ml
index e437ad0a..ef4998ec 100644
--- a/src/Modules.ml
+++ b/src/Modules.ml
@@ -21,6 +21,20 @@ type cfim_module = {
 }
 (** CFIM module *)
 
+let compute_defs_maps (m : cfim_module) :
+    type_def TypeDefId.Map.t * fun_def FunDefId.Map.t =
+  let types_map =
+    List.fold_left
+      (fun m (def : type_def) -> TypeDefId.Map.add def.def_id def m)
+      TypeDefId.Map.empty m.types
+  in
+  let funs_map =
+    List.fold_left
+      (fun m (def : fun_def) -> FunDefId.Map.add def.def_id def m)
+      FunDefId.Map.empty m.functions
+  in
+  (types_map, funs_map)
+
 (** Split a module's declarations between types and functions *)
 let split_declarations (decls : declaration_group list) :
     type_declaration_group list * fun_declaration_group list =
diff --git a/src/Print.ml b/src/Print.ml
index 282b2860..94546675 100644
--- a/src/Print.ml
+++ b/src/Print.ml
@@ -481,20 +481,20 @@ module Contexts = struct
   let ctx_to_rtype_formatter (fmt : ctx_formatter) : PT.rtype_formatter =
     PV.value_to_rtype_formatter fmt
 
-  let type_ctx_to_adt_variant_to_string_fun (ctx : T.type_def list) :
-      T.TypeDefId.id -> T.VariantId.id -> string =
+  let type_ctx_to_adt_variant_to_string_fun (ctx : T.type_def T.TypeDefId.Map.t)
+      : T.TypeDefId.id -> T.VariantId.id -> string =
    fun def_id variant_id ->
-    let def = T.TypeDefId.nth ctx def_id in
+    let def = T.TypeDefId.Map.find def_id ctx in
     match def.kind with
     | Struct _ -> failwith "Unreachable"
     | Enum variants ->
         let variant = T.VariantId.nth variants variant_id in
         name_to_string def.name ^ "::" ^ variant.variant_name
 
-  let type_ctx_to_adt_field_names_fun (ctx : T.type_def list) :
+  let type_ctx_to_adt_field_names_fun (ctx : T.type_def T.TypeDefId.Map.t) :
       T.TypeDefId.id -> T.VariantId.id option -> string list option =
    fun def_id opt_variant_id ->
-    let def = T.TypeDefId.nth ctx def_id in
+    let def = T.TypeDefId.Map.find def_id ctx in
     let fields = TU.type_def_get_fields def opt_variant_id in
     (* TODO: the field name should be optional?? *)
     let fields = List.map (fun f -> f.T.field_name) fields in
@@ -510,7 +510,7 @@ module Contexts = struct
       v.name
     in
     let type_def_id_to_string def_id =
-      let def = T.TypeDefId.nth ctx.type_context.type_defs def_id in
+      let def = C.ctx_lookup_type_def ctx def_id in
       name_to_string def.name
     in
     let adt_variant_to_string =
@@ -615,10 +615,10 @@ module CfimAst = struct
       PT.type_def_id_to_string = fmt.type_def_id_to_string;
     }
 
-  let type_ctx_to_adt_field_to_string_fun (ctx : T.type_def list) :
+  let type_ctx_to_adt_field_to_string_fun (ctx : T.type_def T.TypeDefId.Map.t) :
       T.TypeDefId.id -> T.VariantId.id option -> T.FieldId.id -> string =
    fun def_id opt_variant_id field_id ->
-    let def = T.TypeDefId.nth ctx def_id in
+    let def = T.TypeDefId.Map.find def_id ctx in
     let fields = TU.type_def_get_fields def opt_variant_id in
     let field = T.FieldId.nth fields field_id in
     field.T.field_name
@@ -629,7 +629,7 @@ module CfimAst = struct
       type_ctx_to_adt_field_to_string_fun ctx.type_context.type_defs
     in
     let fun_def_id_to_string def_id =
-      let def = A.FunDefId.nth ctx.fun_context def_id in
+      let def = C.ctx_lookup_fun_def ctx def_id in
       name_to_string def.name
     in
     {
@@ -904,18 +904,19 @@ module PA = CfimAst (* local module *)
 module Module = struct
   (** This function pretty-prints a type definition by using a definition
       context *)
-  let type_def_to_string (type_context : T.type_def list) (def : T.type_def) :
-      string =
+  let type_def_to_string (type_context : T.type_def T.TypeDefId.Map.t)
+      (def : T.type_def) : string =
     let type_def_id_to_string (id : T.TypeDefId.id) : string =
-      let def = T.TypeDefId.nth type_context id in
+      let def = T.TypeDefId.Map.find id type_context in
       name_to_string def.name
     in
     PT.type_def_to_string type_def_id_to_string def
 
   (** Generate an [ast_formatter] by using a definition context in combination
       with the variables local to a function's definition *)
-  let def_ctx_to_ast_formatter (type_context : T.type_def list)
-      (fun_context : A.fun_def list) (def : A.fun_def) : PA.ast_formatter =
+  let def_ctx_to_ast_formatter (type_context : T.type_def T.TypeDefId.Map.t)
+      (fun_context : A.fun_def A.FunDefId.Map.t) (def : A.fun_def) :
+      PA.ast_formatter =
     let rvar_to_string vid =
       let var = T.RegionVarId.nth def.signature.region_params vid in
       PT.region_var_to_string var
@@ -929,11 +930,11 @@ module Module = struct
       PT.type_var_to_string var
     in
     let type_def_id_to_string def_id =
-      let def = T.TypeDefId.nth type_context def_id in
+      let def = T.TypeDefId.Map.find def_id type_context in
       name_to_string def.name
     in
     let fun_def_id_to_string def_id =
-      let def = A.FunDefId.nth fun_context def_id in
+      let def = A.FunDefId.Map.find def_id fun_context in
       name_to_string def.name
     in
     let var_id_to_string vid =
@@ -961,18 +962,20 @@ module Module = struct
 
   (** This function pretty-prints a function definition by using a definition
       context *)
-  let fun_def_to_string (type_context : T.type_def list)
-      (fun_context : A.fun_def list) (def : A.fun_def) : string =
+  let fun_def_to_string (type_context : T.type_def T.TypeDefId.Map.t)
+      (fun_context : A.fun_def A.FunDefId.Map.t) (def : A.fun_def) : string =
     let fmt = def_ctx_to_ast_formatter type_context fun_context def in
     PA.fun_def_to_string fmt "" "  " def
 
   let module_to_string (m : M.cfim_module) : string =
+    let types_defs_map, funs_defs_map = M.compute_defs_maps m in
+
     (* The types *)
-    let type_defs = List.map (type_def_to_string m.M.types) m.M.types in
+    let type_defs = List.map (type_def_to_string types_defs_map) m.M.types in
 
     (* The functions *)
     let fun_defs =
-      List.map (fun_def_to_string m.M.types m.M.functions) m.M.functions
+      List.map (fun_def_to_string types_defs_map funs_defs_map) m.M.functions
     in
 
     (* Put everything together *)
diff --git a/src/TypesAnalysis.ml b/src/TypesAnalysis.ml
index 80748f8b..f49b49ac 100644
--- a/src/TypesAnalysis.ml
+++ b/src/TypesAnalysis.ml
@@ -5,11 +5,12 @@ type subtype_info = {
   under_borrow : bool;  (** Are we inside a borrow? *)
   under_nested_borrows : bool;  (** Are we inside nested borrows? *)
 }
+[@@deriving show]
 
-type type_param_info = subtype_info
+type type_param_info = subtype_info [@@deriving show]
 (** See [type_def_info] *)
 
-type expl_info = subtype_info
+type expl_info = subtype_info [@@deriving show]
 
 type 'p g_type_info = {
   contains_static : bool;
@@ -20,20 +21,26 @@ type 'p g_type_info = {
       (** Does the type (transitively) contains nested borrows? *)
   param_infos : 'p;  (** Gives information about the type parameters *)
 }
+[@@deriving show]
 (** Generic definition *)
 
-type type_def_info = type_param_info list g_type_info
+type type_def_info = type_param_info list g_type_info [@@deriving show]
 (** Information about a type definition. *)
 
-type ty_info = unit g_type_info
+type ty_info = unit g_type_info [@@deriving show]
 (** Information about a type. *)
 
 type partial_type_info = type_param_info list option g_type_info
+[@@deriving show]
 (** Helper definition.
 
     Allows us to factorize code: [analyze_full_ty] is used both to analyze
     type definitions and types. *)
 
+type type_infos = type_def_info TypeDefId.Map.t [@@deriving show]
+
+let expl_info_init = { under_borrow = false; under_nested_borrows = false }
+
 let initialize_g_type_info (param_infos : 'p) : 'p g_type_info =
   {
     contains_static = false;
@@ -73,10 +80,6 @@ let partial_type_info_to_ty_info (info : partial_type_info) : ty_info =
     param_infos = ();
   }
 
-type type_infos = type_def_info TypeDefId.Map.t
-
-let expl_info_init = { under_borrow = false; under_nested_borrows = false }
-
 let analyze_full_ty (updated : bool ref) (infos : type_infos)
     (ty_info : partial_type_info) (ty : 'r gr_ty) : partial_type_info =
   (* Small utility *)