Skip to content

Commit

Permalink
Factorize mem/find_all/find_opt under one function.
Browse files Browse the repository at this point in the history
  • Loading branch information
@lyrm
lyrm committed Feb 20, 2024
1 parent ea9de3a commit 816bb66
Show file tree
Hide file tree
Showing 2 changed files with 56 additions and 95 deletions.
79 changes: 29 additions & 50 deletions src_lockfree/lf_htbl.ml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -241,7 +241,8 @@ let[@tail_mod_cons] rec remove_first_occ key removed = function
| x :: xs -> x :: remove_first_occ key removed xs
| [] -> []

let[@inline] rec try_remove t key hkey =
let[@inline] rec try_remove : ('k, 'v) t -> 'k -> int -> bool =
fun t key hkey ->
let mask = t.mask in
let bucket_index = hkey land mask in
let bucket = Array.get t.buckets bucket_index in
Expand Down Expand Up @@ -328,68 +329,46 @@ let try_remove t key =
let hkey = hash key in
try_remove t key hkey

let mem t key =
let hashed_key = hash key in
let mask = t.mask in
let bucket_index = hashed_key land mask in
let bucket = Array.get t.buckets bucket_index in
if bucket_index == hashed_key then (
let (Link ({ bindings; _ } as before)) = Atomic.get bucket in
type ('a, _) poly =
| Mem : ('a, bool) poly
| All : ('a, 'a list) poly
| Option : ('a, 'a option) poly

if before.incr != Size.used_once then begin
Size.update_once t.size before.incr;
before.incr <- Size.used_once
end;

List.mem_assoc key bindings)
else
let found, _, _, next =
find_node t.size mask bucket (Key.reverse hashed_key)
in
if found == 0 then List.mem_assoc key next.bindings else false

let find_all t key =
let[@inline] find_as : type a r. ('k, a) t -> 'k -> (a, r) poly -> r =
fun t key poly ->
let hashed_key = hash key in
let mask = t.mask in
let bucket_index = hashed_key land mask in
let bucket = Array.get t.buckets bucket_index in
if bucket_index == hashed_key then (
if bucket_index == hashed_key then begin
let (Link ({ bindings; _ } as before)) = Atomic.get bucket in

if before.incr != Size.used_once then begin
Size.update_once t.size before.incr;
before.incr <- Size.used_once
end;

List.fold_right
(fun (k, v) acc -> if k = key then v :: acc else acc)
bindings [])
match poly with
| All ->
List.fold_right
(fun (k, v) acc -> if k = key then v :: acc else acc)
bindings []
| Option -> List.assoc_opt key bindings
| Mem -> List.mem_assoc key bindings
end
else
let found, _, _, next =
find_node t.size mask bucket (Key.reverse hashed_key)
in
if found == 0 then
List.fold_right
(fun (k, v) acc -> if k = key then v :: acc else acc)
next.bindings []
else []

let find_opt t key =
let hashed_key = hash key in
let mask = t.mask in
let bucket_index = hashed_key land mask in
let bucket = Array.get t.buckets bucket_index in
if bucket_index == hashed_key then (
let (Link ({ bindings; _ } as before)) = Atomic.get bucket in

if before.incr != Size.used_once then begin
Size.update_once t.size before.incr;
before.incr <- Size.used_once
end;

List.assoc_opt key bindings)
else
let found, _, _, next =
find_node t.size mask bucket (Key.reverse hashed_key)
in
if found == 0 then List.assoc_opt key next.bindings else None
match poly with
| All ->
List.fold_right
(fun (k, v) acc -> if k = key then v :: acc else acc)
next.bindings []
| Option -> List.assoc_opt key next.bindings
| Mem -> List.mem_assoc key next.bindings
else match poly with All -> [] | Option -> None | Mem -> false

let find_all t key = find_as t key All
let find_opt t key = find_as t key Option
let mem t key = find_as t key Mem
72 changes: 27 additions & 45 deletions src_lockfree/lf_resizable_htbl.ml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -404,62 +404,44 @@ let try_remove t key =
let hkey = hash key in
try_remove t key hkey

let mem t key =
let hkey = hash key in
let mask = Atomic.get t.mask in
let buckets = Atomic.fenceless_get t.buckets in
let bucket_index, bucket = get_bucket t buckets mask hkey in
if bucket_index == hkey then (
let (Link ({ bindings; _ } as before)) = Atomic.get bucket in

if before.incr != Size.used_once then begin
Size.update_once t.size before.incr;
before.incr <- Size.used_once
end;

List.mem_assoc key bindings)
else
let found, _, _, next = find_node t.size mask bucket (Key.reverse hkey) in
if found == 0 then List.mem_assoc key next.bindings else false
type ('a, _) poly =
| Mem : ('a, bool) poly
| All : ('a, 'a list) poly
| Option : ('a, 'a option) poly

let find_all t key =
let[@inline] find_as : type a r. ('k, a) t -> 'k -> (a, r) poly -> r =
fun t key poly ->
let hkey = hash key in
let mask = Atomic.get t.mask in
let buckets = Atomic.fenceless_get t.buckets in
let bucket_index, bucket = get_bucket t buckets mask hkey in
if bucket_index == hkey then (
if bucket_index == hkey then begin
let (Link ({ bindings; _ } as before)) = Atomic.get bucket in

if before.incr != Size.used_once then begin
Size.update_once t.size before.incr;
before.incr <- Size.used_once
end;

List.fold_right
(fun (k, v) acc -> if k = key then v :: acc else acc)
bindings [])
match poly with
| All ->
List.fold_right
(fun (k, v) acc -> if k = key then v :: acc else acc)
bindings []
| Option -> List.assoc_opt key bindings
| Mem -> List.mem_assoc key bindings
end
else
let found, _, _, next = find_node t.size mask bucket (Key.reverse hkey) in
if found == 0 then
List.fold_right
(fun (k, v) acc -> if k = key then v :: acc else acc)
next.bindings []
else []

let find_opt t key =
let hkey = hash key in
let mask = Atomic.get t.mask in
let buckets = Atomic.fenceless_get t.buckets in
let bucket_index, bucket = get_bucket t buckets mask hkey in
if bucket_index == hkey then (
let (Link ({ bindings; _ } as before)) = Atomic.get bucket in

if before.incr != Size.used_once then begin
Size.update_once t.size before.incr;
before.incr <- Size.used_once
end;

List.assoc_opt key bindings)
else
let found, _, _, next = find_node t.size mask bucket (Key.reverse hkey) in
if found == 0 then List.assoc_opt key next.bindings else None
match poly with
| All ->
List.fold_right
(fun (k, v) acc -> if k = key then v :: acc else acc)
next.bindings []
| Option -> List.assoc_opt key next.bindings
| Mem -> List.mem_assoc key next.bindings
else match poly with All -> [] | Option -> None | Mem -> false

let find_all t key = find_as t key All
let find_opt t key = find_as t key Option
let mem t key = find_as t key Mem

0 comments on commit 816bb66

Please sign in to comment.