Function optimizations

CHANGES.md

@@ -2,6 +2,10 @@
 
 ## Features/Changes
 * Compiler: exit-loop-early in more cases (#2077)
+* Compiler/wasm: omit code pointer from closures when not used (#2059)
+
+## Bug fixes
+* Runtime/wasm: add back legacy function caml_string_of_array (#2081)
 
 # 6.1.1 (2025-07-07) - Lille
 

compiler/lib-wasm/call_graph_analysis.ml

@@ -0,0 +1,86 @@
+open! Stdlib
+open Code
+
+let debug = Debug.find "call-graph"
+
+let times = Debug.find "times"
+
+let block_deps ~info ~non_escaping ~ambiguous ~blocks pc =
+  let block = Addr.Map.find pc blocks in
+  List.iter block.body ~f:(fun i ->
+      match i with
+      | Let (_, Apply { f; _ }) -> (
+          try
+            match Var.Tbl.get info.Global_flow.info_approximation f with
+            | Top -> ()
+            | Values { known; others } ->
+                if others || Var.Set.cardinal known > 1
+                then Var.Set.iter (fun x -> Var.Hashtbl.replace ambiguous x ()) known;
+                if debug ()
+                then
+                  Format.eprintf
+                    "CALL others:%b known:%d@."
+                    others
+                    (Var.Set.cardinal known)
+          with Invalid_argument _ -> ())
+      | Let (x, Closure _) -> (
+          match Var.Tbl.get info.Global_flow.info_approximation x with
+          | Top -> ()
+          | Values { known; others } ->
+              if Var.Set.cardinal known = 1 && (not others) && Var.Set.mem x known
+              then (
+                let may_escape = Var.ISet.mem info.Global_flow.info_may_escape x in
+                if debug () then Format.eprintf "CLOSURE may-escape:%b@." may_escape;
+                if not may_escape then Var.Hashtbl.replace non_escaping x ()))
+      | Let (_, (Prim _ | Block _ | Constant _ | Field _ | Special _))
+      | Event _ | Assign _ | Set_field _ | Offset_ref _ | Array_set _ -> ())
+
+type t =
+  { unambiguous_non_escaping : unit Var.Hashtbl.t
+  ; has_tail_calls : unit Var.Hashtbl.t
+  }
+
+let direct_calls_only info f =
+  Config.Flag.optcall () && Var.Hashtbl.mem info.unambiguous_non_escaping f
+
+let has_tail_calls info f = Var.Hashtbl.mem info.has_tail_calls f
+
+let f p info =
+  let t = Timer.make () in
+  let non_escaping = Var.Hashtbl.create 128 in
+  let ambiguous = Var.Hashtbl.create 128 in
+  let has_tail_calls = Var.Hashtbl.create 128 in
+  fold_closures
+    p
+    (fun name_opt _ (pc, _) _ () ->
+      traverse
+        { fold = Code.fold_children }
+        (fun pc () -> block_deps ~info ~non_escaping ~ambiguous ~blocks:p.blocks pc)
+        pc
+        p.blocks
+        ();
+      Option.iter
+        ~f:(fun f ->
+          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 _)) when Code.Var.equal x x' ->
+                      Var.Hashtbl.replace has_tail_calls f ()
+                  | _ -> ())
+              | _ -> ())
+            pc
+            p.blocks
+            ())
+        name_opt)
+    ();
+  if debug ()
+  then Format.eprintf "SUMMARY non-escaping:%d" (Var.Hashtbl.length non_escaping);
+  Var.Hashtbl.iter (fun x () -> Var.Hashtbl.remove non_escaping x) ambiguous;
+  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; has_tail_calls }

compiler/lib-wasm/call_graph_analysis.mli

@@ -0,0 +1,7 @@
+type t
+
+val direct_calls_only : t -> Code.Var.t -> bool
+
+val has_tail_calls : t -> Code.Var.t -> bool
+
+val f : Code.program -> Global_flow.info -> t

compiler/lib-wasm/code_generation.ml

@@ -59,6 +59,7 @@ type context =
   ; mutable globalized_variables : Var.Set.t
   ; value_type : W.value_type
   ; mutable unit_name : string option
+  ; mutable no_tail_call : unit Var.Hashtbl.t
   }
 
 let make_context ~value_type =
@@ -82,6 +83,7 @@ let make_context ~value_type =
   ; globalized_variables = Var.Set.empty
   ; value_type
   ; unit_name = None
+  ; no_tail_call = Var.Hashtbl.create 16
   }
 
 type var =
@@ -224,27 +226,30 @@ let get_global name =
     | Some { init; _ } -> init
     | _ -> None)
 
-let register_import ?(import_module = "env") ~name typ st =
-  ( (try
-       let x, typ' =
-         StringMap.find name (StringMap.find import_module st.context.imports)
-       in
-       (*ZZZ error message*)
-       assert (Poly.equal typ typ');
-       x
-     with Not_found ->
-       let x = Var.fresh_n name in
-       st.context.imports <-
-         StringMap.update
-           import_module
-           (fun m ->
-             Some
-               (match m with
-               | None -> StringMap.singleton name (x, typ)
-               | Some m -> StringMap.add name (x, typ) m))
-           st.context.imports;
-       x)
-  , st )
+let register_import ?(allow_tail_call = true) ?(import_module = "env") ~name typ st =
+  let x =
+    try
+      let x, typ' =
+        StringMap.find name (StringMap.find import_module st.context.imports)
+      in
+      (*ZZZ error message*)
+      assert (Poly.equal typ typ');
+      x
+    with Not_found ->
+      let x = Var.fresh_n name in
+      st.context.imports <-
+        StringMap.update
+          import_module
+          (fun m ->
+            Some
+              (match m with
+              | None -> StringMap.singleton name (x, typ)
+              | Some m -> StringMap.add name (x, typ) m))
+          st.context.imports;
+      x
+  in
+  if not allow_tail_call then Var.Hashtbl.replace st.context.no_tail_call x ();
+  x, st
 
 let register_init_code code st =
   let st' = { var_count = 0; vars = Var.Map.empty; instrs = []; context = st.context } in
@@ -368,6 +373,7 @@ module Arith = struct
       (match e, e' with
       | W.Const (I32 n), W.Const (I32 n') when Int32.(n' < 31l) ->
           W.Const (I32 (Int32.shift_left n (Int32.to_int n')))
+      | _, W.Const (I32 0l) -> e
       | _ -> W.BinOp (I32 Shl, e, e'))
 
   let ( lsr ) = binary (Shr U)
@@ -715,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 body in
+  let body = Tail_call.f ~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

@@ -42,6 +42,7 @@ type context =
   ; mutable globalized_variables : Code.Var.Set.t
   ; value_type : Wasm_ast.value_type
   ; mutable unit_name : string option
+  ; mutable no_tail_call : unit Code.Var.Hashtbl.t
   }
 
 val make_context : value_type:Wasm_ast.value_type -> context
@@ -156,7 +157,11 @@ val register_type : string -> (unit -> type_def t) -> Wasm_ast.var t
 val heap_type_sub : Wasm_ast.heap_type -> Wasm_ast.heap_type -> bool t
 
 val register_import :
-  ?import_module:string -> name:string -> Wasm_ast.import_desc -> Wasm_ast.var t
+     ?allow_tail_call:bool
+  -> ?import_module:string
+  -> name:string
+  -> Wasm_ast.import_desc
+  -> Wasm_ast.var t
 
 val register_global :
      Wasm_ast.var

compiler/lib-wasm/curry.ml

@@ -293,14 +293,7 @@ module Make (Target : Target_sig.S) = struct
         (fun ~typ closure ->
           let* l = expression_list load l in
           call ?typ ~cps:true ~arity closure l)
-        (let* args =
-           (* We don't need the deadcode sentinal when the tag is 0 *)
-           Memory.allocate
-             ~tag:0
-             ~deadcode_sentinal:(Code.Var.fresh ())
-             ~load
-             (List.map ~f:(fun x -> `Var x) (List.tl l))
-         in
+        (let* args = Memory.allocate ~tag:0 (expression_list load (List.tl l)) in
          let* make_iterator =
            register_import ~name:"caml_apply_continuation" (Fun (Type.primitive_type 1))
          in