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

@@ -452,7 +452,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
 
@@ -575,11 +576,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
 
@@ -142,6 +143,8 @@ val event : Parse_info.t -> unit t
 
 val no_event : unit t
 
+val hidden_location : Parse_info.t
+
 type type_def =
   { supertype : Wa_ast.var option
   ; final : bool

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,19 +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:(`Catch x) ~context:(`Catch x :: context))
-         [ ( ocaml_tag
-           , let* () = no_event in
-             store ~always:true x (return (W.Pop Value.value)) )
-         ; ( js_tag
-           , let* () = no_event in
-             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:`Skip
+                         ~context:(`Skip :: `Skip :: `Catch :: context))
+                      [ ocaml_tag, 1, Value.value; js_tag, 0, externref ]
+                  in
+                  instr (W.Push e))
+             in
+             instr (W.CallInstr (f, [ exn ]))))
      in
      let* () = no_event in
      exn_handler ~result_typ ~fall_through ~context)

compiler/lib/wasm/wa_generate.ml

@@ -160,15 +160,15 @@ module Generate (Target : Wa_target_sig.S) = struct
     let rec index_rec context pc i =
       match context with
       | `Block pc' :: _ when pc = pc' -> i
-      | (`Block _ | `Skip | `Catch _) :: rem -> index_rec rem pc (i + 1)
+      | (`Block _ | `Skip | `Catch) :: rem -> index_rec rem pc (i + 1)
       | [] -> assert false
     in
     index_rec context pc 0
 
   let catch_index context =
     let rec index_rec context i =
       match context with
-      | `Catch x :: _ -> Some (x, i)
+      | `Catch :: _ -> Some i
       | (`Block _ | `Skip | `Return) :: rem -> index_rec rem (i + 1)
       | [] -> None
     in
@@ -759,7 +759,7 @@ module Generate (Target : Wa_target_sig.S) = struct
 
   let extend_context fall_through context =
     match fall_through with
-    | (`Block _ | `Catch _) as b -> b :: context
+    | (`Block _ | `Catch | `Skip) as b -> b :: context
     | `Return -> `Skip :: context
 
   let needed_handlers (p : program) pc =
@@ -911,7 +911,7 @@ module Generate (Target : Wa_target_sig.S) = struct
               let* e = load x in
               match fall_through with
               | `Return -> instr (Push e)
-              | `Block _ | `Catch _ -> instr (Return (Some e)))
+              | `Block _ | `Catch | `Skip -> instr (Return (Some e)))
           | Cond (x, cont1, cont2) ->
               let context' = extend_context fall_through context in
               if_
@@ -923,7 +923,7 @@ module Generate (Target : Wa_target_sig.S) = struct
               let* e = Value.unit in
               match fall_through with
               | `Return -> instr (Push e)
-              | `Block _ | `Catch _ -> instr (Return (Some e)))
+              | `Block _ | `Catch | `Skip -> instr (Return (Some e)))
           | Switch (x, a) ->
               let len = Array.length a in
               let l = Array.to_list (Array.sub a ~pos:0 ~len:(len - 1)) in
@@ -937,12 +937,10 @@ module Generate (Target : Wa_target_sig.S) = struct
               let* e = load x in
               let* tag = register_import ~name:exception_name (Tag Value.value) in
               match fall_through with
-              | `Catch x -> store ~always:true x (return e)
-              | `Block _ | `Return -> (
+              | `Catch -> instr (Push e)
+              | `Block _ | `Return | `Skip -> (
                   match catch_index context with
-                  | Some (x, i) ->
-                      let* () = store ~always:true x (return e) in
-                      instr (Br (i, None))
+                  | Some i -> instr (Br (i, Some e))
                   | None -> instr (Throw (tag, e))))
           | Pushtrap (cont, x, cont') ->
               handle_exceptions

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 | Event _ -> ()
   | ArraySet (_, e, e', e'') ->

compiler/lib/wasm/wa_tail_call.ml

@@ -37,12 +37,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_target_sig.ml

@@ -252,9 +252,9 @@ module type S = sig
   val handle_exceptions :
        result_typ:Wa_ast.value_type list
     -> fall_through:'a
-    -> context:([> `Catch of Code.Var.t ] as 'b) list
+    -> context:([> `Catch | `Skip ] as 'b) list
     -> (   result_typ:Wa_ast.value_type list
-        -> fall_through:[> `Catch of Code.Var.t ]
+        -> fall_through:[> `Skip ]
         -> context:'b list
         -> unit Wa_code_generation.t)
     -> Wa_ast.var

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

@@ -442,6 +442,21 @@ 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.Event Wa_code_generation.hidden_location)
+                             @ instruction (Wa_ast.Br (i + 1, Some (Pop ty))))))
+                      catches))
+        ]
   and instruction i =
     match i with
     | Drop e -> [ List (Atom "drop" :: expression e) ]
@@ -460,20 +475,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"