feat(core): allow ensuring inclusion by boxId

fleet-sdk · Dec 12, 2022 · d24494f · d24494f
1 parent 26824dd
commit d24494f
Show file tree

Hide file tree

Showing 2 changed files with 94 additions and 10 deletions.
diff --git a/packages/core/src/builder/selector/boxSelector.spec.ts b/packages/core/src/builder/selector/boxSelector.spec.ts
@@ -133,13 +133,67 @@ describe("Ensure input inclusion", () => {
     const selector = new BoxSelector(regularBoxesMock).ensureInclusion(
       (input) => input.boxId === arbitraryBoxId
     );
-    const boxes = selector.select(target);
+    const boxes = selector.select({ nanoErgs: 10000n });
 
     expect(boxes.some((x) => x.boxId === arbitraryBoxId)).toBe(true);
     expect(boxes).toHaveLength(1);
     expect(sumBy(boxes, (x) => x.value)).toBeGreaterThanOrEqual(target.nanoErgs);
   });
 
+  it("Should forcedly include inputs by boxId", () => {
+    const arbitraryBoxId = "2555e34138d276905fe0bc19240bbeca10f388a71f7b4d2f65a7d0bfd23c846d";
+    const selector = new BoxSelector(regularBoxesMock).ensureInclusion(arbitraryBoxId);
+    const boxes = selector.select({ nanoErgs: 10000n });
+
+    expect(boxes).toHaveLength(1);
+    expect(boxes[0].boxId).toBe(arbitraryBoxId);
+  });
+
+  it("Should forcedly include inputs by multiple boxId", () => {
+    const selector = new BoxSelector(regularBoxesMock).ensureInclusion([
+      "e56847ed19b3dc6b72828fcfb992fdf7310828cf291221269b7ffc72fd66706e",
+      "2555e34138d276905fe0bc19240bbeca10f388a71f7b4d2f65a7d0bfd23c846d"
+    ]);
+    const boxes = selector.select({ nanoErgs: 10000n });
+
+    expect(boxes).toHaveLength(2);
+    expect(boxes[0].boxId).toBe("e56847ed19b3dc6b72828fcfb992fdf7310828cf291221269b7ffc72fd66706e");
+    expect(boxes[1].boxId).toBe("2555e34138d276905fe0bc19240bbeca10f388a71f7b4d2f65a7d0bfd23c846d");
+  });
+
+  it("Should forcedly include inputs by multiple boxId and filter criteria", () => {
+    const boxId1 = "e56847ed19b3dc6b72828fcfb992fdf7310828cf291221269b7ffc72fd66706e";
+    const boxId2 = "2555e34138d276905fe0bc19240bbeca10f388a71f7b4d2f65a7d0bfd23c846d";
+    const boxId3 = "467b6867c6726cc5484be3cbddbf55c30c0a71594a20c1ac28d35b5049632444";
+    const boxId4 = "a2c9821f5c2df9c320f17136f043b33f7716713ab74c84d687885f9dd39d2c8a";
+
+    const selector = new BoxSelector(regularBoxesMock)
+      .ensureInclusion((box) => box.boxId === boxId1 || box.boxId === boxId3)
+      .ensureInclusion([boxId1, boxId2, boxId1])
+      .ensureInclusion(boxId4);
+
+    const boxes = selector.select({ nanoErgs: 10000n });
+
+    expect(boxes).toHaveLength(4);
+    expect(boxes[0].boxId).toBe(boxId1);
+    expect(boxes[1].boxId).toBe(boxId4);
+    expect(boxes[2].boxId).toBe(boxId3);
+    expect(boxes[3].boxId).toBe(boxId2);
+  });
+
+  it("Should forcedly include inputs by boxId and ignore duplicates", () => {
+    const selector = new BoxSelector(regularBoxesMock)
+      .ensureInclusion("e56847ed19b3dc6b72828fcfb992fdf7310828cf291221269b7ffc72fd66706e")
+      .ensureInclusion("2555e34138d276905fe0bc19240bbeca10f388a71f7b4d2f65a7d0bfd23c846d")
+      .ensureInclusion("2555e34138d276905fe0bc19240bbeca10f388a71f7b4d2f65a7d0bfd23c846d");
+
+    const boxes = selector.select({ nanoErgs: 10000n });
+
+    expect(boxes).toHaveLength(2);
+    expect(boxes[0].boxId).toBe("e56847ed19b3dc6b72828fcfb992fdf7310828cf291221269b7ffc72fd66706e");
+    expect(boxes[1].boxId).toBe("2555e34138d276905fe0bc19240bbeca10f388a71f7b4d2f65a7d0bfd23c846d");
+  });
+
   it("Should forcedly include inputs that attends to filter criteria and collect additional inputs until target is reached", () => {
     const arbitraryBoxId = "2555e34138d276905fe0bc19240bbeca10f388a71f7b4d2f65a7d0bfd23c846d";
     const tokenId = "0cd8c9f416e5b1ca9f986a7f10a84191dfb85941619e49e53c0dc30ebf83324b";

diff --git a/packages/core/src/builder/selector/boxSelector.ts b/packages/core/src/builder/selector/boxSelector.ts
@@ -2,9 +2,11 @@ import {
   Amount,
   Box,
   BoxCandidate,
+  BoxId,
   FilterPredicate,
   first,
   isUndefined,
+  OneOrMore,
   SortingDirection,
   SortingSelector,
   TokenTargetAmount
@@ -34,6 +36,7 @@ export class BoxSelector<T extends Box<bigint>> {
   private _ensureFilterPredicate?: FilterPredicate<Box<bigint>>;
   private _inputsSortSelector?: SortingSelector<Box<bigint>>;
   private _inputsSortDir?: SortingDirection;
+  private _ensureInclusionBoxIds?: Set<BoxId>;
 
   constructor(inputs: T[]) {
     this._inputs = inputs;
@@ -56,12 +59,23 @@ export class BoxSelector<T extends Box<bigint>> {
 
     const remaining = this._deepCloneTarget(target);
     let unselected = [...this._inputs];
-    let selected!: Box<bigint>[];
+    let selected: Box<bigint>[] = [];
 
-    if (isDefined(this._ensureFilterPredicate)) {
-      const predicate = this._ensureFilterPredicate;
-      selected = unselected.filter(predicate);
-      unselected = unselected.filter((input) => !predicate(input));
+    const predicate = this._ensureFilterPredicate;
+    const inclusion = this._ensureInclusionBoxIds;
+
+    if (isDefined(predicate)) {
+      if (isDefined(inclusion)) {
+        selected = unselected.filter((box) => predicate(box) || inclusion.has(box.boxId));
+      } else {
+        selected = unselected.filter(predicate);
+      }
+    } else if (isDefined(inclusion)) {
+      selected = unselected.filter((box) => inclusion.has(box.boxId));
+    }
+
+    if (isDefined(selected)) {
+      unselected = unselected.filter((box) => !selected.some((sel) => sel.boxId === box.boxId));
 
       if (isDefined(remaining.nanoErgs)) {
         remaining.nanoErgs -= sumBy(selected, (input) => input.value);
@@ -74,8 +88,6 @@ export class BoxSelector<T extends Box<bigint>> {
           }
         }
       }
-    } else {
-      selected = [];
     }
 
     unselected = this._sort(unselected);
@@ -143,8 +155,26 @@ export class BoxSelector<T extends Box<bigint>> {
     return orderBy(inputs, this._inputsSortSelector, this._inputsSortDir || "asc");
   }
 
-  public ensureInclusion(predicate: FilterPredicate<Box<bigint>>): BoxSelector<T> {
-    this._ensureFilterPredicate = predicate;
+  public ensureInclusion(predicate: FilterPredicate<Box<bigint>>): BoxSelector<T>;
+  public ensureInclusion(boxIds: OneOrMore<BoxId>): BoxSelector<T>;
+  public ensureInclusion(
+    predicateOrBoxIds: FilterPredicate<Box<bigint>> | OneOrMore<BoxId>
+  ): BoxSelector<T> {
+    if (typeof predicateOrBoxIds === "function") {
+      this._ensureFilterPredicate = predicateOrBoxIds;
+    } else {
+      if (isUndefined(this._ensureInclusionBoxIds)) {
+        this._ensureInclusionBoxIds = new Set();
+      }
+
+      if (Array.isArray(predicateOrBoxIds)) {
+        for (const boxId of predicateOrBoxIds) {
+          this._ensureInclusionBoxIds.add(boxId);
+        }
+      } else {
+        this._ensureInclusionBoxIds.add(predicateOrBoxIds);
+      }
+    }
 
     return this;
   }