tail calls

Author: vouillon

compiler/lib-wasm/code_generation.ml

@@ -709,7 +709,7 @@ let need_dummy_fun ~cps ~arity st =
 
 let init_code context = instrs context.init_code
 
-let function_body ~context ~param_names ~body =
+let function_body ~context ~return_exn ~param_names ~body =
   let st = { var_count = 0; vars = Var.Map.empty; instrs = []; context } in
   let (), st = body st in
   let local_count, body = st.var_count, List.rev st.instrs in
@@ -721,7 +721,7 @@ let function_body ~context ~param_names ~body =
       | Local (i, x, typ) -> local_types.(i) <- x, typ
       | Expr _ -> ())
     st.vars;
-  let body = Tail_call.f ~no_tail_call:context.no_tail_call body in
+  let body = Tail_call.f ~return_exn ~no_tail_call:context.no_tail_call body in
   let param_count = List.length param_names in
   let locals =
     local_types

compiler/lib-wasm/code_generation.mli

@@ -199,6 +199,7 @@ val need_dummy_fun : cps:bool -> arity:int -> Code.Var.t t
 
 val function_body :
      context:context
+  -> return_exn:bool
   -> param_names:Code.Var.t list
   -> body:unit t
   -> (Wasm_ast.var * Wasm_ast.value_type) list * Wasm_ast.instruction list

compiler/lib-wasm/curry.ml

@@ -95,7 +95,7 @@ module Make (Target : Target_sig.S) = struct
       loop m [] f None
     in
     let param_names = args @ [ f ] in
-    let locals, body = function_body ~context ~param_names ~body in
+    let locals, body = function_body ~context ~return_exn:false ~param_names ~body in
     W.Function
       { name
       ; exported_name = None
@@ -130,7 +130,7 @@ module Make (Target : Target_sig.S) = struct
       push (Closure.curry_allocate ~cps:false ~arity m ~f:name' ~closure:f ~arg:x)
     in
     let param_names = [ x; f ] in
-    let locals, body = function_body ~context ~param_names ~body in
+    let locals, body = function_body ~context ~return_exn:false ~param_names ~body in
     W.Function
       { name
       ; exported_name = None
@@ -181,7 +181,7 @@ module Make (Target : Target_sig.S) = struct
       loop m [] f None
     in
     let param_names = args @ [ f ] in
-    let locals, body = function_body ~context ~param_names ~body in
+    let locals, body = function_body ~context ~return_exn:false ~param_names ~body in
     W.Function
       { name
       ; exported_name = None
@@ -220,7 +220,7 @@ module Make (Target : Target_sig.S) = struct
       instr (W.Return (Some c))
     in
     let param_names = [ x; cont; f ] in
-    let locals, body = function_body ~context ~param_names ~body in
+    let locals, body = function_body ~context ~return_exn:false ~param_names ~body in
     W.Function
       { name
       ; exported_name = None
@@ -264,7 +264,7 @@ module Make (Target : Target_sig.S) = struct
          build_applies (load f) l)
     in
     let param_names = l @ [ f ] in
-    let locals, body = function_body ~context ~param_names ~body in
+    let locals, body = function_body ~context ~return_exn:false ~param_names ~body in
     W.Function
       { name
       ; exported_name = None
@@ -305,7 +305,7 @@ module Make (Target : Target_sig.S) = struct
          push (call ~cps:true ~arity:2 (load f) [ x; iterate ]))
     in
     let param_names = l @ [ f ] in
-    let locals, body = function_body ~context ~param_names ~body in
+    let locals, body = function_body ~context ~return_exn:false ~param_names ~body in
     W.Function
       { name
       ; exported_name = None
@@ -340,7 +340,7 @@ module Make (Target : Target_sig.S) = struct
       instr (W.Return (Some e))
     in
     let param_names = l @ [ f ] in
-    let locals, body = function_body ~context ~param_names ~body in
+    let locals, body = function_body ~context ~return_exn:false ~param_names ~body in
     W.Function
       { name
       ; exported_name = None

compiler/lib-wasm/generate.ml

@@ -1969,6 +1969,7 @@ module Generate (Target : Target_sig.S) = struct
     let locals, body =
       function_body
         ~context:ctx.global_context
+        ~return_exn
         ~param_names
         ~body:
           (let* () =
@@ -2037,6 +2038,7 @@ module Generate (Target : Target_sig.S) = struct
     let locals, body =
       function_body
         ~context
+        ~return_exn:false
         ~param_names:[]
         ~body:
           (List.fold_right
@@ -2067,7 +2069,7 @@ module Generate (Target : Target_sig.S) = struct
 
   let entry_point context toplevel_fun entry_name =
     let signature, param_names, body = entry_point ~toplevel_fun in
-    let locals, body = function_body ~context ~param_names ~body in
+    let locals, body = function_body ~context ~return_exn:false ~param_names ~body in
     W.Function
       { name = Var.fresh_n "entry_point"
       ; exported_name = Some entry_name

compiler/lib-wasm/tail_call.ml

@@ -24,30 +24,46 @@ let get_return ~tail i =
   | Push (LocalGet y) when tail -> Some y
   | _ -> None
 
-let rewrite_tail_call ~no_tail_call ~y i =
+let rewrite_tail_call ~return_exn ~no_tail_call ~y i =
   match i with
   | Wasm_ast.LocalSet (x, Call (symb, l))
     when Code.Var.equal x y && not (Code.Var.Hashtbl.mem no_tail_call symb) ->
       Some (Wasm_ast.Return_call (symb, l))
   | LocalSet (x, Call_ref (ty, e, l)) when Code.Var.equal x y ->
       Some (Return_call_ref (ty, e, l))
+  | LocalSet (x, Br_on_null (_, Call (symb, l))) when return_exn && Code.Var.equal x y ->
+      Some (Wasm_ast.Return_call (symb, l))
+  | LocalSet (x, Br_on_null (_, Call_ref (ty, e, l)))
+    when return_exn && Code.Var.equal x y -> Some (Return_call_ref (ty, e, l))
   | _ -> None
 
-let rec instruction ~no_tail_call ~tail i =
+let rec instruction ~return_exn ~no_tail_call ~tail i =
   match i with
-  | Wasm_ast.Loop (ty, l) -> Wasm_ast.Loop (ty, instructions ~no_tail_call ~tail l)
-  | Block (ty, l) -> Block (ty, instructions ~no_tail_call ~tail l)
+  | Wasm_ast.Loop (ty, l) ->
+      Wasm_ast.Loop (ty, instructions ~return_exn ~no_tail_call ~tail l)
+  | Block (ty, l) -> Block (ty, instructions ~return_exn ~no_tail_call ~tail l)
   | If (ty, e, l1, l2) ->
-      If (ty, e, instructions ~no_tail_call ~tail l1, instructions ~no_tail_call ~tail l2)
+      If
+        ( ty
+        , e
+        , instructions ~return_exn ~no_tail_call ~tail l1
+        , instructions ~return_exn ~no_tail_call ~tail l2 )
   | Return (Some (Call (symb, l))) when not (Code.Var.Hashtbl.mem no_tail_call symb) ->
       Return_call (symb, l)
   | Return (Some (Call_ref (ty, e, l))) -> Return_call_ref (ty, e, l)
   | Push (Call (symb, l)) when tail && not (Code.Var.Hashtbl.mem no_tail_call symb) ->
       Return_call (symb, l)
   | Push (Call_ref (ty, e, l)) when tail -> Return_call_ref (ty, e, l)
   | Push (Call_ref _) -> i
+  | Return (Some (Br_on_null (_, Call (symb, l)))) when return_exn -> Return_call (symb, l)
+  | Return (Some (Br_on_null (_, Call_ref (ty, e, l)))) when return_exn ->
+      Return_call_ref (ty, e, l)
+  | Push (Br_on_null (_, Call (symb, l))) when return_exn && tail -> Return_call (symb, l)
+  | Push (Br_on_null (_, Call_ref (ty, e, l))) when return_exn && tail ->
+      Return_call_ref (ty, e, l)
+  | Push (Br_on_null (_, Call_ref _)) when return_exn -> i
   | Drop (BlockExpr (typ, l)) ->
-      Drop (BlockExpr (typ, instructions ~no_tail_call ~tail:false l))
+      Drop (BlockExpr (typ, instructions ~return_exn ~no_tail_call ~tail:false l))
   | Drop _
   | LocalSet _
   | GlobalSet _
@@ -67,28 +83,29 @@ let rec instruction ~no_tail_call ~tail i =
   | Unreachable
   | Event _ -> i
 
-and instructions ~no_tail_call ~tail l =
+and instructions ~return_exn ~no_tail_call ~tail l =
   match l with
   | [] -> []
-  | [ i ] -> [ instruction ~no_tail_call ~tail i ]
-  | i :: Nop :: rem -> instructions ~no_tail_call ~tail (i :: rem)
-  | i :: i' :: Nop :: rem -> instructions ~no_tail_call ~tail (i :: i' :: rem)
+  | [ i ] -> [ instruction ~return_exn ~no_tail_call ~tail i ]
+  | i :: Nop :: rem -> instructions ~return_exn ~no_tail_call ~tail (i :: rem)
+  | i :: i' :: Nop :: rem -> instructions ~return_exn ~no_tail_call ~tail (i :: i' :: rem)
   | i :: i' :: (([] | [ Event _ ]) as event_opt) -> (
       (* There can be an event at the end of the function, which we
          should keep. *)
       match get_return ~tail i' with
       | None ->
-          instruction ~no_tail_call ~tail:false i
-          :: instruction ~no_tail_call ~tail i'
+          instruction ~return_exn ~no_tail_call ~tail:false i
+          :: instruction ~return_exn ~no_tail_call ~tail i'
           :: event_opt
       | Some y -> (
-          match rewrite_tail_call ~no_tail_call ~y i with
+          match rewrite_tail_call ~return_exn ~no_tail_call ~y i with
           | None ->
-              instruction ~no_tail_call ~tail:false i
-              :: instruction ~no_tail_call ~tail i'
+              instruction ~return_exn ~no_tail_call ~tail:false i
+              :: instruction ~return_exn ~no_tail_call ~tail i'
               :: event_opt
           | Some i'' -> i'' :: event_opt))
   | i :: rem ->
-      instruction ~no_tail_call ~tail:false i :: instructions ~no_tail_call ~tail rem
+      instruction ~return_exn ~no_tail_call ~tail:false i
+      :: instructions ~return_exn ~no_tail_call ~tail rem
 
-let f ~no_tail_call l = instructions ~no_tail_call ~tail:true l
+let f ~return_exn ~no_tail_call l = instructions ~return_exn ~no_tail_call ~tail:true l

compiler/lib-wasm/tail_call.mli

@@ -17,6 +17,7 @@
  *)
 
 val f :
-     no_tail_call:unit Code.Var.Hashtbl.t
+     return_exn:bool
+  -> no_tail_call:unit Code.Var.Hashtbl.t
   -> Wasm_ast.instruction list
   -> Wasm_ast.instruction list