Skip to content

Commit 5807d1a

Browse files
EunovoEunovo
committed
Add huffman taptree constructor
1 parent 52559f8 commit 5807d1a

11 files changed

+638
-4
lines changed

‎src/psbt/bip371.d.ts

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

‎src/psbt/bip371.js

+31
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
'use strict';
22
Object.defineProperty(exports, '__esModule', { value: true });
33
exports.checkTaprootInputForSigs =
4+
exports.tapTreeUsingHuffmanConstructor =
45
exports.tapTreeFromList =
56
exports.tapTreeToList =
67
exports.tweakInternalPubKey =
@@ -18,6 +19,7 @@ const psbtutils_1 = require('./psbtutils');
1819
const bip341_1 = require('../payments/bip341');
1920
const payments_1 = require('../payments');
2021
const psbtutils_2 = require('./psbtutils');
22+
const sorted_list_1 = require('../sorted_list');
2123
const toXOnly = pubKey => (pubKey.length === 32 ? pubKey : pubKey.slice(1, 33));
2224
exports.toXOnly = toXOnly;
2325
/**
@@ -155,6 +157,35 @@ function tapTreeFromList(leaves = []) {
155157
return instertLeavesInTree(leaves);
156158
}
157159
exports.tapTreeFromList = tapTreeFromList;
160+
/**
161+
* Construct a Taptree where the leaves with the highest likelihood of use are closer to the root.
162+
* @param nodes A list of nodes where each element contains a weight (likelihood of use) and
163+
* a node which could be a Tapleaf or a branch in a Taptree
164+
*/
165+
function tapTreeUsingHuffmanConstructor(nodes) {
166+
const sortedNodes = new sorted_list_1.SortedList(
167+
nodes,
168+
(a, b) => a.weight - b.weight,
169+
); // Create a list sorted in ascending order of weight
170+
let newNode;
171+
let nodeA, nodeB;
172+
while (sortedNodes.length() > 1) {
173+
// Construct a new node from the two nodes with the least weight
174+
nodeA = sortedNodes.pop(); // There will always be an element to pop
175+
nodeB = sortedNodes.pop(); // because loop ends when length <= 1
176+
newNode = {
177+
weight: nodeA.weight + nodeB.weight,
178+
node: [nodeA.node, nodeB.node],
179+
};
180+
// Place newNode back into sorted list
181+
sortedNodes.insert(newNode);
182+
}
183+
// Last node in sortedNodes is the root node
184+
const root = sortedNodes.pop();
185+
if (!root) throw new Error('Cannot create taptree from empty list.');
186+
return root.node;
187+
}
188+
exports.tapTreeUsingHuffmanConstructor = tapTreeUsingHuffmanConstructor;
158189
function checkTaprootInputForSigs(input, action) {
159190
const sigs = extractTaprootSigs(input);
160191
return sigs.some(sig =>

‎src/sorted_list.d.ts

+9
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
export declare class SortedList<T> {
2+
private array;
3+
private compare;
4+
constructor(array: Array<T>, compare: (a: T, b: T) => number);
5+
pop(): T | undefined;
6+
insert(element: T): number;
7+
length(): number;
8+
toArray(): T[];
9+
}

‎src/sorted_list.js

+66
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
'use strict';
2+
Object.defineProperty(exports, '__esModule', { value: true });
3+
exports.SortedList = void 0;
4+
class SortedList {
5+
constructor(array, compare) {
6+
this.array = array;
7+
this.compare = compare;
8+
this.array.sort(compare);
9+
}
10+
pop() {
11+
return this.array.shift();
12+
}
13+
insert(element) {
14+
let high = this.array.length - 1;
15+
let low = 0;
16+
let mid;
17+
let highElement, lowElement, midElement;
18+
let compareHigh, compareLow, compareMid;
19+
let targetIndex;
20+
while (targetIndex === undefined) {
21+
if (high < low) {
22+
targetIndex = low;
23+
continue;
24+
}
25+
mid = Math.floor((low + high) / 2);
26+
highElement = this.array[high];
27+
lowElement = this.array[low];
28+
midElement = this.array[mid];
29+
compareHigh = this.compare(element, highElement);
30+
compareLow = this.compare(element, lowElement);
31+
compareMid = this.compare(element, midElement);
32+
if (low === high) {
33+
// Target index is either to the left or right of element at low
34+
if (compareLow <= 0) targetIndex = low;
35+
else targetIndex = low + 1;
36+
continue;
37+
}
38+
if (compareHigh >= 0) {
39+
// Target index is to the right of high
40+
low = high;
41+
continue;
42+
}
43+
if (compareLow <= 0) {
44+
// Target index is to the left of low
45+
high = low;
46+
continue;
47+
}
48+
if (compareMid <= 0) {
49+
// Target index is to the left of mid
50+
high = mid;
51+
continue;
52+
}
53+
// Target index is to the right of mid
54+
low = mid + 1;
55+
}
56+
this.array.splice(targetIndex, 0, element);
57+
return targetIndex;
58+
}
59+
length() {
60+
return this.array.length;
61+
}
62+
toArray() {
63+
return this.array;
64+
}
65+
}
66+
exports.SortedList = SortedList;

‎src/types.d.ts

+7
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,13 @@ export declare function isTapleaf(o: any): o is Tapleaf;
2626
* The tree has no balancing requirements.
2727
*/
2828
export type Taptree = [Taptree | Tapleaf, Taptree | Tapleaf] | Tapleaf;
29+
export interface HuffmanTapTreeNode {
30+
/**
31+
* weight is the sum of the weight of all children under this node
32+
*/
33+
weight: number;
34+
node: Taptree;
35+
}
2936
export declare function isTaptree(scriptTree: any): scriptTree is Taptree;
3037
export interface TinySecp256k1Interface {
3138
isXOnlyPoint(p: Uint8Array): boolean;

‎test/huffman.spec.ts

+229
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,229 @@
1+
import * as assert from 'assert';
2+
import { describe, it } from 'mocha';
3+
import { HuffmanTapTreeNode, Taptree } from '../src/types';
4+
import { tapTreeUsingHuffmanConstructor } from '../src/psbt/bip371';
5+
6+
describe('Taptree using Huffman Constructor', () => {
7+
const scriptBuff = Buffer.from('');
8+
9+
it('test empty array', () => {
10+
assert.throws(
11+
() => tapTreeUsingHuffmanConstructor([]),
12+
'Cannot create taptree from empty list.',
13+
);
14+
});
15+
16+
it(
17+
'should return only one node for a single leaf',
18+
testLeafDistances([{ weight: 1, node: { output: scriptBuff } }], [0]),
19+
);
20+
21+
it(
22+
'should return a balanced tree for a list of scripts with equal weights',
23+
testLeafDistances(
24+
[
25+
{
26+
weight: 1,
27+
node: {
28+
output: scriptBuff,
29+
},
30+
},
31+
{
32+
weight: 1,
33+
node: {
34+
output: scriptBuff,
35+
},
36+
},
37+
{
38+
weight: 1,
39+
node: {
40+
output: scriptBuff,
41+
},
42+
},
43+
{
44+
weight: 1,
45+
node: {
46+
output: scriptBuff,
47+
},
48+
},
49+
],
50+
[2, 2, 2, 2],
51+
),
52+
);
53+
54+
it(
55+
'should return an optimal binary tree for a list of scripts with weights [1, 2, 3, 4, 5]',
56+
testLeafDistances(
57+
[
58+
{
59+
weight: 1,
60+
node: {
61+
output: scriptBuff,
62+
},
63+
},
64+
{
65+
weight: 2,
66+
node: {
67+
output: scriptBuff,
68+
},
69+
},
70+
{
71+
weight: 3,
72+
node: {
73+
output: scriptBuff,
74+
},
75+
},
76+
{
77+
weight: 4,
78+
node: {
79+
output: scriptBuff,
80+
},
81+
},
82+
{
83+
weight: 5,
84+
node: {
85+
output: scriptBuff,
86+
},
87+
},
88+
],
89+
[3, 3, 2, 2, 2],
90+
),
91+
);
92+
93+
it(
94+
'should return an optimal binary tree for a list of scripts with weights [1, 2, 3, 3]',
95+
testLeafDistances(
96+
[
97+
{
98+
weight: 1,
99+
node: {
100+
output: scriptBuff,
101+
},
102+
},
103+
{
104+
weight: 2,
105+
node: {
106+
output: scriptBuff,
107+
},
108+
},
109+
{
110+
weight: 3,
111+
node: {
112+
output: scriptBuff,
113+
},
114+
},
115+
{
116+
weight: 3,
117+
node: {
118+
output: scriptBuff,
119+
},
120+
},
121+
],
122+
[3, 3, 2, 1],
123+
),
124+
);
125+
126+
it(
127+
'should return an optimal binary tree for a list of scripts with some negative weights: [1, 2, 3, -3]',
128+
testLeafDistances(
129+
[
130+
{
131+
weight: 1,
132+
node: {
133+
output: scriptBuff,
134+
},
135+
},
136+
{
137+
weight: 2,
138+
node: {
139+
output: scriptBuff,
140+
},
141+
},
142+
{
143+
weight: 3,
144+
node: {
145+
output: scriptBuff,
146+
},
147+
},
148+
{
149+
weight: -3,
150+
node: {
151+
output: scriptBuff,
152+
},
153+
},
154+
],
155+
[3, 2, 1, 3],
156+
),
157+
);
158+
159+
it(
160+
'should return an optimal binary tree for a list of scripts with some weights specified as infinity',
161+
testLeafDistances(
162+
[
163+
{
164+
weight: 1,
165+
node: {
166+
output: scriptBuff,
167+
},
168+
},
169+
{
170+
weight: Number.POSITIVE_INFINITY,
171+
node: {
172+
output: scriptBuff,
173+
},
174+
},
175+
{
176+
weight: 3,
177+
node: {
178+
output: scriptBuff,
179+
},
180+
},
181+
{
182+
weight: Number.NEGATIVE_INFINITY,
183+
node: {
184+
output: scriptBuff,
185+
},
186+
},
187+
],
188+
[3, 1, 2, 3],
189+
),
190+
);
191+
});
192+
193+
function testLeafDistances(
194+
input: HuffmanTapTreeNode[],
195+
expectedDistances: number[],
196+
) {
197+
return () => {
198+
const tree = tapTreeUsingHuffmanConstructor(input);
199+
200+
if (!Array.isArray(tree)) {
201+
// tree is just one node
202+
assert.deepEqual([0], expectedDistances);
203+
return;
204+
}
205+
206+
const leaves = input.map(value => value.node);
207+
208+
const map = new Map<Taptree, number>(); // Map of leaf to actual distance
209+
let currentDistance = 1;
210+
let currentArray: Array<Taptree[] | Taptree> = tree as any;
211+
let nextArray: Array<Taptree[] | Taptree> = [];
212+
while (currentArray.length > 0) {
213+
currentArray.forEach(value => {
214+
if (Array.isArray(value)) {
215+
nextArray = nextArray.concat(value);
216+
return;
217+
}
218+
map.set(value, currentDistance);
219+
});
220+
221+
currentDistance += 1; // New level
222+
currentArray = nextArray;
223+
nextArray = [];
224+
}
225+
226+
const actualDistances = leaves.map(value => map.get(value));
227+
assert.deepEqual(actualDistances, expectedDistances);
228+
};
229+
}

0 commit comments

Comments
 (0)