add ensureInclusion param

fleet-sdk · Sep 15, 2024 · 85f049d · 85f049d
1 parent dd7e43f
commit 85f049d
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 12 deletions.
diff --git a/packages/core/src/builder/pluginContext.ts b/packages/core/src/builder/pluginContext.ts
@@ -58,12 +58,8 @@ export function createPluginContext(
   return {
     addInputs: (inputs) =>
       transactionBuilder
-        .from(inputs)
-        .configureSelector((selector) =>
-          selector.ensureInclusion(
-            Array.isArray(inputs) ? inputs.map((input) => input.boxId) : inputs.boxId
-          )
-        ).inputs.length,
+        .from(inputs, { ensureInclusion: true })
+        .inputs.length,
     addOutputs: (outputs, options) =>
       transactionBuilder.to(outputs, options).outputs.length,
     addDataInputs: (dataInputs, options) =>

diff --git a/packages/core/src/builder/transactionBuilder.ts b/packages/core/src/builder/transactionBuilder.ts
@@ -72,12 +72,12 @@ type InputsInclusionOptions = {
 
 export class TransactionBuilder {
   private readonly _inputs!: InputsCollection;
-  private readonly _preservedInputs!: InputsCollection;
   private readonly _dataInputs!: InputsCollection;
   private readonly _outputs!: OutputsCollection;
   private readonly _settings!: TransactionBuilderSettings;
   private readonly _creationHeight!: number;
 
+  private _ensureInclusion?: Set<string>;
   private _selectorCallbacks?: SelectorCallback[];
   private _changeAddress?: ErgoAddress;
   private _feeAmount?: bigint;
@@ -143,15 +143,22 @@ export class TransactionBuilder {
     options: InputsInclusionOptions = { ensureInclusion: false }
   ): TransactionBuilder {
     const items = isCollectionLike(inputs) ? inputs.toArray() : inputs;
-    if (options.ensureInclusion) {
-      this._preservedInputs.add(items);
-    } else {
-      this._inputs.add(items);
-    }
+    if (options.ensureInclusion) this.#ensureInclusionOf(items);
 
+    this._inputs.add(items);
     return this;
   }
 
+  #ensureInclusionOf(inputs: OneOrMore<Box<Amount>>): void {
+    if (!this._ensureInclusion) this._ensureInclusion = new Set();
+
+    if (Array.isArray(inputs)) {
+      for (const input of inputs) this._ensureInclusion.add(input.boxId);
+    } else {
+      this._ensureInclusion.add(inputs.boxId);
+    }
+  }
+
   public to(
     outputs: OneOrMore<OutputBuilder>,
     options?: CollectionAddOptions
@@ -274,6 +281,10 @@ export class TransactionBuilder {
     }
 
     const selector = new BoxSelector(this.inputs.toArray());
+    if (this._ensureInclusion?.size) { 
+      selector.ensureInclusion(Array.from(this._ensureInclusion));
+    }
+
     if (some(this._selectorCallbacks)) {
       for (const selectorCallBack of this._selectorCallbacks) {
         selectorCallBack(selector);