sync spec & code (#107)

* sync spec & code

* move primitives (merkletree) into the backend

* comment on the ops spec and link to issue #108

* typo

* fix github-ci mdbook-publish pages

src/backends/plonky2/mock_signed.rs

@@ -2,12 +2,12 @@ use anyhow::Result;
 use std::any::Any;
 use std::collections::HashMap;
 
+use super::primitives::merkletree::MerkleTree;
 use crate::constants::MAX_DEPTH;
 use crate::middleware::{
     containers::Dictionary, hash_str, AnchoredKey, Hash, Params, Pod, PodId, PodSigner, PodType,
     Statement, Value, KEY_SIGNER, KEY_TYPE,
 };
-use crate::primitives::merkletree::MerkleTree;
 
 pub struct MockSigner {
     pub pk: String,

src/backends/plonky2/mod.rs

@@ -1,3 +1,4 @@
 pub mod basetypes;
 pub mod mock_main;
 pub mod mock_signed;
+pub mod primitives;

src/backends/plonky2/primitives/merkletree.rs

@@ -0,0 +1,644 @@
+//! Module that implements the MerkleTree specified at
+//! https://0xparc.github.io/pod2/merkletree.html .
+use anyhow::{anyhow, Result};
+use plonky2::field::goldilocks_field::GoldilocksField;
+use plonky2::hash::poseidon::PoseidonHash;
+use plonky2::plonk::config::Hasher;
+use std::collections::HashMap;
+use std::fmt;
+use std::iter::IntoIterator;
+
+use crate::backends::plonky2::basetypes::{Hash, Value, F, NULL};
+
+/// Implements the MerkleTree specified at
+/// https://0xparc.github.io/pod2/merkletree.html
+#[derive(Clone, Debug)]
+pub struct MerkleTree {
+    max_depth: usize,
+    root: Node,
+}
+
+impl MerkleTree {
+    /// builds a new `MerkleTree` where the leaves contain the given key-values
+    pub fn new(max_depth: usize, kvs: &HashMap<Value, Value>) -> Result<Self> {
+        let mut root = Node::Intermediate(Intermediate::empty());
+
+        for (k, v) in kvs.iter() {
+            let leaf = Leaf::new(max_depth, *k, *v)?;
+            root.add_leaf(0, max_depth, leaf)?;
+        }
+
+        let _ = root.compute_hash();
+        Ok(Self { max_depth, root })
+    }
+
+    /// returns the root of the tree
+    pub fn root(&self) -> Hash {
+        self.root.hash()
+    }
+
+    /// returns the max_depth parameter from the tree
+    pub fn max_depth(&self) -> usize {
+        self.max_depth
+    }
+
+    /// returns the value at the given key
+    pub fn get(&self, key: &Value) -> Result<Value> {
+        let path = keypath(self.max_depth, *key)?;
+        let key_resolution = self.root.down(0, self.max_depth, path, None)?;
+        match key_resolution {
+            Some((k, v)) if &k == key => Ok(v),
+            _ => Err(anyhow!("key not found")),
+        }
+    }
+
+    /// returns a boolean indicating whether the key exists in the tree
+    pub fn contains(&self, key: &Value) -> Result<bool> {
+        let path = keypath(self.max_depth, *key)?;
+        match self.root.down(0, self.max_depth, path, None) {
+            Ok(Some((k, _))) => {
+                if &k == key {
+                    Ok(true)
+                } else {
+                    Ok(false)
+                }
+            }
+            _ => Ok(false),
+        }
+    }
+
+    /// returns a proof of existence, which proves that the given key exists in
+    /// the tree. It returns the `value` of the leaf at the given `key`, and the
+    /// `MerkleProof`.
+    pub fn prove(&self, key: &Value) -> Result<(Value, MerkleProof)> {
+        let path = keypath(self.max_depth, *key)?;
+
+        let mut siblings: Vec<Hash> = Vec::new();
+
+        match self
+            .root
+            .down(0, self.max_depth, path, Some(&mut siblings))?
+        {
+            Some((k, v)) if &k == key => Ok((
+                v,
+                MerkleProof {
+                    existence: true,
+                    siblings,
+                    other_leaf: None,
+                },
+            )),
+            _ => Err(anyhow!("key not found")),
+        }
+    }
+
+    /// returns a proof of non-existence, which proves that the given
+    /// `key` does not exist in the tree. The return value specifies
+    /// the key-value pair in the leaf reached as a result of
+    /// resolving `key` as well as a `MerkleProof`.
+    pub fn prove_nonexistence(&self, key: &Value) -> Result<MerkleProof> {
+        let path = keypath(self.max_depth, *key)?;
+
+        let mut siblings: Vec<Hash> = Vec::new();
+
+        // note: non-existence of a key can be in 2 cases:
+        match self
+            .root
+            .down(0, self.max_depth, path, Some(&mut siblings))?
+        {
+            // case i) the expected leaf does not exist
+            None => Ok(MerkleProof {
+                existence: false,
+                siblings,
+                other_leaf: None,
+            }),
+            // case ii) the expected leaf does exist in the tree, but it has a different `key`
+            Some((k, v)) if &k != key => Ok(MerkleProof {
+                existence: false,
+                siblings,
+                other_leaf: Some((k, v)),
+            }),
+            _ => Err(anyhow!("key found")),
+        }
+        // both cases prove that the given key don't exist in the tree. ∎
+    }
+
+    /// verifies an inclusion proof for the given `key` and `value`
+    pub fn verify(
+        max_depth: usize,
+        root: Hash,
+        proof: &MerkleProof,
+        key: &Value,
+        value: &Value,
+    ) -> Result<()> {
+        let h = proof.compute_root_from_leaf(max_depth, key, Some(*value))?;
+
+        if h != root {
+            return Err(anyhow!("proof of inclusion does not verify"));
+        } else {
+            Ok(())
+        }
+    }
+
+    /// verifies a non-inclusion proof for the given `key`, that is, the given
+    /// `key` does not exist in the tree
+    pub fn verify_nonexistence(
+        max_depth: usize,
+        root: Hash,
+        proof: &MerkleProof,
+        key: &Value,
+    ) -> Result<()> {
+        match proof.other_leaf {
+            Some((k, _v)) if &k == key => Err(anyhow!("Invalid non-existence proof.")),
+            _ => {
+                let k = proof.other_leaf.map(|(k, _)| k).unwrap_or(*key);
+                let v: Option<Value> = proof.other_leaf.map(|(_, v)| v);
+                let h = proof.compute_root_from_leaf(max_depth, &k, v)?;
+
+                if h != root {
+                    return Err(anyhow!("proof of exclusion does not verify"));
+                } else {
+                    Ok(())
+                }
+            }
+        }
+    }
+
+    /// returns an iterator over the leaves of the tree
+    pub fn iter(&self) -> Iter {
+        Iter {
+            state: vec![&self.root],
+        }
+    }
+}
+
+impl<'a> IntoIterator for &'a MerkleTree {
+    type Item = (&'a Value, &'a Value);
+    type IntoIter = Iter<'a>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.iter()
+    }
+}
+
+impl fmt::Display for MerkleTree {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        writeln!(
+            f,
+            "\nPaste in GraphViz (https://dreampuf.github.io/GraphvizOnline/):\n-----"
+        )?;
+        writeln!(f, "digraph hierarchy {{")?;
+        writeln!(f, "node [fontname=Monospace,fontsize=10,shape=box]")?;
+        write!(f, "{}", self.root)?;
+        writeln!(f, "\n}}\n-----")
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct MerkleProof {
+    // note: currently we don't use the `_existence` field, we would use if we merge the methods
+    // `verify` and `verify_nonexistence` into a single one
+    #[allow(unused)]
+    existence: bool,
+    siblings: Vec<Hash>,
+    // other_leaf is used for non-existence proofs
+    other_leaf: Option<(Value, Value)>,
+}
+
+impl fmt::Display for MerkleProof {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        for (i, s) in self.siblings.iter().enumerate() {
+            if i > 0 {
+                write!(f, ", ")?;
+            }
+            write!(f, "{}", s)?;
+        }
+        Ok(())
+    }
+}
+
+impl MerkleProof {
+    /// Computes the root of the Merkle tree suggested by a Merkle proof given a
+    /// key & value. If a value is not provided, the terminal node is assumed to
+    /// be empty.
+    fn compute_root_from_leaf(
+        &self,
+        max_depth: usize,
+        key: &Value,
+        value: Option<Value>,
+    ) -> Result<Hash> {
+        if self.siblings.len() >= max_depth {
+            return Err(anyhow!("max depth reached"));
+        }
+
+        let path = keypath(max_depth, *key)?;
+        let mut h = value
+            .map(|v| Hash(PoseidonHash::hash_no_pad(&[key.0, v.0].concat()).elements))
+            .unwrap_or(Hash([GoldilocksField(0); 4]));
+        for (i, sibling) in self.siblings.iter().enumerate().rev() {
+            let input: Vec<F> = if path[i] {
+                [sibling.0, h.0].concat()
+            } else {
+                [h.0, sibling.0].concat()
+            };
+            h = Hash(PoseidonHash::hash_no_pad(&input).elements);
+        }
+        Ok(h)
+    }
+}
+
+#[derive(Clone, Debug)]
+enum Node {
+    None,
+    Leaf(Leaf),
+    Intermediate(Intermediate),
+}
+
+impl fmt::Display for Node {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::Intermediate(n) => {
+                writeln!(
+                    f,
+                    "\"{}\" -> {{ \"{}\" \"{}\" }}",
+                    n.hash(),
+                    n.left.hash(),
+                    n.right.hash()
+                )?;
+                write!(f, "{}", n.left)?;
+                write!(f, "{}", n.right)
+            }
+            Self::Leaf(l) => {
+                writeln!(f, "\"{}\" [style=filled]", l.hash())?;
+                writeln!(f, "\"k:{}\\nv:{}\" [style=dashed]", l.key, l.value)?;
+                writeln!(
+                    f,
+                    "\"{}\" -> {{ \"k:{}\\nv:{}\" }}",
+                    l.hash(),
+                    l.key,
+                    l.value,
+                )
+            }
+            Self::None => Ok(()),
+        }
+    }
+}
+
+impl Node {
+    fn is_empty(&self) -> bool {
+        match self {
+            Self::None => true,
+            Self::Leaf(_l) => false,
+            Self::Intermediate(_n) => false,
+        }
+    }
+    fn compute_hash(&mut self) -> Hash {
+        match self {
+            Self::None => NULL,
+            Self::Leaf(l) => l.compute_hash(),
+            Self::Intermediate(n) => n.compute_hash(),
+        }
+    }
+    fn hash(&self) -> Hash {
+        match self {
+            Self::None => NULL,
+            Self::Leaf(l) => l.hash(),
+            Self::Intermediate(n) => n.hash(),
+        }
+    }
+
+    /// Goes down from the current node until it encounters a terminal node,
+    /// viz. a leaf or empty node, or until it reaches the maximum depth. The
+    /// `siblings` parameter is used to store the siblings while going down to
+    /// the leaf, if the given parameter is set to `None`, then no siblings are
+    /// stored. In this way, the same method `down` can be used by MerkleTree
+    /// methods `get`, `contains`, `prove` and `prove_nonexistence`.
+    ///
+    /// Be aware that this method will return the found leaf at the given path,
+    /// which may contain a different key and value than the expected one.
+    fn down(
+        &self,
+        lvl: usize,
+        max_depth: usize,
+        path: Vec<bool>,
+        mut siblings: Option<&mut Vec<Hash>>,
+    ) -> Result<Option<(Value, Value)>> {
+        if lvl >= max_depth {
+            return Err(anyhow!("max depth reached"));
+        }
+
+        match self {
+            Self::Intermediate(n) => {
+                if path[lvl] {
+                    if let Some(s) = siblings.as_mut() {
+                        s.push(n.left.hash());
+                    }
+                    return n.right.down(lvl + 1, max_depth, path, siblings);
+                } else {
+                    if let Some(s) = siblings.as_mut() {
+                        s.push(n.right.hash());
+                    }
+                    return n.left.down(lvl + 1, max_depth, path, siblings);
+                }
+            }
+            Self::Leaf(Leaf {
+                key,
+                value,
+                path: _p,
+                hash: _h,
+            }) => Ok(Some((key.clone(), value.clone()))),
+            _ => Ok(None),
+        }
+    }
+
+    // adds the leaf at the tree from the current node (self), without computing any hash
+    fn add_leaf(&mut self, lvl: usize, max_depth: usize, leaf: Leaf) -> Result<()> {
+        if lvl >= max_depth {
+            return Err(anyhow!("max depth reached"));
+        }
+
+        match self {
+            Self::Intermediate(n) => {
+                if leaf.path[lvl] {
+                    if n.right.is_empty() {
+                        // empty sub-node, add the leaf here
+                        n.right = Box::new(Node::Leaf(leaf));
+                        return Ok(());
+                    }
+                    n.right.add_leaf(lvl + 1, max_depth, leaf)?;
+                } else {
+                    if n.left.is_empty() {
+                        // empty sub-node, add the leaf here
+                        n.left = Box::new(Node::Leaf(leaf));
+                        return Ok(());
+                    }
+                    n.left.add_leaf(lvl + 1, max_depth, leaf)?;
+                }
+            }
+            Self::Leaf(l) => {
+                // in this case, it means that we found a leaf in the new-leaf
+                // path, thus we need to push both leaves (old-leaf and
+                // new-leaf) down the path till their paths diverge.
+
+                // first check that keys of both leaves are different
+                // (l=old-leaf, leaf=new-leaf)
+                if l.key == leaf.key {
+                    // Note: current approach returns an error when trying to
+                    // add to a leaf where the key already exists. We could also
+                    // ignore it if needed.
+                    return Err(anyhow!("key already exists"));
+                }
+                let old_leaf = l.clone();
+                // set self as an intermediate node
+                *self = Node::Intermediate(Intermediate::empty());
+                return self.down_till_divergence(lvl, max_depth, old_leaf, leaf);
+            }
+            Self::None => {
+                return Err(anyhow!("reached empty node, should not have entered"));
+            }
+        }
+        Ok(())
+    }
+
+    /// goes down through a 'virtual' path till finding a divergence. This
+    /// method is used for when adding a new leaf another already existing leaf
+    /// is found, so that both leaves (new and old) are pushed down the path
+    /// till their keys diverge.
+    fn down_till_divergence(
+        &mut self,
+        lvl: usize,
+        max_depth: usize,
+        old_leaf: Leaf,
+        new_leaf: Leaf,
+    ) -> Result<()> {
+        if lvl >= max_depth {
+            return Err(anyhow!("max depth reached"));
+        }
+
+        if let Node::Intermediate(ref mut n) = self {
+            if old_leaf.path[lvl] != new_leaf.path[lvl] {
+                // reached divergence in next level, set the leaves as children
+                // at the current node
+                if new_leaf.path[lvl] {
+                    n.left = Box::new(Node::Leaf(old_leaf));
+                    n.right = Box::new(Node::Leaf(new_leaf));
+                } else {
+                    n.left = Box::new(Node::Leaf(new_leaf));
+                    n.right = Box::new(Node::Leaf(old_leaf));
+                }
+                return Ok(());
+            }
+
+            // no divergence yet, continue going down
+            if new_leaf.path[lvl] {
+                n.right = Box::new(Node::Intermediate(Intermediate::empty()));
+                return n
+                    .right
+                    .down_till_divergence(lvl + 1, max_depth, old_leaf, new_leaf);
+            } else {
+                n.left = Box::new(Node::Intermediate(Intermediate::empty()));
+                return n
+                    .left
+                    .down_till_divergence(lvl + 1, max_depth, old_leaf, new_leaf);
+            }
+        }
+        Ok(())
+    }
+}
+
+#[derive(Clone, Debug)]
+struct Intermediate {
+    hash: Option<Hash>,
+    left: Box<Node>,
+    right: Box<Node>,
+}
+impl Intermediate {
+    fn empty() -> Self {
+        Self {
+            hash: None,
+            left: Box::new(Node::None),
+            right: Box::new(Node::None),
+        }
+    }
+    fn compute_hash(&mut self) -> Hash {
+        if self.left.clone().is_empty() && self.right.clone().is_empty() {
+            self.hash = Some(NULL);
+            return NULL;
+        }
+        let l_hash = self.left.compute_hash();
+        let r_hash = self.right.compute_hash();
+        let input: Vec<F> = [l_hash.0, r_hash.0].concat();
+        let h = Hash(PoseidonHash::hash_no_pad(&input).elements);
+        self.hash = Some(h);
+        h
+    }
+    fn hash(&self) -> Hash {
+        self.hash.unwrap()
+    }
+}
+
+#[derive(Clone, Debug)]
+struct Leaf {
+    hash: Option<Hash>,
+    path: Vec<bool>,
+    key: Value,
+    value: Value,
+}
+impl Leaf {
+    fn new(max_depth: usize, key: Value, value: Value) -> Result<Self> {
+        Ok(Self {
+            hash: None,
+            path: keypath(max_depth, key)?,
+            key,
+            value,
+        })
+    }
+    fn compute_hash(&mut self) -> Hash {
+        let input: Vec<F> = [self.key.0, self.value.0].concat();
+        let h = Hash(PoseidonHash::hash_no_pad(&input).elements);
+        self.hash = Some(h);
+        h
+    }
+    fn hash(&self) -> Hash {
+        self.hash.unwrap()
+    }
+}
+
+// NOTE 1: think if maybe the length of the returned vector can be <256
+// (8*bytes.len()), so that we can do fewer iterations. For example, if the
+// tree.max_depth is set to 20, we just need 20 iterations of the loop, not 256.
+// NOTE 2: which approach do we take with keys that are longer than the
+// max-depth? ie, what happens when two keys share the same path for more bits
+// than the max_depth?
+/// returns the path of the given key
+fn keypath(max_depth: usize, k: Value) -> Result<Vec<bool>> {
+    let bytes = k.to_bytes();
+    if max_depth > 8 * bytes.len() {
+        // note that our current keys are of Value type, which are 4 Goldilocks
+        // field elements, ie ~256 bits, therefore the max_depth can not be
+        // bigger than 256.
+        return Err(anyhow!(
+            "key to short (key length: {}) for the max_depth: {}",
+            8 * bytes.len(),
+            max_depth
+        ));
+    }
+    Ok((0..max_depth)
+        .map(|n| bytes[n / 8] & (1 << (n % 8)) != 0)
+        .collect())
+}
+
+pub struct Iter<'a> {
+    state: Vec<&'a Node>,
+}
+
+impl<'a> Iterator for Iter<'a> {
+    type Item = (&'a Value, &'a Value);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let node = self.state.pop();
+        match node {
+            Some(Node::None) => self.next(),
+            Some(Node::Leaf(Leaf {
+                hash: _,
+                path: _,
+                key,
+                value,
+            })) => Some((key, value)),
+            Some(Node::Intermediate(Intermediate {
+                hash: _,
+                left,
+                right,
+            })) => {
+                self.state.push(&right);
+                self.state.push(&left);
+                self.next()
+            }
+            _ => None,
+        }
+    }
+}
+
+#[cfg(test)]
+pub mod tests {
+    use itertools::Itertools;
+    use std::cmp::Ordering;
+
+    use super::*;
+
+    #[test]
+    fn test_merkletree() -> Result<()> {
+        let mut kvs = HashMap::new();
+        for i in 0..8 {
+            if i == 1 {
+                continue;
+            }
+            kvs.insert(Value::from(i), Value::from(1000 + i));
+        }
+        let key = Value::from(13);
+        let value = Value::from(1013);
+        kvs.insert(key, value);
+
+        let tree = MerkleTree::new(32, &kvs)?;
+        // when printing the tree, it should print the same tree as in
+        // https://0xparc.github.io/pod2/merkletree.html#example-2
+        println!("{}", tree);
+
+        // Inclusion checks
+        let (v, proof) = tree.prove(&Value::from(13))?;
+        assert_eq!(v, Value::from(1013));
+        println!("{}", proof);
+
+        MerkleTree::verify(32, tree.root(), &proof, &key, &value)?;
+
+        // Exclusion checks
+        let key = Value::from(12);
+        let proof = tree.prove_nonexistence(&key)?;
+        assert_eq!(
+            proof.other_leaf.unwrap(),
+            (Value::from(4), Value::from(1004))
+        );
+        println!("{}", proof);
+
+        MerkleTree::verify_nonexistence(32, tree.root(), &proof, &key)?;
+
+        let key = Value::from(1);
+        let proof = tree.prove_nonexistence(&Value::from(1))?;
+        assert_eq!(proof.other_leaf, None);
+        println!("{}", proof);
+
+        MerkleTree::verify_nonexistence(32, tree.root(), &proof, &key)?;
+
+        // Check iterator
+        let collected_kvs: Vec<_> = tree.into_iter().collect::<Vec<_>>();
+
+        // Expected key ordering
+        let cmp = |max_depth: usize| {
+            move |k1, k2| {
+                let path1 = keypath(max_depth, k1).unwrap();
+                let path2 = keypath(max_depth, k2).unwrap();
+
+                let first_unequal_bits = std::iter::zip(path1, path2).find(|(b1, b2)| b1 != b2);
+
+                match first_unequal_bits {
+                    Some((b1, b2)) => {
+                        if b1 < b2 {
+                            Ordering::Less
+                        } else {
+                            Ordering::Greater
+                        }
+                    }
+                    _ => Ordering::Equal,
+                }
+            }
+        };
+
+        let sorted_kvs = kvs
+            .iter()
+            .sorted_by(|(k1, _), (k2, _)| cmp(32)(**k1, **k2))
+            .collect::<Vec<_>>();
+
+        assert_eq!(collected_kvs, sorted_kvs);
+
+        Ok(())
+    }
+}

src/backends/plonky2/primitives/mod.rs

@@ -0,0 +1 @@
+pub mod merkletree;