1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151

use std::{iter::repeat, sync::Arc};

use anyhow::Context;
use mithril_common::{
    crypto_helper::{Bytes, MKTreeLeafIndexer, MKTreeLeafPosition, MKTreeNode, MKTreeStorer},
    StdResult,
};

/// A Merkle tree store with Sqlite backend
/// * This store is used to store the Merkle tree nodes in a Sqlite database with in memory mode.
/// * This store is not suited for proving, it is only used to store the compute the root of a Merkle tree.
/// * This store is slower than the in memory store but uses less memory which is important to minimize the signer footprint on the SPO infrastructure.
pub struct MKTreeStoreSqlite {
    inner_store: sqlite::ConnectionThreadSafe,
}

impl MKTreeStoreSqlite {
    fn build() -> StdResult<Self> {
        Ok(Self {
            inner_store: Self::create_connection()?,
        })
    }

    fn create_connection() -> StdResult<sqlite::ConnectionThreadSafe> {
        let connection = sqlite::Connection::open_thread_safe(":memory:")?;
        connection.execute("pragma shrink_memory; pragma temp_store = FILE;")?;
        connection.execute(
            "create table merkle_tree (
                position integer, 
                element blob, 
                primary key (position)
            )",
        )?;

        Ok(connection)
    }

    fn get_element_at_position(&self, position: u64) -> StdResult<Option<Arc<MKTreeNode>>> {
        let query = "SELECT element FROM merkle_tree WHERE position = ?";
        let mut statement = self.inner_store.prepare(query)?;
        statement.bind((1, position as i64)).unwrap();
        let result = if let Ok(sqlite::State::Row) = statement.next() {
            Some(Arc::new(MKTreeNode::new(
                statement.read::<Bytes, _>("element")?,
            )))
        } else {
            None
        };

        Ok(result)
    }

    fn insert_elements_from_position(
        &self,
        position: u64,
        elements: Vec<Arc<MKTreeNode>>,
    ) -> StdResult<()> {
        let values_columns: Vec<&str> = repeat("(?, ?)").take(elements.len()).collect();
        let values: Vec<sqlite::Value> = elements
            .into_iter()
            .enumerate()
            .flat_map(|(i, elem)| {
                vec![
                    sqlite::Value::Integer((position + i as u64) as i64),
                    sqlite::Value::Binary((**elem).to_vec()),
                ]
            })
            .collect();
        let query = format!(
            "INSERT INTO merkle_tree(position, element) VALUES {}",
            values_columns.join(", ")
        );
        let mut statement = self.inner_store.prepare(query)?;
        statement.bind::<&[(_, sqlite::Value)]>(
            values
                .into_iter()
                .enumerate()
                .map(|(i, v)| (i + 1, v))
                .collect::<Vec<_>>()
                .as_slice(),
        )?;
        statement.next()?;

        Ok(())
    }
}

impl Clone for MKTreeStoreSqlite {
    fn clone(&self) -> Self {
        unimplemented!("Clone is not implemented for MKTreeStoreSqlite")
    }
}

impl MKTreeLeafIndexer for MKTreeStoreSqlite {
    fn set_leaf_position(&self, _pos: MKTreeLeafPosition, _node: Arc<MKTreeNode>) -> StdResult<()> {
        Ok(())
    }

    fn get_leaf_position(&self, _node: &MKTreeNode) -> Option<MKTreeLeafPosition> {
        unimplemented!("get_leaf_position is not implemented for MKTreeStoreSqlite")
    }

    fn total_leaves(&self) -> usize {
        unimplemented!("total_leaves is not implemented for MKTreeStoreSqlite")
    }

    fn leaves(&self) -> Vec<MKTreeNode> {
        unimplemented!("leaves is not implemented for MKTreeStoreSqlite")
    }
}

impl MKTreeStorer for MKTreeStoreSqlite {
    fn build() -> StdResult<Self> {
        Self::build()
    }

    fn get_elem(&self, pos: u64) -> StdResult<Option<Arc<MKTreeNode>>> {
        self.get_element_at_position(pos).with_context(|| {
            format!("MKTreeStoreSqlite failed to retrieve element at position {pos}")
        })
    }

    fn append(&self, pos: u64, elems: Vec<Arc<MKTreeNode>>) -> StdResult<()> {
        self.insert_elements_from_position(pos, elems)
            .with_context(|| {
                format!("MKTreeStoreSqlite failed to insert elements from position {pos}")
            })
    }
}

#[cfg(test)]
mod tests {
    use mithril_common::crypto_helper::MKTree;

    use super::*;

    #[test]
    fn test_golden_merkle_root() {
        let leaves = vec!["golden-1", "golden-2", "golden-3", "golden-4", "golden-5"];
        let mktree =
            MKTree::<MKTreeStoreSqlite>::new(&leaves).expect("MKTree creation should not fail");
        let mkroot = mktree
            .compute_root()
            .expect("MKRoot generation should not fail");

        assert_eq!(
            "3bbced153528697ecde7345a22e50115306478353619411523e804f2323fd921",
            mkroot.to_hex()
        );
    }
}