From c884e625bfdf142b195ac6c252e5825fc4f0be54 Mon Sep 17 00:00:00 2001
From: Nathan Rebours <nathan.p.rebours@gmail.com>
Date: Mon, 25 Mar 2019 18:14:03 +0800
Subject: [PATCH 1/2] Add omp_config to expansion context

Signed-off-by: Nathan Rebours <nathan.p.rebours@gmail.com>
---
 src/ast_traverse.ml       |  46 +++++-----
 src/ast_traverse.mli      |   2 +-
 src/context_free.ml       | 180 +++++++++++++++++++-------------------
 src/context_free.mli      |   2 +-
 src/driver.ml             |  52 ++++++-----
 src/expansion_context.ml  |  34 +++++--
 src/expansion_context.mli |  33 ++++++-
 7 files changed, 203 insertions(+), 146 deletions(-)

diff --git a/src/ast_traverse.ml b/src/ast_traverse.ml
index 41e166d15..bfbbe15e3 100644
--- a/src/ast_traverse.ml
+++ b/src/ast_traverse.ml
@@ -79,38 +79,44 @@ let var_names_of = object
     | _ -> acc
 end
 
-class map_with_code_path = object (self)
-  inherit [Code_path.t] map_with_context as super
+class map_with_expansion_context = object (self)
+  inherit [Expansion_context.Base.t] map_with_context as super
 
-  method! expression path expr =
-    super#expression (Code_path.enter_expr path) expr
+  method! expression ctxt expr =
+    super#expression (Expansion_context.Base.enter_expr ctxt) expr
 
-  method! module_binding path mb =
-       super#module_binding (Code_path.enter_module ~loc:mb.pmb_loc mb.pmb_name.txt path) mb
+  method! module_binding ctxt mb =
+    super#module_binding
+      (Expansion_context.Base.enter_module ~loc:mb.pmb_loc mb.pmb_name.txt ctxt)
+      mb
 
-  method! module_declaration path md =
-    super#module_declaration (Code_path.enter_module ~loc:md.pmd_loc md.pmd_name.txt path) md
+  method! module_declaration ctxt md =
+    super#module_declaration
+      (Expansion_context.Base.enter_module ~loc:md.pmd_loc md.pmd_name.txt ctxt)
+      md
 
-  method! module_type_declaration path mtd =
+  method! module_type_declaration ctxt mtd =
     super#module_type_declaration
-      (Code_path.enter_module ~loc:mtd.pmtd_loc mtd.pmtd_name.txt path)
+      (Expansion_context.Base.enter_module ~loc:mtd.pmtd_loc mtd.pmtd_name.txt ctxt)
       mtd
 
-  method! value_description path vd =
-    super#value_description (Code_path.enter_value ~loc:vd.pval_loc vd.pval_name.txt path) vd
+  method! value_description ctxt vd =
+    super#value_description
+      (Expansion_context.Base.enter_value ~loc:vd.pval_loc vd.pval_name.txt ctxt)
+      vd
 
-  method! value_binding path {pvb_pat; pvb_expr; pvb_attributes; pvb_loc} =
+  method! value_binding ctxt {pvb_pat; pvb_expr; pvb_attributes; pvb_loc} =
     let all_var_names = var_names_of#pattern pvb_pat [] in
     let var_name = Base.List.last all_var_names in
-    let in_binding_path =
+    let in_binding_ctxt =
       Base.Option.fold var_name
-        ~init:path
-        ~f:(fun path var_name -> Code_path.enter_value ~loc:pvb_loc var_name path)
+        ~init:ctxt
+        ~f:(fun ctxt var_name -> Expansion_context.Base.enter_value ~loc:pvb_loc var_name ctxt)
     in
-    let pvb_pat = self#pattern path pvb_pat in
-    let pvb_expr = self#expression in_binding_path pvb_expr in
-    let pvb_attributes = self#attributes in_binding_path pvb_attributes in
-    let pvb_loc = self#location path pvb_loc in
+    let pvb_pat = self#pattern ctxt pvb_pat in
+    let pvb_expr = self#expression in_binding_ctxt pvb_expr in
+    let pvb_attributes = self#attributes in_binding_ctxt pvb_attributes in
+    let pvb_loc = self#location ctxt pvb_loc in
     { pvb_pat; pvb_expr; pvb_attributes; pvb_loc }
 end
 
diff --git a/src/ast_traverse.mli b/src/ast_traverse.mli
index e41a79e53..0876a044a 100644
--- a/src/ast_traverse.mli
+++ b/src/ast_traverse.mli
@@ -54,7 +54,7 @@ end
 
 class map_with_path : [string] map_with_context
 
-class map_with_code_path : [Code_path.t] map_with_context
+class map_with_expansion_context : [Expansion_context.Base.t] map_with_context
 
 class virtual ['res] lift : object
   inherit ['res] Ppxlib_traverse_builtins.lift
diff --git a/src/context_free.ml b/src/context_free.ml
index 73b2ab22e..00a8acecb 100644
--- a/src/context_free.ml
+++ b/src/context_free.ml
@@ -197,33 +197,33 @@ module Generated_code_hook = struct
     | _ -> t.f context { loc with loc_start = loc.loc_end } x
 end
 
-let rec map_node_rec context ts super_call loc path x =
-  let ctxt = Expansion_context.Extension.make ~extension_point_loc:loc ~code_path:path in
+let rec map_node_rec context ts super_call loc base_ctxt x =
+  let ctxt = Expansion_context.Extension.make ~extension_point_loc:loc ~base:base_ctxt in
   match EC.get_extension context x with
-  | None -> super_call path x
+  | None -> super_call base_ctxt x
   | Some (ext, attrs) ->
     match E.For_context.convert ts ~ctxt ext with
-    | None -> super_call path x
+    | None -> super_call base_ctxt x
     | Some x ->
-      map_node_rec context ts super_call loc path (EC.merge_attributes context x attrs)
+      map_node_rec context ts super_call loc base_ctxt (EC.merge_attributes context x attrs)
 ;;
 
-let map_node context ts super_call loc path x ~hook =
-  let ctxt = Expansion_context.Extension.make ~extension_point_loc:loc ~code_path:path in
+let map_node context ts super_call loc base_ctxt x ~hook =
+  let ctxt = Expansion_context.Extension.make ~extension_point_loc:loc ~base:base_ctxt in
   match EC.get_extension context x with
-  | None -> super_call path x
+  | None -> super_call base_ctxt x
   | Some (ext, attrs) ->
     match E.For_context.convert ts ~ctxt ext with
-    | None -> super_call path x
+    | None -> super_call base_ctxt x
     | Some x ->
       let generated_code =
-        map_node_rec context ts super_call loc path (EC.merge_attributes context x attrs)
+        map_node_rec context ts super_call loc base_ctxt (EC.merge_attributes context x attrs)
       in
       Generated_code_hook.replace hook context loc (Single generated_code);
       generated_code
 ;;
 
-let rec map_nodes context ts super_call get_loc path l ~hook ~in_generated_code =
+let rec map_nodes context ts super_call get_loc base_ctxt l ~hook ~in_generated_code =
   match l with
   | [] -> []
   | x :: l ->
@@ -231,29 +231,29 @@ let rec map_nodes context ts super_call get_loc path l ~hook ~in_generated_code
     | None ->
       (* These two lets force the evaluation order, so that errors are reported in the
          same order as they appear in the source file. *)
-      let x = super_call path x in
-      let l = map_nodes context ts super_call get_loc path l ~hook ~in_generated_code in
+      let x = super_call base_ctxt x in
+      let l = map_nodes context ts super_call get_loc base_ctxt l ~hook ~in_generated_code in
       x :: l
     | Some (ext, attrs) ->
       let extension_point_loc = get_loc x in
-      let ctxt = Expansion_context.Extension.make ~extension_point_loc ~code_path:path in
+      let ctxt = Expansion_context.Extension.make ~extension_point_loc ~base:base_ctxt in
       match E.For_context.convert_inline ts ~ctxt ext with
       | None ->
-        let x = super_call path x in
+        let x = super_call base_ctxt x in
         let l =
-          map_nodes context ts super_call get_loc path l ~hook ~in_generated_code
+          map_nodes context ts super_call get_loc base_ctxt l ~hook ~in_generated_code
         in
         x :: l
       | Some x ->
         assert_no_attributes attrs;
         let generated_code =
-          map_nodes context ts super_call get_loc path x ~hook
+          map_nodes context ts super_call get_loc base_ctxt x ~hook
             ~in_generated_code:true
         in
         if not in_generated_code then
           Generated_code_hook.replace hook context extension_point_loc (Many generated_code);
         generated_code
-        @ map_nodes context ts super_call get_loc path l ~hook ~in_generated_code
+        @ map_nodes context ts super_call get_loc base_ctxt l ~hook ~in_generated_code
 
 let map_nodes = map_nodes ~in_generated_code:false
 
@@ -315,23 +315,23 @@ let sort_attr_inline l =
    This complexity is horrible, but in practice we don't care as [attrs] is always a list
    of one element; it only has [@@deriving].
 *)
-let handle_attr_group_inline attrs rf items ~loc ~path =
+let handle_attr_group_inline attrs rf items ~loc ~base_ctxt =
   List.fold_left attrs ~init:[]
     ~f:(fun acc (Rule.Attr_group_inline.T group) ->
       match get_group group.attribute items with
       | None -> acc
       | Some values ->
-        let ctxt = Expansion_context.Deriver.make ~derived_item_loc:loc ~code_path:path in
+        let ctxt = Expansion_context.Deriver.make ~derived_item_loc:loc ~base:base_ctxt in
         let expect_items = group.expand ~ctxt rf items values in
         expect_items :: acc)
 
-let handle_attr_inline attrs item ~loc ~path =
+let handle_attr_inline attrs item ~loc ~base_ctxt =
   List.fold_left attrs ~init:[]
     ~f:(fun acc (Rule.Attr_inline.T a) ->
       match Attribute.get a.attribute item with
       | None -> acc
       | Some value ->
-        let ctxt = Expansion_context.Deriver.make ~derived_item_loc:loc ~code_path:path in
+        let ctxt = Expansion_context.Deriver.make ~derived_item_loc:loc ~base:base_ctxt in
         let expect_items = a.expand ~ctxt item value in
         expect_items :: acc)
 
@@ -405,124 +405,124 @@ class map_top_down ?(expect_mismatch_handler=Expect_mismatch_handler.nop)
   let map_nodes = map_nodes ~hook in
 
   object(self)
-    inherit Ast_traverse.map_with_code_path as super
+    inherit Ast_traverse.map_with_expansion_context as super
 
     (* No point recursing into every location *)
     method! location _ x = x
 
-    method! core_type path x =
-      map_node EC.core_type core_type super#core_type x.ptyp_loc path x
+    method! core_type base_ctxt x =
+      map_node EC.core_type core_type super#core_type x.ptyp_loc base_ctxt x
 
-    method! pattern path x =
-      map_node EC.pattern pattern super#pattern x.ppat_loc path x
+    method! pattern base_ctxt x =
+      map_node EC.pattern pattern super#pattern x.ppat_loc base_ctxt x
 
-    method! expression path e =
+    method! expression base_ctxt e =
       let e =
         match e.pexp_desc with
         | Pexp_extension _ ->
-          map_node EC.expression expression (fun _ e -> e) e.pexp_loc path e
+          map_node EC.expression expression (fun _ e -> e) e.pexp_loc base_ctxt e
         | _ -> e
       in
       let expand_constant kind char text =
         match Hashtbl.find constants (char,kind) with
-        | None -> super#expression path e
-        | Some expand -> self#expression path (expand e.pexp_loc text)
+        | None -> super#expression base_ctxt e
+        | Some expand -> self#expression base_ctxt (expand e.pexp_loc text)
       in
       match e.pexp_desc with
       | Pexp_apply ({ pexp_desc = Pexp_ident id; _ } as func, args) -> begin
           match Hashtbl.find special_functions id.txt with
           | None ->
-            self#pexp_apply_without_traversing_function path e func args
+            self#pexp_apply_without_traversing_function base_ctxt e func args
           | Some pattern ->
             match pattern e with
             | None ->
-              self#pexp_apply_without_traversing_function path e func args
+              self#pexp_apply_without_traversing_function base_ctxt e func args
             | Some e ->
-              self#expression path e
+              self#expression base_ctxt e
         end
       | Pexp_ident id -> begin
           match Hashtbl.find special_functions id.txt with
           | None ->
-            super#expression path e
+            super#expression base_ctxt e
           | Some pattern ->
             match pattern e with
             | None ->
-              super#expression path e
+              super#expression base_ctxt e
             | Some e ->
-              self#expression path e
+              self#expression base_ctxt e
         end
       | Pexp_constant (Pconst_integer (s, Some c)) -> expand_constant Integer c s
       | Pexp_constant (Pconst_float   (s, Some c)) -> expand_constant Float   c s
       | _ ->
-        super#expression path e
+        super#expression base_ctxt e
 
     (* Pre-conditions:
        - e.pexp_desc = Pexp_apply(func, args)
        - func.pexp_desc = Pexp_ident _
     *)
-    method private pexp_apply_without_traversing_function path e func args =
+    method private pexp_apply_without_traversing_function base_ctxt e func args =
       let { pexp_desc = _; pexp_loc; pexp_attributes } = e in
       let func =
         let { pexp_desc; pexp_loc; pexp_attributes } = func in
-        let pexp_attributes = self#attributes path pexp_attributes in
+        let pexp_attributes = self#attributes base_ctxt pexp_attributes in
         { pexp_desc
         ; pexp_loc (* location doesn't need to be traversed *)
         ; pexp_attributes
         }
       in
-      let args = List.map args ~f:(fun (lab, exp) -> (lab, self#expression path exp)) in
-      let pexp_attributes = self#attributes path pexp_attributes in
+      let args = List.map args ~f:(fun (lab, exp) -> (lab, self#expression base_ctxt exp)) in
+      let pexp_attributes = self#attributes base_ctxt pexp_attributes in
       { pexp_loc
       ; pexp_attributes
       ; pexp_desc = Pexp_apply (func, args)
       }
 
-    method! class_type path x =
-      map_node EC.class_type class_type super#class_type x.pcty_loc path x
+    method! class_type base_ctxt x =
+      map_node EC.class_type class_type super#class_type x.pcty_loc base_ctxt x
 
-    method! class_type_field path x =
+    method! class_type_field base_ctxt x =
       map_node EC.class_type_field class_type_field super#class_type_field
-        x.pctf_loc path x
+        x.pctf_loc base_ctxt x
 
-    method! class_expr path x =
-      map_node EC.class_expr class_expr super#class_expr x.pcl_loc path x
+    method! class_expr base_ctxt x =
+      map_node EC.class_expr class_expr super#class_expr x.pcl_loc base_ctxt x
 
-    method! class_field path x =
-      map_node EC.class_field class_field super#class_field x.pcf_loc path x
+    method! class_field base_ctxt x =
+      map_node EC.class_field class_field super#class_field x.pcf_loc base_ctxt x
 
-    method! module_type path x =
-      map_node EC.module_type module_type super#module_type x.pmty_loc path x
+    method! module_type base_ctxt x =
+      map_node EC.module_type module_type super#module_type x.pmty_loc base_ctxt x
 
-    method! module_expr path x =
-      map_node EC.module_expr module_expr super#module_expr x.pmod_loc path x
+    method! module_expr base_ctxt x =
+      map_node EC.module_expr module_expr super#module_expr x.pmod_loc base_ctxt x
 
-    method! structure_item path x =
-      map_node EC.structure_item structure_item super#structure_item x.pstr_loc path x
+    method! structure_item base_ctxt x =
+      map_node EC.structure_item structure_item super#structure_item x.pstr_loc base_ctxt x
 
-    method! signature_item path x =
-      map_node EC.signature_item signature_item super#signature_item x.psig_loc path x
+    method! signature_item base_ctxt x =
+      map_node EC.signature_item signature_item super#signature_item x.psig_loc base_ctxt x
 
-    method! class_structure path { pcstr_self; pcstr_fields } =
-      let pcstr_self = self#pattern path pcstr_self in
+    method! class_structure base_ctxt { pcstr_self; pcstr_fields } =
+      let pcstr_self = self#pattern base_ctxt pcstr_self in
       let pcstr_fields =
         map_nodes EC.class_field class_field super#class_field
-          (fun x -> x.pcf_loc) path pcstr_fields
+          (fun x -> x.pcf_loc) base_ctxt pcstr_fields
       in
       { pcstr_self; pcstr_fields }
 
-    method! class_signature path { pcsig_self; pcsig_fields } =
-      let pcsig_self = self#core_type path pcsig_self in
+    method! class_signature base_ctxt { pcsig_self; pcsig_fields } =
+      let pcsig_self = self#core_type base_ctxt pcsig_self in
       let pcsig_fields =
         map_nodes EC.class_type_field class_type_field super#class_type_field
-          (fun x -> x.pctf_loc) path pcsig_fields
+          (fun x -> x.pctf_loc) base_ctxt pcsig_fields
       in
       { pcsig_self; pcsig_fields }
 
     (* TODO: try to factorize #structure and #signature without meta-programming *)
     (*$*)
-    method! structure path st =
+    method! structure base_ctxt st =
       let rec with_extra_items item ~extra_items ~expect_items ~rest ~in_generated_code =
-        let item = super#structure_item path item in
+        let item = super#structure_item base_ctxt item in
         let extra_items = loop (rev_concat extra_items) ~in_generated_code:true in
         if not in_generated_code then
           Generated_code_hook.insert_after hook Structure_item item.pstr_loc
@@ -545,11 +545,11 @@ class map_top_down ?(expect_mismatch_handler=Expect_mismatch_handler.nop)
           match item.pstr_desc with
           | Pstr_extension (ext, attrs) -> begin
               let extension_point_loc = item.pstr_loc in
-              let ctxt = Expansion_context.Extension.make ~extension_point_loc ~code_path:path in
+              let ctxt = Expansion_context.Extension.make ~extension_point_loc ~base:base_ctxt in
               match E.For_context.convert_inline structure_item ~ctxt ext with
               | None ->
-                let item = super#structure_item path item in
-                let rest = self#structure path rest in
+                let item = super#structure_item base_ctxt item in
+                let rest = self#structure base_ctxt rest in
                 item :: rest
               | Some items ->
                 assert_no_attributes attrs;
@@ -562,38 +562,38 @@ class map_top_down ?(expect_mismatch_handler=Expect_mismatch_handler.nop)
 
           | Pstr_type(rf, tds) ->
             let extra_items =
-              handle_attr_group_inline attr_str_type_decls rf tds ~loc ~path
+              handle_attr_group_inline attr_str_type_decls rf tds ~loc ~base_ctxt
             in
             let expect_items =
-              handle_attr_group_inline attr_str_type_decls_expect rf tds ~loc ~path
+              handle_attr_group_inline attr_str_type_decls_expect rf tds ~loc ~base_ctxt
             in
             with_extra_items item ~extra_items ~expect_items ~rest ~in_generated_code
 
           | Pstr_typext te ->
-            let extra_items = handle_attr_inline attr_str_type_exts te ~loc ~path in
+            let extra_items = handle_attr_inline attr_str_type_exts te ~loc ~base_ctxt in
             let expect_items =
-              handle_attr_inline attr_str_type_exts_expect te ~loc ~path
+              handle_attr_inline attr_str_type_exts_expect te ~loc ~base_ctxt
             in
             with_extra_items item ~extra_items ~expect_items ~rest ~in_generated_code
 
           | Pstr_exception ec ->
-            let extra_items = handle_attr_inline attr_str_exceptions ec ~loc ~path in
+            let extra_items = handle_attr_inline attr_str_exceptions ec ~loc ~base_ctxt in
             let expect_items =
-              handle_attr_inline attr_str_exceptions_expect ec ~loc ~path
+              handle_attr_inline attr_str_exceptions_expect ec ~loc ~base_ctxt
             in
             with_extra_items item ~extra_items ~expect_items ~rest ~in_generated_code
 
           | _ ->
-            let item = self#structure_item path item in
-            let rest = self#structure path rest in
+            let item = self#structure_item base_ctxt item in
+            let rest = self#structure base_ctxt rest in
             item :: rest
       in
       loop st ~in_generated_code:false
 
     (*$ str_to_sig _last_text_block *)
-    method! signature path sg =
+    method! signature base_ctxt sg =
       let rec with_extra_items item ~extra_items ~expect_items ~rest ~in_generated_code =
-        let item = super#signature_item path item in
+        let item = super#signature_item base_ctxt item in
         let extra_items = loop (rev_concat extra_items) ~in_generated_code:true in
         if not in_generated_code then
           Generated_code_hook.insert_after hook Signature_item item.psig_loc
@@ -616,11 +616,11 @@ class map_top_down ?(expect_mismatch_handler=Expect_mismatch_handler.nop)
           match item.psig_desc with
           | Psig_extension (ext, attrs) -> begin
               let extension_point_loc = item.psig_loc in
-              let ctxt = Expansion_context.Extension.make ~extension_point_loc ~code_path:path in
+              let ctxt = Expansion_context.Extension.make ~extension_point_loc ~base:base_ctxt in
               match E.For_context.convert_inline signature_item ~ctxt ext with
               | None ->
-                let item = super#signature_item path item in
-                let rest = self#signature path rest in
+                let item = super#signature_item base_ctxt item in
+                let rest = self#signature base_ctxt rest in
                 item :: rest
               | Some items ->
                 assert_no_attributes attrs;
@@ -633,30 +633,30 @@ class map_top_down ?(expect_mismatch_handler=Expect_mismatch_handler.nop)
 
           | Psig_type(rf, tds) ->
             let extra_items =
-              handle_attr_group_inline attr_sig_type_decls rf tds ~loc ~path
+              handle_attr_group_inline attr_sig_type_decls rf tds ~loc ~base_ctxt
             in
             let expect_items =
-              handle_attr_group_inline attr_sig_type_decls_expect rf tds ~loc ~path
+              handle_attr_group_inline attr_sig_type_decls_expect rf tds ~loc ~base_ctxt
             in
             with_extra_items item ~extra_items ~expect_items ~rest ~in_generated_code
 
           | Psig_typext te ->
-            let extra_items = handle_attr_inline attr_sig_type_exts te ~loc ~path in
+            let extra_items = handle_attr_inline attr_sig_type_exts te ~loc ~base_ctxt in
             let expect_items =
-              handle_attr_inline attr_sig_type_exts_expect te ~loc ~path
+              handle_attr_inline attr_sig_type_exts_expect te ~loc ~base_ctxt
             in
             with_extra_items item ~extra_items ~expect_items ~rest ~in_generated_code
 
           | Psig_exception ec ->
-            let extra_items = handle_attr_inline attr_sig_exceptions ec ~loc ~path in
+            let extra_items = handle_attr_inline attr_sig_exceptions ec ~loc ~base_ctxt in
             let expect_items =
-              handle_attr_inline attr_sig_exceptions_expect ec ~loc ~path
+              handle_attr_inline attr_sig_exceptions_expect ec ~loc ~base_ctxt
             in
             with_extra_items item ~extra_items ~expect_items ~rest ~in_generated_code
 
           | _ ->
-            let item = self#signature_item path item in
-            let rest = self#signature path rest in
+            let item = self#signature_item base_ctxt item in
+            let rest = self#signature base_ctxt rest in
             item :: rest
       in
       loop sg ~in_generated_code:false
diff --git a/src/context_free.mli b/src/context_free.mli
index 49236044c..889eea8b1 100644
--- a/src/context_free.mli
+++ b/src/context_free.mli
@@ -134,4 +134,4 @@ class map_top_down
     -> ?generated_code_hook:Generated_code_hook.t
     (* default: Generated_code_hook.nop *)
     -> Rule.t list
-    -> Ast_traverse.map_with_code_path
+    -> Ast_traverse.map_with_expansion_context
diff --git a/src/driver.ml b/src/driver.ml
index a75012ca4..c06618763 100644
--- a/src/driver.ml
+++ b/src/driver.ml
@@ -132,7 +132,7 @@ module Transform = struct
       Some { first with loc_end = last.loc_end }
   ;;
 
-  let merge_into_generic_mappers t ~hook ~expect_mismatch_handler =
+  let merge_into_generic_mappers t ~hook ~expect_mismatch_handler ~omp_config =
     let { rules; enclose_impl; enclose_intf; impl; intf; _ } = t in
     let map =
       new Context_free.map_top_down rules
@@ -173,7 +173,8 @@ module Transform = struct
             gen_header_and_footer Structure_item whole_loc f
         in
         let file_path = File_path.get_default_path_str st in
-        let st = map#structure (Code_path.top_level ~file_path) st in
+        let base_ctxt = Expansion_context.Base.top_level ~omp_config ~file_path in
+        let st = map#structure base_ctxt st in
         match header, footer with
         | [], [] -> st
         | _      -> List.concat [ header; st; footer ]
@@ -192,7 +193,8 @@ module Transform = struct
             gen_header_and_footer Signature_item whole_loc f
         in
         let file_path = File_path.get_default_path_sig sg in
-        let sg = map#signature (Code_path.top_level ~file_path) sg in
+        let base_ctxt = Expansion_context.Base.top_level ~omp_config ~file_path in
+        let sg = map#signature base_ctxt sg in
         match header, footer with
         | [], [] -> sg
         | _      -> List.concat [ header; sg; footer ]
@@ -301,7 +303,7 @@ let debug_dropped_attribute name ~old_dropped ~new_dropped =
   print_diff "reappeared"  old_dropped new_dropped
 ;;
 
-let get_whole_ast_passes ~hook ~expect_mismatch_handler =
+let get_whole_ast_passes ~hook ~expect_mismatch_handler ~omp_config =
   let cts =
     match !apply_list with
     | None -> List.rev !Transform.all
@@ -320,7 +322,7 @@ let get_whole_ast_passes ~hook ~expect_mismatch_handler =
   end;
   let cts =
     if !no_merge then
-      List.map cts ~f:(Transform.merge_into_generic_mappers ~hook
+      List.map cts ~f:(Transform.merge_into_generic_mappers ~hook ~omp_config
                          ~expect_mismatch_handler)
     else begin
       let get_enclosers ~f =
@@ -357,6 +359,7 @@ let get_whole_ast_passes ~hook ~expect_mismatch_handler =
         Transform.builtin_of_context_free_rewriters ~rules ~hook ~expect_mismatch_handler
           ~enclose_impl:(merge_encloser impl_enclosers)
           ~enclose_intf:(merge_encloser intf_enclosers)
+          ~omp_config
         :: cts
     end
   in linters @ preprocess @ List.filter cts ~f:(fun (ct : Transform.t) ->
@@ -365,24 +368,9 @@ let get_whole_ast_passes ~hook ~expect_mismatch_handler =
     | _          -> true)
 ;;
 
-let print_passes () =
-  let cts =
-    get_whole_ast_passes ~hook:Context_free.Generated_code_hook.nop
-      ~expect_mismatch_handler:Context_free.Expect_mismatch_handler.nop
-  in
-  if !perform_checks then
-    printf "<builtin:freshen-and-collect-attributes>\n";
-  List.iter cts ~f:(fun ct -> printf "%s\n" ct.Transform.name);
-  if !perform_checks then
-    begin
-      printf "<builtin:check-unused-attributes>\n";
-      if !perform_checks_on_extensions
-      then printf "<builtin:check-unused-extensions>\n"
-    end
-;;
-
-let apply_transforms ~field ~lint_field ~dropped_so_far ~hook ~expect_mismatch_handler x =
-  let cts = get_whole_ast_passes ~hook ~expect_mismatch_handler in
+let apply_transforms
+      ~omp_config ~field ~lint_field ~dropped_so_far ~hook ~expect_mismatch_handler x =
+  let cts = get_whole_ast_passes ~omp_config ~hook ~expect_mismatch_handler in
   let x, _dropped, lint_errors =
   List.fold_left cts ~init:(x, [], [])
     ~f:(fun (x, dropped, lint_errors) (ct : Transform.t) ->
@@ -456,6 +444,22 @@ let as_ppx_config () =
     ~debug:!Ocaml_common.Clflags.debug
     ?for_package:!Ocaml_common.Clflags.for_package
 
+let print_passes () =
+  let hook = Context_free.Generated_code_hook.nop in
+  let expect_mismatch_handler = Context_free.Expect_mismatch_handler.nop in
+  let omp_config = config ~hook ~expect_mismatch_handler in
+  let cts = get_whole_ast_passes ~hook ~expect_mismatch_handler ~omp_config in
+  if !perform_checks then
+    printf "<builtin:freshen-and-collect-attributes>\n";
+  List.iter cts ~f:(fun ct -> printf "%s\n" ct.Transform.name);
+  if !perform_checks then
+    begin
+      printf "<builtin:check-unused-attributes>\n";
+      if !perform_checks_on_extensions
+      then printf "<builtin:check-unused-extensions>\n"
+    end
+;;
+
 (*$*)
 let real_map_structure config cookies st =
   let { C. hook; expect_mismatch_handler } = C.find config in
@@ -466,6 +470,7 @@ let real_map_structure config cookies st =
   end;
   let st, lint_errors =
     apply_transforms st
+      ~omp_config:config
       ~field:(fun (ct : Transform.t) -> ct.impl)
       ~lint_field:(fun (ct : Transform.t) -> ct.lint_impl)
       ~dropped_so_far:Attribute.dropped_so_far_structure ~hook ~expect_mismatch_handler
@@ -508,6 +513,7 @@ let real_map_signature config cookies sg =
   end;
   let sg, lint_errors =
     apply_transforms sg
+      ~omp_config:config
       ~field:(fun (ct : Transform.t) -> ct.intf)
       ~lint_field:(fun (ct : Transform.t) -> ct.lint_intf)
       ~dropped_so_far:Attribute.dropped_so_far_signature ~hook ~expect_mismatch_handler
diff --git a/src/expansion_context.ml b/src/expansion_context.ml
index 797df510f..926f2a9ec 100644
--- a/src/expansion_context.ml
+++ b/src/expansion_context.ml
@@ -1,29 +1,47 @@
+module Base = struct
+  type t =
+    { omp_config : Migrate_parsetree.Driver.config
+    ; code_path : Code_path.t
+    }
+
+  let top_level ~omp_config ~file_path =
+    let code_path = Code_path.top_level ~file_path in
+    {omp_config; code_path}
+
+  let enter_expr t = {t with code_path = Code_path.enter_expr t.code_path}
+  let enter_module ~loc name t = {t with code_path = Code_path.enter_module ~loc name t.code_path}
+  let enter_value ~loc name t = {t with code_path = Code_path.enter_value ~loc name t.code_path}
+end
+
 module Extension = struct
   type t =
     { extension_point_loc : Location.t
-    ; code_path : Code_path.t
+    ; base : Base.t
     }
 
-  let make ~extension_point_loc ~code_path = {extension_point_loc; code_path}
+  let make ~extension_point_loc ~base = {extension_point_loc; base}
 
   let extension_point_loc t = t.extension_point_loc
-  let code_path t = t.code_path
+  let code_path t = t.base.code_path
+  let omp_config t = t.base.omp_config
 
   let with_loc_and_path f =
-    fun ~ctxt -> f ~loc:ctxt.extension_point_loc ~path:(Code_path.to_string_path ctxt.code_path)
+    fun ~ctxt ->
+      f ~loc:ctxt.extension_point_loc ~path:(Code_path.to_string_path ctxt.base.code_path)
 end
 
 module Deriver = struct
   type t =
     { derived_item_loc : Location.t
-    ; code_path : Code_path.t
+    ; base : Base.t
     }
 
-  let make ~derived_item_loc ~code_path = {derived_item_loc; code_path}
+  let make ~derived_item_loc ~base = {derived_item_loc; base}
 
   let derived_item_loc t = t.derived_item_loc
-  let code_path t = t.code_path
+  let code_path t = t.base.code_path
+  let omp_config t = t.base.omp_config
 
   let with_loc_and_path f =
-    fun ~ctxt -> f ~loc:ctxt.derived_item_loc ~path:(Code_path.to_string_path ctxt.code_path)
+    fun ~ctxt -> f ~loc:ctxt.derived_item_loc ~path:(Code_path.to_string_path ctxt.base.code_path)
 end
diff --git a/src/expansion_context.mli b/src/expansion_context.mli
index c94e32a9e..70bbd3d04 100644
--- a/src/expansion_context.mli
+++ b/src/expansion_context.mli
@@ -1,3 +1,24 @@
+module Base : sig
+  (** Type for the location independent parts of the expansion context *)
+  type t
+
+  (**/*)
+  (** Undocumented section *)
+
+  (** Build a new base context at the top level of the given file with the given
+      ocaml-mirgate-parsetree configuration.
+  *)
+  val top_level :
+    omp_config:Migrate_parsetree.Driver.config ->
+    file_path:string ->
+    t
+
+  (** Proxy functions to update the wrapped code path. See code_path.mli for details. *)
+  val enter_expr : t -> t
+  val enter_module : loc:Location.t -> string -> t -> t
+  val enter_value : loc:Location.t -> string -> t -> t
+end
+
 module Extension : sig
   (** Type of expansion contexts for extensions *)
   type t
@@ -8,14 +29,17 @@ module Extension : sig
   (** Return the code path for the given context *)
   val code_path : t -> Code_path.t
 
+  (** Return the ocaml-migrate-parsetree configuration for the given expansion context *)
+  val omp_config : t -> Migrate_parsetree.Driver.config
+
   (** Wrap a [fun ~loc ~path] into a [fun ~ctxt] *)
   val with_loc_and_path : (loc:Location.t -> path:string -> 'a) -> (ctxt:t -> 'a)
 
   (**/**)
   (** Undocumented section *)
 
-  (** Build a new expansion context with the given extension point location and code path *)
-  val make : extension_point_loc:Location.t -> code_path:Code_path.t -> t
+  (** Build a new expansion context with the given extension point location and base context *)
+  val make : extension_point_loc:Location.t -> base:Base.t -> t
 end
 
 module Deriver : sig
@@ -28,6 +52,9 @@ module Deriver : sig
   (** Return the code path for the given context *)
   val code_path : t -> Code_path.t
 
+  (** Return the ocaml-migrate-parsetree configuration for the given expansion context *)
+  val omp_config : t -> Migrate_parsetree.Driver.config
+
   (** Wrap a [fun ~loc ~path] into a [fun ~ctxt] *)
   val with_loc_and_path : (loc:Location.t -> path:string -> 'a) -> (ctxt:t -> 'a)
 
@@ -35,5 +62,5 @@ module Deriver : sig
   (** Undocumented section *)
 
   (** Build a new expansion context with the given item location and code path *)
-  val make : derived_item_loc:Location.t -> code_path:Code_path.t -> t
+  val make : derived_item_loc:Location.t -> base:Base.t -> t
 end

From 6a6eba4a40a083393fe12a00d3e711446d45ab03 Mon Sep 17 00:00:00 2001
From: Nathan Rebours <nathan.p.rebours@gmail.com>
Date: Tue, 2 Apr 2019 18:59:37 +0200
Subject: [PATCH 2/2] Add a unit argument to Expansion_context's smart
 constructors

Signed-off-by: Nathan Rebours <nathan.p.rebours@gmail.com>
---
 src/context_free.ml       | 14 +++++++-------
 src/expansion_context.ml  |  4 ++--
 src/expansion_context.mli |  4 ++--
 3 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/src/context_free.ml b/src/context_free.ml
index 00a8acecb..0c35258e5 100644
--- a/src/context_free.ml
+++ b/src/context_free.ml
@@ -198,7 +198,7 @@ module Generated_code_hook = struct
 end
 
 let rec map_node_rec context ts super_call loc base_ctxt x =
-  let ctxt = Expansion_context.Extension.make ~extension_point_loc:loc ~base:base_ctxt in
+  let ctxt = Expansion_context.Extension.make ~extension_point_loc:loc ~base:base_ctxt () in
   match EC.get_extension context x with
   | None -> super_call base_ctxt x
   | Some (ext, attrs) ->
@@ -209,7 +209,7 @@ let rec map_node_rec context ts super_call loc base_ctxt x =
 ;;
 
 let map_node context ts super_call loc base_ctxt x ~hook =
-  let ctxt = Expansion_context.Extension.make ~extension_point_loc:loc ~base:base_ctxt in
+  let ctxt = Expansion_context.Extension.make ~extension_point_loc:loc ~base:base_ctxt () in
   match EC.get_extension context x with
   | None -> super_call base_ctxt x
   | Some (ext, attrs) ->
@@ -236,7 +236,7 @@ let rec map_nodes context ts super_call get_loc base_ctxt l ~hook ~in_generated_
       x :: l
     | Some (ext, attrs) ->
       let extension_point_loc = get_loc x in
-      let ctxt = Expansion_context.Extension.make ~extension_point_loc ~base:base_ctxt in
+      let ctxt = Expansion_context.Extension.make ~extension_point_loc ~base:base_ctxt () in
       match E.For_context.convert_inline ts ~ctxt ext with
       | None ->
         let x = super_call base_ctxt x in
@@ -321,7 +321,7 @@ let handle_attr_group_inline attrs rf items ~loc ~base_ctxt =
       match get_group group.attribute items with
       | None -> acc
       | Some values ->
-        let ctxt = Expansion_context.Deriver.make ~derived_item_loc:loc ~base:base_ctxt in
+        let ctxt = Expansion_context.Deriver.make ~derived_item_loc:loc ~base:base_ctxt () in
         let expect_items = group.expand ~ctxt rf items values in
         expect_items :: acc)
 
@@ -331,7 +331,7 @@ let handle_attr_inline attrs item ~loc ~base_ctxt =
       match Attribute.get a.attribute item with
       | None -> acc
       | Some value ->
-        let ctxt = Expansion_context.Deriver.make ~derived_item_loc:loc ~base:base_ctxt in
+        let ctxt = Expansion_context.Deriver.make ~derived_item_loc:loc ~base:base_ctxt () in
         let expect_items = a.expand ~ctxt item value in
         expect_items :: acc)
 
@@ -545,7 +545,7 @@ class map_top_down ?(expect_mismatch_handler=Expect_mismatch_handler.nop)
           match item.pstr_desc with
           | Pstr_extension (ext, attrs) -> begin
               let extension_point_loc = item.pstr_loc in
-              let ctxt = Expansion_context.Extension.make ~extension_point_loc ~base:base_ctxt in
+              let ctxt = Expansion_context.Extension.make ~extension_point_loc ~base:base_ctxt () in
               match E.For_context.convert_inline structure_item ~ctxt ext with
               | None ->
                 let item = super#structure_item base_ctxt item in
@@ -616,7 +616,7 @@ class map_top_down ?(expect_mismatch_handler=Expect_mismatch_handler.nop)
           match item.psig_desc with
           | Psig_extension (ext, attrs) -> begin
               let extension_point_loc = item.psig_loc in
-              let ctxt = Expansion_context.Extension.make ~extension_point_loc ~base:base_ctxt in
+              let ctxt = Expansion_context.Extension.make ~extension_point_loc ~base:base_ctxt () in
               match E.For_context.convert_inline signature_item ~ctxt ext with
               | None ->
                 let item = super#signature_item base_ctxt item in
diff --git a/src/expansion_context.ml b/src/expansion_context.ml
index 926f2a9ec..a6d848664 100644
--- a/src/expansion_context.ml
+++ b/src/expansion_context.ml
@@ -19,7 +19,7 @@ module Extension = struct
     ; base : Base.t
     }
 
-  let make ~extension_point_loc ~base = {extension_point_loc; base}
+  let make ~extension_point_loc ~base () = {extension_point_loc; base}
 
   let extension_point_loc t = t.extension_point_loc
   let code_path t = t.base.code_path
@@ -36,7 +36,7 @@ module Deriver = struct
     ; base : Base.t
     }
 
-  let make ~derived_item_loc ~base = {derived_item_loc; base}
+  let make ~derived_item_loc ~base () = {derived_item_loc; base}
 
   let derived_item_loc t = t.derived_item_loc
   let code_path t = t.base.code_path
diff --git a/src/expansion_context.mli b/src/expansion_context.mli
index 70bbd3d04..46744eae2 100644
--- a/src/expansion_context.mli
+++ b/src/expansion_context.mli
@@ -39,7 +39,7 @@ module Extension : sig
   (** Undocumented section *)
 
   (** Build a new expansion context with the given extension point location and base context *)
-  val make : extension_point_loc:Location.t -> base:Base.t -> t
+  val make : extension_point_loc:Location.t -> base:Base.t -> unit -> t
 end
 
 module Deriver : sig
@@ -62,5 +62,5 @@ module Deriver : sig
   (** Undocumented section *)
 
   (** Build a new expansion context with the given item location and code path *)
-  val make : derived_item_loc:Location.t -> base:Base.t -> t
+  val make : derived_item_loc:Location.t -> base:Base.t -> unit -> t
 end