merkle_tree.h

// Copyright 2025 Google LLC.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
 
#ifndef PRIVACY_PROOFS_ZK_LIB_MERKLE_MERKLE_TREE_H_
#define PRIVACY_PROOFS_ZK_LIB_MERKLE_MERKLE_TREE_H_
 
#include <stddef.h>
#include <stdint.h>
 
#include <cstring>
#include <vector>
 
#include "util/crypto.h"
#include "util/panic.h"
 
namespace proofs {
 
// This package computes and verifies Merkle Tree inclusion claims.
// The standard Merkle tree algorithm has been implemented.
 
// A digest of a Merkle tree.
struct Digest {
  static constexpr size_t kLength = kSHA256DigestSize;
  uint8_t data[kLength];
 
  bool operator==(const Digest& y) const {
    return memcmp(data, y.data, kLength) == 0;
  }
 
  static Digest hash2(const Digest& L, const Digest& R) {
    SHA256 sha;
    sha.Update(L.data, kLength);
    sha.Update(R.data, kLength);
    Digest output;
    sha.DigestData(output.data);
    return output;
  }
};
 
// Return the length of the proof for N leaves.
// Mimic the code in generate_proof() without actually
// computing the proof
inline size_t merkle_tree_len(size_t n) {
  size_t r = 1;
  size_t pos = (n - 1);  // maximum possible value
  for (pos += n; pos > 1; pos >>= 1) {
    ++r;
  }
  return r;
}
 
// compute the set of all nodes on the path from the
// root to any leaf in POS.
inline std::vector<bool> compressed_merkle_proof_tree(size_t n,
                                                      const size_t pos[/*np*/],
                                                      size_t np) {
  check(np > 0, "A Merkle proof with 0 leaves is not defined.");
  std::vector<bool> tree(2 * n, false);
 
  // leaves are in TREE
  for (size_t ip = 0; ip < np; ++ip) {
    check(pos[ip] < n, "Invalid position for leaf in Merkle tree");
    tree[pos[ip] + n] = true;
  }
 
  // If a child of an inner node is in TREE, then the parent is in TREE.
  for (size_t i = n; i-- > 1;) {
    tree[i] = (tree[2 * i] || tree[2 * i + 1]);
  }
 
  // Assert that the root is in TREE.
  check(tree[1], "tree[1]");
 
  return tree;
}
 
class MerkleTree {
 public:
  explicit MerkleTree(size_t n) : n_(n), layers_(2 * n) {}
 
  void set_leaf(size_t pos, const Digest& leaf) {
    check(pos < n_, "Invalid position for leaf in Merkle tree");
    layers_[pos + n_] = leaf;
  }
 
  Digest build_tree() {
    for (size_t i = n_; i-- > 1;) {
      layers_[i] = Digest::hash2(layers_[2 * i], layers_[2 * i + 1]);
    }
    return layers_[1];
  }
 
  // The generate_proof method writes a Merkle tree proof for the leaf
  // at position pos into the proof array and returns the size of the proof
  // in number of Digests.
  size_t generate_proof(Digest proof[/*logn+1*/], size_t pos) const {
    Digest* begin = proof;
    *proof++ = layers_[pos + n_];
    for (pos += n_; pos > 1; pos >>= 1) {
      *proof++ = layers_[pos ^ 1];
    }
    return (proof - begin);
  }
 
  // Compressed Merkle proofs over a set POS[NP] of leaves.
  //
  // We first compute the set TREE of all nodes that are on the path
  // from the root to any leaf in POS.  Then, for each inner node in
  // TREE, we include in the proof the child that is not in TREE, if
  // any.
  size_t generate_compressed_proof(std::vector<Digest>& proof,
                                   const size_t pos[/*np*/], size_t np) {
    std::vector<bool> tree = compressed_merkle_proof_tree(n_, pos, np);
 
    // For each TREE node, include in the proof the
    // child that is not TREE, if any.
    size_t sz = 0;
    for (size_t i = n_; i-- > 1;) {
      if (tree[i]) {
        size_t child = 2 * i;
        if (tree[child]) {
          // try the other child
          child = 2 * i + 1;
        }
        if (!tree[child]) {
          proof.push_back(layers_[child]);
          ++sz;
        }
      }
    }
    return sz;
  }
 
  size_t n_;
  // layers_[n, 2 * n) stores the leaves (nodes at layer 0).
  // layers_[n/2, n) stores nodes at layer 1.
  // layers_[n/4, n/2) stores nodes at layer 2, etc.
  // The root is at layers_[1] where layers_[0] is not used.
  std::vector<Digest> layers_;
};
 
class MerkleTreeVerifier {
 public:
  explicit MerkleTreeVerifier(size_t n, const Digest& root)
      : n_(n), root_(root) {}
 
  bool verify_proof(const Digest* proof, size_t pos) const {
    Digest t = *proof++;
    for (pos += n_; pos > 1; pos >>= 1) {
      t = (pos & 1) ? Digest::hash2(*proof++, t) : Digest::hash2(t, *proof++);
    }
    return t == root_;
  }
 
  bool verify_compressed_proof(const Digest* proof, size_t proof_len,
                               const Digest leaves[/*np*/],
                               const size_t pos[/*np*/], size_t np) const {
    // Reconstructed layers_, where only the DEFINED subset is
    // defined.
    std::vector<Digest> layers(2 * n_, Digest{});
    std::vector<bool> defined(2 * n_, false);
 
    /*scope for TREE */ {
      std::vector<bool> tree = compressed_merkle_proof_tree(n_, pos, np);
 
      // read the proof
      size_t sz = 0;
      for (size_t i = n_; i-- > 1;) {
        if (tree[i]) {
          size_t child = 2 * i;
          if (tree[child]) {
            // try the other child
            child = 2 * i + 1;
          }
          if (!tree[child]) {
            if (sz >= proof_len) {
              return false;
            }
            layers[child] = proof[sz++];
            defined[child] = true;
          }
        }
      }
    }
 
    // set LAYERS at all leaves in POS
    for (size_t ip = 0; ip < np; ++ip) {
      size_t l = pos[ip] + n_;
      layers[l] = leaves[ip];
      defined[l] = true;
    }
 
    // Recompute as many inner nodes as we can
    for (size_t i = n_; i-- > 1;) {
      if (defined[2 * i] && defined[2 * i + 1]) {
        layers[i] = Digest::hash2(layers[2 * i], layers[2 * i + 1]);
        defined[i] = true;
      }
    }
 
    return (defined[1] && (root_ == layers[1]));
  }
 
 private:
  size_t n_;
  Digest root_;
};
 
}  // namespace proofs
 
#endif  // PRIVACY_PROOFS_ZK_LIB_MERKLE_MERKLE_TREE_H_