feat(mutate): accept updates: AwaitIterable<Update | Update[]>

Author: tabcat

src/mutate.ts

@@ -1,4 +1,4 @@
-import { firstElement } from "@tabcat/ith-element";
+import { firstElement, lastElement } from "@tabcat/ith-element";
 import { union } from "@tabcat/ordered-sets/union";
 import { pairwiseTraversal } from "@tabcat/ordered-sets/util";
 import { Blockstore } from "interface-blockstore";
@@ -74,21 +74,33 @@ const handleUpdate = (
   }
 };
 
-async function takeOne<T>(it: AwaitIterable<T>): Promise<T | void> {
-  for await (const v of it) return v;
+async function takeOneUpdate(
+  updates: AwaitIterable<Update | Update[]>,
+): Promise<Update | void> {
+  for await (const u of updates) return Array.isArray(u) ? u[0] : u;
 }
 
 async function populateUpdts(
-  updates: AwaitIterable<Update>,
+  updates: AwaitIterable<Update | Update[]>,
   updts: Update[],
   updatee: Bucket,
   isHead: boolean,
 ): Promise<void> {
-  for await (const update of updates) {
-    updts.push(update);
-    const boundary = updatee.getBoundary();
+  for await (let _updates of updates) {
+    if (!Array.isArray(_updates)) {
+      _updates = [_updates];
+    }
+
+    for (const u of _updates) {
+      updts.push(u);
+    }
 
-    if (boundary != null && !isHead && compareTuples(update, boundary) >= 0) {
+    const boundary = updatee.getBoundary();
+    if (
+      boundary != null &&
+      !isHead &&
+      compareTuples(lastElement(updts), boundary) >= 0
+    ) {
       break;
     }
   }
@@ -106,13 +118,14 @@ async function populateUpdts(
 export async function* mutate(
   blockstore: Blockstore,
   tree: ProllyTree,
-  updates: AwaitIterable<Update>,
+  updates: AwaitIterable<Update | Update[]>,
 ): AsyncGenerator<ProllyTreeDiff> {
   // whole function should be rewritten around updates async iterator, too complicated right now
   if (Array.isArray(updates)) {
     updates = updates[Symbol.iterator]();
   }
-  const firstUpdate = await takeOne(updates);
+
+  const firstUpdate = await takeOneUpdate(updates);
 
   if (firstUpdate == null) {
     return tree;
@@ -270,7 +283,7 @@ export async function* mutate(
       break;
     }
 
-    let nextUpdt = updts[0] ?? (await takeOne(updates));
+    let nextUpdt = updts[0] ?? (await takeOneUpdate(updates));
     let nextLevel = level;
 
     if (nextUpdt == null) {

test/mutate.test.ts

@@ -3,7 +3,7 @@ import { pairwiseTraversal } from "@tabcat/ordered-sets/util";
 import { describe, expect, it } from "vitest";
 import { compareBuckets, compareBytes, compareTuples } from "../src/compare.js";
 import { BucketDiff, NodeDiff } from "../src/diff.js";
-import { cloneTree } from "../src/index.js";
+import { cloneTree, createEmptyTree } from "../src/index.js";
 import { Node, ProllyTree } from "../src/interface.js";
 import { Update, mutate } from "../src/mutate.js";
 import { nodeToTuple } from "../src/utils.js";
@@ -162,4 +162,21 @@ describe("mutate", () => {
       });
     }
   }
+
+  it("accepts updates with type AwaitIterable<Update | Update[]>", async () => {
+    const tree1 = createEmptyTree();
+    const tree2 = createEmptyTree();
+
+    const updates: Update[] = [
+      { timestamp: 1, hash: new Uint8Array(32), message: new Uint8Array() },
+    ];
+
+    for await (const _ of mutate(blockstore, tree1, updates)) {
+    }
+
+    for await (const _ of mutate(blockstore, tree2, [updates])) {
+    }
+
+    expect(tree1).to.deep.equal(tree2);
+  });
 });