Extend ast_mapper to allow mapping locations.

git-svn-id: http://caml.inria.fr/svn/ocaml/trunk@13210 f963ae5c-01c2-4b8c-9fe0-0dff7051ff02
thelema · Jan 7, 2013 · 42a339e · 42a339e
1 parent f7721f0
commit 42a339e
Showing 1 changed file with 93 additions and 62 deletions.
diff --git a/tools/ast_mapper.ml b/tools/ast_mapper.ml
@@ -24,6 +24,8 @@ let map_snd f (x, y) = (x, f y)
 let map_tuple f1 f2 (x, y) = (f1 x, f2 y)
 let map_opt f = function None -> None | Some x -> Some (f x)
 
+let map_loc sub {loc; txt} = {loc = sub # location loc; txt}
+
 module T = struct
   (* Type expressions for the core language *)
 
@@ -52,42 +54,46 @@ module T = struct
     field_type ?loc (Pfield (s, t))
   let field_var ?loc () = field_type ?loc Pfield_var
 
-  let core_field_type sub = function
-    | {pfield_desc = Pfield (s, d); pfield_loc = loc} -> field ~loc s (sub # typ d)
-    | x -> x
+  let core_field_type sub {pfield_desc = desc; pfield_loc = loc} =
+    let loc = sub # location loc in
+    match desc with
+    | Pfield (s, d) -> field ~loc:(sub # location loc) s (sub # typ d)
+    | Pfield_var -> field_var ~loc ()
 
   let row_field sub = function
     | Rtag (l, b, tl) -> Rtag (l, b, List.map (sub # typ) tl)
     | Rinherit t -> Rinherit (sub # typ t)
 
   let map sub {ptyp_desc = desc; ptyp_loc = loc} =
+    let loc = sub # location loc in
     match desc with
     | Ptyp_any -> any ~loc ()
     | Ptyp_var s -> var ~loc s
     | Ptyp_arrow (lab, t1, t2) -> arrow ~loc lab (sub # typ t1) (sub # typ t2)
     | Ptyp_tuple tyl -> tuple ~loc (List.map (sub # typ) tyl)
-    | Ptyp_constr (lid, tl) -> constr ~loc lid (List.map (sub # typ) tl)
+    | Ptyp_constr (lid, tl) -> constr ~loc (map_loc sub lid) (List.map (sub # typ) tl)
     | Ptyp_object l -> object_ ~loc (List.map (core_field_type sub) l)
-    | Ptyp_class (lid, tl, ll) -> class_ ~loc lid (List.map (sub # typ) tl) ll
+    | Ptyp_class (lid, tl, ll) -> class_ ~loc (map_loc sub lid) (List.map (sub # typ) tl) ll
     | Ptyp_alias (t, s) -> alias ~loc (sub # typ t) s
     | Ptyp_variant (rl, b, ll) -> variant ~loc (List.map (row_field sub) rl) b ll
     | Ptyp_poly (sl, t) -> poly ~loc sl (sub # typ t)
-    | Ptyp_package (lid, l) -> package ~loc lid (List.map (map_snd (sub # typ)) l)
+    | Ptyp_package (lid, l) -> package ~loc (map_loc sub lid) (List.map (map_tuple (map_loc sub) (sub # typ)) l)
 
   let map_type_declaration sub td =
     {td with
      ptype_cstrs =
      List.map
-       (fun (ct1, ct2, loc) -> sub # typ ct1, sub # typ ct2, loc)
+       (fun (ct1, ct2, loc) -> sub # typ ct1, sub # typ ct2, sub # location loc)
        td.ptype_cstrs;
      ptype_kind = sub # type_kind td.ptype_kind;
      ptype_manifest = map_opt (sub # typ) td.ptype_manifest;
+     ptype_loc = sub # location td.ptype_loc;
     }
 
   let map_type_kind sub = function
     | Ptype_abstract -> Ptype_abstract
-    | Ptype_variant l -> Ptype_variant (List.map (fun (s, tl, t, loc) -> (s, List.map (sub # typ) tl, map_opt (sub # typ) t, loc)) l)
-    | Ptype_record l -> Ptype_record (List.map (fun (s, flags, t, loc) -> (s, flags, sub # typ t, loc)) l)
+    | Ptype_variant l -> Ptype_variant (List.map (fun (s, tl, t, loc) -> (map_loc sub s, List.map (sub # typ) tl, map_opt (sub # typ) t, sub # location loc)) l)
+    | Ptype_record l -> Ptype_record (List.map (fun (s, flags, t, loc) -> (map_loc sub s, flags, sub # typ t, sub # location loc)) l)
 end
 
 module CT = struct
@@ -100,8 +106,9 @@ module CT = struct
   let fun_ ?loc a b c = mk ?loc (Pcty_fun (a, b, c))
 
   let map sub {pcty_loc = loc; pcty_desc = desc} =
+    let loc = sub # location loc in
     match desc with
-    | Pcty_constr (lid, tys) -> constr ~loc lid (List.map (sub # typ) tys)
+    | Pcty_constr (lid, tys) -> constr ~loc (map_loc sub lid) (List.map (sub # typ) tys)
     | Pcty_signature x -> signature ~loc (sub # class_signature x)
     | Pcty_fun (lab, t, ct) ->
         fun_ ~loc lab
@@ -117,6 +124,7 @@ module CT = struct
   let cstr ?loc a b = mk_field ?loc (Pctf_cstr (a, b))
 
   let map_field sub {pctf_desc = desc; pctf_loc = loc} =
+    let loc = sub # location loc in
     match desc with
     | Pctf_inher ct -> inher ~loc (sub # class_type ct)
     | Pctf_val (s, m, v, t) -> val_ ~loc s m v (sub # typ t)
@@ -128,7 +136,7 @@ module CT = struct
     {
      pcsig_self = sub # typ pcsig_self;
      pcsig_fields = List.map (sub # class_type_field) pcsig_fields;
-     pcsig_loc;
+     pcsig_loc = sub # location pcsig_loc ;
     }
 end
 
@@ -143,18 +151,19 @@ module MT = struct
   let typeof_ ?loc a = mk ?loc (Pmty_typeof a)
 
   let map sub {pmty_desc = desc; pmty_loc = loc} =
+    let loc = sub # location loc in
     match desc with
-    | Pmty_ident s -> ident ~loc s
+    | Pmty_ident s -> ident ~loc (map_loc sub s)
     | Pmty_signature sg -> signature ~loc (sub # signature sg)
-    | Pmty_functor (s, mt1, mt2) -> functor_ ~loc s (sub # module_type mt1) (sub # module_type mt2)
-    | Pmty_with (mt, l) -> with_ ~loc (sub # module_type mt) (List.map (map_snd (sub # with_constraint)) l)
+    | Pmty_functor (s, mt1, mt2) -> functor_ ~loc (map_loc sub s) (sub # module_type mt1) (sub # module_type mt2)
+    | Pmty_with (mt, l) -> with_ ~loc (sub # module_type mt) (List.map (map_tuple (map_loc sub) (sub # with_constraint)) l)
     | Pmty_typeof me -> typeof_ ~loc (sub # module_expr me)
 
   let map_with_constraint sub = function
     | Pwith_type d -> Pwith_type (sub # type_declaration d)
-    | Pwith_module s -> Pwith_module s
+    | Pwith_module s -> Pwith_module (map_loc sub s)
     | Pwith_typesubst d -> Pwith_typesubst (sub # type_declaration d)
-    | Pwith_modsubst s -> Pwith_modsubst s
+    | Pwith_modsubst s -> Pwith_modsubst (map_loc sub s)
 
   let mk_item ?(loc = Location.none) x = {psig_desc = x; psig_loc = loc}
 
@@ -170,15 +179,16 @@ module MT = struct
   let class_type ?loc a = mk_item ?loc (Psig_class_type a)
 
   let map_signature_item sub {psig_desc = desc; psig_loc = loc} =
+    let loc = sub # location loc in
     match desc with
-    | Psig_value (s, vd) -> value ~loc s (sub # value_description vd)
-    | Psig_type l -> type_ ~loc (List.map (map_snd (sub # type_declaration)) l)
-    | Psig_exception (s, ed) -> exception_ ~loc s (sub # exception_declaration ed)
-    | Psig_module (s, mt) -> module_ ~loc s (sub # module_type mt)
-    | Psig_recmodule l -> rec_module ~loc (List.map (map_snd (sub # module_type)) l)
-    | Psig_modtype (s, Pmodtype_manifest mt) -> modtype ~loc s (Pmodtype_manifest  (sub # module_type mt))
-    | Psig_modtype (s, Pmodtype_abstract) -> modtype ~loc s Pmodtype_abstract
-    | Psig_open s -> open_ ~loc s
+    | Psig_value (s, vd) -> value ~loc (map_loc sub s) (sub # value_description vd)
+    | Psig_type l -> type_ ~loc (List.map (map_tuple (map_loc sub) (sub # type_declaration)) l)
+    | Psig_exception (s, ed) -> exception_ ~loc (map_loc sub s) (sub # exception_declaration ed)
+    | Psig_module (s, mt) -> module_ ~loc (map_loc sub s) (sub # module_type mt)
+    | Psig_recmodule l -> rec_module ~loc (List.map (map_tuple (map_loc sub) (sub # module_type)) l)
+    | Psig_modtype (s, Pmodtype_manifest mt) -> modtype ~loc (map_loc sub s) (Pmodtype_manifest  (sub # module_type mt))
+    | Psig_modtype (s, Pmodtype_abstract) -> modtype ~loc (map_loc sub s) Pmodtype_abstract
+    | Psig_open s -> open_ ~loc (map_loc sub s)
     | Psig_include mt -> include_ ~loc (sub # module_type mt)
     | Psig_class l -> class_ ~loc (List.map (sub # class_description) l)
     | Psig_class_type l -> class_type ~loc (List.map (sub # class_type_declaration) l)
@@ -198,10 +208,11 @@ module M = struct
   let unpack ?loc e = mk ?loc (Pmod_unpack e)
 
   let map sub {pmod_loc = loc; pmod_desc = desc} =
+    let loc = sub # location loc in
     match desc with
-    | Pmod_ident x -> ident ~loc x
+    | Pmod_ident x -> ident ~loc (map_loc sub x)
     | Pmod_structure str -> structure ~loc (sub # structure str)
-    | Pmod_functor (arg, arg_ty, body) -> functor_ ~loc arg (sub # module_type arg_ty) (sub # module_expr body)
+    | Pmod_functor (arg, arg_ty, body) -> functor_ ~loc (map_loc sub arg) (sub # module_type arg_ty) (sub # module_expr body)
     | Pmod_apply (m1, m2) -> apply ~loc (sub # module_expr m1) (sub # module_expr m2)
     | Pmod_constraint (m, mty) -> constraint_ ~loc (sub # module_expr m) (sub # module_type mty)
     | Pmod_unpack e -> unpack ~loc (sub # expr e)
@@ -222,17 +233,18 @@ module M = struct
   let include_ ?loc a = mk_item ?loc (Pstr_include a)
 
   let map_structure_item sub {pstr_loc = loc; pstr_desc = desc} =
+    let loc = sub # location loc in
     match desc with
     | Pstr_eval x -> eval ~loc (sub # expr x)
     | Pstr_value (r, pel) -> value ~loc r (List.map (map_tuple (sub # pat) (sub # expr)) pel)
-    | Pstr_primitive (name, vd) -> primitive ~loc name (sub # value_description vd)
-    | Pstr_type l -> type_ ~loc (List.map (fun (s, d) -> (s, sub # type_declaration d)) l)
-    | Pstr_exception (name, ed) -> exception_ ~loc name (sub # exception_declaration ed)
-    | Pstr_exn_rebind (s, lid) -> exn_rebind ~loc s lid
-    | Pstr_module (s, m) -> module_ ~loc s (sub # module_expr m)
-    | Pstr_recmodule l -> rec_module ~loc (List.map (fun (s, mty, me) -> (s, sub # module_type mty, sub # module_expr me)) l)
-    | Pstr_modtype (s, mty) -> modtype ~loc s (sub # module_type mty)
-    | Pstr_open lid -> open_ ~loc lid
+    | Pstr_primitive (name, vd) -> primitive ~loc (map_loc sub name) (sub # value_description vd)
+    | Pstr_type l -> type_ ~loc (List.map (map_tuple (map_loc sub) (sub # type_declaration)) l)
+    | Pstr_exception (name, ed) -> exception_ ~loc (map_loc sub name) (sub # exception_declaration ed)
+    | Pstr_exn_rebind (s, lid) -> exn_rebind ~loc (map_loc sub s) (map_loc sub lid)
+    | Pstr_module (s, m) -> module_ ~loc (map_loc sub s) (sub # module_expr m)
+    | Pstr_recmodule l -> rec_module ~loc (List.map (fun (s, mty, me) -> (map_loc sub s, sub # module_type mty, sub # module_expr me)) l)
+    | Pstr_modtype (s, mty) -> modtype ~loc (map_loc sub s) (sub # module_type mty)
+    | Pstr_open lid -> open_ ~loc (map_loc sub lid)
     | Pstr_class l -> class_ ~loc (List.map (sub # class_declaration) l)
     | Pstr_class_type l -> class_type ~loc (List.map (sub # class_type_declaration) l)
     | Pstr_include e -> include_ ~loc (sub # module_expr e)
@@ -282,40 +294,41 @@ module E = struct
   let strconst ?loc x = constant ?loc (Const_string x)
 
   let map sub {pexp_loc = loc; pexp_desc = desc} =
+    let loc = sub # location loc in
     match desc with
-    | Pexp_ident x -> ident ~loc x
+    | Pexp_ident x -> ident ~loc (map_loc sub x)
     | Pexp_constant x -> constant ~loc x
     | Pexp_let (r, pel, e) -> let_ ~loc r (List.map (map_tuple (sub # pat) (sub # expr)) pel) (sub # expr e)
     | Pexp_function (lab, def, pel) -> function_ ~loc lab (map_opt (sub # expr) def) (List.map (map_tuple (sub # pat) (sub # expr)) pel)
     | Pexp_apply (e, l) -> apply ~loc (sub # expr e) (List.map (map_snd (sub # expr)) l)
     | Pexp_match (e, l) -> match_ ~loc (sub # expr e) (List.map (map_tuple (sub # pat) (sub # expr)) l)
     | Pexp_try (e, l) -> try_ ~loc (sub # expr e) (List.map (map_tuple (sub # pat) (sub # expr)) l)
     | Pexp_tuple el -> tuple ~loc (List.map (sub # expr) el)
-    | Pexp_construct (lid, arg, b) -> construct ~loc lid (map_opt (sub # expr) arg) b
+    | Pexp_construct (lid, arg, b) -> construct ~loc (map_loc sub lid) (map_opt (sub # expr) arg) b
     | Pexp_variant (lab, eo) -> variant ~loc lab (map_opt (sub # expr) eo)
-    | Pexp_record (l, eo) -> record ~loc (List.map (fun (id, e) -> (id, sub # expr e)) l) (map_opt (sub # expr) eo)
-    | Pexp_field (e, lid) -> field ~loc (sub # expr e) lid
-    | Pexp_setfield (e1, lid, e2) -> setfield ~loc (sub # expr e1) lid (sub # expr e2)
+    | Pexp_record (l, eo) -> record ~loc (List.map (map_tuple (map_loc sub) (sub # expr)) l) (map_opt (sub # expr) eo)
+    | Pexp_field (e, lid) -> field ~loc (sub # expr e) (map_loc sub lid)
+    | Pexp_setfield (e1, lid, e2) -> setfield ~loc (sub # expr e1) (map_loc sub lid) (sub # expr e2)
     | Pexp_array el -> array ~loc (List.map (sub # expr) el)
     | Pexp_ifthenelse (e1, e2, e3) -> ifthenelse ~loc (sub # expr e1) (sub # expr e2) (map_opt (sub # expr) e3)
     | Pexp_sequence (e1, e2) -> sequence ~loc (sub # expr e1) (sub # expr e2)
     | Pexp_while (e1, e2) -> while_ ~loc (sub # expr e1) (sub # expr e2)
-    | Pexp_for (id, e1, e2, d, e3) -> for_ ~loc id (sub # expr e1) (sub # expr e2) d (sub # expr e3)
+    | Pexp_for (id, e1, e2, d, e3) -> for_ ~loc (map_loc sub id) (sub # expr e1) (sub # expr e2) d (sub # expr e3)
     | Pexp_constraint (e, t1, t2) -> constraint_ ~loc (sub # expr e) (map_opt (sub # typ) t1) (map_opt (sub # typ) t2)
     | Pexp_when (e1, e2) -> when_ ~loc (sub # expr e1) (sub # expr e2)
     | Pexp_send (e, s) -> send ~loc (sub # expr e) s
-    | Pexp_new lid -> new_ ~loc lid
-    | Pexp_setinstvar (s, e) -> setinstvar ~loc s (sub # expr e)
-    | Pexp_override sel -> override ~loc (List.map (map_snd (sub # expr)) sel)
-    | Pexp_letmodule (s, me, e) -> letmodule ~loc (s, sub # module_expr me, sub # expr e)
+    | Pexp_new lid -> new_ ~loc (map_loc sub lid)
+    | Pexp_setinstvar (s, e) -> setinstvar ~loc (map_loc sub s) (sub # expr e)
+    | Pexp_override sel -> override ~loc (List.map (map_tuple (map_loc sub) (sub # expr)) sel)
+    | Pexp_letmodule (s, me, e) -> letmodule ~loc (map_loc sub s, sub # module_expr me, sub # expr e)
     | Pexp_assert e -> assert_ ~loc (sub # expr e)
     | Pexp_assertfalse -> assertfalse ~loc ()
     | Pexp_lazy e -> lazy_ ~loc (sub # expr e)
     | Pexp_poly (e, t) -> poly ~loc (sub # expr e) (map_opt (sub # typ) t)
     | Pexp_object cls -> object_ ~loc (sub # class_structure cls)
     | Pexp_newtype (s, e) -> newtype ~loc s (sub # expr e)
     | Pexp_pack me -> pack ~loc (sub # module_expr me)
-    | Pexp_open (lid, e) -> open_ ~loc lid (sub # expr e)
+    | Pexp_open (lid, e) -> open_ ~loc (map_loc sub lid) (sub # expr e)
 end
 
 module P = struct
@@ -338,24 +351,23 @@ module P = struct
   let unpack ?loc a = mk ?loc (Ppat_unpack a)
 
   let map sub {ppat_desc = desc; ppat_loc = loc} =
+    let loc = sub # location loc in
     match desc with
     | Ppat_any -> any ~loc ()
-    | Ppat_var s -> var ~loc s
-    | Ppat_alias (p, s) -> alias ~loc (sub # pat p) s
+    | Ppat_var s -> var ~loc (map_loc sub s)
+    | Ppat_alias (p, s) -> alias ~loc (sub # pat p) (map_loc sub s)
     | Ppat_constant c -> constant ~loc c
     | Ppat_tuple pl -> tuple ~loc (List.map (sub # pat) pl)
-    | Ppat_construct (l, p, b) -> construct ~loc l (map_opt (sub # pat) p) b
+    | Ppat_construct (l, p, b) -> construct ~loc (map_loc sub l) (map_opt (sub # pat) p) b
     | Ppat_variant (l, p) -> variant ~loc l (map_opt (sub # pat) p)
     | Ppat_record (lpl, cf) ->
-        (*record ~loc (List.map (map_snd (sub # pat)) lpl) cf*)
-        record ~loc
-          (List.map (fun (s, p) -> (s, sub # pat p)) lpl) cf
+        record ~loc (List.map (map_tuple (map_loc sub) (sub # pat)) lpl) cf
     | Ppat_array pl -> array ~loc (List.map (sub # pat) pl)
     | Ppat_or (p1, p2) -> or_ ~loc (sub # pat p1) (sub # pat p2)
     | Ppat_constraint (p, t) -> constraint_ ~loc (sub # pat p) (sub # typ t)
-    | Ppat_type s -> type_ ~loc s
+    | Ppat_type s -> type_ ~loc (map_loc sub s)
     | Ppat_lazy p -> lazy_ ~loc (sub # pat p)
-    | Ppat_unpack s -> unpack ~loc s
+    | Ppat_unpack s -> unpack ~loc (map_loc sub s)
 end
 
 module CE = struct
@@ -371,8 +383,9 @@ module CE = struct
   let constraint_ ?loc a b = mk ?loc (Pcl_constraint (a, b))
 
   let map sub {pcl_loc = loc; pcl_desc = desc} =
+    let loc = sub # location loc in
     match desc with
-    | Pcl_constr (lid, tys) -> constr ~loc lid (List.map (sub # typ) tys)
+    | Pcl_constr (lid, tys) -> constr ~loc (map_loc sub lid) (List.map (sub # typ) tys)
     | Pcl_structure s ->
         structure ~loc (sub # class_structure s)
     | Pcl_fun (lab, e, p, ce) ->
@@ -401,12 +414,13 @@ module CE = struct
   let init ?loc a = mk_field ?loc (Pcf_init a)
 
   let map_field sub {pcf_desc = desc; pcf_loc = loc} =
+    let loc = sub # location loc in
     match desc with
     | Pcf_inher (o, ce, s) -> inher ~loc o (sub # class_expr ce) s
-    | Pcf_valvirt (s, m, t) -> valvirt ~loc s m (sub # typ t)
-    | Pcf_val (s, m, o, e) -> val_ ~loc s m o (sub # expr e)
-    | Pcf_virt (s, p, t) -> virt ~loc s p (sub # typ t)
-    | Pcf_meth (s, p, o, e) -> meth ~loc s p o (sub # expr e)
+    | Pcf_valvirt (s, m, t) -> valvirt ~loc (map_loc sub s) m (sub # typ t)
+    | Pcf_val (s, m, o, e) -> val_ ~loc (map_loc sub s) m o (sub # expr e)
+    | Pcf_virt (s, p, t) -> virt ~loc (map_loc sub s) p (sub # typ t)
+    | Pcf_meth (s, p, o, e) -> meth ~loc (map_loc sub s) p o (sub # expr e)
     | Pcf_constr (t1, t2) -> constr ~loc (sub # typ t1) (sub # typ t2)
     | Pcf_init e -> init ~loc (sub # expr e)
 
@@ -415,6 +429,16 @@ module CE = struct
      pcstr_pat = sub # pat pcstr_pat;
      pcstr_fields = List.map (sub # class_field) pcstr_fields;
     }
+
+  let class_infos sub f {pci_virt; pci_params = (pl, ploc); pci_name; pci_expr; pci_variance; pci_loc} =
+    {
+     pci_virt;
+     pci_params = List.map (map_loc sub) pl, sub # location ploc;
+     pci_name = map_loc sub pci_name;
+     pci_expr = f pci_expr;
+     pci_variance;
+     pci_loc = sub # location pci_loc;
+    }
 end
 
 (* Now, a generic AST mapper class, to be extended to cover all kinds
@@ -467,7 +491,7 @@ class create =
     method module_type = MT.map this
     method with_constraint c = MT.map_with_constraint this c
 
-    method class_declaration decl = {decl with pci_expr = this # class_expr decl.pci_expr}
+    method class_declaration = CE.class_infos this (this # class_expr)
     method class_expr = CE.map this
     method class_field = CE.map_field this
     method class_structure = CE.map_structure this
@@ -476,18 +500,25 @@ class create =
     method class_type_field = CT.map_field this
     method class_signature = CT.map_signature this
 
-    method class_type_declaration decl = {decl with pci_expr = this # class_type decl.pci_expr}
-    method class_description decl = {decl with pci_expr = this # class_type decl.pci_expr}
+    method class_type_declaration = CE.class_infos this (this # class_type)
+    method class_description = CE.class_infos this (this # class_type)
 
     method type_declaration = T.map_type_declaration this
     method type_kind = T.map_type_kind this
     method typ = T.map this
 
-    method value_description vd = {vd with pval_type = this # typ vd.pval_type}
+    method value_description {pval_type; pval_prim; pval_loc} =
+      {
+       pval_type = this # typ pval_type;
+       pval_prim;
+       pval_loc = this # location pval_loc;
+      }
     method pat = P.map this
     method expr = E.map this
 
     method exception_declaration tl = List.map (this # typ) tl
+
+    method location l = l
   end