Start working on greedy symbolic value expansion and expansion before

assignment

1 files changed, 44 insertions, 0 deletions

diff --git a/src/Modules.ml b/src/Modules.ml
index bf5e9835..e4cb06c1 100644
--- a/src/Modules.ml
+++ b/src/Modules.ml
@@ -14,3 +14,47 @@ type cfim_module = {
   functions : fun_def list;
 }
 (** CFIM module *)
+
+type 'id decl_group = NonRec of 'id | Rec of 'id list [@@deriving show]
+
+type types_decl_group = TypeDefId.id decl_group [@@deriving show]
+
+type funs_decl_group = FunDefId.id decl_group [@@deriving show]
+
+(** Split a module's declarations between types and functions *)
+let split_declarations (decls : declaration list) :
+    types_decl_group list * funs_decl_group list =
+  let rec split decls =
+    match decls with
+    | [] -> ([], [])
+    | d :: decls' -> (
+        let types, funs = split decls' in
+        match d with
+        | Type id -> (NonRec id :: types, funs)
+        | Fun id -> (types, NonRec id :: funs)
+        | RecTypes ids -> (Rec ids :: types, funs)
+        | RecFuns ids -> (types, Rec ids :: funs))
+  in
+  split decls
+
+(** Split a module's declarations into two maps from type/fun ids to
+    declaration groups.
+ *)
+let split_declarations_to_group_maps (decls : declaration list) :
+    types_decl_group TypeDefId.Map.t * funs_decl_group FunDefId.Map.t =
+  let module G (M : Map.S) = struct
+    let add_group (map : M.key decl_group M.t) (group : M.key decl_group) :
+        M.key decl_group M.t =
+      match group with
+      | NonRec id -> M.add id group map
+      | Rec ids -> List.fold_left (fun map id -> M.add id group map) map ids
+
+    let create_map (groups : M.key decl_group list) : M.key decl_group M.t =
+      List.fold_left add_group M.empty groups
+  end in
+  let types, funs = split_declarations decls in
+  let module TG = G (TypeDefId.Map) in
+  let types = TG.create_map types in
+  let module FG = G (FunDefId.Map) in
+  let funs = FG.create_map funs in
+  (types, funs)