add sibling proofs to relevant blocks

bluesky-social · Feb 7, 2025 · c3eb038 · c3eb038
1 parent 4f1d84d
commit c3eb038
Show file tree

Hide file tree

Showing 3 changed files with 83 additions and 13 deletions.
diff --git a/packages/repo/src/block-map.ts b/packages/repo/src/block-map.ts
@@ -66,6 +66,7 @@ export class BlockMap {
     toAdd.forEach((bytes, cid) => {
       this.set(cid, bytes)
     })
+    return this
   }
 
   get size(): number {

diff --git a/packages/repo/src/mst/mst.ts b/packages/repo/src/mst/mst.ts
@@ -755,20 +755,90 @@ export class MST {
     return cids
   }
 
-  async addBlocksForPath(key: string, blocks: BlockMap) {
+  async addSelf(blocks: BlockMap) {
     const serialized = await this.serialize()
     blocks.set(serialized.cid, serialized.bytes)
+    return blocks
+  }
+
+  async proofForLeftSib(key: string): Promise<BlockMap> {
     const index = await this.findGtOrEqualLeafIndex(key)
-    const found = await this.atIndex(index)
-    if (found && found.isLeaf() && found.key === key) {
-      return
-    }
     const prev = await this.atIndex(index - 1)
-    if (prev && prev.isTree()) {
-      await prev.addBlocksForPath(key, blocks)
+    let blocks: BlockMap
+    if (!prev || prev.isLeaf()) {
+      blocks = new BlockMap()
+    } else {
+      blocks = await prev.proofForLeftSib(key)
+    }
+    return this.addSelf(blocks)
+  }
+
+  async proofForRightSib(key: string): Promise<BlockMap> {
+    const index = await this.findGtOrEqualLeafIndex(key)
+    const found = await this.atIndex(index)
+    let blocks: BlockMap
+    if (!found) {
+      // shouldn't ever hit, null case
+      blocks = new BlockMap()
+    } else if (found.isTree()) {
+      blocks = await found.proofForRightSib(key)
+      // recurse down
+    } else {
+      const node =
+        found.key === key
+          ? await this.atIndex(index + 1)
+          : await this.atIndex(index - 1)
+      if (!node || node.isLeaf()) {
+        blocks = new BlockMap()
+      } else {
+        blocks = await node.proofForRightSib(key)
+      }
     }
+    return this.addSelf(blocks)
   }
 
+  async proofForKey(key: string): Promise<BlockMap> {
+    const index = await this.findGtOrEqualLeafIndex(key)
+    const found = await this.atIndex(index)
+    let blocks: BlockMap
+    if (found && found.isLeaf() && found.key === key) {
+      blocks = new BlockMap()
+    } else {
+      const prev = await this.atIndex(index - 1)
+      if (!prev || prev.isLeaf()) {
+        return new BlockMap()
+      } else {
+        blocks = await prev.proofForKey(key)
+      }
+    }
+    return this.addSelf(blocks)
+  }
+
+  async getCoveringProof(key: string): Promise<BlockMap> {
+    const [self, left, right] = await Promise.all([
+      this.proofForKey(key),
+      this.proofForLeftSib(key),
+      this.proofForRightSib(key),
+    ])
+    return self.addMap(left).addMap(right)
+  }
+
+  // @TODO delete
+
+  // async addBlocksForPath(key: string, blocks: BlockMap) {
+  //   const serialized = await this.serialize()
+  //   blocks.set(serialized.cid, serialized.bytes)
+  //   const index = await this.findGtOrEqualLeafIndex(key)
+  //   const found = await this.atIndex(index)
+  //   if (found && found.isLeaf() && found.key === key) {
+  //     return
+  //   }
+  //   const prev = await this.atIndex(index - 1)
+  //   if (prev && prev.isTree()) {
+  //     await prev.addBlocksForPath(key, blocks)
+  //   }
+  // }
+
   // Matching Leaf interface
   // -------------------
 

diff --git a/packages/repo/src/repo.ts b/packages/repo/src/repo.ts
@@ -142,15 +142,14 @@ export class Repo extends ReadableRepo {
     const newBlocks = diff.newMstBlocks
     const removedCids = diff.removedCids
 
-    const relevantBlocks = new BlockMap()
-    await Promise.all(
+    const proofs = await Promise.all(
       writes.map((op) =>
-        data.addBlocksForPath(
-          util.formatDataKey(op.collection, op.rkey),
-          relevantBlocks,
-        ),
+        data.getCoveringProof(util.formatDataKey(op.collection, op.rkey)),
       ),
     )
+    const relevantBlocks = proofs.reduce((acc, cur) => {
+      return acc.addMap(cur)
+    }, new BlockMap())
 
     const addedLeaves = leaves.getMany(diff.newLeafCids.toList())
     if (addedLeaves.missing.length > 0) {