Add huffman taptree constructor

src/psbt/bip371.d.ts

@@ -1,5 +1,5 @@
 /// <reference types="node" />
-import { Taptree } from '../types';
+import { Taptree, HuffmanTapTreeNode } from '../types';
 import { PsbtInput, PsbtOutput, TapLeaf } from 'bip174/src/lib/interfaces';
 export declare const toXOnly: (pubKey: Buffer) => Buffer;
 /**
@@ -38,4 +38,10 @@ export declare function tapTreeToList(tree: Taptree): TapLeaf[];
  * @returns the corresponding taptree, or throws an exception if the tree cannot be reconstructed
  */
 export declare function tapTreeFromList(leaves?: TapLeaf[]): Taptree;
+/**
+ * Construct a Taptree where the leaves with the highest likelihood of use are closer to the root.
+ * @param nodes A list of nodes where each element contains a weight (likelihood of use) and
+ * a node which could be a Tapleaf or a branch in a Taptree
+ */
+export declare function tapTreeUsingHuffmanConstructor(nodes: HuffmanTapTreeNode[]): Taptree;
 export declare function checkTaprootInputForSigs(input: PsbtInput, action: string): boolean;

src/psbt/bip371.js

@@ -1,6 +1,7 @@
 'use strict';
 Object.defineProperty(exports, '__esModule', { value: true });
 exports.checkTaprootInputForSigs =
+  exports.tapTreeUsingHuffmanConstructor =
   exports.tapTreeFromList =
   exports.tapTreeToList =
   exports.tweakInternalPubKey =
@@ -18,6 +19,7 @@ const psbtutils_1 = require('./psbtutils');
 const bip341_1 = require('../payments/bip341');
 const payments_1 = require('../payments');
 const psbtutils_2 = require('./psbtutils');
+const sorted_list_1 = require('../sorted_list');
 const toXOnly = pubKey => (pubKey.length === 32 ? pubKey : pubKey.slice(1, 33));
 exports.toXOnly = toXOnly;
 /**
@@ -155,6 +157,35 @@ function tapTreeFromList(leaves = []) {
   return instertLeavesInTree(leaves);
 }
 exports.tapTreeFromList = tapTreeFromList;
+/**
+ * Construct a Taptree where the leaves with the highest likelihood of use are closer to the root.
+ * @param nodes A list of nodes where each element contains a weight (likelihood of use) and
+ * a node which could be a Tapleaf or a branch in a Taptree
+ */
+function tapTreeUsingHuffmanConstructor(nodes) {
+  const sortedNodes = new sorted_list_1.SortedList(
+    nodes,
+    (a, b) => a.weight - b.weight,
+  ); // Create a list sorted in ascending order of weight
+  let newNode;
+  let nodeA, nodeB;
+  while (sortedNodes.length() > 1) {
+    // Construct a new node from the two nodes with the least weight
+    nodeA = sortedNodes.pop(); // There will always be an element to pop
+    nodeB = sortedNodes.pop(); // because loop ends when length <= 1
+    newNode = {
+      weight: nodeA.weight + nodeB.weight,
+      node: [nodeA.node, nodeB.node],
+    };
+    // Place newNode back into sorted list
+    sortedNodes.insert(newNode);
+  }
+  // Last node in sortedNodes is the root node
+  const root = sortedNodes.pop();
+  if (!root) throw new Error('Cannot create taptree from empty list.');
+  return root.node;
+}
+exports.tapTreeUsingHuffmanConstructor = tapTreeUsingHuffmanConstructor;
 function checkTaprootInputForSigs(input, action) {
   const sigs = extractTaprootSigs(input);
   return sigs.some(sig =>

src/sorted_list.d.ts

@@ -0,0 +1,9 @@
+export declare class SortedList<T> {
+    private array;
+    private compare;
+    constructor(array: Array<T>, compare: (a: T, b: T) => number);
+    pop(): T | undefined;
+    insert(element: T): number;
+    length(): number;
+    toArray(): T[];
+}

src/sorted_list.js

@@ -0,0 +1,66 @@
+'use strict';
+Object.defineProperty(exports, '__esModule', { value: true });
+exports.SortedList = void 0;
+class SortedList {
+  constructor(array, compare) {
+    this.array = array;
+    this.compare = compare;
+    this.array.sort(compare);
+  }
+  pop() {
+    return this.array.shift();
+  }
+  insert(element) {
+    let high = this.array.length - 1;
+    let low = 0;
+    let mid;
+    let highElement, lowElement, midElement;
+    let compareHigh, compareLow, compareMid;
+    let targetIndex;
+    while (targetIndex === undefined) {
+      if (high < low) {
+        targetIndex = low;
+        continue;
+      }
+      mid = Math.floor((low + high) / 2);
+      highElement = this.array[high];
+      lowElement = this.array[low];
+      midElement = this.array[mid];
+      compareHigh = this.compare(element, highElement);
+      compareLow = this.compare(element, lowElement);
+      compareMid = this.compare(element, midElement);
+      if (low === high) {
+        // Target index is either to the left or right of element at low
+        if (compareLow <= 0) targetIndex = low;
+        else targetIndex = low + 1;
+        continue;
+      }
+      if (compareHigh >= 0) {
+        // Target index is to the right of high
+        low = high;
+        continue;
+      }
+      if (compareLow <= 0) {
+        // Target index is to the left of low
+        high = low;
+        continue;
+      }
+      if (compareMid <= 0) {
+        // Target index is to the left of mid
+        high = mid;
+        continue;
+      }
+      // Target index is to the right of mid
+      low = mid + 1;
+    }
+    this.array.splice(targetIndex, 0, element);
+    return targetIndex;
+  }
+  length() {
+    return this.array.length;
+  }
+  toArray() {
+    return this.array;
+  }
+}
+exports.SortedList = SortedList;

src/types.d.ts

@@ -26,6 +26,13 @@ export declare function isTapleaf(o: any): o is Tapleaf;
  * The tree has no balancing requirements.
  */
 export type Taptree = [Taptree | Tapleaf, Taptree | Tapleaf] | Tapleaf;
+export interface HuffmanTapTreeNode {
+    /**
+     * weight is the sum of the weight of all children under this node
+     */
+    weight: number;
+    node: Taptree;
+}
 export declare function isTaptree(scriptTree: any): scriptTree is Taptree;
 export interface TinySecp256k1Interface {
     isXOnlyPoint(p: Uint8Array): boolean;

test/huffman.spec.ts

@@ -0,0 +1,229 @@
+import * as assert from 'assert';
+import { describe, it } from 'mocha';
+import { HuffmanTapTreeNode, Taptree } from '../src/types';
+import { tapTreeUsingHuffmanConstructor } from '../src/psbt/bip371';
+
+describe('Taptree using Huffman Constructor', () => {
+  const scriptBuff = Buffer.from('');
+
+  it('test empty array', () => {
+    assert.throws(
+      () => tapTreeUsingHuffmanConstructor([]),
+      'Cannot create taptree from empty list.',
+    );
+  });
+
+  it(
+    'should return only one node for a single leaf',
+    testLeafDistances([{ weight: 1, node: { output: scriptBuff } }], [0]),
+  );
+
+  it(
+    'should return a balanced tree for a list of scripts with equal weights',
+    testLeafDistances(
+      [
+        {
+          weight: 1,
+          node: {
+            output: scriptBuff,
+          },
+        },
+        {
+          weight: 1,
+          node: {
+            output: scriptBuff,
+          },
+        },
+        {
+          weight: 1,
+          node: {
+            output: scriptBuff,
+          },
+        },
+        {
+          weight: 1,
+          node: {
+            output: scriptBuff,
+          },
+        },
+      ],
+      [2, 2, 2, 2],
+    ),
+  );
+
+  it(
+    'should return an optimal binary tree for a list of scripts with weights [1, 2, 3, 4, 5]',
+    testLeafDistances(
+      [
+        {
+          weight: 1,
+          node: {
+            output: scriptBuff,
+          },
+        },
+        {
+          weight: 2,
+          node: {
+            output: scriptBuff,
+          },
+        },
+        {
+          weight: 3,
+          node: {
+            output: scriptBuff,
+          },
+        },
+        {
+          weight: 4,
+          node: {
+            output: scriptBuff,
+          },
+        },
+        {
+          weight: 5,
+          node: {
+            output: scriptBuff,
+          },
+        },
+      ],
+      [3, 3, 2, 2, 2],
+    ),
+  );
+
+  it(
+    'should return an optimal binary tree for a list of scripts with weights [1, 2, 3, 3]',
+    testLeafDistances(
+      [
+        {
+          weight: 1,
+          node: {
+            output: scriptBuff,
+          },
+        },
+        {
+          weight: 2,
+          node: {
+            output: scriptBuff,
+          },
+        },
+        {
+          weight: 3,
+          node: {
+            output: scriptBuff,
+          },
+        },
+        {
+          weight: 3,
+          node: {
+            output: scriptBuff,
+          },
+        },
+      ],
+      [3, 3, 2, 1],
+    ),
+  );
+
+  it(
+    'should return an optimal binary tree for a list of scripts with some negative weights: [1, 2, 3, -3]',
+    testLeafDistances(
+      [
+        {
+          weight: 1,
+          node: {
+            output: scriptBuff,
+          },
+        },
+        {
+          weight: 2,
+          node: {
+            output: scriptBuff,
+          },
+        },
+        {
+          weight: 3,
+          node: {
+            output: scriptBuff,
+          },
+        },
+        {
+          weight: -3,
+          node: {
+            output: scriptBuff,
+          },
+        },
+      ],
+      [3, 2, 1, 3],
+    ),
+  );
+
+  it(
+    'should return an optimal binary tree for a list of scripts with some weights specified as infinity',
+    testLeafDistances(
+      [
+        {
+          weight: 1,
+          node: {
+            output: scriptBuff,
+          },
+        },
+        {
+          weight: Number.POSITIVE_INFINITY,
+          node: {
+            output: scriptBuff,
+          },
+        },
+        {
+          weight: 3,
+          node: {
+            output: scriptBuff,
+          },
+        },
+        {
+          weight: Number.NEGATIVE_INFINITY,
+          node: {
+            output: scriptBuff,
+          },
+        },
+      ],
+      [3, 1, 2, 3],
+    ),
+  );
+});
+
+function testLeafDistances(
+  input: HuffmanTapTreeNode[],
+  expectedDistances: number[],
+) {
+  return () => {
+    const tree = tapTreeUsingHuffmanConstructor(input);
+
+    if (!Array.isArray(tree)) {
+      // tree is just one node
+      assert.deepEqual([0], expectedDistances);
+      return;
+    }
+
+    const leaves = input.map(value => value.node);
+
+    const map = new Map<Taptree, number>(); // Map of leaf to actual distance
+    let currentDistance = 1;
+    let currentArray: Array<Taptree[] | Taptree> = tree as any;
+    let nextArray: Array<Taptree[] | Taptree> = [];
+    while (currentArray.length > 0) {
+      currentArray.forEach(value => {
+        if (Array.isArray(value)) {
+          nextArray = nextArray.concat(value);
+          return;
+        }
+        map.set(value, currentDistance);
+      });
+
+      currentDistance += 1; // New level
+      currentArray = nextArray;
+      nextArray = [];
+    }
+
+    const actualDistances = leaves.map(value => map.get(value));
+    assert.deepEqual(actualDistances, expectedDistances);
+  };
+}