Wasm_of_ocaml: small changes (#1964)

* Improve placement of Wasm block statements

Place Wasm block statements around control instructions rather than
around whole blocks of instructions. This results in less variable being
uninitialized (variables initialized within a Wasm block are no longer
considered as initialized after the block).

* Clean-up regarding function and primitive types

* Wasm AST: add br_on_null

Author: vouillon

compiler/lib-wasm/code_generation.ml

@@ -474,6 +474,7 @@ let rec is_smi e =
   | RefNull _
   | Br_on_cast _
   | Br_on_cast_fail _
+  | Br_on_null _
   | Try _
   | ExternConvertAny _ -> false
   | BinOp ((F32 _ | F64 _), _, _) | RefTest _ | RefEq _ -> true

compiler/lib-wasm/curry.ml

@@ -24,11 +24,6 @@ open Code_generation
 module Make (Target : Target_sig.S) = struct
   open Target
 
-  let func_type n =
-    { W.params = List.init ~len:(n + 1) ~f:(fun _ -> Value.value)
-    ; result = [ Value.value ]
-    }
-
   let bind_parameters l =
     List.fold_left
       ~f:(fun l x ->
@@ -105,7 +100,7 @@ module Make (Target : Target_sig.S) = struct
       { name
       ; exported_name = None
       ; typ = None
-      ; signature = func_type 1
+      ; signature = Type.func_type 1
       ; param_names
       ; locals
       ; body
@@ -140,7 +135,7 @@ module Make (Target : Target_sig.S) = struct
       { name
       ; exported_name = None
       ; typ = None
-      ; signature = func_type 1
+      ; signature = Type.func_type 1
       ; param_names
       ; locals
       ; body
@@ -191,7 +186,7 @@ module Make (Target : Target_sig.S) = struct
       { name
       ; exported_name = None
       ; typ = None
-      ; signature = func_type 2
+      ; signature = Type.func_type 2
       ; param_names
       ; locals
       ; body
@@ -230,7 +225,7 @@ module Make (Target : Target_sig.S) = struct
       { name
       ; exported_name = None
       ; typ = None
-      ; signature = func_type 2
+      ; signature = Type.func_type 2
       ; param_names
       ; locals
       ; body
@@ -274,7 +269,7 @@ module Make (Target : Target_sig.S) = struct
       { name
       ; exported_name = None
       ; typ = None
-      ; signature = func_type arity
+      ; signature = Type.primitive_type (arity + 1)
       ; param_names
       ; locals
       ; body
@@ -306,7 +301,7 @@ module Make (Target : Target_sig.S) = struct
              (List.map ~f:(fun x -> `Var x) (List.tl l))
          in
          let* make_iterator =
-           register_import ~name:"caml_apply_continuation" (Fun (func_type 0))
+           register_import ~name:"caml_apply_continuation" (Fun (Type.primitive_type 1))
          in
          let iterate = Var.fresh_n "iterate" in
          let* () = store iterate (return (W.Call (make_iterator, [ args ]))) in
@@ -321,7 +316,7 @@ module Make (Target : Target_sig.S) = struct
       { name
       ; exported_name = None
       ; typ = None
-      ; signature = func_type arity
+      ; signature = Type.primitive_type (arity + 1)
       ; param_names
       ; locals
       ; body
@@ -356,7 +351,7 @@ module Make (Target : Target_sig.S) = struct
       { name
       ; exported_name = None
       ; typ = None
-      ; signature = func_type arity
+      ; signature = Type.func_type arity
       ; param_names
       ; locals
       ; body

compiler/lib-wasm/gc_target.ml

@@ -202,8 +202,10 @@ module Type = struct
                 ]
           })
 
-  let func_type n =
-    { W.params = List.init ~len:(n + 1) ~f:(fun _ -> value); result = [ value ] }
+  let primitive_type n =
+    { W.params = List.init ~len:n ~f:(fun _ -> value); result = [ value ] }
+
+  let func_type n = primitive_type (n + 1)
 
   let function_type ~cps n =
     let n = if cps then n + 1 else n in
@@ -423,8 +425,6 @@ module Type = struct
 end
 
 module Value = struct
-  let value = Type.value
-
   let block_type =
     let* t = Type.block_type in
     return (W.Ref { nullable = false; typ = Type t })
@@ -528,6 +528,7 @@ module Value = struct
     | Call_ref _
     | Br_on_cast _
     | Br_on_cast_fail _
+    | Br_on_null _
     | Try _ -> false
     | IfExpr (_, e1, e2, e3) -> effect_free e1 && effect_free e2 && effect_free e3
     | ArrayNewFixed (_, l) | StructNew (_, l) -> List.for_all ~f:effect_free l
@@ -743,13 +744,13 @@ module Memory = struct
     let a = Code.Var.fresh_n "a" in
     let i = Code.Var.fresh_n "i" in
     block_expr
-      { params = []; result = [ Value.value ] }
+      { params = []; result = [ Type.value ] }
       (let* () = store a e in
        let* () = store ~typ:I32 i (Value.int_val e') in
        let* () =
          drop
            (block_expr
-              { params = []; result = [ Value.value ] }
+              { params = []; result = [ Type.value ] }
               (let* block = Type.block_type in
                let* a = load a in
                let* e =
@@ -779,7 +780,7 @@ module Memory = struct
       (let* () =
          drop
            (block_expr
-              { params = []; result = [ Value.value ] }
+              { params = []; result = [ Type.value ] }
               (let* block = Type.block_type in
                let* a = load a in
                let* () =
@@ -840,7 +841,7 @@ module Memory = struct
     let* () =
       drop
         (block_expr
-           { params = []; result = [ Value.value ] }
+           { params = []; result = [ Type.value ] }
            (let* e =
               if_match
                 ~typ:(Some (W.Ref { nullable = false; typ = Type fun_ty }))
@@ -1406,7 +1407,7 @@ let internal_primitives =
     let arity = List.length args in
     (* [Type.func_type] counts one additional argument for the closure environment (absent
        here) *)
-    let* f = register_import ~name (Fun (Type.func_type (arity - 1))) in
+    let* f = register_import ~name (Fun (Type.primitive_type arity)) in
     let args = List.map ~f:transl_prim_arg args in
     let* args = expression_list Fun.id args in
     return (W.Call (f, args))
@@ -1675,19 +1676,19 @@ let externref = W.Ref { nullable = true; typ = Extern }
 
 let handle_exceptions ~result_typ ~fall_through ~context body x exn_handler =
   let* js_tag = register_import ~name:"javascript_exception" (Tag externref) in
-  let* ocaml_tag = register_import ~name:"ocaml_exception" (Tag Value.value) in
+  let* ocaml_tag = register_import ~name:"ocaml_exception" (Tag Type.value) in
   let* f =
     register_import
       ~name:"caml_wrap_exception"
-      (Fun { params = [ externref ]; result = [ Value.value ] })
+      (Fun { params = [ externref ]; result = [ Type.value ] })
   in
   block
     { params = []; result = result_typ }
     (let* () =
        store
          x
          (block_expr
-            { params = []; result = [ Value.value ] }
+            { params = []; result = [ Type.value ] }
             (let* exn =
                block_expr
                  { params = []; result = [ externref ] }
@@ -1698,7 +1699,7 @@ let handle_exceptions ~result_typ ~fall_through ~context body x exn_handler =
                          ~result_typ:[ externref ]
                          ~fall_through:`Skip
                          ~context:(`Skip :: `Skip :: `Catch :: context))
-                      [ ocaml_tag, 1, Value.value; js_tag, 0, externref ]
+                      [ ocaml_tag, 1, Type.value; js_tag, 0, externref ]
                   in
                   instr (W.Push e))
              in

compiler/lib-wasm/generate.ml

@@ -72,13 +72,13 @@ module Generate (Target : Target_sig.S) = struct
 
   let repr_type r =
     match r with
-    | Value -> Value.value
+    | Value -> Type.value
     | Float -> F64
     | Int32 -> I32
     | Nativeint -> I32
     | Int64 -> I64
 
-  let specialized_func_type (_, params, result) =
+  let specialized_primitive_type (_, params, result) =
     { W.params = List.map ~f:repr_type params; result = [ repr_type result ] }
 
   let box_value r e =
@@ -127,9 +127,6 @@ module Generate (Target : Target_sig.S) = struct
       ];
     h
 
-  let func_type n =
-    { W.params = List.init ~len:n ~f:(fun _ -> Value.value); result = [ Value.value ] }
-
   let float_bin_op' op f g =
     Memory.box_float (op (Memory.unbox_float f) (Memory.unbox_float g))
 
@@ -711,7 +708,7 @@ module Generate (Target : Target_sig.S) = struct
                   let ((_, arg_typ, res_typ) as typ) =
                     Hashtbl.find specialized_primitives name
                   in
-                  let* f = register_import ~name (Fun (specialized_func_type typ)) in
+                  let* f = register_import ~name (Fun (specialized_primitive_type typ)) in
                   let rec loop acc arg_typ l =
                     match arg_typ, l with
                     | [], [] -> box_value res_typ (return (W.Call (f, List.rev acc)))
@@ -722,7 +719,9 @@ module Generate (Target : Target_sig.S) = struct
                   in
                   loop [] arg_typ l
                 with Not_found ->
-                  let* f = register_import ~name (Fun (func_type (List.length l))) in
+                  let* f =
+                    register_import ~name (Fun (Type.primitive_type (List.length l)))
+                  in
                   let rec loop acc l =
                     match l with
                     | [] -> return (W.Call (f, List.rev acc))
@@ -921,6 +920,8 @@ module Generate (Target : Target_sig.S) = struct
         | _ -> Structure.is_merge_node g pc'
       in
       let code ~context =
+        let block = Addr.Map.find pc ctx.blocks in
+        let* () = translate_instrs ctx context block.body in
         translate_node_within
           ~result_typ
           ~fall_through
@@ -965,7 +966,6 @@ module Generate (Target : Target_sig.S) = struct
           translate_tree result_typ fall_through pc' context
       | [] -> (
           let block = Addr.Map.find pc ctx.blocks in
-          let* () = translate_instrs ctx context block.body in
           let branch = block.branch in
           match branch with
           | Branch cont -> translate_branch result_typ fall_through pc cont context
@@ -997,7 +997,7 @@ module Generate (Target : Target_sig.S) = struct
               instr (Br_table (e, List.map ~f:dest l, dest a.(len - 1)))
           | Raise (x, _) -> (
               let* e = load x in
-              let* tag = register_import ~name:exception_name (Tag Value.value) in
+              let* tag = register_import ~name:exception_name (Tag Type.value) in
               match fall_through with
               | `Catch -> instr (Push e)
               | `Block _ | `Return | `Skip -> (
@@ -1082,7 +1082,7 @@ module Generate (Target : Target_sig.S) = struct
              wrap_with_handlers
                p
                pc
-               ~result_typ:[ Value.value ]
+               ~result_typ:[ Type.value ]
                ~fall_through:`Return
                ~context:[]
                (fun ~result_typ ~fall_through ~context ->
@@ -1103,7 +1103,10 @@ module Generate (Target : Target_sig.S) = struct
           | None -> Option.map ~f:(fun name -> name ^ ".init") unit_name
           | Some _ -> None)
       ; typ = None
-      ; signature = func_type param_count
+      ; signature =
+          (match name_opt with
+          | None -> Type.primitive_type param_count
+          | Some _ -> Type.func_type (param_count - 1))
       ; param_names
       ; locals
       ; body
@@ -1112,7 +1115,7 @@ module Generate (Target : Target_sig.S) = struct
 
   let init_function ~context ~to_link =
     let name = Code.Var.fresh_n "initialize" in
-    let signature = { W.params = []; result = [ Value.value ] } in
+    let signature = { W.params = []; result = [ Type.value ] } in
     let locals, body =
       function_body
         ~context
@@ -1288,7 +1291,7 @@ module G = Generate (Gc_target)
 
 let init = G.init
 
-let start () = make_context ~value_type:Gc_target.Value.value
+let start () = make_context ~value_type:Gc_target.Type.value
 
 let f ~context ~unit_name p ~live_vars ~in_cps ~deadcode_sentinal =
   let t = Timer.make () in

compiler/lib-wasm/initialize_locals.ml

@@ -47,6 +47,7 @@ let rec scan_expression ctx e =
   | RefTest (_, e')
   | Br_on_cast (_, _, _, e')
   | Br_on_cast_fail (_, _, _, e')
+  | Br_on_null (_, e') -> scan_expression ctx e'
   | ExternConvertAny e' -> scan_expression ctx e'
   | BinOp (_, e', e'')
   | ArrayNew (_, e', e'')

compiler/lib-wasm/target_sig.ml

@@ -96,9 +96,15 @@ module type S = sig
     val unbox_nativeint : expression -> expression
   end
 
-  module Value : sig
+  module Type : sig
     val value : Wasm_ast.value_type
 
+    val func_type : int -> Wasm_ast.func_type
+
+    val primitive_type : int -> Wasm_ast.func_type
+  end
+
+  module Value : sig
     val unit : expression
 
     val val_int : expression -> expression

compiler/lib-wasm/wasm_ast.ml

@@ -167,6 +167,7 @@ type expression =
   | RefNull of heap_type
   | Br_on_cast of int * ref_type * ref_type * expression
   | Br_on_cast_fail of int * ref_type * ref_type * expression
+  | Br_on_null of int * expression
   | IfExpr of value_type * expression * expression * expression
   | Try of func_type * instruction list * (var * int * value_type) list
   | ExternConvertAny of expression

compiler/lib-wasm/wasm_output.ml

@@ -645,6 +645,11 @@ end = struct
         output_uint ch i;
         output_heaptype st.type_names ch typ1.typ;
         output_heaptype st.type_names ch typ2.typ
+    | Br_on_null (i, e') ->
+        Feature.require gc;
+        output_expression st ch e';
+        output_byte ch 0xD5;
+        output_uint ch i
     | IfExpr (typ, e1, e2, e3) ->
         output_expression st ch e1;
         output_byte ch 0x04;
@@ -885,6 +890,7 @@ end = struct
     | RefTest (_, e')
     | Br_on_cast (_, _, _, e')
     | Br_on_cast_fail (_, _, _, e')
+    | Br_on_null (_, e') -> expr_function_references e' set
     | ExternConvertAny e' -> expr_function_references e' set
     | BinOp (_, e', e'')
     | ArrayNew (_, e', e'')

compiler/lib-wasm/wat_output.ml

@@ -452,6 +452,8 @@ let expression_or_instructions ctx st in_function =
             :: ref_type st ty'
             :: expression e)
         ]
+    | Br_on_null (i, e) ->
+        [ List (Atom "br_on_null" :: Atom (string_of_int i) :: expression e) ]
     | IfExpr (ty, cond, ift, iff) ->
         [ List
             ((Atom "if" :: block_type st { params = []; result = [ ty ] })

runtime/wasm/domain.wat

@@ -23,8 +23,6 @@
       (func $caml_ml_mutex_unlock (param (ref eq)) (result (ref eq))))
 
    (type $block (array (mut (ref eq))))
-   (type $function_1 (func (param (ref eq) (ref eq)) (result (ref eq))))
-   (type $closure (sub (struct (;(field i32);) (field (ref $function_1)))))
 
    (func (export "caml_atomic_cas")
       (param $ref (ref eq)) (param $o (ref eq)) (param $n (ref eq))