Use a try_table like instruction internally

Author: vouillon

compiler/lib/wasm/wa_asm_output.ml

@@ -343,7 +343,17 @@ module Output () = struct
     | RefEq _
     | RefNull _
     | Br_on_cast _
-    | Br_on_cast_fail _ -> assert false (* Not supported *)
+    | Br_on_cast_fail _-> assert false (* Not supported *)
+    | Try (ty, body, catches) ->
+        Feature.require exception_handling;
+        line (string "try" ^^ block_type ty)
+        ^^ indent (concat_map (instruction m) body)
+        ^^ concat_map
+             (fun (tag, i, ty) ->
+               line (string "catch " ^^ index tag)
+               ^^ indent (instruction m (Wa_ast.Br (i, Some (Pop ty)))))
+             catches
+        ^^ line (string "end_try")
 
   and instruction m i =
     match i with
@@ -394,19 +404,6 @@ module Output () = struct
     | CallInstr (x, l) -> concat_map (expression m) l ^^ line (string "call " ^^ index x)
     | Nop -> empty
     | Push e -> expression m e
-    | Try (ty, body, catches, catch_all) ->
-        Feature.require exception_handling;
-        line (string "try" ^^ block_type ty)
-        ^^ indent (concat_map (instruction m) body)
-        ^^ concat_map
-             (fun (tag, l) ->
-               line (string "catch " ^^ index tag)
-               ^^ indent (concat_map (instruction m) l))
-             catches
-        ^^ (match catch_all with
-           | None -> empty
-           | Some l -> line (string "catch_all") ^^ indent (concat_map (instruction m) l))
-        ^^ line (string "end_try")
     | Throw (i, e) ->
         Feature.require exception_handling;
         expression m e ^^ line (string "throw " ^^ index i)

compiler/lib/wasm/wa_ast.ml

@@ -174,6 +174,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
@@ -191,11 +192,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

@@ -458,7 +458,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
 
@@ -583,11 +584,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_core_target.ml

@@ -650,13 +650,17 @@ let internal_primitives = Hashtbl.create 0
 
 let handle_exceptions ~result_typ ~fall_through ~context body x exn_handler =
   let* ocaml_tag = register_import ~name:"ocaml_exception" (Tag Value.value) in
-  try_
+  block
     { params = []; result = result_typ }
-    (body ~result_typ ~fall_through:(`Block (-1)) ~context)
-    [ ( ocaml_tag
-      , let* () = store ~always:true x (return (W.Pop Value.value)) in
-        exn_handler ~result_typ ~fall_through ~context )
-    ]
+    (let* () =
+       store
+         x
+         (try_expr
+            { params = []; result = [ Value.value ] }
+            (body ~result_typ ~fall_through:(`Block (-1)) ~context:(`Skip :: context))
+            [ ocaml_tag, 0, Value.value ])
+     in
+     exn_handler ~result_typ ~fall_through ~context)
 
 let post_process_function_body ~param_names:_ ~locals:_ instrs = instrs
 

compiler/lib/wasm/wa_gc_target.ml

@@ -530,7 +530,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
 
@@ -1701,16 +1702,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

@@ -71,6 +71,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
 
@@ -95,10 +96,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

@@ -41,12 +41,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_indirect (ty, e, l))) -> Return_call_indirect (ty, e, l)
   | Return (Some (Call_ref (ty, e, l))) -> Return_call_ref (ty, e, l)

compiler/lib/wasm/wa_wat_output.ml

@@ -476,6 +476,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) ]
@@ -509,20 +522,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"