flambda-backend: Support stack_ exp syntax (#2540)

Author: riaqn

boot/menhir/parser.mli

@@ -19,6 +19,7 @@ type token =
   | STRUCT
   | STRING of (string * Location.t * string option)
   | STAR
+  | STACK
   | SIG
   | SEMISEMI
   | SEMI

lambda/translcore.ml

@@ -290,7 +290,7 @@ let fuse_method_arity (parent : fusable_function) : fusable_function =
         (function (Texp_poly _, _, _) -> true | _ -> false)
         exp_extra
     ->
-      begin match transl_alloc_mode_r method_.alloc_mode with
+      begin match transl_alloc_mode method_.alloc_mode with
       | Alloc_heap -> ()
       | Alloc_local ->
           (* If we support locally-allocated objects, we'll also have to
@@ -466,7 +466,7 @@ and transl_exp0 ~in_new_scope ~scopes sort e =
         Lconst(Const_block(0, List.map extract_constant ll))
       with Not_constant ->
         Lprim(Pmakeblock(0, Immutable, Some shape,
-                         transl_alloc_mode_r alloc_mode),
+                         transl_alloc_mode alloc_mode),
               ll,
               (of_location ~scopes e.exp_loc))
       end
@@ -508,7 +508,7 @@ and transl_exp0 ~in_new_scope ~scopes sort e =
           begin match const_block with
           | Some const_block -> const_block
           | None ->
-            let alloc_mode = transl_alloc_mode_r (Option.get alloc_mode) in
+            let alloc_mode = transl_alloc_mode (Option.get alloc_mode) in
             let makeblock =
               match cstr.cstr_shape with
               | Constructor_uniform_value ->
@@ -534,7 +534,7 @@ and transl_exp0 ~in_new_scope ~scopes sort e =
                that out by checking that the sort list is empty *)
             lam)
           else
-            let alloc_mode = transl_alloc_mode_r (Option.get alloc_mode) in
+            let alloc_mode = transl_alloc_mode (Option.get alloc_mode) in
             let makeblock =
               match cstr.cstr_shape with
               | Constructor_uniform_value ->
@@ -569,13 +569,13 @@ and transl_exp0 ~in_new_scope ~scopes sort e =
                                    extract_constant lam]))
           with Not_constant ->
             Lprim(Pmakeblock(0, Immutable, None,
-                             transl_alloc_mode_r alloc_mode),
+                             transl_alloc_mode alloc_mode),
                   [Lconst(const_int tag); lam],
                   of_location ~scopes e.exp_loc)
       end
   | Texp_record {fields; representation; extended_expression; alloc_mode} ->
       transl_record ~scopes e.exp_loc e.exp_env
-        (Option.map transl_alloc_mode_r alloc_mode)
+        (Option.map transl_alloc_mode alloc_mode)
         fields representation extended_expression
   | Texp_field(arg, id, lbl, float) ->
       let targ = transl_exp ~scopes Jkind.Sort.for_record arg in
@@ -595,7 +595,7 @@ and transl_exp0 ~in_new_scope ~scopes sort e =
             | Boxing (alloc_mode, _) -> alloc_mode
             | Non_boxing _ -> assert false
           in
-          let mode = transl_alloc_mode_r alloc_mode in
+          let mode = transl_alloc_mode alloc_mode in
           Lprim (Pfloatfield (lbl.lbl_pos, sem, mode), [targ],
                  of_location ~scopes e.exp_loc)
         | Record_ufloat ->
@@ -614,7 +614,7 @@ and transl_exp0 ~in_new_scope ~scopes sort e =
                 | Float_boxed ->
                   (match float with
                     | Boxing (mode, _) ->
-                        flat_read_float_boxed (transl_alloc_mode_r mode)
+                        flat_read_float_boxed (transl_alloc_mode mode)
                     | Non_boxing _ ->
                         Misc.fatal_error
                           "expected typechecking to make [float] boxing mode\
@@ -668,7 +668,7 @@ and transl_exp0 ~in_new_scope ~scopes sort e =
                      transl_exp ~scopes lbl_sort newval],
             of_location ~scopes e.exp_loc)
   | Texp_array (amut, element_sort, expr_list, alloc_mode) ->
-      let mode = transl_alloc_mode_r alloc_mode in
+      let mode = transl_alloc_mode alloc_mode in
       let kind = array_kind e element_sort in
       let ll =
         transl_list ~scopes
@@ -1578,7 +1578,7 @@ and transl_function ~in_new_scope ~scopes e params body
       ~alloc_mode ~ret_mode:sreturn_mode ~ret_sort:sreturn_sort ~region:sregion
       ~zero_alloc =
   let attrs = e.exp_attributes in
-  let mode = transl_alloc_mode_r alloc_mode in
+  let mode = transl_alloc_mode alloc_mode in
   let assume_zero_alloc =
     Builtin_attributes.assume_zero_alloc ~is_check_allowed:true zero_alloc
   in
@@ -1981,7 +1981,7 @@ and transl_match ~scopes ~arg_sort ~return_sort e arg pat_expr_list partial =
     match arg, exn_cases with
     | {exp_desc = Texp_tuple (argl, alloc_mode)}, [] ->
       assert (static_handlers = []);
-      let mode = transl_alloc_mode_r alloc_mode in
+      let mode = transl_alloc_mode alloc_mode in
       let argl =
         List.map (fun (_, a) -> (a, Jkind.Sort.for_tuple_element)) argl
       in
@@ -2000,7 +2000,7 @@ and transl_match ~scopes ~arg_sort ~return_sort e arg pat_expr_list partial =
             argl
           |> List.split
         in
-        let mode = transl_alloc_mode_r alloc_mode in
+        let mode = transl_alloc_mode alloc_mode in
         static_catch (transl_list ~scopes argl) val_ids
           (Matching.for_multiple_match ~scopes ~return_layout e.exp_loc
              lvars mode val_cases partial)

lambda/translmode.ml

@@ -35,6 +35,9 @@ let transl_alloc_mode_r mode =
   (* we only take the locality axis *)
   Alloc.proj (Comonadic Areality) mode |> transl_locality_mode_r
 
+let transl_alloc_mode (mode : Typedtree.alloc_mode) =
+  transl_alloc_mode_r mode.mode
+
 let transl_modify_mode locality =
   match Locality.zap_to_floor locality with
   | Global -> modify_heap

lambda/translmode.mli

@@ -20,4 +20,6 @@ val transl_alloc_mode_l : (allowed * 'r) Alloc.t -> Lambda.alloc_mode
 
 val transl_alloc_mode_r : ('l * allowed) Alloc.t -> Lambda.alloc_mode
 
+val transl_alloc_mode : Typedtree.alloc_mode -> Lambda.alloc_mode
+
 val transl_modify_mode : (allowed * 'r) Locality.t -> Lambda.modify_mode

parsing/ast_helper.ml

@@ -223,6 +223,7 @@ module Exp = struct
     mk ?loc ?attrs (Pexp_letop {let_; ands; body})
   let extension ?loc ?attrs a = mk ?loc ?attrs (Pexp_extension a)
   let unreachable ?loc ?attrs () = mk ?loc ?attrs Pexp_unreachable
+  let stack ?loc ?attrs e = mk ?loc ?attrs (Pexp_stack e)
 
   let case lhs ?guard rhs =
     {

parsing/ast_helper.mli

@@ -201,6 +201,7 @@ module Exp:
                -> binding_op list -> expression -> expression
     val extension: ?loc:loc -> ?attrs:attrs -> extension -> expression
     val unreachable: ?loc:loc -> ?attrs:attrs -> unit -> expression
+    val stack : ?loc:loc -> ?attrs:attrs -> expression -> expression
 
     val case: pattern -> ?guard:expression -> expression -> case
     val binding_op: str -> pattern -> expression -> loc -> binding_op

parsing/ast_iterator.ml

@@ -660,6 +660,7 @@ module E = struct
         sub.expr sub body
     | Pexp_extension x -> sub.extension sub x
     | Pexp_unreachable -> ()
+    | Pexp_stack e -> sub.expr sub e
 
   let iter_binding_op sub {pbop_op; pbop_pat; pbop_exp; pbop_loc} =
     iter_loc sub pbop_op;

parsing/ast_mapper.ml

@@ -780,6 +780,7 @@ module E = struct
           (List.map (sub.binding_op sub) ands) (sub.expr sub body)
     | Pexp_extension x -> extension ~loc ~attrs (sub.extension sub x)
     | Pexp_unreachable -> unreachable ~loc ~attrs ()
+    | Pexp_stack e -> stack ~loc ~attrs (sub.expr sub e)
 
   let map_binding_op sub {pbop_op; pbop_pat; pbop_exp; pbop_loc} =
     let open Exp in

parsing/depend.ml

@@ -328,6 +328,7 @@ let rec add_expr bv exp =
       | Ok { arg; _ } -> add_expr bv arg
       end
   | Pexp_extension e -> handle_extension e
+  | Pexp_stack e -> add_expr bv e
   | Pexp_unreachable -> ()
 
 and add_expr_jane_syntax bv : Jane_syntax.Expression.t -> _ = function

parsing/lexer.mll

@@ -84,6 +84,7 @@ let keyword_table =
     "private", PRIVATE;
     "rec", REC;
     "sig", SIG;
+    "stack_", STACK;
     "struct", STRUCT;
     "then", THEN;
     "to", TO;

parsing/parser.mly

@@ -1042,6 +1042,7 @@ let unboxed_type sloc lident tys =
 %token HASH_SUFFIX            "# "
 %token <string> HASHOP        "##" (* just an example *)
 %token SIG                    "sig"
+%token STACK                  "stack_"
 %token STAR                   "*"
 %token <string * Location.t * string option>
        STRING                 "\"hello\"" (* just an example *)
@@ -1133,7 +1134,7 @@ The precedences must be listed from low to high.
 %nonassoc BACKQUOTE BANG BEGIN CHAR FALSE FLOAT HASH_FLOAT INT HASH_INT OBJECT
           LBRACE LBRACELESS LBRACKET LBRACKETBAR LBRACKETCOLON LIDENT LPAREN
           NEW PREFIXOP STRING TRUE UIDENT
-          LBRACKETPERCENT QUOTED_STRING_EXPR
+          LBRACKETPERCENT QUOTED_STRING_EXPR STACK
 
 
 /* Entry points */
@@ -2811,6 +2812,8 @@ fun_expr:
 %inline expr_:
   | simple_expr nonempty_llist(labeled_simple_expr)
       { mkexp ~loc:$sloc (Pexp_apply($1, $2)) }
+  | STACK simple_expr
+      { mkexp ~loc:$sloc (Pexp_stack $2) }
   | labeled_tuple %prec below_COMMA
       { pexp_ltuple $sloc $1 }
   | mkrhs(constr_longident) simple_expr %prec below_HASH

parsing/parsetree.mli

@@ -440,6 +440,7 @@ and expression_desc =
             - [let* P0 = E00 and* P1 = E01 in E1] *)
   | Pexp_extension of extension  (** [[%id]] *)
   | Pexp_unreachable  (** [.] *)
+  | Pexp_stack of expression (** stack_ exp *)
 
 and case =
     {

parsing/pprintast.ml

@@ -933,6 +933,9 @@ and expression ?(jane_syntax_parens = false) ctxt f x =
                 end (e,l)
         end
 
+    | Pexp_stack e ->
+        (* Similar to the common case of [Pexp_apply] *)
+        pp f "@[<hov2>stack_@ %a@]" (expression2 reset_ctxt)  e
     | Pexp_construct (li, Some eo)
       when not (is_simple_construct (view_expr x))-> (* Not efficient FIXME*)
         (match view_expr x with

parsing/printast.ml

@@ -384,6 +384,9 @@ and expression i ppf x =
       payload i ppf arg
   | Pexp_unreachable ->
       line i ppf "Pexp_unreachable"
+  | Pexp_stack e ->
+      line i ppf "Pexp_stack\n";
+      expression i ppf e
 
 and value_description i ppf x =
   line i ppf "value_description %a %a\n" fmt_string_loc

testsuite/tests/language-extensions/pprintast_unconditional.ml

@@ -34,6 +34,11 @@ module Example = struct
      end"
 
   let local_exp = parse expression "let x = foo (local_ x) in local_ y"
+  let stack_exp = parse expression
+    "let x = stack_ 42 in \
+     let y = stack_ (f x) in \
+     let z = foo (stack_ 42) in \
+     foo (stack_ (f x))"
 
   let modal_kind_struct =
     parse module_expr "struct \
@@ -178,6 +183,7 @@ end = struct
   let modality_val = test "modality_val" module_type Example.modality_val
 
   let local_exp = test "local_exp" expression Example.local_exp
+  let stack_exp = test "stack_exp" expression Example.stack_exp
 
   let longident = test "longident" longident Example.longident
   let expression = test "expression" expression Example.expression

testsuite/tests/language-extensions/pprintast_unconditional.reference

@@ -17,6 +17,10 @@ modality_val: sig val t : string -> local_ string @@ foo bar end
 
 local_exp: let x = foo (local_ x) in local_ y
 
+stack_exp:
+  let x = stack_ 42 in
+  let y = stack_ (f x) in let z = foo (stack_ 42) in foo (stack_ (f x))
+
 longident: No.Longidents.Require.extensions
 
 expression: [x for x = 1 to 10]
@@ -118,6 +122,10 @@ modality_val: sig val t : string -> local_ string @@ foo bar end
 
 local_exp: let x = foo (local_ x) in local_ y
 
+stack_exp:
+  let x = stack_ 42 in
+  let y = stack_ (f x) in let z = foo (stack_ 42) in foo (stack_ (f x))
+
 longident: No.Longidents.Require.extensions
 
 expression: [x for x = 1 to 10]