Optimize exceptions

compiler/lib-wasm/call_graph_analysis.ml

@@ -42,6 +42,156 @@ type t = { unambiguous_non_escaping : unit Var.Hashtbl.t }
 let direct_calls_only info f =
   Config.Flag.optcall () && Var.Hashtbl.mem info.unambiguous_non_escaping f
 
+let callee_if_known info call_info exact f =
+  match get_approx info f with
+  | Top -> None
+  | Values { known; others } ->
+      if
+        exact
+        && (not others)
+        && Var.Set.for_all (fun f -> direct_calls_only call_info f) known
+      then Some (Var.Set.choose known)
+      else None
+
+let propagate nodes edges eligible =
+  let rec propagate n =
+    List.iter
+      ~f:(fun n' ->
+        if (not (Var.Hashtbl.mem nodes n')) && eligible n'
+        then (
+          Var.Hashtbl.add nodes n' ();
+          propagate n'))
+      (Var.Hashtbl.find_all edges n)
+  in
+  Var.Hashtbl.iter (fun n () -> propagate n) nodes
+
+let call_graph p info call_info eligible =
+  let under_handler = Var.Hashtbl.create 16 in
+  let callees = Var.Hashtbl.create 16 in
+  let callers = Var.Hashtbl.create 16 in
+  let has_tail_calls = Var.Hashtbl.create 16 in
+  let tail_callers = Var.Hashtbl.create 16 in
+  let rec traverse name_opt pc visited nesting =
+    if not (Addr.Set.mem pc visited)
+    then (
+      let visited = Addr.Set.add pc visited in
+      let block = Addr.Map.find pc p.blocks in
+      List.iter block.body ~f:(fun i ->
+          match i with
+          | Let (_, Apply { f; exact; _ }) -> (
+              match get_approx info f with
+              | Top -> ()
+              | Values { known; others } ->
+                  if
+                    exact
+                    && (not others)
+                    && Var.Set.for_all (fun f -> direct_calls_only call_info f) known
+                  then
+                    if nesting > 0
+                    then
+                      Var.Set.iter
+                        (fun f ->
+                          (*                        Format.eprintf "BBB %a@." Code.Var.print f; *)
+                          Var.Hashtbl.replace under_handler f ())
+                        known
+                    else
+                      Option.iter
+                        ~f:(fun f ->
+                          Var.Set.iter
+                            (fun g ->
+                              Var.Hashtbl.add callees f g;
+                              Var.Hashtbl.add callers g f)
+                            known)
+                        name_opt)
+          | Let (_, (Closure _ | Prim _ | Block _ | Constant _ | Field _ | Special _))
+          | Event _ | Assign _ | Set_field _ | Offset_ref _ | Array_set _ -> ());
+      if nesting = 0
+      then
+        Option.iter
+          ~f:(fun f ->
+            Code.traverse
+              { fold = Code.fold_children }
+              (fun pc () ->
+                let block = Addr.Map.find pc p.blocks in
+                match block.branch with
+                | Return x -> (
+                    match last_instr block.body with
+                    | Some (Let (x', Apply { f = g; exact; _ })) when Code.Var.equal x x'
+                      -> (
+                        match callee_if_known info call_info exact g with
+                        | None -> Var.Hashtbl.replace has_tail_calls f ()
+                        | Some g -> Var.Hashtbl.add tail_callers g f)
+                    | _ -> ())
+                | _ -> ())
+              pc
+              p.blocks
+              ())
+          name_opt;
+      Code.fold_children
+        p.blocks
+        pc
+        (fun pc' visited ->
+          let nesting =
+            match block.branch with
+            | Pushtrap ((body_pc, _), _, _) when pc' = body_pc -> nesting + 1
+            | Poptrap _ -> nesting - 1
+            | _ -> nesting
+          in
+          traverse name_opt pc' visited nesting)
+        visited)
+    else visited
+  in
+  fold_closures
+    p
+    (fun name_opt _ (pc, _) _ () -> ignore (traverse name_opt pc Addr.Set.empty 0))
+    ();
+  propagate has_tail_calls tail_callers eligible;
+  under_handler, callers, callees, has_tail_calls
+
+let function_do_raise p pc =
+  Code.traverse
+    { fold = Code.fold_children_skip_try_body }
+    (fun pc do_raise ->
+      let block = Addr.Map.find pc p.blocks in
+      do_raise
+      ||
+      match block.branch with
+      | Raise _ -> true
+      | _ -> false)
+    pc
+    p.blocks
+    false
+
+let raising_functions p info call_info eligible =
+  let under_handler, callers, callees, has_tail_calls =
+    call_graph p info call_info eligible
+  in
+  propagate under_handler callees (fun f ->
+      eligible f && not (Var.Hashtbl.mem has_tail_calls f));
+  let h = Var.Hashtbl.create 16 in
+  let eligible f =
+    eligible f
+    && Var.Hashtbl.mem under_handler f
+    && not (Var.Hashtbl.mem has_tail_calls f)
+  in
+  Code.fold_closures
+    p
+    (fun name_opt _params (pc, _) _ () ->
+      match name_opt with
+      | None -> ()
+      | Some name ->
+          if direct_calls_only call_info name && eligible name && function_do_raise p pc
+          then Var.Hashtbl.add h name ())
+    ();
+  propagate h callers eligible;
+  if false
+  then
+    Var.Hashtbl.iter
+      (fun name () ->
+        Format.eprintf "ZZZ %a %b@." Var.print name (Var.Hashtbl.mem under_handler name))
+      h;
+  h
+
 let f p info =
   let t = Timer.make () in
   let non_escaping = Var.Hashtbl.create 128 in
@@ -62,4 +212,12 @@ let f p info =
   if debug ()
   then Format.eprintf " unambiguous-non-escaping:%d@." (Var.Hashtbl.length non_escaping);
   if times () then Format.eprintf "  call graph analysis: %a@." Timer.print t;
-  { unambiguous_non_escaping = non_escaping }
+  (*
+  Var.Hashtbl.iter (fun f _ -> Format.eprintf "AAA %a@." Code.Var.print f) non_escaping;
+*)
+  let call_info = { unambiguous_non_escaping = non_escaping } in
+  call_info
+
+(*
+- Optimize tail-calls
+*)

compiler/lib-wasm/call_graph_analysis.mli

@@ -2,4 +2,7 @@ type t
 
 val direct_calls_only : t -> Code.Var.t -> bool
 
+val raising_functions :
+  Code.program -> Global_flow.info -> t -> (Code.Var.t -> bool) -> unit Code.Var.Hashtbl.t
+
 val f : Code.program -> Global_flow.info -> t

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/gc_target.ml

@@ -25,6 +25,8 @@ type expression = Wasm_ast.expression Code_generation.t
 module Type = struct
   let value = W.Ref { nullable = false; typ = Eq }
 
+  let value_or_exn = W.Ref { nullable = true; typ = Eq }
+
   let block_type =
     register_type "block" (fun () ->
         return