Merge pull request #12131 from NickBarnes/nick-get-copy5.1

Simplify weak hash sets

(cherry picked from commit 088fc769eab02fcc016c296bbe6ea3337d8f680e)

2 files changed, 23 insertions, 61 deletions

diff --git a/Changes b/Changes
index 9e9c36aa4d..adf554dcda 100644
--- a/Changes
+++ b/Changes
@@ -132,6 +132,10 @@ OCaml 5.1.0
   (B. Szilvasy, Gabriel Scherer and Xavier Leroy, review by
    Stefan Muenzel, Guillaume Munch-Maccagnoni and Damien Doligez)
 
+- #12131: Simplify implementation of weak hash sets, fixing a
+  performance regression. (Nick Barnes, review by François Bobot,
+  Alain Frisch and Damien Doligez).
+
 ### Type system:
 
 - #6941, #11187: prohibit using classes through recursive modules
diff --git a/stdlib/weak.ml b/stdlib/weak.ml
index 354015bebd..e1e0a20966 100644
--- a/stdlib/weak.ml
+++ b/stdlib/weak.ml
@@ -271,77 +271,40 @@ module Make (H : Hashtbl.HashedType) : (S with type data = H.t) = struct
     let h = H.hash d in
     add_aux t set (Some d) h (get_index t h)
 
+  (* General auxiliary function for searching for a particular value
+   * in a hash-set, and acting according to whether or not it's found *)
 
-  let find_or t d ifnotfound =
+  let find_aux t d found notfound =
     let h = H.hash d in
     let index = get_index t h in
     let bucket = t.table.(index) in
     let hashes = t.hashes.(index) in
     let sz = length bucket in
     let rec loop i =
-      if i >= sz then ifnotfound h index
+      if i >= sz then notfound h index
       else if h = hashes.(i) then begin
-        match get_copy bucket i with
-        | Some v when H.equal v d
-           -> begin match get bucket i with
-              | Some v -> v
-              | None -> loop (i + 1)
-              end
+        match get bucket i with
+        | Some v as opt when H.equal v d -> found bucket i opt v
         | _ -> loop (i + 1)
       end else loop (i + 1)
     in
     loop 0
 
+  let find_opt t d = find_aux t d (fun _b _i  o _v -> o)
+                                  (fun _h _i -> None)
 
-  let merge t d =
-    find_or t d (fun h index -> add_aux t set (Some d) h index; d)
+  let merge t d    = find_aux t d (fun _b _i _o  v -> v)
+                                  (fun  h  i ->
+                                        add_aux t set (Some d) h i; d)
 
+  let find t d     = find_aux t d (fun _b _i _o  v -> v)
+                                  (fun _h _i -> raise Not_found)
 
-  let find t d = find_or t d (fun _h _index -> raise Not_found)
-
-  let find_opt t d =
-    let h = H.hash d in
-    let index = get_index t h in
-    let bucket = t.table.(index) in
-    let hashes = t.hashes.(index) in
-    let sz = length bucket in
-    let rec loop i =
-      if i >= sz then None
-      else if h = hashes.(i) then begin
-        match get_copy bucket i with
-        | Some v when H.equal v d
-           -> begin match get bucket i with
-              | Some _ as v -> v
-              | None -> loop (i + 1)
-              end
-        | _ -> loop (i + 1)
-      end else loop (i + 1)
-    in
-    loop 0
-
-
-  let find_shadow t d iffound ifnotfound =
-    let h = H.hash d in
-    let index = get_index t h in
-    let bucket = t.table.(index) in
-    let hashes = t.hashes.(index) in
-    let sz = length bucket in
-    let rec loop i =
-      if i >= sz then ifnotfound
-      else if h = hashes.(i) then begin
-        match get_copy bucket i with
-        | Some v when H.equal v d -> iffound bucket i
-        | _ -> loop (i + 1)
-      end else loop (i + 1)
-    in
-    loop 0
-
-
-  let remove t d = find_shadow t d (fun w i -> set w i None) ()
-
-
-  let mem t d = find_shadow t d (fun _w _i -> true) false
+  let remove t d   = find_aux t d (fun  b  i _o _v -> set b i None)
+                                  (fun _h _i -> ())
 
+  let mem t d      = find_aux t d (fun _b _i _o _v -> true)
+                                  (fun _h _i -> false)
 
   let find_all t d =
     let h = H.hash d in
@@ -352,18 +315,13 @@ module Make (H : Hashtbl.HashedType) : (S with type data = H.t) = struct
     let rec loop i accu =
       if i >= sz then accu
       else if h = hashes.(i) then begin
-        match get_copy bucket i with
-        | Some v when H.equal v d
-           -> begin match get bucket i with
-              | Some v -> loop (i + 1) (v :: accu)
-              | None -> loop (i + 1) accu
-              end
+        match get bucket i with
+        | Some v when H.equal v d -> loop (i + 1) (v :: accu)
         | _ -> loop (i + 1) accu
       end else loop (i + 1) accu
     in
     loop 0 []
 
-
   let stats t =
     let len = Array.length t.table in
     let lens = Array.map length t.table in