flambda-backend: Fix ghost locations for modular extensions (#1348)

* Fix ghost locations in modular extension AST nodes

* Add missing ghostification

Thank you, Carl!

* Comment update about ghostiness (+ word-wrapping)

* Add ghostify function (#1)

* Add `Location.ghostify`

* Update the parser's `make_ghost` to save an allocation in some cases

* Promote parser.ml

* Mark the inner mutable arrays for iarrays as ghost

* Add comment about ghostification for comprehensions

* Explain that ppxlib is where the ghostiness requirement is enforced

* Use `Ast_helper.default_loc` to default the generated locations

* Restore propagating the location, now via `Ast_helper.default_loc`

* Drop obsolete comment

* Rewrite advisory comment about locations

Author: antalsz

boot/menhir/parser.ml

@@ -17,8 +17,8 @@ open Extensions_parsing
    expression to translate.  So we just check for the immutable arrays extension
    when processing a comprehension expression for an immutable array.
 
-   Note [Wrapping with make_extension]
-   ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+   Note [Wrapping with make_entire_extension]
+   ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
    The topmost node in the encoded AST must always look like e.g.
    [%extension.comprehensions]. This allows the decoding machinery to know
@@ -31,12 +31,12 @@ open Extensions_parsing
    structurally impossible/hard to forget taking this final step.
 
    However, the final step is only one line of code (a call to
-   [make_extension]), but yet the name of the extension varies, as does the type
-   of the payload. It would thus take several lines of code to execute this
-   command otherwise, along with dozens of lines to create the structure in the
-   first place. And so instead we just manually call [make_extension] and refer
-   to this Note as a reminder to authors of future extensions to remember to do
-   this wrapping.
+   [make_entire_extension]), but yet the name of the extension varies, as does
+   the type of the payload. It would thus take several lines of code to execute
+   this command otherwise, along with dozens of lines to create the structure in
+   the first place. And so instead we just manually call [make_entire_extension]
+   and refer to this Note as a reminder to authors of future extensions to
+   remember to do this wrapping.
 *)
 
 (** List and array comprehensions *)
@@ -89,77 +89,61 @@ module Comprehensions = struct
      v}
   *)
 
-  let comprehension_expr ~loc names x =
-    Expression.wrap_desc ~loc ~attrs:[] @@ Expression.make_extension ~loc (extension_string :: names) x
+  let comprehension_expr names x =
+    Expression.wrap_desc ~attrs:[] @@
+    Expression.make_extension (extension_string :: names) x
 
   (** First, we define how to go from the nice AST to the OCaml AST; this is
       the [expr_of_...] family of expressions, culminating in
       [expr_of_comprehension_expr]. *)
 
-  let expr_of_iterator ~loc = function
+  let expr_of_iterator = function
     | Range { start; stop; direction } ->
         comprehension_expr
-          ~loc
           [ "for"
           ; "range"
           ; match direction with
             | Upto   -> "upto"
             | Downto -> "downto" ]
           (Ast_helper.Exp.tuple [start; stop])
     | In seq ->
-        comprehension_expr ~loc ["for"; "in"] seq
+        comprehension_expr ["for"; "in"] seq
 
-  let expr_of_clause_binding ~loc { pattern; iterator; attributes } =
-    Ast_helper.Vb.mk
-      ~loc
-      ~attrs:attributes
-      pattern
-      (expr_of_iterator ~loc iterator)
+  let expr_of_clause_binding { pattern; iterator; attributes } =
+    Ast_helper.Vb.mk ~attrs:attributes pattern (expr_of_iterator iterator)
 
-  let expr_of_clause ~loc clause rest = match clause with
+  let expr_of_clause clause rest = match clause with
     | For iterators ->
         comprehension_expr
-          ~loc
           ["for"]
           (Ast_helper.Exp.let_
-             Nonrecursive
-             (List.map (expr_of_clause_binding ~loc) iterators)
+             Nonrecursive (List.map expr_of_clause_binding iterators)
              rest)
     | When cond ->
-        comprehension_expr
-          ~loc
-          ["when"]
-          (Ast_helper.Exp.sequence cond rest)
+        comprehension_expr ["when"] (Ast_helper.Exp.sequence cond rest)
 
-  let expr_of_comprehension ~loc ~type_ { body; clauses } =
+  let expr_of_comprehension ~type_ { body; clauses } =
     comprehension_expr
-      ~loc
       type_
       (List.fold_right
-         (expr_of_clause ~loc)
+         expr_of_clause
          clauses
-         (comprehension_expr ~loc ["body"] body))
+         (comprehension_expr ["body"] body))
 
   let expr_of ~loc eexpr =
-    let ghost_loc = { loc with Location.loc_ghost = true } in
-    let expr_of_comprehension_type type_ =
-      expr_of_comprehension ~loc:ghost_loc ~type_
-    in
-    (* See Note [Wrapping with make_extension] *)
-    Expression.make_extension ~loc [extension_string] @@
-    match eexpr with
-    | Cexp_list_comprehension comp ->
-        expr_of_comprehension_type ["list"]  comp
-    | Cexp_array_comprehension (amut, comp) ->
-        expr_of_comprehension_type
-          [ "array"
-          ; match amut with
-            | Mutable   ->
-                "mutable"
-            | Immutable ->
-                "immutable"
-          ]
-          comp
+    (* See Note [Wrapping with make_entire_extension] *)
+    Expression.make_entire_extension ~loc extension_string (fun () ->
+      match eexpr with
+      | Cexp_list_comprehension comp ->
+          expr_of_comprehension ~type_:["list"] comp
+      | Cexp_array_comprehension (amut, comp) ->
+          expr_of_comprehension
+            ~type_:[ "array"
+                   ; match amut with
+                     | Mutable   -> "mutable"
+                     | Immutable -> "immutable"
+                   ]
+            comp)
 
   (** Then, we define how to go from the OCaml AST to the nice AST; this is
       the [..._of_expr] family of expressions, culminating in
@@ -280,19 +264,19 @@ module Immutable_arrays = struct
 
   let expr_of ~loc = function
     | Iaexp_immutable_array elts ->
-      (* See Note [Wrapping with make_extension] *)
-      Expression.make_extension ~loc [extension_string] @@
-      Ast_helper.Exp.array ~loc elts
+      (* See Note [Wrapping with make_entire_extension] *)
+      Expression.make_entire_extension ~loc extension_string (fun () ->
+        Ast_helper.Exp.array elts)
 
   let of_expr expr = match expr.pexp_desc with
     | Pexp_array elts -> Iaexp_immutable_array elts
     | _ -> failwith "Malformed immutable array expression"
 
   let pat_of ~loc = function
     | Iapat_immutable_array elts ->
-      (* See Note [Wrapping with make_extension] *)
-      Pattern.make_extension ~loc [extension_string] @@
-      Ast_helper.Pat.array ~loc elts
+      (* See Note [Wrapping with make_entire_extension] *)
+      Pattern.make_entire_extension ~loc extension_string (fun () ->
+        Ast_helper.Pat.array elts)
 
   let of_pat expr = match expr.ppat_desc with
     | Ppat_array elts -> Iapat_immutable_array elts
@@ -311,10 +295,10 @@ module Strengthen = struct
      [(module M)] can be the inferred type for [M], so this should be fine. *)
 
   let mty_of ~loc { mty; mod_id } =
-    (* See Note [Wrapping with make_extension] *)
-    Module_type.make_extension ~loc [extension_string] @@
+    (* See Note [Wrapping with make_entire_extension] *)
+    Module_type.make_entire_extension ~loc extension_string (fun () ->
       Ast_helper.Mty.functor_ (Named (Location.mknoloc None, mty))
-        (Ast_helper.Mty.alias mod_id)
+        (Ast_helper.Mty.alias mod_id))
 
   let of_mty mty = match mty.pmty_desc with
     | Pmty_functor(Named(_, mty), {pmty_desc = Pmty_alias mod_id}) ->

parsing/extensions.ml

@@ -136,14 +136,17 @@ module type AST_parameters = sig
       [fun tm -> tm.pCAT_loc] for the appropriate syntactic category [CAT]. *)
   val location : ast -> Location.t
 
-  (** Turn an [ast_desc] into an [ast] by adding the appropriate metadata *)
+  (** Turn an [ast_desc] into an [ast] by adding the appropriate metadata.  When
+      creating [ast] nodes afresh for an extension, the location should be
+      omitted; in this case, it will default to [!Ast_helper.default_loc], which
+      should be [ghost]. *)
   val wrap_desc :
-    loc:Location.t -> attrs:Parsetree.attributes -> ast_desc -> ast
+    ?loc:Location.t -> attrs:Parsetree.attributes -> ast_desc -> ast
 
   (** How to construct an extension node for this AST (something of the shape
       [[%name]] or [[%%name]], depending on the AST).  Should just be
-      [Ast_helper.CAT.extension] for the appropriate syntactic category
-      [CAT]. *)
+      [Ast_helper.CAT.extension] for the appropriate syntactic category [CAT].
+      (This means that [?loc] should default to [!Ast_helper.default_loc.].) *)
   val make_extension_node :
     ?loc:Location.t -> ?attrs:attributes -> extension -> ast
 
@@ -172,9 +175,12 @@ module type AST = sig
   val location : ast -> Location.t
 
   val wrap_desc :
-    loc:Location.t -> attrs:Parsetree.attributes -> ast_desc -> ast
+    ?loc:Location.t -> attrs:Parsetree.attributes -> ast_desc -> ast
 
-  val make_extension : loc:Location.t -> string list -> ast -> ast_desc
+  val make_extension : string list -> ast -> ast_desc
+
+  val make_entire_extension :
+    loc:Location.t -> string -> (unit -> ast) -> ast_desc
 
   val match_extension : ast -> (string list * ast) option
 end
@@ -195,12 +201,17 @@ module Make_AST (AST_parameters : AST_parameters) :
   struct
     include AST_parameters
 
-    let make_extension ~loc names =
+    let make_extension names =
       make_extension_use
         ~extension_node:
           (make_extension_node
-             ~loc
-             ({ txt = String.concat "." ("extension" :: names); loc }, PStr []))
+             ({ txt = String.concat "." ("extension" :: names);
+                loc = !Ast_helper.default_loc },
+              PStr []))
+
+    let make_entire_extension ~loc name ast =
+      make_extension [name]
+        (Ast_helper.with_default_loc (Location.ghostify loc) ast)
 
     (* This raises an error if the language extension node is malformed.
        Malformed means either:
@@ -237,7 +248,7 @@ module Expression = Make_AST(struct
 
   let location expr = expr.pexp_loc
 
-  let wrap_desc ~loc ~attrs = Ast_helper.Exp.mk ~loc ~attrs
+  let wrap_desc ?loc ~attrs = Ast_helper.Exp.mk ?loc ~attrs
 
   let make_extension_node = Ast_helper.Exp.extension
 
@@ -262,7 +273,7 @@ module Pattern = Make_AST(struct
 
   let location pat = pat.ppat_loc
 
-  let wrap_desc ~loc ~attrs = Ast_helper.Pat.mk ~loc ~attrs
+  let wrap_desc ?loc ~attrs = Ast_helper.Pat.mk ?loc ~attrs
 
   let make_extension_node = Ast_helper.Pat.extension
 
@@ -286,7 +297,7 @@ module Module_type = Make_AST(struct
 
     let location mty = mty.pmty_loc
 
-    let wrap_desc ~loc ~attrs = Ast_helper.Mty.mk ~loc ~attrs
+    let wrap_desc ?loc ~attrs = Ast_helper.Mty.mk ?loc ~attrs
 
     let make_extension_node = Ast_helper.Mty.extension
 

parsing/extensions_parsing.ml

@@ -36,20 +36,20 @@
 
     b. We define an *overall auxiliary AST* for each syntactic category that's
        just for our language extensions; for expressions, it's called
-       [Extensions.Expression.t].  It contains one constructor for each of the AST types
-       defined as described in design point (1).  This addresses concern (2); we
-       can now match on actual OCaml constructors, as long as we can get a hold
-       of them.  And to do that…
-
-    c. We define a general scheme for how we represent language extensions in terms
-       of the existing ASTs, and provide a few primitives for consuming/creating
-       AST nodes of this form, for each syntactic category.  There's not a lot
-       of abstraction to be done, or at least it's not (yet) apparent what
-       abstraction there is to do, so most of this remains manual.  (Setting up
-       a full lens-based/otherwise bidirectional approach sounds like a great
-       opportunity for yak-shaving, but not *actually* a good idea.)  This
-       solves concern (3), and by doing it uniformly helps us address multiple
-       cases at one stroke.
+       [Extensions.Expression.t].  It contains one constructor for each of the
+       AST types defined as described in design point (1).  This addresses
+       concern (2); we can now match on actual OCaml constructors, as long as we
+       can get a hold of them.  And to do that…
+
+    c. We define a general scheme for how we represent language extensions in
+       terms of the existing ASTs, and provide a few primitives for
+       consuming/creating AST nodes of this form, for each syntactic category.
+       There's not a lot of abstraction to be done, or at least it's not (yet)
+       apparent what abstraction there is to do, so most of this remains manual.
+       (Setting up a full lens-based/otherwise bidirectional approach sounds
+       like a great opportunity for yak-shaving, but not *actually* a good
+       idea.)  This solves concern (3), and by doing it uniformly helps us
+       address multiple cases at one stroke.
 
     Then, for each syntactic category, we define a module (in extensions.ml)
     that contains functions for converting between the Parsetree representation
@@ -68,7 +68,17 @@
     writing out extension points or attributes directly, we write the result of
     [Some_ast.make_extension ~loc [name1; name2; ...; NameN] a] as the special
     syntax [{% 'name1.name2.....nameN' | a %}] in the BNF.  Other pieces of the
-    OCaml AST are used as normal. *)
+    OCaml AST are used as normal.
+
+    One detail which we hide as much as possible is locations: whenever
+    constructing an OCaml AST node -- whether with [wrap_desc], the functions in
+    [Ast_helper], or some other way -- the location should be left to be
+    defaulted (and the default, [!Ast_helper.make_default], should be ghost).
+    The [make_entire_extension] function will handle making sure this default
+    location is set appropriately.  If this isn't done and any locations on
+    subterms aren't marked as ghost, the compiler will work fine, but ppxlib may
+    detect that you've violated its well-formedness constraints and fail to
+    parse the resulting AST. *)
 
 (** Errors around the extension representation.  These should mostly just be
     fatal, but they're needed for one test case
@@ -109,14 +119,28 @@ module type AST = sig
   (** How to get the location attached to an AST node *)
   val location : ast -> Location.t
 
-  (** Turn an [ast_desc] into an [ast] by adding the appropriate metadata *)
+  (** Turn an [ast_desc] into an [ast] by adding the appropriate metadata.  When
+      creating [ast] nodes afresh for an extension, the location should be
+      omitted; in this case, it will default to [!Ast_helper.default_loc], which
+      should be [ghost]. *)
   val wrap_desc :
-    loc:Location.t -> attrs:Parsetree.attributes -> ast_desc -> ast
-
-  (** Embed a language extension term in the AST with the given name
-      and body (the [ast]).  The name will be joined with dots
-      and preceded by [extension.].  Partial inverse of [match_extension]. *)
-  val make_extension  : loc:Location.t -> string list -> ast -> ast_desc
+    ?loc:Location.t -> attrs:Parsetree.attributes -> ast_desc -> ast
+
+  (** Embed a language extension term in the AST with the given name and body
+      (the [ast]).  The name will be joined with dots and preceded by
+      [extension.].  Any locations in the generated AST will be set to
+      [!Ast_helper.default_loc], which should be [ghost].  Partial inverse of
+      [match_extension]. *)
+  val make_extension  : string list -> ast -> ast_desc
+
+  (** As [make_extension], but specifically for the AST node corresponding to
+      the entire piece of extension syntax (e.g., for a list comprehension, the
+      whole [[x for x in xs]], and not a subterm like [for x in xs]).  This sets
+      [Ast_helper.default_loc] locally to the [ghost] version of the provided
+      location, which is why the [ast] is generated from a function call; it is
+      during this call that the location is so set. *)
+  val make_entire_extension :
+    loc:Location.t -> string -> (unit -> ast) -> ast_desc
 
   (** Given an AST node, check if it's a language extension term; if it is,
       split it back up into its name (the [string list]) and the body (the

parsing/extensions_parsing.mli

@@ -41,6 +41,11 @@ let init lexbuf fname =
   }
 ;;
 
+let ghostify l =
+  if l.loc_ghost
+  then l
+  else { l with loc_ghost = true }
+
 let symbol_rloc () = {
   loc_start = Parsing.symbol_start_pos ();
   loc_end = Parsing.symbol_end_pos ();

parsing/location.ml

@@ -68,6 +68,9 @@ val init : Lexing.lexbuf -> string -> unit
 val curr : Lexing.lexbuf -> t
 (** Get the location of the current token from the [lexbuf]. *)
 
+val ghostify : t -> t
+(** Return a version of the location with [loc_ghost = true] *)
+
 val symbol_rloc: unit -> t
 val symbol_gloc: unit -> t
 

parsing/location.mli

@@ -504,7 +504,10 @@ let lapply ~loc p1 p2 =
 let loc_map (f : 'a -> 'b) (x : 'a Location.loc) : 'b Location.loc =
   { x with txt = f x.txt }
 
-let make_ghost x = { x with loc = { x.loc with loc_ghost = true }}
+let make_ghost x =
+  if x.loc.loc_ghost
+  then x (* Save an allocation *)
+  else { x with loc = Location.ghostify x.loc }
 
 let loc_last (id : Longident.t Location.loc) : string Location.loc =
   loc_map Longident.last id

parsing/parser.mly

@@ -1916,8 +1916,7 @@ module Conv = struct
             | Mutable   -> Ppat_array pats
             | Immutable ->
                 Extensions.Immutable_arrays.pat_of
-                  ~loc:pat.pat_loc
-                  (Iapat_immutable_array pats)
+                  ~loc:pat.pat_loc (Iapat_immutable_array pats)
           in
           mkpat ppat
       | Tpat_lazy p ->