From 174200aa352b0aaf46af96dd3909337230109095 Mon Sep 17 00:00:00 2001
From: Zolisa Bleki <zolisa.bleki@gmail.com>
Date: Thu, 29 Aug 2024 05:07:10 +0200
Subject: [PATCH] Optimize MemoryStore abstract functions.

Minimize the number of Deferred.return calls.
---
 zarr/src/storage/memory.ml | 55 +++++++++++++++++---------------------
 1 file changed, 24 insertions(+), 31 deletions(-)

diff --git a/zarr/src/storage/memory.ml b/zarr/src/storage/memory.ml
index adeabaf..4952175 100644
--- a/zarr/src/storage/memory.ml
+++ b/zarr/src/storage/memory.ml
@@ -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