From 444ccd79e8d99fcc06b7528344aab683dcc3341c Mon Sep 17 00:00:00 2001
From: Alan Shaw <alan.shaw@protocol.ai>
Date: Tue, 19 Mar 2024 09:39:47 +0000
Subject: [PATCH] feat: filter entries with keys greater than this string (#37)

---
 src/index.js         | 37 ++++++++++++++++++++----
 test/entries.test.js | 68 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 99 insertions(+), 6 deletions(-)

diff --git a/src/index.js b/src/index.js
index aaa08d51..0c199f48 100644
--- a/src/index.js
+++ b/src/index.js
@@ -222,11 +222,13 @@ export const del = async (blocks, root, key) => {
  * @param {API.BlockFetcher} blocks Bucket block storage.
  * @param {API.ShardLink} root CID of the root node of the bucket.
  * @param {object} [options]
- * @param {string} [options.prefix]
+ * @param {string} [options.prefix] Filter results to entries with keys prefixed with this string.
+ * @param {string} [options.gt] Filter results to entries with keys greater than this string.
+ * @param {string} [options.gte] Filter results to entries with keys greater than or equal to this string.
  * @returns {AsyncIterableIterator<API.ShardValueEntry>}
  */
 export const entries = async function * (blocks, root, options = {}) {
-  const { prefix } = options
+  const { prefix, gt, gte } = options
   const shards = new ShardFetcher(blocks)
   const rshard = await shards.get(root)
 
@@ -238,7 +240,12 @@ export const entries = async function * (blocks, root, options = {}) {
 
         if (Array.isArray(entry[1])) {
           if (entry[1][1]) {
-            if (!prefix || (prefix && key.startsWith(prefix))) {
+            if (
+              (prefix && key.startsWith(prefix)) ||
+              (gt && key > gt) ||
+              (gte && key >= gte) ||
+              (!prefix && !gt && !gte)
+            ) {
               yield [key, entry[1][1]]
             }
           }
@@ -250,13 +257,31 @@ export const entries = async function * (blocks, root, options = {}) {
             if (prefix.length > key.length && !prefix.startsWith(key)) {
               continue
             }
+          } else if (gt) {
+            if (gt.length <= key.length && !key.startsWith(gt)) {
+              continue
+            }
+            if (gt.length > key.length && !gt.startsWith(key)) {
+              continue
+            }
+          } else if (gte) {
+            if (gte.length <= key.length && !key.startsWith(gte)) {
+              continue
+            }
+            if (gte.length > key.length && !gte.startsWith(key)) {
+              continue
+            }
           }
           yield * ents(await shards.get(entry[1][0], key))
         } else {
-          if (prefix && !key.startsWith(prefix)) {
-            continue
+          if (
+            (prefix && key.startsWith(prefix)) ||
+            (gt && key > gt) ||
+            (gte && key >= gte) ||
+            (!prefix && !gt && !gte)
+          ) {
+            yield [key, entry[1]]
           }
-          yield [key, entry[1]]
         }
       }
     }
diff --git a/test/entries.test.js b/test/entries.test.js
index ca167b83..60c446c6 100644
--- a/test/entries.test.js
+++ b/test/entries.test.js
@@ -73,4 +73,72 @@ describe('entries', () => {
       assert.equal(results[i][0], key)
     }
   })
+
+  it('lists entries by key greater than string', async () => {
+    const empty = await ShardBlock.create()
+    const blocks = new Blockstore()
+    await blocks.put(empty.cid, empty.bytes)
+
+    /** @type {Array<[string, API.UnknownLink]>} */
+    const testdata = [
+      ['cccc', await randomCID(32)],
+      ['deee', await randomCID(32)],
+      ['dooo', await randomCID(32)],
+      ['beee', await randomCID(32)]
+    ]
+
+    /** @type {API.ShardLink} */
+    let root = empty.cid
+    for (const [k, v] of testdata) {
+      const res = await put(blocks, root, k, v)
+      for (const b of res.additions) {
+        await blocks.put(b.cid, b.bytes)
+      }
+      root = res.root
+    }
+
+    const gt = 'beee'
+    const results = []
+    for await (const entry of entries(blocks, root, { gt })) {
+      results.push(entry)
+    }
+
+    for (const [i, key] of testdata.map(d => d[0]).filter(k => k > gt).sort().entries()) {
+      assert.equal(results[i][0], key)
+    }
+  })
+
+  it('lists entries by key greater than or equal to string', async () => {
+    const empty = await ShardBlock.create()
+    const blocks = new Blockstore()
+    await blocks.put(empty.cid, empty.bytes)
+
+    /** @type {Array<[string, API.UnknownLink]>} */
+    const testdata = [
+      ['cccc', await randomCID(32)],
+      ['deee', await randomCID(32)],
+      ['dooo', await randomCID(32)],
+      ['beee', await randomCID(32)]
+    ]
+
+    /** @type {API.ShardLink} */
+    let root = empty.cid
+    for (const [k, v] of testdata) {
+      const res = await put(blocks, root, k, v)
+      for (const b of res.additions) {
+        await blocks.put(b.cid, b.bytes)
+      }
+      root = res.root
+    }
+
+    const gte = 'beee'
+    const results = []
+    for await (const entry of entries(blocks, root, { gte })) {
+      results.push(entry)
+    }
+
+    for (const [i, key] of testdata.map(d => d[0]).filter(k => k >= gte).sort().entries()) {
+      assert.equal(results[i][0], key)
+    }
+  })
 })