Refactoring changes after @altgr's suggestions

compiler/desugared/ast.ml

@@ -24,6 +24,7 @@ open Shared_ast
 (** Inside a scope, a definition can refer either to a scope def, or a subscope
     def *)
 module ScopeDef = struct
+  module Base = struct
     type t =
       | Var of ScopeVar.t * StateName.t option
       | SubScopeVar of SubScopeName.t * ScopeVar.t * Pos.t
@@ -37,7 +38,9 @@ module ScopeDef = struct
         | 0 -> Option.compare StateName.compare stx sty
         | n -> n)
       | SubScopeVar (x', x, _), SubScopeVar (y', y, _) -> (
-      match SubScopeName.compare x' y' with 0 -> ScopeVar.compare x y | n -> n)
+        match SubScopeName.compare x' y' with
+        | 0 -> ScopeVar.compare x y
+        | n -> n)
       | Var _, _ -> -1
       | _, Var _ -> 1
 
@@ -61,10 +64,12 @@ module ScopeDef = struct
       | Var (v, Some sv) -> Int.logxor (ScopeVar.hash v) (StateName.hash sv)
       | SubScopeVar (w, v, _) ->
         Int.logxor (SubScopeName.hash w) (ScopeVar.hash v)
-end
+  end
 
-module ScopeDefMap : Map.S with type key = ScopeDef.t = Map.Make (ScopeDef)
-module ScopeDefSet : Set.S with type elt = ScopeDef.t = Set.Make (ScopeDef)
+  include Base
+  module Map = Map.Make (Base)
+  module Set = Set.Make (Base)
+end
 
 (** {1 AST} *)
 
@@ -195,7 +200,7 @@ type scope = {
   scope_vars : var_or_states ScopeVar.Map.t;
   scope_sub_scopes : ScopeName.t SubScopeName.Map.t;
   scope_uid : ScopeName.t;
-  scope_defs : scope_def ScopeDefMap.t;
+  scope_defs : scope_def ScopeDef.Map.t;
   scope_assertions : assertion list;
   scope_options : catala_option Marked.pos list;
   scope_meta_assertions : meta_assertion list;
@@ -218,9 +223,9 @@ let rec locations_used e : LocationSet.t =
       (fun e -> LocationSet.union (locations_used e))
       e LocationSet.empty
 
-let free_variables (def : rule RuleName.Map.t) : Pos.t ScopeDefMap.t =
-  let add_locs (acc : Pos.t ScopeDefMap.t) (locs : LocationSet.t) :
-      Pos.t ScopeDefMap.t =
+let free_variables (def : rule RuleName.Map.t) : Pos.t ScopeDef.Map.t =
+  let add_locs (acc : Pos.t ScopeDef.Map.t) (locs : LocationSet.t) :
+      Pos.t ScopeDef.Map.t =
     LocationSet.fold
       (fun (loc, loc_pos) acc ->
         let usage =
@@ -235,7 +240,9 @@ let free_variables (def : rule RuleName.Map.t) : Pos.t ScopeDefMap.t =
                    Marked.get_mark sub_index ))
           | ToplevelVar _ -> None
         in
-        match usage with Some u -> ScopeDefMap.add u loc_pos acc | None -> acc)
+        match usage with
+        | Some u -> ScopeDef.Map.add u loc_pos acc
+        | None -> acc)
       locs acc
   in
   RuleName.Map.fold
@@ -246,14 +253,14 @@ let free_variables (def : rule RuleName.Map.t) : Pos.t ScopeDefMap.t =
           (locations_used (Expr.unbox rule.rule_cons))
       in
       add_locs acc locs)
-    def ScopeDefMap.empty
+    def ScopeDef.Map.empty
 
 let fold_exprs ~(f : 'a -> expr -> 'a) ~(init : 'a) (p : program) : 'a =
   let acc =
     ScopeName.Map.fold
       (fun _ scope acc ->
         let acc =
-          ScopeDefMap.fold
+          ScopeDef.Map.fold
             (fun _ scope_def acc ->
               RuleName.Map.fold
                 (fun _ rule acc ->

compiler/desugared/ast.mli

@@ -30,10 +30,10 @@ module ScopeDef : sig
   val get_position : t -> Pos.t
   val format_t : Format.formatter -> t -> unit
   val hash : t -> int
-end
 
-module ScopeDefMap : Map.S with type key = ScopeDef.t
-module ScopeDefSet : Set.S with type elt = ScopeDef.t
+  module Map : Map.S with type key = t
+  module Set : Set.S with type elt = t
+end
 
 (** {1 AST} *)
 
@@ -118,7 +118,7 @@ type scope = {
   scope_vars : var_or_states ScopeVar.Map.t;
   scope_sub_scopes : ScopeName.t SubScopeName.Map.t;
   scope_uid : ScopeName.t;
-  scope_defs : scope_def ScopeDefMap.t;
+  scope_defs : scope_def ScopeDef.Map.t;
   scope_assertions : assertion list;
   scope_options : catala_option Marked.pos list;
   scope_meta_assertions : meta_assertion list;
@@ -133,7 +133,7 @@ type program = {
 (** {1 Helpers} *)
 
 val locations_used : expr -> LocationSet.t
-val free_variables : rule RuleName.Map.t -> Pos.t ScopeDefMap.t
+val free_variables : rule RuleName.Map.t -> Pos.t ScopeDef.Map.t
 
 val fold_exprs : f:('a -> expr -> 'a) -> init:'a -> program -> 'a
 (** Usage: [fold_exprs ~f ~init program] applies ~f to all the expressions

compiler/desugared/dependency.ml

@@ -173,11 +173,11 @@ let build_scope_dependencies (scope : Ast.scope) : ScopeDependencies.t =
       scope.scope_sub_scopes g
   in
   let g =
-    Ast.ScopeDefMap.fold
+    Ast.ScopeDef.Map.fold
       (fun def_key scope_def g ->
         let def = scope_def.Ast.scope_def_rules in
         let fv = Ast.free_variables def in
-        Ast.ScopeDefMap.fold
+        Ast.ScopeDef.Map.fold
           (fun fv_def fv_def_pos g ->
             match def_key, fv_def with
             | ( Ast.ScopeDef.Var (v_defined, s_defined),
@@ -246,7 +246,10 @@ let build_scope_dependencies (scope : Ast.scope) : ScopeDependencies.t =
 module ExceptionVertex = struct
   type t = { rules : Pos.t RuleName.Map.t; label : LabelName.t }
 
-  let compare x y = RuleName.Map.compare compare x.rules y.rules
+  let compare x y =
+    RuleName.Map.compare
+      (fun _ _ -> 0 (* we don't care about positions here*))
+      x.rules y.rules
 
   let hash (x : t) : int =
     RuleName.Map.fold

compiler/desugared/disambiguate.ml

@@ -45,7 +45,7 @@ let rule ctx env rule =
 let scope ctx env scope =
   let env = Typing.Env.open_scope scope.scope_uid env in
   let scope_defs =
-    ScopeDefMap.map
+    ScopeDef.Map.map
       (fun def ->
         let scope_def_rules =
           (* Note: ordering in file order might be better for error reporting ?
@@ -75,7 +75,7 @@ let program prg =
     ScopeName.Map.fold
       (fun scope_name scope env ->
         let vars =
-          ScopeDefMap.fold
+          ScopeDef.Map.fold
             (fun var def vars ->
               match var with
               | Var (v, _states) -> ScopeVar.Map.add v def.scope_def_typ vars

compiler/desugared/from_surface.ml

@@ -1019,7 +1019,7 @@ let process_def
       (Marked.get_mark def.definition_name)
   in
   let scope_def_ctxt =
-    Ast.ScopeDefMap.find def_key scope_ctxt.scope_defs_contexts
+    Ast.ScopeDef.Map.find def_key scope_ctxt.scope_defs_contexts
   in
   (* We add to the name resolution context the name of the parameter variable *)
   let new_ctxt, param_uids =
@@ -1028,7 +1028,7 @@ let process_def
       def
   in
   let scope_updated =
-    let scope_def = Ast.ScopeDefMap.find def_key scope.scope_defs in
+    let scope_def = Ast.ScopeDef.Map.find def_key scope.scope_defs in
     let rule_name = def.definition_id in
     let label_situation =
       match def.definition_label with
@@ -1075,7 +1075,7 @@ let process_def
     in
     {
       scope with
-      scope_defs = Ast.ScopeDefMap.add def_key scope_def scope.scope_defs;
+      scope_defs = Ast.ScopeDef.Map.add def_key scope_def scope.scope_defs;
     }
   in
   {
@@ -1204,7 +1204,7 @@ let check_unlabeled_exception
       (* should not happen *)
     in
     let scope_def_ctxt =
-      Ast.ScopeDefMap.find def_key scope_ctxt.scope_defs_contexts
+      Ast.ScopeDef.Map.find def_key scope_ctxt.scope_defs_contexts
     in
     match exception_to with
     | Surface.Ast.NotAnException | Surface.Ast.ExceptionToLabel _ -> ()
@@ -1296,7 +1296,7 @@ let attribute_to_io (attr : Surface.Ast.scope_decl_context_io) : Ast.io =
 let init_scope_defs
     (ctxt : Name_resolution.context)
     (scope_idmap : Name_resolution.scope_var_or_subscope IdentName.Map.t) :
-    Ast.scope_def Ast.ScopeDefMap.t =
+    Ast.scope_def Ast.ScopeDef.Map.t =
   (* Initializing the definitions of all scopes and subscope vars, with no rules
      yet inside *)
   let add_def _ v scope_def_map =
@@ -1306,7 +1306,7 @@ let init_scope_defs
       match v_sig.var_sig_states_list with
       | [] ->
         let def_key = Ast.ScopeDef.Var (v, None) in
-        Ast.ScopeDefMap.add def_key
+        Ast.ScopeDef.Map.add def_key
           {
             Ast.scope_def_rules = RuleName.Map.empty;
             Ast.scope_def_typ = v_sig.var_sig_typ;
@@ -1344,7 +1344,7 @@ let init_scope_defs
                      { io_input; io_output });
                 }
               in
-              Ast.ScopeDefMap.add def_key def acc, i + 1)
+              Ast.ScopeDef.Map.add def_key def acc, i + 1)
             (scope_def_map, 0) states
         in
         scope_def)
@@ -1364,7 +1364,7 @@ let init_scope_defs
               Ast.ScopeDef.SubScopeVar
                 (v0, v, Marked.get_mark (ScopeVar.get_info v))
             in
-            Ast.ScopeDefMap.add def_key
+            Ast.ScopeDef.Map.add def_key
               {
                 Ast.scope_def_rules = RuleName.Map.empty;
                 Ast.scope_def_typ = v_sig.var_sig_typ;
@@ -1375,7 +1375,7 @@ let init_scope_defs
               scope_def_map)
         sub_scope_def.Name_resolution.var_idmap scope_def_map
   in
-  IdentName.Map.fold add_def scope_idmap Ast.ScopeDefMap.empty
+  IdentName.Map.fold add_def scope_idmap Ast.ScopeDef.Map.empty
 
 (** Main function of this module *)
 let translate_program

compiler/desugared/linting.ml

@@ -22,7 +22,7 @@ open Catala_utils
 let detect_empty_definitions (p : program) : unit =
   ScopeName.Map.iter
     (fun (scope_name : ScopeName.t) scope ->
-      ScopeDefMap.iter
+      ScopeDef.Map.iter
         (fun scope_def_key scope_def ->
           if
             (match scope_def_key with ScopeDef.Var _ -> true | _ -> false)
@@ -59,7 +59,7 @@ let detect_unused_scope_vars (p : program) : unit =
   in
   ScopeName.Map.iter
     (fun (scope_name : ScopeName.t) scope ->
-      ScopeDefMap.iter
+      ScopeDef.Map.iter
         (fun scope_def_key scope_def ->
           match scope_def_key with
           | ScopeDef.Var (v, _)

compiler/desugared/name_resolution.ml

@@ -39,7 +39,7 @@ type scope_var_or_subscope =
 type scope_context = {
   var_idmap : scope_var_or_subscope IdentName.Map.t;
       (** All variables, including scope variables and subscopes *)
-  scope_defs_contexts : scope_def_context Ast.ScopeDefMap.t;
+  scope_defs_contexts : scope_def_context Ast.ScopeDef.Map.t;
       (** What is the default rule to refer to for unnamed exceptions, if any *)
   sub_scopes : ScopeName.Set.t;
       (** Other scopes referred to by this scope. Used for dependency analysis *)
@@ -624,7 +624,7 @@ let process_name_item (ctxt : context) (item : Surface.Ast.code_item Marked.pos)
         ScopeName.Map.add scope_uid
           {
             var_idmap = IdentName.Map.empty;
-            scope_defs_contexts = Ast.ScopeDefMap.empty;
+            scope_defs_contexts = Ast.ScopeDef.Map.empty;
             sub_scopes = ScopeName.Set.empty;
           }
           ctxt.scopes;
@@ -853,7 +853,7 @@ let process_definition
               {
                 s_ctxt with
                 scope_defs_contexts =
-                  Ast.ScopeDefMap.update def_key
+                  Ast.ScopeDef.Map.update def_key
                     (fun def_key_ctx ->
                       Some
                         (update_def_key_ctx d

compiler/desugared/name_resolution.mli

@@ -39,7 +39,7 @@ type scope_var_or_subscope =
 type scope_context = {
   var_idmap : scope_var_or_subscope IdentName.Map.t;
       (** All variables, including scope variables and subscopes *)
-  scope_defs_contexts : scope_def_context Ast.ScopeDefMap.t;
+  scope_defs_contexts : scope_def_context Ast.ScopeDef.Map.t;
       (** What is the default rule to refer to for unnamed exceptions, if any *)
   sub_scopes : ScopeName.Set.t;
       (** Other scopes referred to by this scope. Used for dependency analysis *)

compiler/desugared/print.ml

@@ -88,13 +88,13 @@ let build_exception_tree exc_graph =
 
 let print_exceptions_graph
     (scope : ScopeName.t)
-    (var : DesugaredVarName.t)
+    (var : Ast.ScopeDef.t)
     (g : Dependency.ExceptionsDependencies.t) =
   Cli.result_format
     "Printing the tree of exceptions for the definitions of variable %a of \
      scope %a."
     (Cli.format_with_style [ANSITerminal.yellow])
-    (Format.asprintf "\"%a\"" DesugaredVarName.format var)
+    (Format.asprintf "\"%a\"" Ast.ScopeDef.format_t var)
     (Cli.format_with_style [ANSITerminal.yellow])
     (Format.asprintf "\"%a\"" ScopeName.format_t scope);
   Dependency.ExceptionsDependencies.iter_vertex

compiler/desugared/print.mli

@@ -16,7 +16,7 @@
 
 val print_exceptions_graph :
   Shared_ast.ScopeName.t ->
-  Shared_ast.DesugaredVarName.t ->
+  Ast.ScopeDef.t ->
   Dependency.ExceptionsDependencies.t ->
   unit
 (** Prints the exception graph of a variable to the terminal *)

compiler/driver.ml

@@ -21,6 +21,127 @@ open Catala_utils
     string representation. *)
 let extensions = [".catala_fr", "fr"; ".catala_en", "en"; ".catala_pl", "pl"]
 
+let get_scope_uid
+    (options : Cli.options)
+    (backend : Plugin.t Cli.backend_option)
+    (ctxt : Desugared.Name_resolution.context) =
+  match options.ex_scope, backend with
+  | None, `Interpret ->
+    Errors.raise_error "No scope was provided for execution."
+  | None, _ ->
+    let _, scope =
+      try
+        Shared_ast.IdentName.Map.filter_map
+          (fun _ -> function
+            | Desugared.Name_resolution.TScope (uid, _) -> Some uid
+            | _ -> None)
+          ctxt.typedefs
+        |> Shared_ast.IdentName.Map.choose
+      with Not_found ->
+        Errors.raise_error "There isn't any scope inside the program."
+    in
+    scope
+  | Some name, _ -> (
+    match Shared_ast.IdentName.Map.find_opt name ctxt.typedefs with
+    | Some (Desugared.Name_resolution.TScope (uid, _)) -> uid
+    | _ ->
+      Errors.raise_error "There is no scope %a inside the program."
+        (Cli.format_with_style [ANSITerminal.yellow])
+        ("\"" ^ name ^ "\""))
+
+let get_variable_uid
+    (options : Cli.options)
+    (backend : Plugin.t Cli.backend_option)
+    (ctxt : Desugared.Name_resolution.context)
+    (scope_uid : Shared_ast.ScopeName.t) =
+  match options.ex_variable, backend with
+  | None, `Exceptions ->
+    Errors.raise_error
+      "Please specify a variable with the -v option to print its exception \
+       tree."
+  | None, _ -> None
+  | Some name, _ -> (
+    (* Sometimes the variable selected is of the form [a.b]*)
+    let first_part, second_part =
+      match
+        Re.(
+          exec_opt
+            (compile
+            @@ whole_string
+            @@ seq [group (rep1 (compl [char '.'])); char '.'; group (rep1 any)]
+            )
+            name)
+      with
+      | None -> name, None
+      | Some groups -> Re.Group.get groups 1, Some (Re.Group.get groups 2)
+    in
+    match
+      Shared_ast.IdentName.Map.find_opt first_part
+        (Shared_ast.ScopeName.Map.find scope_uid ctxt.scopes).var_idmap
+    with
+    | None ->
+      Errors.raise_error "Variable %a not found inside scope %a"
+        (Cli.format_with_style [ANSITerminal.yellow])
+        ("\"" ^ name ^ "\"")
+        (Cli.format_with_style [ANSITerminal.yellow])
+        (Format.asprintf "\"%a\"" Shared_ast.ScopeName.format_t scope_uid)
+    | Some
+        (Desugared.Name_resolution.SubScope (subscope_var_name, subscope_name))
+      -> (
+      match second_part with
+      | None ->
+        Errors.raise_error
+          "Subscope %a of scope %a cannot be selected by itself, please add \
+           \".<var>\" where <var> is a subscope variable."
+          (Cli.format_with_style [ANSITerminal.yellow])
+          (Format.asprintf "\"%a\"" Shared_ast.SubScopeName.format_t
+             subscope_var_name)
+          (Cli.format_with_style [ANSITerminal.yellow])
+          (Format.asprintf "\"%a\"" Shared_ast.ScopeName.format_t scope_uid)
+      | Some second_part -> (
+        match
+          Shared_ast.IdentName.Map.find_opt second_part
+            (Shared_ast.ScopeName.Map.find subscope_name ctxt.scopes).var_idmap
+        with
+        | Some (Desugared.Name_resolution.ScopeVar v) ->
+          Some
+            (Desugared.Ast.ScopeDef.SubScopeVar
+               (subscope_var_name, v, Pos.no_pos))
+        | _ ->
+          Errors.raise_error
+            "Var %a of subscope %a in scope %a does not exist, please check \
+             your command line arguments."
+            (Cli.format_with_style [ANSITerminal.yellow])
+            ("\"" ^ second_part ^ "\"")
+            (Cli.format_with_style [ANSITerminal.yellow])
+            (Format.asprintf "\"%a\"" Shared_ast.SubScopeName.format_t
+               subscope_var_name)
+            (Cli.format_with_style [ANSITerminal.yellow])
+            (Format.asprintf "\"%a\"" Shared_ast.ScopeName.format_t scope_uid)))
+    | Some (Desugared.Name_resolution.ScopeVar v) ->
+      Some
+        (Desugared.Ast.ScopeDef.Var
+           ( v,
+             Option.map
+               (fun second_part ->
+                 let var_sig = Shared_ast.ScopeVar.Map.find v ctxt.var_typs in
+                 match
+                   Shared_ast.IdentName.Map.find_opt second_part
+                     var_sig.var_sig_states_idmap
+                 with
+                 | Some state -> state
+                 | None ->
+                   Errors.raise_error
+                     "State %a is not found for variable %a of scope %a"
+                     (Cli.format_with_style [ANSITerminal.yellow])
+                     ("\"" ^ second_part ^ "\"")
+                     (Cli.format_with_style [ANSITerminal.yellow])
+                     ("\"" ^ first_part ^ "\"")
+                     (Cli.format_with_style [ANSITerminal.yellow])
+                     (Format.asprintf "\"%a\"" Shared_ast.ScopeName.format_t
+                        scope_uid))
+               second_part )))
+
 (** Entry function for the executable. Returns a negative number in case of
     error. Usage: [driver source_file options]*)
 let driver source_file (options : Cli.options) : int =
@@ -144,129 +265,9 @@ let driver source_file (options : Cli.options) : int =
       backend -> (
       Cli.debug_print "Name resolution...";
       let ctxt = Desugared.Name_resolution.form_context prgm in
-      let scope_uid =
-        match options.ex_scope, backend with
-        | None, `Interpret ->
-          Errors.raise_error "No scope was provided for execution."
-        | None, _ ->
-          let _, scope =
-            try
-              Shared_ast.IdentName.Map.filter_map
-                (fun _ -> function
-                  | Desugared.Name_resolution.TScope (uid, _) -> Some uid
-                  | _ -> None)
-                ctxt.typedefs
-              |> Shared_ast.IdentName.Map.choose
-            with Not_found ->
-              Errors.raise_error "There isn't any scope inside the program."
-          in
-          scope
-        | Some name, _ -> (
-          match Shared_ast.IdentName.Map.find_opt name ctxt.typedefs with
-          | Some (Desugared.Name_resolution.TScope (uid, _)) -> uid
-          | _ ->
-            Errors.raise_error "There is no scope %a inside the program."
-              (Cli.format_with_style [ANSITerminal.yellow])
-              ("\"" ^ name ^ "\""))
-      in
+      let scope_uid = get_scope_uid options backend ctxt in
       (* This uid is a Desugared identifier *)
-      let variable_uid =
-        match options.ex_variable, backend with
-        | None, `Exceptions ->
-          Errors.raise_error
-            "Please specify a variable with the -v option to print its \
-             exception tree."
-        | None, _ -> None
-        | Some name, _ -> (
-          (* Sometimes the variable selected is of the form [a.b]*)
-          let first_part, second_part =
-            match
-              Re.(
-                exec_opt
-                  (compile
-                  @@ whole_string
-                  @@ seq
-                       [
-                         group (rep1 (compl [char '.']));
-                         char '.';
-                         group (rep1 any);
-                       ])
-                  name)
-            with
-            | None -> name, None
-            | Some groups -> Re.Group.get groups 1, Some (Re.Group.get groups 2)
-          in
-          match
-            Shared_ast.IdentName.Map.find_opt first_part
-              (Shared_ast.ScopeName.Map.find scope_uid ctxt.scopes).var_idmap
-          with
-          | None ->
-            Errors.raise_error "Variable %a not found inside scope %a"
-              (Cli.format_with_style [ANSITerminal.yellow])
-              ("\"" ^ name ^ "\"")
-              (Cli.format_with_style [ANSITerminal.yellow])
-              (Format.asprintf "\"%a\"" Shared_ast.ScopeName.format_t scope_uid)
-          | Some
-              (Desugared.Name_resolution.SubScope
-                (subscope_var_name, subscope_name)) -> (
-            match second_part with
-            | None ->
-              Errors.raise_error
-                "Subscope %a of scope %a cannot be selected by itself, please \
-                 add \".<var>\" where <var> is a subscope variable."
-                (Cli.format_with_style [ANSITerminal.yellow])
-                (Format.asprintf "\"%a\"" Shared_ast.SubScopeName.format_t
-                   subscope_var_name)
-                (Cli.format_with_style [ANSITerminal.yellow])
-                (Format.asprintf "\"%a\"" Shared_ast.ScopeName.format_t
-                   scope_uid)
-            | Some second_part -> (
-              match
-                Shared_ast.IdentName.Map.find_opt second_part
-                  (Shared_ast.ScopeName.Map.find subscope_name ctxt.scopes)
-                    .var_idmap
-              with
-              | Some (Desugared.Name_resolution.ScopeVar v) ->
-                Some
-                  (Shared_ast.DesugaredVarName.SubScopeVar (subscope_var_name, v))
-              | _ ->
-                Errors.raise_error
-                  "Var %a of subscope %a in scope %a does not exist, please \
-                   check your command line arguments."
-                  (Cli.format_with_style [ANSITerminal.yellow])
-                  ("\"" ^ second_part ^ "\"")
-                  (Cli.format_with_style [ANSITerminal.yellow])
-                  (Format.asprintf "\"%a\"" Shared_ast.SubScopeName.format_t
-                     subscope_var_name)
-                  (Cli.format_with_style [ANSITerminal.yellow])
-                  (Format.asprintf "\"%a\"" Shared_ast.ScopeName.format_t
-                     scope_uid)))
-          | Some (Desugared.Name_resolution.ScopeVar v) ->
-            Some
-              (Shared_ast.DesugaredVarName.ScopeVar
-                 ( v,
-                   Option.map
-                     (fun second_part ->
-                       let var_sig =
-                         Shared_ast.ScopeVar.Map.find v ctxt.var_typs
-                       in
-                       match
-                         Shared_ast.IdentName.Map.find_opt second_part
-                           var_sig.var_sig_states_idmap
-                       with
-                       | Some state -> state
-                       | None ->
-                         Errors.raise_error
-                           "State %a is not found for variable %a of scope %a"
-                           (Cli.format_with_style [ANSITerminal.yellow])
-                           ("\"" ^ second_part ^ "\"")
-                           (Cli.format_with_style [ANSITerminal.yellow])
-                           ("\"" ^ first_part ^ "\"")
-                           (Cli.format_with_style [ANSITerminal.yellow])
-                           (Format.asprintf "\"%a\""
-                              Shared_ast.ScopeName.format_t scope_uid))
-                     second_part )))
-      in
+      let variable_uid = get_variable_uid options backend ctxt scope_uid in
       Cli.debug_print "Desugaring...";
       let prgm = Desugared.From_surface.translate_program ctxt prgm in
       Cli.debug_print "Disambiguating...";
@@ -287,7 +288,7 @@ let driver source_file (options : Cli.options) : int =
               "Please provide a scope variable to analyze with the -v option."
         in
         Desugared.Print.print_exceptions_graph scope_uid variable_uid
-          (Shared_ast.DesugaredVarName.Map.find variable_uid exceptions_graphs)
+          (Desugared.Ast.ScopeDef.Map.find variable_uid exceptions_graphs)
       | `Scopelang ->
         let _output_file, with_output = get_output_format () in
         with_output

compiler/scopelang/from_desugared.ml

@@ -435,7 +435,7 @@ let translate_def
 let translate_rule ctx (scope : Desugared.Ast.scope) = function
   | Desugared.Dependency.Vertex.Var (var, state) -> (
     let scope_def =
-      Desugared.Ast.ScopeDefMap.find
+      Desugared.Ast.ScopeDef.Map.find
         (Desugared.Ast.ScopeDef.Var (var, state))
         scope.scope_defs
     in
@@ -455,7 +455,7 @@ let translate_rule ctx (scope : Desugared.Ast.scope) = function
              (RuleName.Map.bindings var_def))
         "It is impossible to give a definition to a scope variable tagged as \
          input."
-    | OnlyInput -> [], DesugaredVarName.Map.empty
+    | OnlyInput -> [], Desugared.Ast.ScopeDef.Map.empty
     (* we do not provide any definition for an input-only variable *)
     | _ ->
       let expr_def, exc_graph =
@@ -479,8 +479,8 @@ let translate_rule ctx (scope : Desugared.Ast.scope) = function
               scope_def.Desugared.Ast.scope_def_io,
               Expr.unbox expr_def );
         ],
-        DesugaredVarName.Map.singleton
-          (DesugaredVarName.ScopeVar (var, state))
+        Desugared.Ast.ScopeDef.Map.singleton
+          (Desugared.Ast.ScopeDef.Var (var, state))
           exc_graph ))
   | Desugared.Dependency.Vertex.SubScope sub_scope_index ->
     (* Before calling the sub_scope, we need to include all the re-definitions
@@ -489,7 +489,7 @@ let translate_rule ctx (scope : Desugared.Ast.scope) = function
       SubScopeName.Map.find sub_scope_index scope.scope_sub_scopes
     in
     let sub_scope_vars_redefs_candidates =
-      Desugared.Ast.ScopeDefMap.filter
+      Desugared.Ast.ScopeDef.Map.filter
         (fun def_key scope_def ->
           match def_key with
           | Desugared.Ast.ScopeDef.Var _ -> false
@@ -507,7 +507,7 @@ let translate_rule ctx (scope : Desugared.Ast.scope) = function
         scope.scope_defs
     in
     let sub_scope_vars_redefs =
-      Desugared.Ast.ScopeDefMap.mapi
+      Desugared.Ast.ScopeDef.Map.mapi
         (fun def_key scope_def ->
           let def = scope_def.Desugared.Ast.scope_def_rules in
           let def_typ = scope_def.scope_def_typ in
@@ -573,11 +573,11 @@ let translate_rule ctx (scope : Desugared.Ast.scope) = function
                   def_typ,
                   scope_def.Desugared.Ast.scope_def_io,
                   Expr.unbox expr_def ),
-              (exc_graph, sub_scope_var) ))
+              (exc_graph, sub_scope_var, var_pos) ))
         sub_scope_vars_redefs_candidates
     in
     let sub_scope_vars_redefs_and_exc_graphs =
-      List.map snd (Desugared.Ast.ScopeDefMap.bindings sub_scope_vars_redefs)
+      List.map snd (Desugared.Ast.ScopeDef.Map.bindings sub_scope_vars_redefs)
     in
     let sub_scope_vars_redefs =
       List.map fst sub_scope_vars_redefs_and_exc_graphs
@@ -593,17 +593,19 @@ let translate_rule ctx (scope : Desugared.Ast.scope) = function
                 } );
         ],
       List.fold_left
-        (fun exc_graphs (new_exc_graph, subscope_var) ->
-          DesugaredVarName.Map.add
-            (DesugaredVarName.SubScopeVar (sub_scope_index, subscope_var))
+        (fun exc_graphs (new_exc_graph, subscope_var, var_pos) ->
+          Desugared.Ast.ScopeDef.Map.add
+            (Desugared.Ast.ScopeDef.SubScopeVar
+               (sub_scope_index, subscope_var, var_pos))
             new_exc_graph exc_graphs)
-        DesugaredVarName.Map.empty
+        Desugared.Ast.ScopeDef.Map.empty
         (List.map snd sub_scope_vars_redefs_and_exc_graphs) )
 
 (** Translates a scope *)
 let translate_scope (ctx : ctx) (scope : Desugared.Ast.scope) :
     untyped Ast.scope_decl
-    * Desugared.Dependency.ExceptionsDependencies.t DesugaredVarName.Map.t =
+    * Desugared.Dependency.ExceptionsDependencies.t Desugared.Ast.ScopeDef.Map.t
+    =
   let scope_dependencies =
     Desugared.Dependency.build_scope_dependencies scope
   in
@@ -618,10 +620,10 @@ let translate_scope (ctx : ctx) (scope : Desugared.Ast.scope) :
           translate_rule ctx scope scope_def_key
         in
         ( scope_decl_rules @ new_rules,
-          DesugaredVarName.Map.union
+          Desugared.Ast.ScopeDef.Map.union
             (fun _ _ _ -> assert false (* there should not be key conflicts *))
             new_exceptions_graphs exceptions_graphs ))
-      ([], DesugaredVarName.Map.empty)
+      ([], Desugared.Ast.ScopeDef.Map.empty)
       scope_ordering
   in
   (* Then, after having computed all the scopes variables, we add the
@@ -641,7 +643,7 @@ let translate_scope (ctx : ctx) (scope : Desugared.Ast.scope) :
         match states with
         | WholeVar ->
           let scope_def =
-            Desugared.Ast.ScopeDefMap.find
+            Desugared.Ast.ScopeDef.Map.find
               (Desugared.Ast.ScopeDef.Var (var, None))
               scope.scope_defs
           in
@@ -659,7 +661,7 @@ let translate_scope (ctx : ctx) (scope : Desugared.Ast.scope) :
           List.fold_left
             (fun acc (state : StateName.t) ->
               let scope_def =
-                Desugared.Ast.ScopeDefMap.find
+                Desugared.Ast.ScopeDef.Map.find
                   (Desugared.Ast.ScopeDef.Var (var, Some state))
                   scope.scope_defs
               in
@@ -686,7 +688,8 @@ let translate_scope (ctx : ctx) (scope : Desugared.Ast.scope) :
 
 let translate_program (pgrm : Desugared.Ast.program) :
     untyped Ast.program
-    * Desugared.Dependency.ExceptionsDependencies.t DesugaredVarName.Map.t =
+    * Desugared.Dependency.ExceptionsDependencies.t Desugared.Ast.ScopeDef.Map.t
+    =
   (* First we give mappings to all the locations between Desugared and This
      involves creating a new Scopelang scope variable for every state of a
      Desugared variable. *)
@@ -748,11 +751,11 @@ let translate_program (pgrm : Desugared.Ast.program) :
           translate_scope ctx scope
         in
         ( ScopeName.Map.add scope_name new_program_scope new_program_scopes,
-          DesugaredVarName.Map.union
+          Desugared.Ast.ScopeDef.Map.union
             (fun _ _ _ -> assert false (* key conflicts should not happen*))
             new_exceptions_graphs exceptions_graph ))
       pgrm.program_scopes
-      (ScopeName.Map.empty, DesugaredVarName.Map.empty)
+      (ScopeName.Map.empty, Desugared.Ast.ScopeDef.Map.empty)
   in
   ( {
       Ast.program_topdefs =

compiler/scopelang/from_desugared.mli

@@ -19,7 +19,6 @@
 val translate_program :
   Desugared.Ast.program ->
   Shared_ast.untyped Ast.program
-  * Desugared.Dependency.ExceptionsDependencies.t
-    Shared_ast.DesugaredVarName.Map.t
+  * Desugared.Dependency.ExceptionsDependencies.t Desugared.Ast.ScopeDef.Map.t
 (** This functions returns the translated program as well as all the graphs of
     exceptions inferred for each scope variable of the program. *)

compiler/shared_ast/definitions.ml

@@ -45,67 +45,6 @@ module SubScopeName = Uid.Gen ()
 module StateName = Uid.Gen ()
 
 (** {1 Abstract syntax tree} *)
-module DesugaredVarName : sig
-  type t =
-    | ScopeVar of ScopeVar.t * StateName.t option
-    | SubScopeVar of SubScopeName.t * ScopeVar.t
-
-  val hash : t -> int
-  val compare : t -> t -> int
-  val equal : t -> t -> bool
-  val format : Format.formatter -> t -> unit
-
-  module Map : Map.S with type key = t
-  module Set : Set.S with type elt = t
-end = struct
-  module Ordering = struct
-    type t =
-      | ScopeVar of ScopeVar.t * StateName.t option
-      | SubScopeVar of SubScopeName.t * ScopeVar.t
-
-    let hash x =
-      match x with
-      | ScopeVar (x, None) -> ScopeVar.hash x
-      | ScopeVar (x, Some sx) ->
-        Int.logxor (ScopeVar.hash x) (StateName.hash sx)
-      | SubScopeVar (x, y) -> Int.logxor (SubScopeName.hash x) (ScopeVar.hash y)
-
-    let compare x y =
-      match x, y with
-      | ScopeVar (x, xst), ScopeVar (y, yst) -> (
-        match ScopeVar.compare x y with
-        | 0 -> Option.compare StateName.compare xst yst
-        | n -> n)
-      | SubScopeVar (x, xv), SubScopeVar (y, yv) -> (
-        match SubScopeName.compare x y with
-        | 0 -> ScopeVar.compare xv yv
-        | n -> n)
-      | ScopeVar _, _ -> -1
-      | _, ScopeVar _ -> 1
-      | SubScopeVar _, _ -> .
-      | _, SubScopeVar _ -> .
-
-    let equal x y =
-      match x, y with
-      | ScopeVar (x, sx), ScopeVar (y, sy) ->
-        ScopeVar.equal x y && Option.equal StateName.equal sx sy
-      | SubScopeVar (x, xv), SubScopeVar (y, yv) ->
-        SubScopeName.equal x y && ScopeVar.equal xv yv
-      | (ScopeVar _ | SubScopeVar _), _ -> false
-
-    let format fmt x =
-      match x with
-      | ScopeVar (v, None) -> ScopeVar.format_t fmt v
-      | ScopeVar (v, Some st) ->
-        Format.fprintf fmt "%a.%a" ScopeVar.format_t v StateName.format_t st
-      | SubScopeVar (ss, v) ->
-        Format.fprintf fmt "%a.%a" SubScopeName.format_t ss ScopeVar.format_t v
-  end
-
-  include Ordering
-  module Map = Map.Make (Ordering)
-  module Set = Set.Make (Ordering)
-end
 
 (** Define a common base type for the expressions in most passes of the compiler *)