From 13128c5e192d8cb78a1ea5e7e1f8fec705516683 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=B4me=20Vouillon?= <jerome.vouillon@gmail.com>
Date: Mon, 23 Sep 2024 11:23:37 +0200
Subject: [PATCH] Fix performance issue when assigning short names to variables

---
 compiler/lib/wasm/wa_wasm_output.ml | 19 ++++++++++++++-----
 compiler/lib/wasm/wa_wat_output.ml  | 17 +++++++++++++----
 2 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/compiler/lib/wasm/wa_wasm_output.ml b/compiler/lib/wasm/wa_wasm_output.ml
index f90543581..114ce0edc 100644
--- a/compiler/lib/wasm/wa_wasm_output.ml
+++ b/compiler/lib/wasm/wa_wasm_output.ml
@@ -1043,14 +1043,23 @@ end = struct
     output_byte ch id;
     with_size f ch x
 
-  let rec find_available_name used name i =
-    let nm = Printf.sprintf "%s$%d" name i in
-    if StringSet.mem nm used then find_available_name used name (i + 1) else nm
-
   let assign_names f tbl =
     let names = Hashtbl.fold (fun name idx rem -> (idx, name) :: rem) tbl [] in
     let names = List.sort ~cmp:(fun (idx, _) (idx', _) -> compare idx idx') names in
     let used = ref StringSet.empty in
+    let counts = Hashtbl.create 101 in
+    let rec find_available_name used name =
+      let i =
+        try Hashtbl.find counts name
+        with Not_found ->
+          let i = ref 0 in
+          Hashtbl.replace counts name i;
+          i
+      in
+      incr i;
+      let nm = Printf.sprintf "%s$%d" name !i in
+      if StringSet.mem nm used then find_available_name used name else nm
+    in
     let names =
       List.map
         ~f:(fun (idx, x) ->
@@ -1058,7 +1067,7 @@ end = struct
           | None -> idx, None
           | Some nm ->
               let nm =
-                if StringSet.mem nm !used then find_available_name !used nm 1 else nm
+                if StringSet.mem nm !used then find_available_name !used nm else nm
               in
               used := StringSet.add nm !used;
               idx, Some nm)
diff --git a/compiler/lib/wasm/wa_wat_output.ml b/compiler/lib/wasm/wa_wat_output.ml
index 3e72493be..a01874b4e 100644
--- a/compiler/lib/wasm/wa_wat_output.ml
+++ b/compiler/lib/wasm/wa_wat_output.ml
@@ -23,9 +23,18 @@ let target = `Binaryen (*`Reference*)
 
 let assign_names ?(reversed = true) f names =
   let used = ref StringSet.empty in
-  let rec find_available_name used name i =
-    let nm = Printf.sprintf "%s$%d" name i in
-    if StringSet.mem nm used then find_available_name used name (i + 1) else nm
+  let counts = Hashtbl.create 101 in
+  let rec find_available_name used name =
+    let i =
+      try Hashtbl.find counts name
+      with Not_found ->
+        let i = ref 0 in
+        Hashtbl.replace counts name i;
+        i
+    in
+    incr i;
+    let nm = Printf.sprintf "%s$%d" name !i in
+    if StringSet.mem nm used then find_available_name used name else nm
   in
   let names = if reversed then List.rev names else names in
   let names =
@@ -35,7 +44,7 @@ let assign_names ?(reversed = true) f names =
         | None -> x, None
         | Some nm ->
             let nm =
-              if StringSet.mem nm !used then find_available_name !used nm 1 else nm
+              if StringSet.mem nm !used then find_available_name !used nm else nm
             in
             used := StringSet.add nm !used;
             x, Some nm)