From 21326c51033063068d058e4aa07118520809e812 Mon Sep 17 00:00:00 2001
From: SChernykh <15806605+SChernykh@users.noreply.github.com>
Date: Tue, 14 May 2024 19:37:38 +0200
Subject: [PATCH] Merkle trees: added `get_position_from_path`

Also double check against Monero's Merkle tree path code.
---
 src/merkle.cpp             | 30 +++++++++++++++++++
 src/merkle.h               |  1 +
 tests/src/merkle_tests.cpp | 60 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 91 insertions(+)

diff --git a/src/merkle.cpp b/src/merkle.cpp
index 40628cb..0be9008 100644
--- a/src/merkle.cpp
+++ b/src/merkle.cpp
@@ -338,4 +338,34 @@ bool find_aux_nonce(const std::vector<hash>& aux_id, uint32_t& nonce, uint32_t m
 	}
 }
 
+size_t get_position_from_path(size_t count, uint32_t path)
+{
+	if (count <= 1) {
+		return 0;
+	}
+
+	size_t depth = 0;
+	size_t k = 1;
+
+	while (k < count) {
+		++depth;
+		k <<= 1;
+	}
+
+	k -= count;
+
+	size_t pos = 0;
+
+	for (size_t i = 1; i < depth; ++i) {
+		pos = (pos << 1) | (path & 1);
+		path >>= 1;
+	}
+
+	if (pos < k) {
+		return pos;
+	}
+
+	return (((pos - k) << 1) | (path & 1)) + k;
+}
+
 } // namespace p2pool
diff --git a/src/merkle.h b/src/merkle.h
index 1f36c62..77b1027 100644
--- a/src/merkle.h
+++ b/src/merkle.h
@@ -31,5 +31,6 @@ bool verify_merkle_proof(hash h, const std::vector<hash>& proof, uint32_t path,
 
 uint32_t get_aux_slot(const hash &id, uint32_t nonce, uint32_t n_aux_chains);
 bool find_aux_nonce(const std::vector<hash>& aux_id, uint32_t& nonce, uint32_t max_nonce = 0xFFFF);
+size_t get_position_from_path(size_t count, uint32_t path);
 
 } // namespace p2pool
diff --git a/tests/src/merkle_tests.cpp b/tests/src/merkle_tests.cpp
index 31ce7a2..53d04c8 100644
--- a/tests/src/merkle_tests.cpp
+++ b/tests/src/merkle_tests.cpp
@@ -24,6 +24,60 @@
 
 namespace p2pool {
 
+// Original Monero's tree_path function to test against
+
+size_t tree_hash_cnt(size_t count) {
+	size_t pow = 2;
+	while(pow < count) pow <<= 1;
+	return pow >> 1;
+}
+
+bool tree_path(size_t count, size_t idx, uint32_t* path)
+{
+	if (count == 0)
+		return false;
+
+	if (count == 1) {
+		*path = 0;
+	}
+	else if (count == 2) {
+		*path = idx == 0 ? 0 : 1;
+	}
+	else {
+		size_t i, j;
+
+		*path = 0;
+		size_t cnt = tree_hash_cnt(count);
+
+		for (i = 2 * cnt - count, j = 2 * cnt - count; j < cnt; i += 2, ++j) {
+			if (idx == i || idx == i + 1)
+			{
+				*path = (*path << 1) | (idx == i ? 0 : 1);
+				idx = j;
+			}
+		}
+		assert(i == count);
+
+		while (cnt > 2) {
+			cnt >>= 1;
+			for (i = 0, j = 0; j < cnt; i += 2, ++j) {
+				if (idx == i || idx == i + 1)
+				{
+					*path = (*path << 1) | (idx == i ? 0 : 1);
+					idx = j;
+				}
+			}
+		}
+
+		if (idx == 0 || idx == 1)
+		{
+			*path = (*path << 1) | (idx == 0 ? 0 : 1);
+			idx = 0;
+		}
+	}
+	return true;
+}
+
 TEST(merkle, tree)
 {
 	hash input[10];
@@ -78,8 +132,14 @@ TEST(merkle, tree)
 			uint32_t path;
 
 			ASSERT_TRUE(get_merkle_proof(tree, h, proof, path));
+
+			uint32_t path_monero;
+			ASSERT_TRUE(tree_path(n, i, &path_monero));
+			ASSERT_EQ(path, path_monero);
+
 			ASSERT_TRUE(verify_merkle_proof(h, proof, i, n, root));
 			ASSERT_TRUE(verify_merkle_proof(h, proof, path, root));
+			ASSERT_EQ(get_position_from_path(n, path), i);
 		}
 	};