Skip to content

Commit

Permalink
Optimize MemoryStore abstract functions.
Browse files Browse the repository at this point in the history
Minimize the number of Deferred.return calls.
  • Loading branch information
zoj613 committed Aug 29, 2024
1 parent 62466f4 commit 174200a
Showing 1 changed file with 24 additions and 31 deletions.
55 changes: 24 additions & 31 deletions zarr/src/storage/memory.ml
Original file line number Diff line number Diff line change
Expand Up @@ -49,45 +49,38 @@ module Make (Deferred : Types.Deferred) = struct
let get_partial_values t key ranges =
get t key >>= fun v ->
let size = String.length v in
Deferred.fold_left
(fun acc (rs, len) ->
Deferred.return
(match len with
| None -> String.sub v rs (size - rs) :: acc
| Some l -> String.sub v rs l :: acc)) [] @@ List.rev ranges
Deferred.return @@
List.fold_left
(fun acc (ofs, len) ->
match len with
| None -> String.sub v ofs (size - ofs) :: acc
| Some l -> String.sub v ofs l :: acc) [] @@ List.rev ranges

let rec set_partial_values t key ?(append=false) rv =
let f =
if append then
fun acc (_, v) ->
Deferred.return @@ acc ^ v
else
fun acc (rs, v) ->
let s = Bytes.of_string acc in
String.(length v |> Bytes.blit_string v 0 s rs);
Deferred.return @@ Bytes.to_string s
in
let f = if append then fun acc (_, v) -> acc ^ v else
fun acc (rs, v) ->
let s = Bytes.of_string acc in
String.(length v |> Bytes.blit_string v 0 s rs);
Bytes.to_string s in
let m = Atomic.get t in
let ov = StrMap.find key m in
Deferred.fold_left f ov rv >>= fun nv ->
let m' = StrMap.add key nv m in
let m' = StrMap.add key (List.fold_left f ov rv) m in
let success = Atomic.compare_and_set t m m' in
if not success then set_partial_values t key ~append rv
else Deferred.return_unit

let list_dir t prefix =
let module StrSet = Util.StrSet in
list t >>= fun xs ->
let module S = Util.StrSet in
list t >>| fun xs ->
let n = String.length prefix in
Deferred.fold_left
(fun ((l, r) as a) key ->
let pred = String.starts_with ~prefix key in
match key with
| k when pred && String.contains_from k n '/' ->
let l' = StrSet.add String.(sub k 0 @@ 1 + index_from k n '/') l in
Deferred.return (l', r)
| k when pred -> Deferred.return (l, k :: r)
| _ -> Deferred.return a) (StrSet.empty, []) xs
>>| fun (prefs, keys) ->
keys, StrSet.elements prefs
let prefs, keys =
List.fold_left
(fun ((l, r) as a) key ->
let pred = String.starts_with ~prefix key in
match key with
| k when pred && String.contains_from k n '/' ->
let l' = S.add String.(sub k 0 @@ 1 + index_from k n '/') l in l', r
| k when pred -> l, k :: r
| _ -> a) (S.empty, []) xs
in keys, S.elements prefs
end

0 comments on commit 174200a

Please sign in to comment.