Wasm code generation: separate functions to allocate regular blocks and float arrays

Author: vouillon

compiler/lib-wasm/curry.ml

@@ -293,14 +293,7 @@ module Make (Target : Target_sig.S) = struct
         (fun ~typ closure ->
           let* l = expression_list load l in
           call ?typ ~cps:true ~arity closure l)
-        (let* args =
-           (* We don't need the deadcode sentinal when the tag is 0 *)
-           Memory.allocate
-             ~tag:0
-             ~deadcode_sentinal:(Code.Var.fresh ())
-             ~load
-             (List.map ~f:(fun x -> `Var x) (List.tl l))
-         in
+        (let* args = Memory.allocate ~tag:0 (expression_list load (List.tl l)) in
          let* make_iterator =
            register_import ~name:"caml_apply_continuation" (Fun (Type.primitive_type 1))
          in

compiler/lib-wasm/gc_target.ml

@@ -666,33 +666,23 @@ module Memory = struct
     let* ty = Type.float_type in
     wasm_struct_get ty (wasm_cast ty e) 0
 
-  let allocate ~tag ~deadcode_sentinal ~load l =
-    if tag = 254
-    then
-      let* l =
-        expression_list
-          (fun v ->
-            match v with
-            | `Var y ->
-                if Code.Var.equal y deadcode_sentinal
-                then return (W.Const (F64 0.))
-                else unbox_float (load y)
-            | `Expr e -> unbox_float (return e))
-          l
-      in
-      let* ty = Type.float_array_type in
-      return (W.ArrayNewFixed (ty, l))
-    else
-      let* l =
-        expression_list
-          (fun v ->
-            match v with
-            | `Var y -> load y
-            | `Expr e -> return e)
-          l
-      in
-      let* ty = Type.block_type in
-      return (W.ArrayNewFixed (ty, RefI31 (Const (I32 (Int32.of_int tag))) :: l))
+  let allocate ~tag l =
+    assert (tag <> 254);
+    let* l = l in
+    let* ty = Type.block_type in
+    return (W.ArrayNewFixed (ty, RefI31 (Const (I32 (Int32.of_int tag))) :: l))
+
+  let allocate_float_array ~deadcode_sentinal ~load l =
+    let* l =
+      expression_list
+        (fun y ->
+          if Code.Var.equal y deadcode_sentinal
+          then return (W.Const (F64 0.))
+          else unbox_float (load y))
+        l
+    in
+    let* ty = Type.float_array_type in
+    return (W.ArrayNewFixed (ty, l))
 
   let tag e = wasm_array_get e (Arith.const 0l)
 

compiler/lib-wasm/generate.ml

@@ -773,16 +773,7 @@ module Generate (Target : Target_sig.S) = struct
       ~ty:(Int Normalized)
       (fun i j -> Arith.((j < i) - (i < j)));
     register_prim "%js_array" `Pure (fun ctx _ l ->
-        let* l =
-          List.fold_right
-            ~f:(fun x acc ->
-              let* x = transl_prim_arg ctx x in
-              let* acc = acc in
-              return (`Expr x :: acc))
-            l
-            ~init:(return [])
-        in
-        Memory.allocate ~tag:0 ~deadcode_sentinal:ctx.deadcode_sentinal ~load l)
+        Memory.allocate ~tag:0 (expression_list (fun x -> transl_prim_arg ctx x) l))
 
   let rec translate_expr ctx context x e =
     match e with
@@ -822,11 +813,16 @@ module Generate (Target : Target_sig.S) = struct
           in
           return (W.Call (apply, args @ [ closure ]))
     | Block (tag, a, _, _) ->
-        Memory.allocate
-          ~deadcode_sentinal:ctx.deadcode_sentinal
-          ~tag
-          ~load:(fun x -> load_and_box ctx x)
-          (List.map ~f:(fun x -> `Var x) (Array.to_list a))
+        if tag = 254
+        then
+          Memory.allocate_float_array
+            ~deadcode_sentinal:ctx.deadcode_sentinal
+            ~load
+            (Array.to_list a)
+        else
+          Memory.allocate
+            ~tag
+            (expression_list (fun x -> load_and_box ctx x) (Array.to_list a))
     | Field (x, n, Non_float) -> Memory.field (load_and_box ctx x) n
     | Field (x, n, Float) ->
         Memory.float_array_get

compiler/lib-wasm/target_sig.ml

@@ -20,11 +20,12 @@ module type S = sig
   type expression = Code_generation.expression
 
   module Memory : sig
-    val allocate :
-         tag:int
-      -> deadcode_sentinal:Code.Var.t
+    val allocate : tag:int -> Wasm_ast.expression list Code_generation.t -> expression
+
+    val allocate_float_array :
+         deadcode_sentinal:Code.Var.t
       -> load:(Code.Var.t -> expression)
-      -> [ `Expr of Wasm_ast.expression | `Var of Wasm_ast.var ] list
+      -> Wasm_ast.var list
       -> expression
 
     val load_function_pointer :