From 19d9d31b0cc2f067278c4816323349d33a397d1e Mon Sep 17 00:00:00 2001
From: peilun-conflux <48905552+peilun-conflux@users.noreply.github.com>
Date: Sun, 28 Jan 2024 16:52:28 +0800
Subject: [PATCH] Fix proof data update in corner cases. (#13)

* Fix proof data update in corner cases.

* Fix fmt.
---
 common/append_merkle/src/lib.rs   | 15 ++++++++++-----
 common/append_merkle/src/proof.rs | 20 ++++++++++++++++++--
 2 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/common/append_merkle/src/lib.rs b/common/append_merkle/src/lib.rs
index 52ded90..f9d7fb9 100644
--- a/common/append_merkle/src/lib.rs
+++ b/common/append_merkle/src/lib.rs
@@ -213,6 +213,7 @@ impl<E: HashElement, A: Algorithm<E>> AppendMerkleTree<E, A> {
         mut tx_merkle_nodes: Vec<(usize, E)>,
         start_index: u64,
     ) -> Result<()> {
+        let tx_merkle_nodes_size = tx_merkle_nodes.len();
         if self.leaf_height != 0 {
             tx_merkle_nodes = tx_merkle_nodes
                 .into_iter()
@@ -228,7 +229,8 @@ impl<E: HashElement, A: Algorithm<E>> AppendMerkleTree<E, A> {
         if tx_merkle_nodes.is_empty() {
             return Ok(());
         }
-        let mut position_and_data = proof.file_proof_nodes_in_tree(tx_merkle_nodes);
+        let mut position_and_data =
+            proof.file_proof_nodes_in_tree(tx_merkle_nodes, tx_merkle_nodes_size);
         let start_index = (start_index >> self.leaf_height) as usize;
         for (i, (position, _)) in position_and_data.iter_mut().enumerate() {
             *position += start_index >> i;
@@ -253,10 +255,13 @@ impl<E: HashElement, A: Algorithm<E>> AppendMerkleTree<E, A> {
                 continue;
             }
             if layer[position] != E::null() && layer[position] != data {
-                // This is possible for a valid file proof only when the file proof node is an intermediate node,
-                // so the correct proof node in the flow merkle tree must have been computed as we pad rear data.
-                // Thus, it's okay to skip this case directly.
-                continue;
+                bail!(
+                    "conflict data layer={} position={} tree_data={:?} proof_data={:?}",
+                    i,
+                    position,
+                    layer[position],
+                    data
+                );
             }
             layer[position] = data;
         }
diff --git a/common/append_merkle/src/proof.rs b/common/append_merkle/src/proof.rs
index cc53941..77d983b 100644
--- a/common/append_merkle/src/proof.rs
+++ b/common/append_merkle/src/proof.rs
@@ -103,22 +103,38 @@ impl<T: HashElement> Proof<T> {
         r
     }
 
-    pub fn file_proof_nodes_in_tree(&self, tx_merkle_nodes: Vec<(usize, T)>) -> Vec<(usize, T)> {
+    pub fn file_proof_nodes_in_tree(
+        &self,
+        tx_merkle_nodes: Vec<(usize, T)>,
+        tx_merkle_nodes_size: usize,
+    ) -> Vec<(usize, T)> {
         let mut r = Vec::with_capacity(self.lemma.len());
         let mut subtree_pos = 0;
         let mut root_pos = 0;
-        let mut in_subtree = tx_merkle_nodes.len() == 1;
+        let mut in_subtree = tx_merkle_nodes_size == 1;
         for (i, is_left) in self.path.iter().rev().enumerate() {
             if !in_subtree {
                 if *is_left {
                     in_subtree = true;
+                    root_pos >>= tx_merkle_nodes[i].0;
                 } else if i < tx_merkle_nodes.len() {
                     root_pos += 1 << tx_merkle_nodes[i].0;
+
+                    // This is the last node, so there is no more intermediate node.
+                    if i == tx_merkle_nodes_size - 2 {
+                        if i + 1 >= tx_merkle_nodes.len() {
+                            break;
+                        } else {
+                            in_subtree = true;
+                            root_pos >>= tx_merkle_nodes[i + 1].0;
+                        }
+                    }
                 } else {
                     break;
                 }
             } else {
                 subtree_pos <<= 1;
+                root_pos <<= 1;
                 if !*is_left {
                     subtree_pos += 1;
                 }