From ce129faed3c3914358e5b697a4d1e4d612e711e7 Mon Sep 17 00:00:00 2001
From: michael1011 <me@michael1011.at>
Date: Fri, 22 Mar 2024 15:51:33 +0100
Subject: [PATCH] feat: add SwapTree compare function

---
 lib/Boltz.ts                           |  2 +
 lib/swap/SwapTreeCompare.ts            | 26 +++++++++
 test/unit/swap/SwapTreeCompare.spec.ts | 77 ++++++++++++++++++++++++++
 3 files changed, 105 insertions(+)
 create mode 100644 test/unit/swap/SwapTreeCompare.spec.ts

diff --git a/lib/Boltz.ts b/lib/Boltz.ts
index 8105833..9f220e5 100644
--- a/lib/Boltz.ts
+++ b/lib/Boltz.ts
@@ -28,6 +28,7 @@ import swapTree, {
   extractClaimPublicKeyFromSwapTree,
   extractRefundPublicKeyFromSwapTree,
 } from './swap/SwapTree';
+import { compareTrees } from './swap/SwapTreeCompare';
 import * as SwapTreeSerializer from './swap/SwapTreeSerializer';
 import * as SwapUtils from './swap/SwapUtils';
 import * as TaprootUtils from './swap/TaprootUtils';
@@ -57,6 +58,7 @@ export {
   targetFee,
   swapScript,
   detectSwap,
+  compareTrees,
   detectPreimage,
   reverseSwapTree,
   reverseSwapScript,
diff --git a/lib/swap/SwapTreeCompare.ts b/lib/swap/SwapTreeCompare.ts
index e69de29..f78179d 100644
--- a/lib/swap/SwapTreeCompare.ts
+++ b/lib/swap/SwapTreeCompare.ts
@@ -0,0 +1,26 @@
+import { Taptree } from 'bitcoinjs-lib/src/types';
+import { SwapTree, Tapleaf } from '../consts/Types';
+
+const compareLeaf = (leaf: Tapleaf, compareLeaf: Tapleaf) =>
+  leaf.version === compareLeaf.version &&
+  leaf.output.equals(compareLeaf.output);
+
+const compareTree = (tree: Taptree, compare: Taptree) => {
+  if (Array.isArray(tree) !== Array.isArray(compare)) {
+    return false;
+  }
+
+  if (Array.isArray(tree) && Array.isArray(compare)) {
+    return (
+      tree.length === compare.length &&
+      tree.every((leaf, i) => compareTree(leaf, compare[i]))
+    );
+  } else {
+    return compareLeaf(tree as Tapleaf, compare as Tapleaf);
+  }
+};
+
+export const compareTrees = <T extends SwapTree>(
+  tree: T,
+  compare: T,
+): boolean => compareTree(tree.tree, compare.tree);
diff --git a/test/unit/swap/SwapTreeCompare.spec.ts b/test/unit/swap/SwapTreeCompare.spec.ts
new file mode 100644
index 0000000..5e146b3
--- /dev/null
+++ b/test/unit/swap/SwapTreeCompare.spec.ts
@@ -0,0 +1,77 @@
+import { randomBytes } from 'crypto';
+import swapTree from '../../../lib/swap/SwapTree';
+import { compareTrees } from '../../../lib/swap/SwapTreeCompare';
+import { ECPair } from '../Utils';
+
+describe('SwapTreeCompare', () => {
+  test('should compare SwapTrees', () => {
+    const preimageHash = randomBytes(32);
+    const claimKey = ECPair.makeRandom().publicKey;
+    const refundKey = ECPair.makeRandom().publicKey;
+    const timeoutBlockHeight = 123;
+
+    expect(
+      compareTrees(
+        swapTree(false, preimageHash, claimKey, refundKey, timeoutBlockHeight),
+        swapTree(false, preimageHash, claimKey, refundKey, timeoutBlockHeight),
+      ),
+    ).toEqual(true);
+  });
+
+  test('should return false on version mismatch', () => {
+    const preimageHash = randomBytes(32);
+    const claimKey = ECPair.makeRandom().publicKey;
+    const refundKey = ECPair.makeRandom().publicKey;
+    const timeoutBlockHeight = 123;
+
+    expect(
+      compareTrees(
+        swapTree(true, preimageHash, claimKey, refundKey, timeoutBlockHeight),
+        swapTree(false, preimageHash, claimKey, refundKey, timeoutBlockHeight),
+      ),
+    ).toEqual(false);
+  });
+
+  test('should return false on script mismatch', () => {
+    const preimageHash = randomBytes(32);
+    const claimKey = ECPair.makeRandom().publicKey;
+    const refundKey = ECPair.makeRandom().publicKey;
+    const timeoutBlockHeight = 123;
+
+    expect(
+      compareTrees(
+        swapTree(true, preimageHash, claimKey, refundKey, timeoutBlockHeight),
+        swapTree(
+          true,
+          randomBytes(32),
+          claimKey,
+          refundKey,
+          timeoutBlockHeight,
+        ),
+      ),
+    ).toEqual(false);
+  });
+
+  test('should return false on depth mismatch', () => {
+    const preimageHash = randomBytes(32);
+    const claimKey = ECPair.makeRandom().publicKey;
+    const refundKey = ECPair.makeRandom().publicKey;
+    const timeoutBlockHeight = 123;
+
+    const compare = swapTree(
+      false,
+      preimageHash,
+      claimKey,
+      refundKey,
+      timeoutBlockHeight,
+    );
+    compare.tree = compare.claimLeaf;
+
+    expect(
+      compareTrees(
+        swapTree(false, preimageHash, claimKey, refundKey, timeoutBlockHeight),
+        compare,
+      ),
+    ).toEqual(false);
+  });
+});