@atproto/repo 0.0.1 → 0.2.0

package/src/mst/mst.ts

@@ -1,13 +1,14 @@
 import z from 'zod'
 import { CID } from 'multiformats'
 
-import IpldStore from '../blockstore/ipld-store'
-import { def, cidForData } from '@atproto/common'
-import { DataDiff } from './diff'
-import { DataStore } from '../types'
+import { ReadableBlockstore } from '../storage'
+import { schema as common, cidForCbor } from '@atproto/common'
 import { BlockWriter } from '@ipld/car/api'
 import * as util from './util'
-import MstWalker from './walker'
+import BlockMap from '../block-map'
+import CidSet from '../cid-set'
+import { MissingBlockError, MissingBlocksError } from '../error'
+import * as parse from '../parse'
 
 /**
  * This is an implementation of a Merkle Search Tree (MST)
@@ -19,6 +20,9 @@ import MstWalker from './walker'
  * This is a merkle tree, so each subtree is referred to by it's hash (CID).
  * When a leaf is changed, ever tree on the path to that leaf is changed as well,
  * thereby updating the root hash.
+ *
+ * For atproto, we use SHA-256 as the key hashing algorithm, and ~4 fanout
+ * (2-bits of zero per layer).
  */
 
 /**
@@ -39,74 +43,78 @@ import MstWalker from './walker'
  * Then the first will be described as `prefix: 0, key: 'bsky/posts/abcdefg'`,
  * and the second will be described as `prefix: 16, key: 'hi'.`
  */
-const subTreePointer = z.nullable(def.cid)
+const subTreePointer = z.nullable(common.cid)
 const treeEntry = z.object({
-  p: z.number(), // prefix count of utf-8 chars that this key shares with the prev key
-  k: z.string(), // the rest of the key outside the shared prefix
-  v: def.cid, // value
+  p: z.number(), // prefix count of ascii chars that this key shares with the prev key
+  k: common.bytes, // the rest of the key outside the shared prefix
+  v: common.cid, // value
   t: subTreePointer, // next subtree (to the right of leaf)
 })
-export const nodeDataDef = z.object({
+const nodeData = z.object({
   l: subTreePointer, // left-most subtree
   e: z.array(treeEntry), //entries
 })
-export type NodeData = z.infer<typeof nodeDataDef>
+export type NodeData = z.infer<typeof nodeData>
+
+export const nodeDataDef = {
+  name: 'mst node',
+  schema: nodeData,
+}
 
 export type NodeEntry = MST | Leaf
 
-const DEFAULT_MST_FANOUT = 16
-export type Fanout = 2 | 8 | 16 | 32 | 64
 export type MstOpts = {
   layer: number
-  fanout: Fanout
 }
 
-export class MST implements DataStore {
-  blockstore: IpldStore
-  fanout: Fanout
+export class MST {
+  storage: ReadableBlockstore
   entries: NodeEntry[] | null
   layer: number | null
   pointer: CID
   outdatedPointer = false
 
   constructor(
-    blockstore: IpldStore,
-    fanout: Fanout,
+    storage: ReadableBlockstore,
     pointer: CID,
     entries: NodeEntry[] | null,
     layer: number | null,
   ) {
-    this.blockstore = blockstore
-    this.fanout = fanout
+    this.storage = storage
     this.entries = entries
     this.layer = layer
     this.pointer = pointer
   }
 
   static async create(
-    blockstore: IpldStore,
+    storage: ReadableBlockstore,
     entries: NodeEntry[] = [],
     opts?: Partial<MstOpts>,
   ): Promise<MST> {
     const pointer = await util.cidForEntries(entries)
-    const { layer = 0, fanout = DEFAULT_MST_FANOUT } = opts || {}
-    return new MST(blockstore, fanout, pointer, entries, layer)
+    const { layer = null } = opts || {}
+    return new MST(storage, pointer, entries, layer)
   }
 
   static async fromData(
-    blockstore: IpldStore,
+    storage: ReadableBlockstore,
     data: NodeData,
     opts?: Partial<MstOpts>,
   ): Promise<MST> {
-    const { layer = null, fanout = DEFAULT_MST_FANOUT } = opts || {}
-    const entries = await util.deserializeNodeData(blockstore, data, opts)
-    const pointer = await cidForData(data)
-    return new MST(blockstore, fanout, pointer, entries, layer)
+    const { layer = null } = opts || {}
+    const entries = await util.deserializeNodeData(storage, data, opts)
+    const pointer = await cidForCbor(data)
+    return new MST(storage, pointer, entries, layer)
   }
 
-  static load(blockstore: IpldStore, cid: CID, opts?: Partial<MstOpts>): MST {
-    const { layer = null, fanout = DEFAULT_MST_FANOUT } = opts || {}
-    return new MST(blockstore, fanout, cid, null, layer)
+  // this is really a *lazy* load, doesn't actually touch storage
+  static load(
+    storage: ReadableBlockstore,
+    cid: CID,
+    opts?: Partial<MstOpts>,
+  ): MST {
+    const { layer = null } = opts || {}
+    return new MST(storage, cid, null, layer)
   }
 
   // Immutability
@@ -114,13 +122,7 @@ export class MST implements DataStore {
 
   // We never mutate an MST, we just return a new MST with updated values
   async newTree(entries: NodeEntry[]): Promise<MST> {
-    const mst = new MST(
-      this.blockstore,
-      this.fanout,
-      this.pointer,
-      entries,
-      this.layer,
-    )
+    const mst = new MST(this.storage, this.pointer, entries, this.layer)
     mst.outdatedPointer = true
     return mst
   }
@@ -132,15 +134,14 @@ export class MST implements DataStore {
   async getEntries(): Promise<NodeEntry[]> {
     if (this.entries) return [...this.entries]
     if (this.pointer) {
-      const data = await this.blockstore.get(this.pointer, nodeDataDef)
+      const data = await this.storage.readObj(this.pointer, nodeDataDef)
       const firstLeaf = data.e[0]
       const layer =
         firstLeaf !== undefined
-          ? await util.leadingZerosOnHash(firstLeaf.k, this.fanout)
+          ? await util.leadingZerosOnHash(firstLeaf.k)
           : undefined
-      this.entries = await util.deserializeNodeData(this.blockstore, data, {
+      this.entries = await util.deserializeNodeData(this.storage, data, {
         layer,
-        fanout: this.fanout,
       })
 
       return this.entries
@@ -178,7 +179,7 @@ export class MST implements DataStore {
   async attemptGetLayer(): Promise<number | null> {
     if (this.layer !== null) return this.layer
     const entries = await this.getEntries()
-    let layer = await util.layerForEntries(entries, this.fanout)
+    let layer = await util.layerForEntries(entries)
     if (layer === null) {
       for (const entry of entries) {
         if (entry.isTree()) {
@@ -197,39 +198,29 @@ export class MST implements DataStore {
   // Core functionality
   // -------------------
 
-  // Persist the MST to the blockstore
-  // If the topmost tree only has one entry and it's a subtree, we can eliminate the topmost tree
-  // However, lower trees with only one entry must be preserved
-  async stage(): Promise<CID> {
-    return this.stageRecurse(true)
-  }
-
-  async stageRecurse(trimTop = false): Promise<CID> {
+  // Return the necessary blocks to persist the MST to repo storage
+  async getUnstoredBlocks(): Promise<{ root: CID; blocks: BlockMap }> {
+    const blocks = new BlockMap()
     const pointer = await this.getPointer()
-    const alreadyHas = await this.blockstore.has(pointer)
-    if (alreadyHas) return pointer
+    const alreadyHas = await this.storage.has(pointer)
+    if (alreadyHas) return { root: pointer, blocks }
     const entries = await this.getEntries()
-    if (entries.length === 1 && trimTop) {
-      const node = entries[0]
-      if (node.isTree()) {
-        return node.stageRecurse(true)
-      }
-    }
     const data = util.serializeNodeData(entries)
-    await this.blockstore.stage(data)
+    await blocks.add(data)
     for (const entry of entries) {
       if (entry.isTree()) {
-        await entry.stageRecurse(false)
+        const subtree = await entry.getUnstoredBlocks()
+        blocks.addMap(subtree.blocks)
       }
     }
-    return pointer
+    return { root: pointer, blocks: blocks }
   }
 
   // Adds a new leaf for the given key/value pair
   // Throws if a leaf with that key already exists
   async add(key: string, value: CID, knownZeros?: number): Promise<MST> {
-    const keyZeros =
-      knownZeros ?? (await util.leadingZerosOnHash(key, this.fanout))
+    util.ensureValidMstKey(key)
+    const keyZeros = knownZeros ?? (await util.leadingZerosOnHash(key))
     const layer = await this.getLayer()
     const newLeaf = new Leaf(key, value)
     if (keyZeros === layer) {
@@ -288,9 +279,8 @@ export class MST implements DataStore {
       if (left) updated.push(left)
       updated.push(new Leaf(key, value))
       if (right) updated.push(right)
-      const newRoot = await MST.create(this.blockstore, updated, {
+      const newRoot = await MST.create(this.storage, updated, {
         layer: keyZeros,
-        fanout: this.fanout,
       })
       newRoot.outdatedPointer = true
       return newRoot
@@ -314,6 +304,7 @@ export class MST implements DataStore {
   // Edits the value at the given key
   // Throws if the given key does not exist
   async update(key: string, value: CID): Promise<MST> {
+    util.ensureValidMstKey(key)
     const index = await this.findGtOrEqualLeafIndex(key)
     const found = await this.atIndex(index)
     if (found && found.isLeaf() && found.key === key) {
@@ -329,6 +320,11 @@ export class MST implements DataStore {
 
   // Deletes the value at the given key
   async delete(key: string): Promise<MST> {
+    const altered = await this.deleteRecurse(key)
+    return altered.trimTop()
+  }
+
+  async deleteRecurse(key: string): Promise<MST> {
     const index = await this.findGtOrEqualLeafIndex(key)
     const found = await this.atIndex(index)
     // if found, remove it on this level
@@ -349,7 +345,7 @@ export class MST implements DataStore {
     // else recurse down to find it
     const prev = await this.atIndex(index - 1)
     if (prev?.isTree()) {
-      const subtree = await prev.delete(key)
+      const subtree = await prev.deleteRecurse(key)
       const subTreeEntries = await subtree.getEntries()
       if (subTreeEntries.length === 0) {
         return this.removeEntry(index - 1)
@@ -361,114 +357,6 @@ export class MST implements DataStore {
     }
   }
 
-  // Walk two MSTs to find the semantic changes
-  async diff(other: MST): Promise<DataDiff> {
-    await this.getPointer()
-    await other.getPointer()
-    const diff = new DataDiff()
-
-    const leftWalker = new MstWalker(this)
-    const rightWalker = new MstWalker(other)
-    while (!leftWalker.status.done || !rightWalker.status.done) {
-      // if one walker is finished, continue walking the other & logging all nodes
-      if (leftWalker.status.done && !rightWalker.status.done) {
-        const node = rightWalker.status.curr
-        if (node.isLeaf()) {
-          diff.recordAdd(node.key, node.value)
-        } else {
-          diff.recordNewCid(node.pointer)
-        }
-        await rightWalker.advance()
-        continue
-      } else if (!leftWalker.status.done && rightWalker.status.done) {
-        const node = leftWalker.status.curr
-        if (node.isLeaf()) {
-          diff.recordDelete(node.key, node.value)
-        }
-        await leftWalker.advance()
-        continue
-      }
-      if (leftWalker.status.done || rightWalker.status.done) break
-      const left = leftWalker.status.curr
-      const right = rightWalker.status.curr
-      if (left === null || right === null) break
-
-      // if both pointers are leaves, record an update & advance both or record the lowest key and advance that pointer
-      if (left.isLeaf() && right.isLeaf()) {
-        if (left.key === right.key) {
-          if (!left.value.equals(right.value)) {
-            diff.recordUpdate(left.key, left.value, right.value)
-          }
-          await leftWalker.advance()
-          await rightWalker.advance()
-        } else if (left.key < right.key) {
-          diff.recordDelete(left.key, left.value)
-          await leftWalker.advance()
-        } else {
-          diff.recordAdd(right.key, right.value)
-          await rightWalker.advance()
-        }
-        continue
-      }
-
-      // next, ensure that we're on the same layer
-      // if one walker is at a higher layer than the other, we need to do one of two things
-      // if the higher walker is pointed at a tree, step into that tree to try to catch up with the lower
-      // if the higher walker is pointed at a leaf, then advance the lower walker to try to catch up the higher
-      if (leftWalker.layer() > rightWalker.layer()) {
-        if (left.isLeaf()) {
-          if (right.isLeaf()) {
-            diff.recordAdd(right.key, right.value)
-          } else {
-            diff.recordNewCid(right.pointer)
-          }
-          await rightWalker.advance()
-        } else {
-          await leftWalker.stepInto()
-        }
-        continue
-      } else if (leftWalker.layer() < rightWalker.layer()) {
-        if (right.isLeaf()) {
-          if (left.isLeaf()) {
-            diff.recordDelete(left.key, left.value)
-          }
-          await leftWalker.advance()
-        } else {
-          diff.recordNewCid(right.pointer)
-          await rightWalker.stepInto()
-        }
-        continue
-      }
-
-      // if we're on the same level, and both pointers are trees, do a comparison
-      // if they're the same, step over. if they're different, step in to find the subdiff
-      if (left.isTree() && right.isTree()) {
-        if (left.pointer.equals(right.pointer)) {
-          await leftWalker.stepOver()
-          await rightWalker.stepOver()
-        } else {
-          diff.recordNewCid(right.pointer)
-          await leftWalker.stepInto()
-          await rightWalker.stepInto()
-        }
-        continue
-      }
-
-      // finally, if one pointer is a tree and the other is a leaf, simply step into the tree
-      if (left.isLeaf() && right.isTree()) {
-        await diff.recordNewCid(right.pointer)
-        await rightWalker.stepInto()
-        continue
-      } else if (left.isTree() && right.isLeaf()) {
-        await leftWalker.stepInto()
-        continue
-      }
-
-      throw new Error('Unidentifiable case in diff walk')
-    }
-    return diff
-  }
-
   // Simple Operations
   // -------------------
 
@@ -543,6 +431,16 @@ export class MST implements DataStore {
     return this.newTree(update)
   }
 
+  // if the topmost node in the tree only points to another tree, trim the top and return the subtree
+  async trimTop(): Promise<MST> {
+    const entries = await this.getEntries()
+    if (entries.length === 1 && entries[0].isTree()) {
+      return entries[0].trimTop()
+    } else {
+      return this
+    }
+  }
+
   // Subtree & Splits
   // -------------------
 
@@ -604,17 +502,15 @@ export class MST implements DataStore {
 
   async createChild(): Promise<MST> {
     const layer = await this.getLayer()
-    return MST.create(this.blockstore, [], {
+    return MST.create(this.storage, [], {
       layer: layer - 1,
-      fanout: this.fanout,
     })
   }
 
   async createParent(): Promise<MST> {
     const layer = await this.getLayer()
-    const parent = await MST.create(this.blockstore, [this], {
+    const parent = await MST.create(this.storage, [this], {
       layer: layer + 1,
-      fanout: this.fanout,
     })
     parent.outdatedPointer = true
     return parent
@@ -660,7 +556,11 @@ export class MST implements DataStore {
     }
   }
 
-  async list(count: number, after?: string, before?: string): Promise<Leaf[]> {
+  async list(
+    count = Number.MAX_SAFE_INTEGER,
+    after?: string,
+    before?: string,
+  ): Promise<Leaf[]> {
     const vals: Leaf[] = []
     for await (const leaf of this.walkLeavesFrom(after || '')) {
       if (leaf.key === after) continue
@@ -726,6 +626,22 @@ export class MST implements DataStore {
     return nodes
   }
 
+  // Walks tree & returns all cids
+  async allCids(): Promise<CidSet> {
+    const cids = new CidSet()
+    const entries = await this.getEntries()
+    for (const entry of entries) {
+      if (entry.isLeaf()) {
+        cids.add(entry.value)
+      } else {
+        const subtreeCids = await entry.allCids()
+        cids.addSet(subtreeCids)
+      }
+    }
+    cids.add(await this.getPointer())
+    return cids
+  }
+
   // Walks tree & returns all leaves
   async leaves() {
     const leaves: Leaf[] = []
@@ -741,17 +657,93 @@ export class MST implements DataStore {
     return leaves.length
   }
 
-  // Sync Protocol
+  // Reachable tree traversal
+  // -------------------
 
-  async writeToCarStream(car: BlockWriter): Promise<void> {
-    for await (const entry of this.walk()) {
+  // Walk reachable branches of tree & emit nodes, consumer can bail at any point by returning false
+  async *walkReachable(): AsyncIterable<NodeEntry> {
+    yield this
+    const entries = await this.getEntries()
+    for (const entry of entries) {
       if (entry.isTree()) {
-        const pointer = await entry.getPointer()
-        await this.blockstore.addToCar(car, pointer)
+        try {
+          for await (const e of entry.walkReachable()) {
+            yield e
+          }
+        } catch (err) {
+          if (err instanceof MissingBlockError) {
+            continue
+          } else {
+            throw err
+          }
+        }
       } else {
-        await this.blockstore.addToCar(car, entry.value)
+        yield entry
+      }
+    }
+  }
+
+  async reachableLeaves(): Promise<Leaf[]> {
+    const leaves: Leaf[] = []
+    for await (const entry of this.walkReachable()) {
+      if (entry.isLeaf()) leaves.push(entry)
+    }
+    return leaves
+  }
+
+  // Sync Protocol
+
+  async writeToCarStream(car: BlockWriter): Promise<void> {
+    const leaves = new CidSet()
+    let toFetch = new CidSet()
+    toFetch.add(await this.getPointer())
+    while (toFetch.size() > 0) {
+      const nextLayer = new CidSet()
+      const fetched = await this.storage.getBlocks(toFetch.toList())
+      if (fetched.missing.length > 0) {
+        throw new MissingBlocksError('mst node', fetched.missing)
+      }
+      for (const cid of toFetch.toList()) {
+        const found = await parse.getAndParseByDef(
+          fetched.blocks,
+          cid,
+          nodeDataDef,
+        )
+        await car.put({ cid, bytes: found.bytes })
+        const entries = await util.deserializeNodeData(this.storage, found.obj)
+
+        for (const entry of entries) {
+          if (entry.isLeaf()) {
+            leaves.add(entry.value)
+          } else {
+            nextLayer.add(await entry.getPointer())
+          }
+        }
       }
+      toFetch = nextLayer
+    }
+    const leafData = await this.storage.getBlocks(leaves.toList())
+    if (leafData.missing.length > 0) {
+      throw new MissingBlocksError('mst leaf', leafData.missing)
+    }
+
+    for (const leaf of leafData.blocks.entries()) {
+      await car.put(leaf)
+    }
+  }
+
+  async cidsForPath(key: string): Promise<CID[]> {
+    const cids: CID[] = [await this.getPointer()]
+    const index = await this.findGtOrEqualLeafIndex(key)
+    const found = await this.atIndex(index)
+    if (found && found.isLeaf() && found.key === key) {
+      return [...cids, found.value]
+    }
+    const prev = await this.atIndex(index - 1)
+    if (prev && prev.isTree()) {
+      return [...cids, ...(await prev.cidsForPath(key))]
     }
+    return cids
   }
 
   // Matching Leaf interface

package/src/mst/util.ts

@@ -1,68 +1,60 @@
 import { CID } from 'multiformats'
 import * as uint8arrays from 'uint8arrays'
-import IpldStore from '../blockstore/ipld-store'
+import { ReadableBlockstore } from '../storage'
 import { sha256 } from '@atproto/crypto'
-import { MST, Leaf, NodeEntry, NodeData, MstOpts, Fanout } from './mst'
-import { cidForData } from '@atproto/common'
+import { MST, Leaf, NodeEntry, NodeData, MstOpts } from './mst'
+import { cidForCbor } from '@atproto/common'
 
-type SupportedBases = 'base2' | 'base8' | 'base16' | 'base32' | 'base64'
-
-export const leadingZerosOnHash = async (
-  key: string,
-  fanout: Fanout,
-): Promise<number> => {
-  if ([2, 8, 16, 32, 64].indexOf(fanout) < 0) {
-    throw new Error(`Not a valid fanout: ${fanout}`)
-  }
-  const base: SupportedBases = `base${fanout}`
-  const zeroChar = uint8arrays.toString(new Uint8Array(1), base)[0]
+export const leadingZerosOnHash = async (key: string | Uint8Array) => {
   const hash = await sha256(key)
-  const encoded = uint8arrays.toString(hash, base)
-  let count = 0
-  for (const char of encoded) {
-    if (char === zeroChar) {
-      count++
+  let leadingZeros = 0
+  for (let i = 0; i < hash.length; i++) {
+    const byte = hash[i]
+    if (byte < 64) leadingZeros++
+    if (byte < 16) leadingZeros++
+    if (byte < 4) leadingZeros++
+    if (byte === 0) {
+      leadingZeros++
     } else {
       break
     }
   }
-  return count
+  return leadingZeros
 }
 
 export const layerForEntries = async (
   entries: NodeEntry[],
-  fanout: Fanout,
 ): Promise<number | null> => {
   const firstLeaf = entries.find((entry) => entry.isLeaf())
   if (!firstLeaf || firstLeaf.isTree()) return null
-  return await leadingZerosOnHash(firstLeaf.key, fanout)
+  return await leadingZerosOnHash(firstLeaf.key)
 }
 
 export const deserializeNodeData = async (
-  blockstore: IpldStore,
+  storage: ReadableBlockstore,
   data: NodeData,
   opts?: Partial<MstOpts>,
 ): Promise<NodeEntry[]> => {
-  const { layer, fanout } = opts || {}
+  const { layer } = opts || {}
   const entries: NodeEntry[] = []
   if (data.l !== null) {
     entries.push(
-      await MST.load(blockstore, data.l, {
+      await MST.load(storage, data.l, {
         layer: layer ? layer - 1 : undefined,
-        fanout,
       }),
     )
   }
   let lastKey = ''
   for (const entry of data.e) {
-    const key = lastKey.slice(0, entry.p) + entry.k
+    const keyStr = uint8arrays.toString(entry.k, 'ascii')
+    const key = lastKey.slice(0, entry.p) + keyStr
+    ensureValidMstKey(key)
     entries.push(new Leaf(key, entry.v))
     lastKey = key
     if (entry.t !== null) {
       entries.push(
-        await MST.load(blockstore, entry.t, {
+        await MST.load(storage, entry.t, {
           layer: layer ? layer - 1 : undefined,
-          fanout,
         }),
       )
     }
@@ -93,10 +85,11 @@ export const serializeNodeData = (entries: NodeEntry[]): NodeData => {
       subtree = next.pointer
       i++
     }
+    ensureValidMstKey(leaf.key)
     const prefixLen = countPrefixLen(lastKey, leaf.key)
     data.e.push({
       p: prefixLen,
-      k: leaf.key.slice(prefixLen),
+      k: uint8arrays.fromString(leaf.key.slice(prefixLen), 'ascii'),
       v: leaf.value,
       t: subtree,
     })
@@ -118,5 +111,35 @@ export const countPrefixLen = (a: string, b: string): number => {
 
 export const cidForEntries = async (entries: NodeEntry[]): Promise<CID> => {
   const data = serializeNodeData(entries)
-  return cidForData(data)
+  return cidForCbor(data)
+}
+
+export const isValidMstKey = (str: string): boolean => {
+  const split = str.split('/')
+  return (
+    str.length <= 256 &&
+    split.length === 2 &&
+    split[0].length > 0 &&
+    split[1].length > 0 &&
+    isValidChars(split[0]) &&
+    isValidChars(split[1])
+  )
+}
+
+export const validCharsRegex = /^[a-zA-Z0-9_\-:.]*$/
+
+export const isValidChars = (str: string): boolean => {
+  return str.match(validCharsRegex) !== null
+}
+
+export const ensureValidMstKey = (str: string) => {
+  if (!isValidMstKey(str)) {
+    throw new InvalidMstKeyError(str)
+  }
+}
+
+export class InvalidMstKeyError extends Error {
+  constructor(public key: string) {
+    super(`Not a valid MST key: ${key}`)
+  }
 }