feat: add to/from Psbt TapTree conversion

motorina0 · motorina0 · commit 6cbac53c1d03 · 2022-11-29T10:41:17.000+02:00
diff --git a/src/payments/taprootutils.d.ts b/src/payments/taprootutils.d.ts
@@ -1,6 +1,7 @@
 /// <reference types="node" />
 import { Tapleaf, Taptree } from '../types';
 export declare const LEAF_VERSION_TAPSCRIPT = 192;
+export declare const MAX_TAPTREE_DEPTH = 128;
 interface HashLeaf {
     hash: Buffer;
 }
diff --git a/src/payments/taprootutils.js b/src/payments/taprootutils.js
@@ -1,12 +1,13 @@
 'use strict';
 Object.defineProperty(exports, '__esModule', { value: true });
-exports.tweakKey = exports.tapTweakHash = exports.tapleafHash = exports.findScriptPath = exports.toHashTree = exports.rootHashFromPath = exports.LEAF_VERSION_TAPSCRIPT = void 0;
+exports.tweakKey = exports.tapTweakHash = exports.tapleafHash = exports.findScriptPath = exports.toHashTree = exports.rootHashFromPath = exports.MAX_TAPTREE_DEPTH = exports.LEAF_VERSION_TAPSCRIPT = void 0;
 const buffer_1 = require('buffer');
 const ecc_lib_1 = require('../ecc_lib');
 const bcrypto = require('../crypto');
 const bufferutils_1 = require('../bufferutils');
 const types_1 = require('../types');
 exports.LEAF_VERSION_TAPSCRIPT = 0xc0;
+exports.MAX_TAPTREE_DEPTH = 128;
 const isHashBranch = ht => 'left' in ht && 'right' in ht;
 function rootHashFromPath(controlBlock, leafHash) {
   const m = (controlBlock.length - 33) / 32;
diff --git a/src/psbt/bip371.d.ts b/src/psbt/bip371.d.ts
@@ -1,5 +1,6 @@
 /// <reference types="node" />
-import { PsbtInput } from 'bip174/src/lib/interfaces';
+import { Taptree } from '../types';
+import { PsbtInput, TapLeaf } from 'bip174/src/lib/interfaces';
 export declare const toXOnly: (pubKey: Buffer) => Buffer;
 /**
  * Default tapscript finalizer. It searches for the `tapLeafHashToFinalize` if provided.
@@ -17,3 +18,21 @@ export declare function serializeTaprootSignature(sig: Buffer, sighashType?: num
 export declare function isTaprootInput(input: PsbtInput): boolean;
 export declare function checkTaprootInputFields(inputData: PsbtInput, newInputData: PsbtInput, action: string): void;
 export declare function tweakInternalPubKey(inputIndex: number, input: PsbtInput): Buffer;
+/**
+ * Convert a binary tree to a BIP371 type list. Each element of the list is (according to BIP371):
+ * One or more tuples representing the depth, leaf version, and script for a leaf in the Taproot tree,
+ * allowing the entire tree to be reconstructed. The tuples must be in depth first search order so that
+ * the tree is correctly reconstructed.
+ * @param tree the binary tap tree
+ * @returns a list of BIP 371 tapleaves
+ */
+export declare function tapTreeToList(tree: Taptree): TapLeaf[];
+/**
+ * Convert a BIP371 TapLeaf list to a TapTree (binary).
+ * @param leaves a list of tapleaves where each element of the list is (according to BIP371):
+ * One or more tuples representing the depth, leaf version, and script for a leaf in the Taproot tree,
+ * allowing the entire tree to be reconstructed. The tuples must be in depth first search order so that
+ * the tree is correctly reconstructed.
+ * @returns the corresponding taptree, or throws an exception if the tree cannot be reconstructed
+ */
+export declare function tapTreeFromList(leaves?: TapLeaf[]): Taptree;
diff --git a/src/psbt/bip371.js b/src/psbt/bip371.js
@@ -1,6 +1,7 @@
 'use strict';
 Object.defineProperty(exports, '__esModule', { value: true });
-exports.tweakInternalPubKey = exports.checkTaprootInputFields = exports.isTaprootInput = exports.serializeTaprootSignature = exports.tapScriptFinalizer = exports.toXOnly = void 0;
+exports.tapTreeFromList = exports.tapTreeToList = exports.tweakInternalPubKey = exports.checkTaprootInputFields = exports.isTaprootInput = exports.serializeTaprootSignature = exports.tapScriptFinalizer = exports.toXOnly = void 0;
+const types_1 = require('../types');
 const psbtutils_1 = require('./psbtutils');
 const taprootutils_1 = require('../payments/taprootutils');
 const toXOnly = pubKey => (pubKey.length === 32 ? pubKey : pubKey.slice(1, 33));
@@ -69,6 +70,76 @@ function tweakInternalPubKey(inputIndex, input) {
   return outputKey.x;
 }
 exports.tweakInternalPubKey = tweakInternalPubKey;
+/**
+ * Convert a binary tree to a BIP371 type list. Each element of the list is (according to BIP371):
+ * One or more tuples representing the depth, leaf version, and script for a leaf in the Taproot tree,
+ * allowing the entire tree to be reconstructed. The tuples must be in depth first search order so that
+ * the tree is correctly reconstructed.
+ * @param tree the binary tap tree
+ * @returns a list of BIP 371 tapleaves
+ */
+function tapTreeToList(tree) {
+  return _tapTreeToList(tree);
+}
+exports.tapTreeToList = tapTreeToList;
+/**
+ * Convert a BIP371 TapLeaf list to a TapTree (binary).
+ * @param leaves a list of tapleaves where each element of the list is (according to BIP371):
+ * One or more tuples representing the depth, leaf version, and script for a leaf in the Taproot tree,
+ * allowing the entire tree to be reconstructed. The tuples must be in depth first search order so that
+ * the tree is correctly reconstructed.
+ * @returns the corresponding taptree, or throws an exception if the tree cannot be reconstructed
+ */
+function tapTreeFromList(leaves = []) {
+  if (leaves.length === 1 && leaves[0].depth === 0)
+    return {
+      output: leaves[0].script,
+      version: leaves[0].leafVersion,
+    };
+  return instertLeavesInTree(leaves);
+}
+exports.tapTreeFromList = tapTreeFromList;
+function _tapTreeToList(tree, leaves = [], depth = 0) {
+  if (depth > taprootutils_1.MAX_TAPTREE_DEPTH)
+    throw new Error('Max taptree depth exceeded.');
+  if (!tree) return [];
+  if ((0, types_1.isTapleaf)(tree)) {
+    leaves.push({
+      depth,
+      leafVersion: tree.version || taprootutils_1.LEAF_VERSION_TAPSCRIPT,
+      script: tree.output,
+    });
+    return leaves;
+  }
+  if (tree[0]) _tapTreeToList(tree[0], leaves, depth + 1);
+  if (tree[1]) _tapTreeToList(tree[1], leaves, depth + 1);
+  return leaves;
+}
+function instertLeavesInTree(leaves) {
+  let tree;
+  for (const leaf of leaves) {
+    tree = instertLeafInTree(leaf, tree);
+    if (!tree) throw new Error(`No room left to insert tapleaf in tree`);
+  }
+  return tree;
+}
+function instertLeafInTree(leaf, tree, depth = 0) {
+  if (depth > taprootutils_1.MAX_TAPTREE_DEPTH)
+    throw new Error('Max taptree depth exceeded.');
+  if (leaf.depth === depth) {
+    if (!tree)
+      return {
+        output: leaf.script,
+        version: leaf.leafVersion,
+      };
+    return;
+  }
+  if ((0, types_1.isTapleaf)(tree)) return;
+  const leftSide = instertLeafInTree(leaf, tree && tree[0], depth + 1);
+  if (leftSide) return [leftSide, tree && tree[1]];
+  const rightSide = instertLeafInTree(leaf, tree && tree[1], depth + 1);
+  if (rightSide) return [tree && tree[0], rightSide];
+}
 function checkMixedTaprootAndNonTaprootFields(inputData, newInputData, action) {
   const isBadTaprootUpdate =
     isTaprootInput(inputData) && hasNonTaprootInputFields(newInputData);
diff --git a/ts_src/payments/taprootutils.ts b/ts_src/payments/taprootutils.ts
@@ -6,6 +6,7 @@ import { varuint } from '../bufferutils';
 import { Tapleaf, Taptree, isTapleaf } from '../types';
 
 export const LEAF_VERSION_TAPSCRIPT = 0xc0;
+export const MAX_TAPTREE_DEPTH = 128;
 
 interface HashLeaf {
   hash: Buffer;
diff --git a/ts_src/psbt/bip371.ts b/ts_src/psbt/bip371.ts
@@ -1,7 +1,9 @@
+import { Taptree, Tapleaf, isTapleaf } from '../types';
 import {
   PsbtInput,
   TapLeafScript,
   TapScriptSig,
+  TapLeaf,
 } from 'bip174/src/lib/interfaces';
 
 import {
@@ -13,6 +15,8 @@ import {
   tweakKey,
   tapleafHash,
   rootHashFromPath,
+  LEAF_VERSION_TAPSCRIPT,
+  MAX_TAPTREE_DEPTH,
 } from '../payments/taprootutils';
 
 export const toXOnly = (pubKey: Buffer) =>
@@ -98,6 +102,94 @@ export function tweakInternalPubKey(
   return outputKey.x;
 }
 
+/**
+ * Convert a binary tree to a BIP371 type list. Each element of the list is (according to BIP371):
+ * One or more tuples representing the depth, leaf version, and script for a leaf in the Taproot tree,
+ * allowing the entire tree to be reconstructed. The tuples must be in depth first search order so that
+ * the tree is correctly reconstructed.
+ * @param tree the binary tap tree
+ * @returns a list of BIP 371 tapleaves
+ */
+export function tapTreeToList(tree: Taptree): TapLeaf[] {
+  return _tapTreeToList(tree);
+}
+
+/**
+ * Convert a BIP371 TapLeaf list to a TapTree (binary).
+ * @param leaves a list of tapleaves where each element of the list is (according to BIP371):
+ * One or more tuples representing the depth, leaf version, and script for a leaf in the Taproot tree,
+ * allowing the entire tree to be reconstructed. The tuples must be in depth first search order so that
+ * the tree is correctly reconstructed.
+ * @returns the corresponding taptree, or throws an exception if the tree cannot be reconstructed
+ */
+export function tapTreeFromList(leaves: TapLeaf[] = []): Taptree {
+  if (leaves.length === 1 && leaves[0].depth === 0)
+    return {
+      output: leaves[0].script,
+      version: leaves[0].leafVersion,
+    };
+
+  return instertLeavesInTree(leaves);
+}
+
+function _tapTreeToList(
+  tree: Taptree,
+  leaves: TapLeaf[] = [],
+  depth = 0,
+): TapLeaf[] {
+  if (depth > MAX_TAPTREE_DEPTH) throw new Error('Max taptree depth exceeded.');
+  if (!tree) return [];
+  if (isTapleaf(tree)) {
+    leaves.push({
+      depth,
+      leafVersion: tree.version || LEAF_VERSION_TAPSCRIPT,
+      script: tree.output,
+    });
+    return leaves;
+  }
+  if (tree[0]) _tapTreeToList(tree[0], leaves, depth + 1);
+  if (tree[1]) _tapTreeToList(tree[1], leaves, depth + 1);
+  return leaves;
+}
+
+// Just like Taptree, but it accepts empty branches
+type PartialTaptree =
+  | [PartialTaptree | Tapleaf, PartialTaptree | Tapleaf]
+  | Tapleaf
+  | undefined;
+function instertLeavesInTree(leaves: TapLeaf[]): Taptree {
+  let tree: PartialTaptree;
+  for (const leaf of leaves) {
+    tree = instertLeafInTree(leaf, tree);
+    if (!tree) throw new Error(`No room left to insert tapleaf in tree`);
+  }
+
+  return tree as Taptree;
+}
+
+function instertLeafInTree(
+  leaf: TapLeaf,
+  tree?: PartialTaptree,
+  depth = 0,
+): PartialTaptree {
+  if (depth > MAX_TAPTREE_DEPTH) throw new Error('Max taptree depth exceeded.');
+  if (leaf.depth === depth) {
+    if (!tree)
+      return {
+        output: leaf.script,
+        version: leaf.leafVersion,
+      };
+    return;
+  }
+
+  if (isTapleaf(tree)) return;
+  const leftSide = instertLeafInTree(leaf, tree && tree[0], depth + 1);
+  if (leftSide) return [leftSide, tree && tree[1]];
+
+  const rightSide = instertLeafInTree(leaf, tree && tree[1], depth + 1);
+  if (rightSide) return [tree && tree[0], rightSide];
+}
+
 function checkMixedTaprootAndNonTaprootFields(
   inputData: PsbtInput,
   newInputData: PsbtInput,

Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,7 @@`
`1`	`1`	`/// <reference types="node" />`
`2`	`2`	`import { Tapleaf, Taptree } from '../types';`
`3`	`3`	`export declare const LEAF_VERSION_TAPSCRIPT = 192;`
	`4`	`+export declare const MAX_TAPTREE_DEPTH = 128;`
`4`	`5`	`interface HashLeaf {`
`5`	`6`	`hash: Buffer;`
`6`	`7`	`}`