flambda-backend: Full blown Jane Syntax for mode exprs (#2335)

* Full blown jane syntax for mode exprs

* Add some comments

* small adjustment to printing of typedtree

* Change the way that value binding elaboration encodes modes on Pexp_constraint to allow Jane Syntax to be stricter in validation

* rename variable

.depend

@@ -1138,8 +1138,7 @@ typing/mode.cmi : \
     typing/mode_intf.cmi
 typing/mode_intf.cmi : \
     typing/solver_intf.cmi \
-    typing/solver.cmi \
-    utils/misc.cmi
+    typing/solver.cmi
 typing/mtype.cmo : \
     typing/types.cmi \
     typing/subst.cmi \
@@ -1991,6 +1990,7 @@ typing/typedtree.cmo : \
     parsing/longident.cmi \
     parsing/location.cmi \
     typing/jkind.cmi \
+    parsing/jane_syntax.cmi \
     parsing/jane_asttypes.cmi \
     typing/ident.cmi \
     typing/env.cmi \
@@ -2006,6 +2006,7 @@ typing/typedtree.cmx : \
     parsing/longident.cmx \
     parsing/location.cmx \
     typing/jkind.cmx \
+    parsing/jane_syntax.cmx \
     parsing/jane_asttypes.cmx \
     typing/ident.cmx \
     typing/env.cmx \
@@ -2021,6 +2022,7 @@ typing/typedtree.cmi : \
     parsing/longident.cmi \
     parsing/location.cmi \
     typing/jkind.cmi \
+    parsing/jane_syntax.cmi \
     parsing/jane_asttypes.cmi \
     typing/ident.cmi \
     typing/env.cmi \
@@ -4176,13 +4178,11 @@ lambda/transl_array_comprehension.cmi : \
     lambda/debuginfo.cmi
 lambda/transl_comprehension_utils.cmo : \
     utils/targetint.cmi \
-    typing/primitive.cmi \
     lambda/lambda.cmi \
     typing/ident.cmi \
     lambda/transl_comprehension_utils.cmi
 lambda/transl_comprehension_utils.cmx : \
     utils/targetint.cmx \
-    typing/primitive.cmx \
     lambda/lambda.cmx \
     typing/ident.cmx \
     lambda/transl_comprehension_utils.cmi

boot/menhir/parser.ml

@@ -379,7 +379,7 @@ let mkexp_with_modes ?(ghost=false) ~loc modes exp =
   let loc =
     if ghost then ghost_loc loc else make_loc loc
   in
-  Mode.expr_of_coerce ~loc modes exp
+  Jane_syntax.Modes.expr_of ~loc (Coerce (modes, exp))
 
 (* For modes-related attributes, no need to call [register_attr] because they
 result from native syntax which is only parsed at proper places that are

parsing/ast_invariants.ml

@@ -147,6 +147,7 @@ let iterator =
     | Jexp_comprehension _
     | Jexp_immutable_array _
     | Jexp_layout _
+    | Jexp_modes _
       -> ()
   in
   let expr self exp =

parsing/ast_iterator.ml

@@ -43,7 +43,6 @@ type iterator = {
   constructor_declaration: iterator -> constructor_declaration -> unit;
   expr: iterator -> expression -> unit;
   expr_jane_syntax: iterator -> Jane_syntax.Expression.t -> unit;
-  expr_mode_syntax: iterator -> Jane_syntax.Mode_expr.t -> expression -> unit;
   extension: iterator -> extension -> unit;
   extension_constructor: iterator -> extension_constructor -> unit;
   include_declaration: iterator -> include_declaration -> unit;
@@ -467,6 +466,7 @@ module E = struct
   module L = Jane_syntax.Layouts
   module N_ary = Jane_syntax.N_ary_functions
   module LT = Jane_syntax.Labeled_tuples
+  module Modes = Jane_syntax.Modes
 
   let iter_iterator sub : C.iterator -> _ = function
     | Range { start; stop; direction = _ } ->
@@ -543,24 +543,22 @@ module E = struct
   let iter_labeled_tuple sub : LT.expression -> _ = function
     | el -> List.iter (iter_snd (sub.expr sub)) el
 
+  let iter_modes_exp sub : Modes.expression -> _ = function
+    | Coerce (modes, expr) ->
+        sub.modes sub modes;
+        sub.expr sub expr
+
   let iter_jst sub : Jane_syntax.Expression.t -> _ = function
     | Jexp_comprehension comp_exp -> iter_comp_exp sub comp_exp
     | Jexp_immutable_array iarr_exp -> iter_iarr_exp sub iarr_exp
     | Jexp_layout layout_exp -> iter_layout_exp sub layout_exp
     | Jexp_n_ary_function n_ary_exp -> iter_n_ary_function sub n_ary_exp
     | Jexp_tuple lt_exp -> iter_labeled_tuple sub lt_exp
-
-  let iter_mode sub modes expr =
-    sub.modes sub modes;
-    sub.expr sub expr
+    | Jexp_modes mode_exp -> iter_modes_exp sub mode_exp
 
   let iter sub
         ({pexp_loc = loc; pexp_desc = desc; pexp_attributes = attrs} as expr)=
     sub.location sub loc;
-    match Jane_syntax.Mode_expr.coerce_of_expr expr with
-    | Some (modes, e) ->
-        sub.expr_mode_syntax sub modes e
-    | None ->
     match Jane_syntax.Expression.of_ast expr with
     | Some (jexp, attrs) ->
         sub.attributes sub attrs;
@@ -827,7 +825,6 @@ let default_iterator =
     pat_mode_syntax = P.iter_mode;
     expr = E.iter;
     expr_jane_syntax = E.iter_jst;
-    expr_mode_syntax = E.iter_mode;
     binding_op = E.iter_binding_op;
 
     module_declaration =

parsing/ast_iterator.mli

@@ -46,7 +46,6 @@ type iterator = {
   constructor_declaration: iterator -> constructor_declaration -> unit;
   expr: iterator -> expression -> unit;
   expr_jane_syntax : iterator -> Jane_syntax.Expression.t -> unit;
-  expr_mode_syntax: iterator -> Jane_syntax.Mode_expr.t -> expression -> unit;
   extension: iterator -> extension -> unit;
   extension_constructor: iterator -> extension_constructor -> unit;
   include_declaration: iterator -> include_declaration -> unit;

parsing/ast_mapper.ml

@@ -105,6 +105,18 @@ let map_loc sub {loc; txt} = {loc = sub.location sub loc; txt}
 let map_loc_txt sub f {loc; txt} =
   {loc = sub.location sub loc; txt = f sub txt}
 
+let map_mode_expr sub (mode_expr : Jane_syntax.Mode_expr.t)
+  : Jane_syntax.Mode_expr.t =
+  map_loc_txt sub
+    (fun sub modes ->
+       List.map
+         (fun (mode : Jane_syntax.Mode_expr.Const.t) ->
+            let { loc; txt } = (mode :> string loc) in
+            let loc = sub.location sub loc in
+            Jane_syntax.Mode_expr.Const.mk txt loc)
+         modes)
+    mode_expr
+
 module C = struct
   (* Constants *)
 
@@ -536,6 +548,7 @@ module E = struct
   module L = Jane_syntax.Layouts
   module N_ary = Jane_syntax.N_ary_functions
   module LT = Jane_syntax.Labeled_tuples
+  module Modes = Jane_syntax.Modes
 
   let map_iterator sub : C.iterator -> C.iterator = function
     | Range { start; stop; direction } ->
@@ -629,13 +642,19 @@ module E = struct
     (* CR labeled tuples: Eventually mappers may want to see the labels. *)
     | el -> List.map (map_snd (sub.expr sub)) el
 
+  let map_modes_exp sub : Modes.expression -> Modes.expression = function
+    (* CR modes: One day mappers might want to see the modes *)
+    | Coerce (modes, exp) ->
+        Coerce (map_mode_expr sub modes, sub.expr sub exp)
+
   let map_jst sub : Jane_syntax.Expression.t -> Jane_syntax.Expression.t =
     function
     | Jexp_comprehension x -> Jexp_comprehension (map_cexp sub x)
     | Jexp_immutable_array x -> Jexp_immutable_array (map_iaexp sub x)
     | Jexp_layout x -> Jexp_layout (map_layout_exp sub x)
     | Jexp_n_ary_function x -> Jexp_n_ary_function (map_n_ary_exp sub x)
     | Jexp_tuple ltexp -> Jexp_tuple (map_ltexp sub ltexp)
+    | Jexp_modes mode_exp -> Jexp_modes (map_modes_exp sub mode_exp)
 
   let map sub
         ({pexp_loc = loc; pexp_desc = desc; pexp_attributes = attrs} as exp) =

parsing/depend.ml

@@ -332,6 +332,11 @@ and add_expr_jane_syntax bv : Jane_syntax.Expression.t -> _ = function
   | Jexp_layout x -> add_layout_expr bv x
   | Jexp_n_ary_function n_ary -> add_n_ary_function bv n_ary
   | Jexp_tuple x -> add_labeled_tuple_expr bv x
+  | Jexp_modes x -> add_modes_expr bv x
+
+and add_modes_expr bv : Jane_syntax.Modes.expression -> _ =
+  function
+  | Coerce (_modes, exp) -> add_expr bv exp
 
 and add_comprehension_expr bv : Jane_syntax.Comprehensions.expression -> _ =
   function

parsing/jane_syntax.ml

@@ -459,8 +459,6 @@ module Mode_expr = struct
 
   let attribute_name = attribute_or_extension_name
 
-  let extension_name = attribute_or_extension_name
-
   let payload_of { txt; _ } =
     match txt with
     | [] -> None
@@ -508,26 +506,38 @@ module Mode_expr = struct
     let loc = { loc with loc_ghost = true } in
     let txt = List.map Const.ghostify txt in
     { loc; txt }
+end
+
+(** Some mode-related constructs *)
+module Modes = struct
+  let feature : Feature.t = Language_extension Mode
+
+  type nonrec expression = Coerce of Mode_expr.t * expression
 
-  let coerce_of_expr { pexp_desc; _ } =
+  let extension_name = Mode_expr.attribute_or_extension_name
+
+  let of_expr ({ pexp_desc; pexp_attributes; _ } as expr) =
     match pexp_desc with
     | Pexp_apply
         ( { pexp_desc = Pexp_extension ({ txt; _ }, payload); pexp_loc; _ },
           [(Nolabel, body)] )
       when txt = extension_name ->
-      let modes = of_payload ~loc:pexp_loc payload in
-      Some (modes, body)
-    | _ -> None
-
-  let expr_of_coerce ~loc modes body =
-    match payload_of modes with
+      let modes = Mode_expr.of_payload ~loc:pexp_loc payload in
+      Coerce (modes, body), pexp_attributes
+    | _ ->
+      Misc.fatal_errorf "Improperly encoded modes expression: %a"
+        (Printast.expression 0) expr
+
+  let expr_of ~loc (Coerce (modes, body)) =
+    match Mode_expr.payload_of modes with
     | None -> body
     | Some payload ->
       let ext =
         Ast_helper.Exp.extension ~loc:modes.loc
           (Location.mknoloc extension_name, payload)
       in
-      Ast_helper.Exp.apply ~loc ext [Nolabel, body]
+      Expression.make_entire_jane_syntax ~loc feature (fun () ->
+          Ast_helper.Exp.apply ~loc ext [Nolabel, body])
 end
 
 (** List and array comprehensions *)
@@ -1911,6 +1921,7 @@ module Expression = struct
     | Jexp_layout of Layouts.expression
     | Jexp_n_ary_function of N_ary_functions.expression
     | Jexp_tuple of Labeled_tuples.expression
+    | Jexp_modes of Modes.expression
 
   let of_ast_internal (feat : Feature.t) expr =
     match feat with
@@ -1930,6 +1941,9 @@ module Expression = struct
     | Language_extension Labeled_tuples ->
       let expr, attrs = Labeled_tuples.of_expr expr in
       Some (Jexp_tuple expr, attrs)
+    | Language_extension Mode ->
+      let expr, attrs = Modes.of_expr expr in
+      Some (Jexp_modes expr, attrs)
     | _ -> None
 
   let of_ast = Expression.make_of_ast ~of_ast_internal
@@ -1942,6 +1956,7 @@ module Expression = struct
       | Jexp_layout x -> Layouts.expr_of ~loc x
       | Jexp_n_ary_function x -> N_ary_functions.expr_of ~loc x
       | Jexp_tuple x -> Labeled_tuples.expr_of ~loc x
+      | Jexp_modes x -> Modes.expr_of ~loc x
     in
     (* Performance hack: save an allocation if [attrs] is empty. *)
     match attrs with

parsing/jane_syntax.mli

@@ -159,20 +159,27 @@ module Mode_expr : sig
       attribute is found. *)
   val of_attrs : Parsetree.attributes -> t * Parsetree.attributes
 
-  (** Decode mode coercion and returns the mode and the body.
-      For example, return [Some (local, expr)] on input [local_ expr].
-      Returns [None] if the given expression is not a mode coercion. *)
-  val coerce_of_expr : Parsetree.expression -> (t * Parsetree.expression) option
-
-  (** Encode a mode coercion like [local_ expr] into an expression *)
-  val expr_of_coerce :
-    loc:Location.t -> t -> Parsetree.expression -> Parsetree.expression
-
   (** In some cases, a single mode expression appears twice in the parsetree;
       one of them needs to be made ghost to make our internal tools happy. *)
   val ghostify : t -> t
 end
 
+(** A subset of the mode-related syntax extensions that is embedded
+    using full-blown Jane Syntax. By "full-blown" Jane Syntax, we
+    mean the [Expression], [Pattern], (etc.) modules below that
+    attempt to create a variant of all possible Jane Street syntax
+    for the syntactic form.
+
+    We avoid full-blown Jane Syntax when it isn't very lightweight to fit the
+    new construct into the (somewhat opinionated) framework. Mode coercions are
+    lightweight to fit into full-blown Jane Syntax.
+*)
+module Modes : sig
+  type expression = Coerce of Mode_expr.t * Parsetree.expression
+
+  val expr_of : loc:Location.t -> expression -> Parsetree.expression
+end
+
 module N_ary_functions : sig
   (** These types use the [P] prefix to match how they are represented in the
       upstream compiler *)
@@ -586,6 +593,7 @@ module Expression : sig
     | Jexp_layout of Layouts.expression
     | Jexp_n_ary_function of N_ary_functions.expression
     | Jexp_tuple of Labeled_tuples.expression
+    | Jexp_modes of Modes.expression
 
   include
     AST

parsing/parser.mly

@@ -154,7 +154,7 @@ let mkexp_with_modes ?(ghost=false) ~loc modes exp =
   let loc =
     if ghost then ghost_loc loc else make_loc loc
   in
-  Mode.expr_of_coerce ~loc modes exp
+  Jane_syntax.Modes.expr_of ~loc (Coerce (modes, exp))
 
 (* For modes-related attributes, no need to call [register_attr] because they
 result from native syntax which is only parsed at proper places that are

parsing/pprintast.ml

@@ -788,10 +788,6 @@ and sugar_expr ctxt f e =
    expressions that aren't already self-delimiting.
 *)
 and expression ?(jane_syntax_parens = false) ctxt f x =
-  match Jane_syntax.Mode_expr.coerce_of_expr x with
-  | Some (m, body) ->
-    pp f "@[<2>%s %a@]" (modes m) (expression ctxt) body
-  | None ->
   match Jane_syntax.Expression.of_ast x with
   | Some (jexpr, attrs) ->
       jane_syntax_expr ctxt attrs f jexpr ~parens:jane_syntax_parens
@@ -1484,9 +1480,6 @@ and payload ctxt f = function
       pp f " when "; expression ctxt f e
 
 and pp_print_pexp_function ctxt sep f x =
-  (* do not print [@jane.erasable.mode] on expressions *)
-  let _, attrs = maybe_modes_of_attrs x.pexp_attributes in
-  let x = { x with pexp_attributes = attrs } in
   (* We go to some trouble to print nested [Pexp_newtype]/[Lexp_newtype] as
      newtype parameters of the same "fun" (rather than printing several nested
      "fun (type a) -> ..."). This isn't necessary for round-tripping -- it just
@@ -1586,15 +1579,35 @@ and binding ctxt f {pvb_pat=p; pvb_expr=x; pvb_constraint = ct; _} =
 (* [in] is not printed *)
 and bindings ctxt f (rf,l) =
   let binding kwd rf f x =
-    let modes, attrs = maybe_modes_of_attrs x.pvb_attributes in
+    let modes_on_binding, attrs =
+      Jane_syntax.Mode_expr.maybe_of_attrs x.pvb_attributes
+    in
     let x =
-      match modes, Jane_syntax.Mode_expr.coerce_of_expr x.pvb_expr with
-      | Some _ , Some (_, sbody) ->
-          {x with pvb_expr = sbody}
+      (* For [let local_ x = e in ...] and [let x @ local = e in ...],
+         the parser puts attributes on both the let-binding and on e.
+
+         The below code is meant to print the modes only in one place,
+         not both. (We print it on the let-binding, not the expression.)
+      *)
+      match modes_on_binding, Jane_syntax.Expression.of_ast x.pvb_expr with
+      | Some modes_on_binding,
+        Some (Jexp_modes (Coerce (modes_on_expr, sbody)), _) ->
+          (* Sanity check: only suppress the printing of one mode expression if
+             the mode expressions are in fact identical.
+          *)
+          let mode_names (modes : Jane_syntax.Mode_expr.t) =
+            List.map Location.get_txt (modes.txt :> string loc list)
+          in
+          if
+            List.equal String.equal
+              (mode_names modes_on_binding)
+              (mode_names modes_on_expr)
+          then {x with pvb_expr = sbody}
+          else x
       | _ -> x
     in
     pp f "@[<2>%s %a%s%a@]%a" kwd rec_flag rf
-      (match modes with Some s -> s ^ " " | None -> "")
+      (match modes_on_binding with Some s -> modes s ^ " " | None -> "")
       (binding ctxt) x (item_attributes ctxt) attrs
   in
   match l with
@@ -2006,6 +2019,11 @@ and jane_syntax_expr ctxt attrs f (jexp : Jane_syntax.Expression.t) ~parens =
       if parens then pp f "(%a)" (n_ary_function_expr reset_ctxt) x
       else n_ary_function_expr ctxt f x
   | Jexp_tuple ltexp        -> labeled_tuple_expr ctxt f ltexp
+  | Jexp_modes mexp -> mode_expr ctxt f mexp
+
+and mode_expr ctxt f (mexp : Jane_syntax.Modes.expression) =
+  match mexp with
+  | Coerce (m, body) -> pp f "@[<2>%s %a@]" (modes m) (expression ctxt) body
 
 and comprehension_expr ctxt f (cexp : Jane_syntax.Comprehensions.expression) =
   let punct, comp = match cexp with