sajidanwar.com
forked from
mary.my.id/atcute
a collection of lightweight TypeScript packages for AT Protocol, the protocol powering Bluesky
mary.my.id
1import * as CBOR from '@atcute/cbor';
2import type { CidLink } from '@atcute/cid';
3import * as CID from '@atcute/cid';
4import { decodeUtf8From, encodeUtf8, toSha256 } from '@atcute/uint8array';
5
6import { assertMstKey } from './key.ts';
7import { isNodeData, type NodeData, type TreeEntry } from './types.ts';
8
9/**
10 * represents a node in a Merkle Search Tree (MST)
11 * stores sorted keys, their associated values (CIDs), and subtree pointers
12 */
13export class MSTNode {
14 /**
15 * cached height of this node in the tree
16 * @internal
17 */
18 _height: number | null | undefined;
19 /**
20 * cached CID for this node
21 * @internal
22 */
23 _cid: CidLink | undefined;
24 /**
25 * cached serialized bytes for this node
26 * @internal
27 */
28 _bytes: Uint8Array<ArrayBuffer> | undefined;
29
30 protected constructor(
31 /** sorted array of keys stored in this node */
32 readonly keys: readonly string[],
33 /** array of value CIDs corresponding to each key */
34 readonly values: readonly CidLink[],
35 /** array of subtree CIDs (length is keys.length + 1) */
36 readonly subtrees: readonly (CidLink | null)[],
37 ) {}
38
39 /**
40 * creates a new MST node with validation
41 * @param keys sorted array of keys
42 * @param values array of value CIDs corresponding to keys
43 * @param subtrees array of subtree CIDs (length must be keys.length + 1)
44 * @returns a new validated MST node
45 * @throws {TypeError} if node structure is invalid or keys have inconsistent heights
46 */
47 static async create(
48 keys: readonly string[],
49 values: readonly CidLink[],
50 subtrees: readonly (CidLink | null)[],
51 ): Promise<MSTNode> {
52 if (subtrees.length !== keys.length + 1) {
53 throw new TypeError(`malformed MST node; invalid subtree count`);
54 }
55
56 if (keys.length !== values.length) {
57 throw new TypeError(`malformed MST node; mismatched keys/values lengths`);
58 }
59
60 let expectedHeight: number | undefined;
61 for (const key of keys) {
62 const height = await getKeyHeight(key);
63 expectedHeight ??= height;
64
65 if (height !== expectedHeight) {
66 throw new TypeError(`malformed MST node; inconsistent key heights`);
67 }
68 }
69
70 return new MSTNode(keys, values, subtrees);
71 }
72
73 /**
74 * creates an empty MST node
75 * @returns a new empty node
76 */
77 static empty(): MSTNode {
78 return new MSTNode([], [], [null]);
79 }
80
81 /**
82 * deserializes an MST node from CBOR-encoded bytes
83 * @param bytes the CBOR-encoded node data
84 * @returns the deserialized MST node
85 * @throws {TypeError} if the bytes don't represent a valid MST node
86 */
87 static async deserialize(bytes: Uint8Array): Promise<MSTNode> {
88 const node = CBOR.decode(bytes);
89 if (!isNodeData(node)) {
90 throw new TypeError(`malformed MST node; invalid structure`);
91 }
92
93 const keys: string[] = [];
94 const values: CidLink[] = [];
95 const subtrees: (CidLink | null)[] = [node.l];
96
97 let prevKey = '';
98
99 for (const entry of node.e) {
100 const prefixLen = entry.p;
101 if (prefixLen > prevKey.length) {
102 throw new TypeError(`malformed MST node; unexpected key prefix length`);
103 }
104
105 const suffix = decodeUtf8From(CBOR.fromBytes(entry.k));
106 if (prevKey[prefixLen] === suffix[0]) {
107 throw new TypeError(`malformed MST node; suboptimal key prefix length`);
108 }
109
110 const key = prevKey.slice(0, prefixLen) + suffix;
111 if (key <= prevKey) {
112 throw new TypeError(`malformed MST node; invalid key sort order`);
113 }
114
115 assertMstKey(key);
116 keys.push(key);
117 values.push(entry.v);
118 subtrees.push(entry.t);
119
120 prevKey = key;
121 }
122
123 return await MSTNode.create(keys, values, subtrees);
124 }
125
126 /**
127 * serializes the node to CBOR-encoded bytes with prefix compression
128 * @returns the CBOR-encoded node data
129 */
130 async serialize(): Promise<Uint8Array<ArrayBuffer>> {
131 let bytes = this._bytes;
132 if (bytes === undefined) {
133 const keys = this.keys;
134 const values = this.values;
135 const subtrees = this.subtrees;
136
137 const e: TreeEntry[] = [];
138
139 let prevKey = '';
140
141 for (let idx = 0, len = keys.length; idx < len; idx++) {
142 const key = keys[idx];
143 const prefixLen = commonPrefixLength(prevKey, key);
144 const suffix = key.slice(prefixLen);
145
146 e.push({
147 k: CBOR.toBytes(encodeUtf8(suffix)),
148 p: prefixLen,
149 t: subtrees[idx + 1],
150 v: values[idx],
151 });
152
153 prevKey = key;
154 }
155
156 const n: NodeData = {
157 l: subtrees[0],
158 e: e,
159 };
160
161 this._bytes = bytes = CBOR.encode(n);
162 }
163
164 return bytes;
165 }
166
167 /**
168 * whether the node is empty (no keys or values)
169 */
170 get isEmpty(): boolean {
171 return this.subtrees.length === 1 && this.subtrees[0] === null;
172 }
173
174 /**
175 * computes the CID for this node
176 * @returns the CID link for this node
177 */
178 async cid(): Promise<CidLink> {
179 let cid = this._cid;
180 if (cid === undefined) {
181 this._cid = cid = CID.toCidLink(await CID.create(0x71, await this.serialize()));
182 }
183
184 return cid;
185 }
186
187 /**
188 * computes the height of this node in the MST
189 * @returns the height, or null if indeterminate (empty intermediate node)
190 */
191 async height(): Promise<number | null> {
192 let height = this._height;
193 if (height === undefined) {
194 const keys = this.keys;
195
196 if (this.isEmpty) {
197 height = 0;
198 } else if (keys.length > 0) {
199 height = await getKeyHeight(keys[0]);
200 } else {
201 height = null;
202 }
203
204 this._height = height;
205 }
206
207 return height;
208 }
209
210 /**
211 * gets the node height, throwing if indeterminate
212 * @returns the height
213 * @throws {Error} if height cannot be determined
214 */
215 async requireHeight(): Promise<number> {
216 const height = await this.height();
217 if (height === null) {
218 throw new Error(`indeterminate node height`);
219 }
220
221 return height;
222 }
223
224 /**
225 * finds the index of the first key >= the given key
226 * @param key the key to search for
227 * @returns the index of the lower bound
228 */
229 lowerBound(key: string): number {
230 const keys = this.keys;
231 const len = keys.length;
232
233 for (let idx = 0; idx < len; idx++) {
234 if (key <= keys[idx]) {
235 return idx;
236 }
237 }
238
239 return len;
240 }
241
242 /**
243 * returns the node's CID if it's not empty, null otherwise
244 * @returns the CID or null
245 * @internal
246 */
247 async _toNullable(): Promise<CidLink | null> {
248 if (this.isEmpty) {
249 return null;
250 }
251 return await this.cid();
252 }
253}
254
255/**
256 * computes the MST height for a given key by counting leading zeros in its hash
257 * @param key the key to compute height for
258 * @returns the height (number of leading zero bits in 2-bit chunks)
259 */
260export const getKeyHeight = async (key: string): Promise<number> => {
261 const hash = await toSha256(encodeUtf8(key));
262
263 let lz = 0;
264 for (let idx = 0, len = hash.length; idx < len; idx++) {
265 const byte = hash[idx];
266
267 if (byte < 64) {
268 lz++;
269 }
270 if (byte < 16) {
271 lz++;
272 }
273 if (byte < 4) {
274 lz++;
275 }
276
277 if (byte === 0) {
278 lz++;
279 } else {
280 break;
281 }
282 }
283
284 return lz;
285};
286
287/**
288 * computes the length of the common prefix between two strings
289 * @param a first string
290 * @param b second string
291 * @returns length of common prefix
292 */
293const commonPrefixLength = (a: string, b: string): number => {
294 let idx = 0;
295 for (let len = Math.min(a.length, b.length); idx < len; idx++) {
296 if (a[idx] !== b[idx]) {
297 break;
298 }
299 }
300
301 return idx;
302};