Use a try_table like instruction internally

Author: vouillon

compiler/lib/wasm/wa_ast.ml

@@ -165,6 +165,7 @@ type expression =
   | Br_on_cast of int * ref_type * ref_type * expression
   | Br_on_cast_fail of int * ref_type * ref_type * expression
   | IfExpr of value_type * expression * expression * expression
+  | Try of func_type * instruction list * (var * int * value_type) list
 
 and instruction =
   | Drop of expression
@@ -180,11 +181,6 @@ and instruction =
   | CallInstr of var * expression list
   | Nop
   | Push of expression
-  | Try of
-      func_type
-      * instruction list
-      * (var * instruction list) list
-      * instruction list option
   | Throw of var * expression
   | Rethrow of int
   | ArraySet of var * expression * expression * expression

compiler/lib/wasm/wa_code_generation.ml

@@ -447,7 +447,8 @@ let rec is_smi e =
   | RefCast _
   | RefNull _
   | Br_on_cast _
-  | Br_on_cast_fail _ -> false
+  | Br_on_cast_fail _
+  | Try _ -> false
   | BinOp ((F32 _ | F64 _), _, _) | RefTest _ | RefEq _ -> true
   | IfExpr (_, _, ift, iff) -> is_smi ift && is_smi iff
 
@@ -572,11 +573,9 @@ let if_ ty e l1 l2 =
   | W.UnOp (I32 Eqz, e') -> instr (If (ty, e', instrs2, instrs1))
   | _ -> instr (If (ty, e, instrs1, instrs2))
 
-let try_ ty body handlers =
+let try_expr ty body handlers =
   let* body = blk body in
-  let tags = List.map ~f:fst handlers in
-  let* handler_bodies = expression_list blk (List.map ~f:snd handlers) in
-  instr (Try (ty, body, List.combine tags handler_bodies, None))
+  return (W.Try (ty, body, handlers))
 
 let need_apply_fun ~cps ~arity st =
   let ctx = st.context in

compiler/lib/wasm/wa_code_generation.mli

@@ -128,7 +128,8 @@ val block_expr : Wa_ast.func_type -> unit t -> expression
 
 val if_ : Wa_ast.func_type -> expression -> unit t -> unit t -> unit t
 
-val try_ : Wa_ast.func_type -> unit t -> (Code.Var.t * unit t) list -> unit t
+val try_expr :
+  Wa_ast.func_type -> unit t -> (Code.Var.t * int * Wa_ast.value_type) list -> expression
 
 val add_var : ?typ:Wa_ast.value_type -> Wa_ast.var -> Wa_ast.var t
 

compiler/lib/wasm/wa_gc_target.ml

@@ -526,7 +526,8 @@ module Value = struct
     | Pop _
     | Call_ref _
     | Br_on_cast _
-    | Br_on_cast_fail _ -> false
+    | Br_on_cast_fail _
+    | 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
 
@@ -1676,16 +1677,25 @@ let handle_exceptions ~result_typ ~fall_through ~context body x exn_handler =
   block
     { params = []; result = result_typ }
     (let* () =
-       try_
-         { params = []; result = [] }
-         (body ~result_typ:[] ~fall_through:(`Block (-1)) ~context:(`Skip :: context))
-         [ ocaml_tag, store ~always:true x (return (W.Pop Value.value))
-         ; ( js_tag
-           , let exn = Code.Var.fresh () in
-             let* () = store ~always:true ~typ:externref exn (return (W.Pop externref)) in
-             let* exn = load exn in
-             store ~always:true x (return (W.Call (f, [ exn ]))) )
-         ]
+       store
+         x
+         (block_expr
+            { params = []; result = [ Value.value ] }
+            (let* exn =
+               block_expr
+                 { params = []; result = [ externref ] }
+                 (let* e =
+                    try_expr
+                      { params = []; result = [ externref ] }
+                      (body
+                         ~result_typ:[ externref ]
+                         ~fall_through:(`Block (-1))
+                         ~context:(`Skip :: `Skip :: `Skip :: context))
+                      [ ocaml_tag, 1, Value.value; js_tag, 0, externref ]
+                  in
+                  instr (W.Push e))
+             in
+             instr (W.CallInstr (f, [ exn ]))))
      in
      exn_handler ~result_typ ~fall_through ~context)
 

compiler/lib/wasm/wa_initialize_locals.ml

@@ -68,6 +68,7 @@ let rec scan_expression ctx e =
       scan_expression ctx cond;
       scan_expression (fork_context ctx) e1;
       scan_expression (fork_context ctx) e2
+  | Try (_, body, _) -> scan_instructions ctx body
 
 and scan_expressions ctx l = List.iter ~f:(fun e -> scan_expression ctx e) l
 
@@ -92,10 +93,6 @@ and scan_instruction ctx i =
       scan_expression ctx e;
       scan_instructions ctx l;
       scan_instructions ctx l'
-  | Try (_, body, catches, catch_all) ->
-      scan_instructions ctx body;
-      List.iter ~f:(fun (_, l) -> scan_instructions ctx l) catches;
-      Option.iter ~f:(fun l -> scan_instructions ctx l) catch_all
   | CallInstr (_, l) | Return_call (_, l) -> scan_expressions ctx l
   | Br (_, None) | Return None | Rethrow _ | Nop -> ()
   | ArraySet (_, e, e', e'') ->

compiler/lib/wasm/wa_tail_call.ml

@@ -39,12 +39,6 @@ let rec instruction ~tail i =
   | Wa_ast.Loop (ty, l) -> Wa_ast.Loop (ty, instructions ~tail l)
   | Block (ty, l) -> Block (ty, instructions ~tail l)
   | If (ty, e, l1, l2) -> If (ty, e, instructions ~tail l1, instructions ~tail l2)
-  | Try (ty, l, catches, catch_all) ->
-      Try
-        ( ty
-        , l
-        , List.map ~f:(fun (tag, l) -> tag, instructions ~tail l) catches
-        , Option.map ~f:(fun l -> instructions ~tail l) catch_all )
   | Return (Some (Call (symb, l))) -> Return_call (symb, l)
   | Return (Some (Call_ref (ty, e, l))) -> Return_call_ref (ty, e, l)
   | Push (Call (symb, l)) when tail -> Return_call (symb, l)

compiler/lib/wasm/wa_wasm_output.ml

@@ -628,6 +628,18 @@ end = struct
         output_byte ch 0x05;
         output_expression st ch e3;
         output_byte ch 0x0B
+    | Try (typ, l, catches) ->
+        Feature.require exception_handling;
+        output_byte ch 0x06;
+        output_blocktype st.type_names ch typ;
+        List.iter ~f:(fun i' -> output_instruction st ch i') l;
+        List.iter
+          ~f:(fun (tag, l, ty) ->
+            output_byte ch 0x07;
+            output_uint ch (Hashtbl.find st.tag_names tag);
+            output_instruction st ch (Br (l + 1, Some (Pop ty))))
+          catches;
+        output_byte ch 0X0B
 
   and output_instruction st ch i =
     match i with
@@ -688,23 +700,6 @@ end = struct
         output_uint ch (Hashtbl.find st.func_names f)
     | Nop -> ()
     | Push e -> output_expression st ch e
-    | Try (typ, l, catches, catchall) ->
-        Feature.require exception_handling;
-        output_byte ch 0x06;
-        output_blocktype st.type_names ch typ;
-        List.iter ~f:(fun i' -> output_instruction st ch i') l;
-        List.iter
-          ~f:(fun (tag, l) ->
-            output_byte ch 0x07;
-            output_uint ch (Hashtbl.find st.tag_names tag);
-            List.iter ~f:(fun i' -> output_instruction st ch i') l)
-          catches;
-        (match catchall with
-        | Some l ->
-            output_byte ch 0x05;
-            List.iter ~f:(fun i' -> output_instruction st ch i') l
-        | None -> ());
-        output_byte ch 0X0B
     | Throw (tag, e) ->
         Feature.require exception_handling;
         output_expression st ch e;
@@ -881,6 +876,8 @@ end = struct
           ~f:(fun set i -> expr_function_references i set)
           ~init:(expr_function_references e' set)
           l
+    | Try (_, l, _) ->
+        List.fold_left ~f:(fun set i -> instr_function_references i set) ~init:set l
 
   and instr_function_references i set =
     match i with
@@ -905,24 +902,6 @@ end = struct
     | Br (_, None) | Return None | Nop | Rethrow _ -> set
     | CallInstr (_, l) ->
         List.fold_left ~f:(fun set i -> expr_function_references i set) ~init:set l
-    | Try (_, l, catches, catchall) ->
-        List.fold_left ~f:(fun set i -> instr_function_references i set) ~init:set l
-        |> (fun init ->
-             List.fold_left
-               ~f:(fun set (_, l) ->
-                 List.fold_left
-                   ~f:(fun set i -> instr_function_references i set)
-                   ~init:set
-                   l)
-               ~init
-               catches)
-        |> fun init ->
-        List.fold_left
-          ~f:(fun set i -> instr_function_references i set)
-          ~init
-          (match catchall with
-          | Some l -> l
-          | None -> [])
     | ArraySet (_, e1, e2, e3) ->
         set
         |> expr_function_references e1

compiler/lib/wasm/wa_wat_output.ml

@@ -428,6 +428,19 @@ let expression_or_instructions ctx st in_function =
             @ [ List (Atom "then" :: expression ift) ]
             @ [ List (Atom "else" :: expression iff) ])
         ]
+    | Try (ty, body, catches) ->
+        [ List
+            (Atom "try"
+            :: (block_type st ty
+               @ List (Atom "do" :: instructions body)
+                 :: List.map
+                      ~f:(fun (tag, i, ty) ->
+                        List
+                          (Atom "catch"
+                          :: index st.tag_names tag
+                          :: instruction (Wa_ast.Br (i + 1, Some (Pop ty)))))
+                      catches))
+        ]
   and instruction i =
     match i with
     | Drop e -> [ List (Atom "drop" :: expression e) ]
@@ -446,20 +459,6 @@ let expression_or_instructions ctx st in_function =
                @ list ~always:true "then" instructions (remove_nops l1)
                @ list "else" instructions (remove_nops l2)))
         ]
-    | Try (ty, body, catches, catch_all) ->
-        [ List
-            (Atom "try"
-            :: (block_type st ty
-               @ List (Atom "do" :: instructions body)
-                 :: (List.map
-                       ~f:(fun (tag, l) ->
-                         List (Atom "catch" :: index st.tag_names tag :: instructions l))
-                       catches
-                    @
-                    match catch_all with
-                    | None -> []
-                    | Some l -> [ List (Atom "catch_all" :: instructions l) ])))
-        ]
     | Br_table (e, l, i) ->
         [ List
             (Atom "br_table"