PCDCollection merge extracted from #1347

proofcarryingdata · Dec 14, 2023 · 66671ce · 66671ce
1 parent 5d1899d
commit 66671ce
Show file tree

Hide file tree

Showing 2 changed files with 149 additions and 0 deletions.
diff --git a/packages/pcd-collection/src/PCDCollection.ts b/packages/pcd-collection/src/PCDCollection.ts
@@ -30,6 +30,8 @@ export type MatchingActionPermission =
 
 type AddPCDOptions = { upsert?: boolean };
 
+export type MergeFilterPredicate = (pcd: PCD, target: PCDCollection) => boolean;
+
 export function matchActionToPermission(
   action: PCDAction,
   permissions: PCDPermission[]
@@ -453,6 +455,36 @@ export class PCDCollection {
 
     return collection;
   }
+
+  /**
+   * Merges another PCD collection into this one.
+   * There are two options:
+   * - `shouldInclude` is a function used to filter out PCDs from the other
+   *   collection during merging, e.g. to filter out duplicates or PCDs of
+   *   a type that should not be copied.
+   */
+  public merge(
+    other: PCDCollection,
+    options?: {
+      shouldInclude?: MergeFilterPredicate;
+    }
+  ): void {
+    let pcds = other.getAll();
+
+    // If the caller has specified a filter function, run that first to filter
+    // out unwanted PCDs from the merge.
+    if (options?.shouldInclude) {
+      pcds = pcds.filter((pcd: PCD) => options.shouldInclude?.(pcd, this));
+    }
+
+    this.addAll(pcds, { upsert: true });
+
+    for (const pcd of pcds) {
+      if (other.folders[pcd.id]) {
+        this.setFolder(pcd.id, other.folders[pcd.id]);
+      }
+    }
+  }
 }
 
 /**

diff --git a/packages/pcd-collection/test/PCDCollection.spec.ts b/packages/pcd-collection/test/PCDCollection.spec.ts
@@ -271,6 +271,123 @@ describe("PCDCollection", async function () {
     );
     expect(await deserializedCollection.getHash()).to.eq(firstHash);
   });
+
+  it("should merge two PCD collections", async function () {
+    const pcdList = await Promise.all([
+      newPCD(),
+      newPCD(),
+      newPCD(),
+      newPCD(),
+      newPCD()
+    ]);
+
+    const firstCollection = new PCDCollection(packages);
+    const secondCollection = new PCDCollection(packages);
+
+    firstCollection.add(pcdList[0]);
+    firstCollection.add(pcdList[1]);
+    firstCollection.add(pcdList[2]);
+    secondCollection.add(pcdList[3]);
+    secondCollection.add(pcdList[4]);
+
+    firstCollection.merge(secondCollection);
+
+    expect(firstCollection.getAll()).to.deep.eq(pcdList);
+  });
+
+  it("should upsert PCDs with the same ID by default", async function () {
+    const pcdList = await Promise.all([
+      newPCD(),
+      newPCD(),
+      newPCD(),
+      newPCD(),
+      newPCD()
+    ]);
+
+    const firstCollection = new PCDCollection(packages);
+    const secondCollection = new PCDCollection(packages);
+
+    firstCollection.add(pcdList[0]);
+    firstCollection.add(pcdList[1]);
+    firstCollection.add(pcdList[2]);
+    firstCollection.add(pcdList[3]);
+    firstCollection.add(pcdList[4]);
+    secondCollection.add(pcdList[3]);
+    secondCollection.add(pcdList[4]);
+
+    firstCollection.merge(secondCollection);
+
+    expect(firstCollection.getAll()).to.deep.eq(pcdList);
+  });
+
+  it("should filter out PCDs during merge", async function () {
+    const pcdList = await Promise.all([
+      newPCD(),
+      newPCD(),
+      newPCD(),
+      newPCD(),
+      newPCD()
+    ]);
+
+    const firstCollection = new PCDCollection(packages);
+    const secondCollection = new PCDCollection(packages);
+
+    firstCollection.add(pcdList[0]);
+    firstCollection.add(pcdList[1]);
+    firstCollection.add(pcdList[2]);
+    secondCollection.add(pcdList[3]);
+    secondCollection.add(pcdList[4]);
+
+    // secondCollection contains two PCDs, but we want to filter one of them
+    // out.
+    firstCollection.merge(secondCollection, {
+      shouldInclude: (pcd) => {
+        return pcd.id !== pcdList[4].id;
+      }
+    });
+
+    expect(firstCollection.getAll()).to.deep.eq([
+      pcdList[0],
+      pcdList[1],
+      pcdList[2],
+      pcdList[3]
+      // pcdList[4] should have been filtered out
+    ]);
+  });
+
+  it("should copy folder details", async function () {
+    const pcdList = await Promise.all([newPCD(), newPCD()]);
+
+    const firstCollection = new PCDCollection(packages);
+    const secondCollection = new PCDCollection(packages);
+
+    firstCollection.add(pcdList[0]);
+    firstCollection.setPCDFolder(pcdList[0].id, "A");
+    secondCollection.add(pcdList[1]);
+    secondCollection.setPCDFolder(pcdList[1].id, "B");
+
+    firstCollection.merge(secondCollection);
+
+    expect(firstCollection.getAll()).to.deep.eq(pcdList);
+    expect(firstCollection.getFolderOfPCD(pcdList[1].id)).to.eq("B");
+  });
+
+  it("should copy folder details when folder exists in both collections", async function () {
+    const pcdList = await Promise.all([newPCD(), newPCD()]);
+
+    const firstCollection = new PCDCollection(packages);
+    const secondCollection = new PCDCollection(packages);
+
+    firstCollection.add(pcdList[0]);
+    firstCollection.setPCDFolder(pcdList[0].id, "A");
+    secondCollection.add(pcdList[1]);
+    secondCollection.setPCDFolder(pcdList[1].id, "A");
+
+    firstCollection.merge(secondCollection);
+
+    expect(firstCollection.getAll()).to.deep.eq(pcdList);
+    expect(firstCollection.getFolderOfPCD(pcdList[1].id)).to.eq("A");
+  });
 });
 
 function waitForNewHash(