miden-client-sqlite-store 0.17.0-rc.3

#![allow(clippy::items_after_statements)]

use std::collections::{BTreeMap, BTreeSet};
use std::num::NonZeroUsize;
use std::vec::Vec;

use miden_client::Word;
use miden_client::block::BlockHeader;
use miden_client::crypto::{Forest, InOrderIndex, MmrPeaks};
use miden_client::note::BlockNumber;
use miden_client::store::{BlockRelevance, PartialBlockchainFilter, StoreError};
use miden_client::utils::{Deserializable, Serializable};
use rusqlite::{Connection, Transaction, params, params_from_iter};

use super::SqliteStore;
use crate::sql_error::SqlResultExt;
use crate::sync::query_sync_height;
use crate::{insert_sql, int_array, subst, with_write_tx};

impl SqliteStore {
    pub(crate) fn get_block_headers(
        conn: &mut Connection,
        block_numbers: &BTreeSet<BlockNumber>,
    ) -> Result<Vec<(BlockHeader, BlockRelevance)>, StoreError> {
        let block_number_list =
            int_array(block_numbers.iter().map(|block_number| u64::from(block_number.as_u32())));

        const QUERY: &str =
            "SELECT header, has_client_notes FROM block_headers WHERE block_num IN rarray(?)";

        conn.prepare(QUERY)
            .into_store_error()?
            .query_map(params![block_number_list], parse_block_headers_columns)
            .into_store_error()?
            .map(|result| parse_block_header(result.into_store_error()?))
            .collect()
    }

    pub(crate) fn get_tracked_block_headers(
        conn: &mut Connection,
    ) -> Result<Vec<BlockHeader>, StoreError> {
        // `idx_block_headers_has_notes` is declared `WHERE has_client_notes = 1`, and SQLite
        // matches a partial index only when the predicate is spelled the same way.
        const QUERY: &str =
            "SELECT header, has_client_notes FROM block_headers WHERE has_client_notes=1";
        conn.prepare(QUERY)
            .into_store_error()?
            .query_map(params![], parse_block_headers_columns)
            .into_store_error()?
            .map(|result| parse_block_header(result.into_store_error()?).map(|(block, _)| block))
            .collect()
    }

    pub(crate) fn get_tracked_block_header_numbers(
        conn: &mut Connection,
    ) -> Result<BTreeSet<usize>, StoreError> {
        const QUERY: &str = "SELECT block_num FROM block_headers WHERE has_client_notes=1";
        conn.prepare(QUERY)
            .into_store_error()?
            .query_map(params![], |row| row.get::<_, u32>(0))
            .into_store_error()?
            .map(|result| {
                let block_num: u32 = result.into_store_error()?;
                Ok(block_num as usize)
            })
            .collect()
    }

    pub(crate) fn get_partial_blockchain_nodes(
        conn: &mut Connection,
        filter: &PartialBlockchainFilter,
    ) -> Result<BTreeMap<InOrderIndex, Word>, StoreError> {
        match filter {
            PartialBlockchainFilter::All => query_partial_blockchain_nodes(
                conn,
                "SELECT id, node FROM partial_blockchain_nodes",
                params![],
            ),

            PartialBlockchainFilter::List(ids) if ids.is_empty() => Ok(BTreeMap::new()),
            PartialBlockchainFilter::List(ids) => {
                let id_values = int_array(ids.iter().map(|id| id.inner() as u64));

                query_partial_blockchain_nodes(
                    conn,
                    "SELECT id, node FROM partial_blockchain_nodes WHERE id IN rarray(?)",
                    params_from_iter([id_values]),
                )
            },

            PartialBlockchainFilter::Forest(forest) if forest.is_empty() => Ok(BTreeMap::new()),
            PartialBlockchainFilter::Forest(forest) => {
                let max_index = i64::try_from(
                    forest.rightmost_in_order_index().expect("forest is not empty").inner(),
                )
                .expect("id is a valid i64");

                query_partial_blockchain_nodes(
                    conn,
                    "SELECT id, node FROM partial_blockchain_nodes WHERE id <= ?",
                    params![max_index],
                )
            },
        }
    }

    pub(crate) fn get_current_blockchain_peaks(
        conn: &mut Connection,
    ) -> Result<MmrPeaks, StoreError> {
        const QUERY: &str =
            "SELECT block_num, partial_blockchain_peaks FROM blockchain_checkpoint LIMIT 1";

        let (block_num, peaks_bytes): (u32, Vec<u8>) = conn
            .prepare(QUERY)
            .into_store_error()?
            .query_row(params![], |row| {
                Ok((row.get("block_num")?, row.get("partial_blockchain_peaks")?))
            })
            .into_store_error()?;

        if peaks_bytes.is_empty() {
            return Ok(MmrPeaks::new(Forest::empty(), vec![])?);
        }
        parse_partial_blockchain_peaks(block_num, &peaks_bytes)
    }

    pub(crate) fn insert_block_header(
        conn: &mut Connection,
        block_header: &BlockHeader,
        nodes: &[(InOrderIndex, Word)],
        has_client_notes: bool,
    ) -> Result<(), StoreError> {
        with_write_tx(conn, |tx| {
            Self::insert_block_header_tx(tx, block_header, has_client_notes)?;
            Self::insert_partial_blockchain_nodes_tx(tx, nodes)
        })
    }

    /// Inserts a list of MMR authentication nodes to the Partial Blockchain nodes table.
    pub(crate) fn insert_partial_blockchain_nodes_tx(
        tx: &Transaction<'_>,
        nodes: &[(InOrderIndex, Word)],
    ) -> Result<(), StoreError> {
        const QUERY: &str = insert_sql!(partial_blockchain_nodes { id, node } | IGNORE);
        let mut stmt = tx.prepare_cached(QUERY).into_store_error()?;
        for (index, node) in nodes {
            stmt.execute(params![index.inner(), node.to_bytes()]).into_store_error()?;
        }
        Ok(())
    }

    /// Inserts a block header using a [`rusqlite::Transaction`].
    ///
    /// If the block header exists and `has_client_notes` is `true` then the `has_client_notes`
    /// column is updated to `true` to signify that the block now contains a relevant note.
    pub(crate) fn insert_block_header_tx(
        tx: &Transaction<'_>,
        block_header: &BlockHeader,
        has_client_notes: bool,
    ) -> Result<(), StoreError> {
        const QUERY: &str =
            insert_sql!(block_headers { block_num, header, has_client_notes } | IGNORE);
        let block_num = block_header.block_num().as_u32();
        tx.execute(QUERY, params![block_num, block_header.to_bytes(), has_client_notes])
            .into_store_error()?;

        set_block_header_has_client_notes(tx, u64::from(block_num), has_client_notes)?;
        Ok(())
    }

    /// Prunes irrelevant block data from the store.
    ///
    /// This performs three operations in a single transaction:
    /// 1. Deletes MMR authentication nodes at the given `node_indices`.
    /// 2. Sets `has_client_notes = false` for `blocks_to_untrack`.
    /// 3. Deletes block headers with `has_client_notes = false` that are not the genesis or
    ///    sync-height block.
    pub(crate) fn untrack_and_prune_irrelevant_blocks(
        conn: &mut Connection,
        blocks_to_untrack: &[BlockNumber],
        node_indices_to_remove: &[InOrderIndex],
    ) -> Result<(), StoreError> {
        with_write_tx(conn, |tx| {
            // 1. Delete stale MMR authentication nodes.
            if !node_indices_to_remove.is_empty() {
                let id_values =
                    int_array(node_indices_to_remove.iter().map(|id| id.inner() as u64));

                tx.execute(
                    "DELETE FROM partial_blockchain_nodes WHERE id IN rarray(?)",
                    params![id_values],
                )
                .into_store_error()?;
            }

            // 2. Mark untracked blocks as irrelevant.
            if !blocks_to_untrack.is_empty() {
                let block_values =
                    int_array(blocks_to_untrack.iter().map(|b| u64::from(b.as_u32())));

                tx.execute(
                    "UPDATE block_headers SET has_client_notes = 0 WHERE block_num IN rarray(?)",
                    params![block_values],
                )
                .into_store_error()?;
            }

            // 3. Delete irrelevant block headers.
            let genesis: u32 = BlockNumber::GENESIS.as_u32();
            let sync_height = query_sync_height(tx)?.as_u32();

            tx.execute(
                "DELETE FROM block_headers \
                 WHERE has_client_notes = 0 \
                 AND block_num > ?1 \
                 AND block_num < ?2",
                rusqlite::params![genesis, sync_height],
            )
            .into_store_error()?;

            Ok(())
        })
    }
}

// HELPERS
// ================================================================================================

fn query_partial_blockchain_nodes<P: rusqlite::Params>(
    conn: &mut Connection,
    sql: &str,
    params: P,
) -> Result<BTreeMap<InOrderIndex, Word>, StoreError> {
    let mut stmt = conn.prepare_cached(sql).into_store_error()?;

    stmt.query_map(params, |row| Ok((row.get::<_, usize>("id")?, row.get::<_, Vec<u8>>("node")?)))
        .into_store_error()?
        .map(|row_res| {
            let (id, node) = row_res.into_store_error()?;
            // An in-order index is never zero. A zero value means the row is corrupt.
            let id = NonZeroUsize::new(id).ok_or_else(|| {
                StoreError::ParsingError(
                    "stored partial blockchain node id must be non-zero".to_string(),
                )
            })?;
            Ok((InOrderIndex::new(id), Word::read_from_bytes(&node)?))
        })
        .collect()
}

fn parse_partial_blockchain_peaks(forest: u32, peaks_nodes: &[u8]) -> Result<MmrPeaks, StoreError> {
    let mmr_peaks_nodes = Vec::<Word>::read_from_bytes(peaks_nodes)?;

    let forest_size = usize::try_from(forest).expect("u64 should fit in usize");
    let forest = Forest::new(forest_size).map_err(|err| {
        StoreError::DataDeserializationError(
            miden_protocol::utils::serde::DeserializationError::InvalidValue(format!(
                "invalid forest size {forest_size}: {err}"
            )),
        )
    })?;
    MmrPeaks::new(forest, mmr_peaks_nodes).map_err(StoreError::MmrError)
}

fn parse_block_headers_columns(
    row: &rusqlite::Row<'_>,
) -> Result<(Vec<u8>, bool), rusqlite::Error> {
    Ok((row.get("header")?, row.get("has_client_notes")?))
}

fn parse_block_header(
    (header, has_client_notes): (Vec<u8>, bool),
) -> Result<(BlockHeader, BlockRelevance), StoreError> {
    Ok((BlockHeader::read_from_bytes(&header)?, has_client_notes.into()))
}

pub(crate) fn set_block_header_has_client_notes(
    tx: &Transaction<'_>,
    block_num: u64,
    has_client_notes: bool,
) -> Result<(), StoreError> {
    // Only update to change has_client_notes to true if it was false previously
    const QUERY: &str = "\
        UPDATE block_headers
        SET has_client_notes=?
        WHERE block_num=? AND has_client_notes=0;";
    tx.execute(QUERY, params![has_client_notes, block_num]).into_store_error()?;
    Ok(())
}

#[cfg(test)]
mod test {
    use std::collections::{BTreeMap, BTreeSet};
    use std::vec::Vec;

    use miden_client::Word;
    use miden_client::block::BlockHeader;
    use miden_client::crypto::{Forest, InOrderIndex, MmrPeaks};
    use miden_client::note::BlockNumber;
    use miden_client::store::{PartialBlockchainFilter, Store};
    use miden_client::utils::Serializable;
    use miden_protocol::crypto::merkle::mmr::Mmr;
    use rusqlite::params;

    use crate::SqliteStore;
    use crate::tests::create_test_store;

    async fn insert_dummy_block_headers(store: &mut SqliteStore) -> Vec<BlockHeader> {
        let block_headers: Vec<BlockHeader> =
            (0..5).map(|block_num| BlockHeader::mock(block_num, None, None, &[])).collect();

        let block_headers_clone = block_headers.clone();
        store
            .interact_with_connection(move |conn| {
                let tx = conn.transaction().unwrap();
                (0..5).for_each(|block_num| {
                    SqliteStore::insert_block_header_tx(
                        &tx,
                        &block_headers_clone[block_num],
                        false,
                    )
                    .unwrap();
                });
                tx.commit().unwrap();
                Ok(())
            })
            .await
            .unwrap();

        block_headers
    }

    #[tokio::test]
    async fn insert_and_get_block_headers_by_number() {
        let mut store = create_test_store().await;
        let block_headers = insert_dummy_block_headers(&mut store).await;

        let block_header = Store::get_block_header_by_num(&store, 3.into()).await.unwrap().unwrap();
        assert_eq!(block_headers[3], block_header.0);
    }

    #[tokio::test]
    async fn insert_and_get_block_headers_by_list() {
        let mut store = create_test_store().await;
        let mock_block_headers = insert_dummy_block_headers(&mut store).await;

        let block_headers: Vec<BlockHeader> =
            Store::get_block_headers(&store, &[1.into(), 3.into()].into_iter().collect())
                .await
                .unwrap()
                .into_iter()
                .map(|(block_header, _has_notes)| block_header)
                .collect();
        assert_eq!(
            &[mock_block_headers[1].clone(), mock_block_headers[3].clone()],
            &block_headers[..]
        );
    }

    /// Tests that `insert_block_header` persists the tracked header and its MMR authentication
    /// nodes in the same call, so both are retrievable afterwards.
    #[tokio::test]
    async fn insert_block_header_stores_header_and_nodes() {
        let store = create_test_store().await;
        const TOTAL_BLOCKS: usize = 8;

        let headers: Vec<BlockHeader> = (0..TOTAL_BLOCKS)
            .map(|n| BlockHeader::mock(u32::try_from(n).unwrap(), None, None, &[]))
            .collect();
        let mut mmr = Mmr::default();
        for header in &headers {
            mmr.add(header.commitment()).expect("valid MMR append");
        }

        let tracked: BTreeSet<usize> = [5].into();
        let auth_nodes = collect_auth_nodes(&mmr, &headers, &tracked);
        let header = headers[5].clone();

        Store::insert_block_header(&store, &header, &auth_nodes, true).await.unwrap();

        // The header is stored and marked as tracked.
        let stored = Store::get_block_headers(&store, &[5.into()].into_iter().collect())
            .await
            .unwrap();
        assert_eq!(stored.len(), 1);
        assert_eq!(stored[0].0, header);

        let tracked = Store::get_tracked_block_header_numbers(&store).await.unwrap();
        assert_eq!(tracked, [5].into());

        // Every authentication node was stored by the same call.
        let stored_nodes =
            Store::get_partial_blockchain_nodes(&store, PartialBlockchainFilter::All)
                .await
                .unwrap();
        let expected: BTreeMap<InOrderIndex, Word> = auth_nodes.iter().copied().collect();
        assert_eq!(stored_nodes, expected);
    }

    /// Tests that a failure inserting the MMR nodes rolls back the block header written in the same
    /// call, proving both land in a single transaction.
    #[tokio::test]
    async fn insert_block_header_rolls_back_header_when_nodes_fail() {
        let store = create_test_store().await;
        let header = BlockHeader::mock(5, None, None, &[]);
        // One node so the node insert actually runs (an empty slice would be a no-op).
        let nodes = [(InOrderIndex::from_leaf_pos(5), header.commitment())];

        // Force the node insert (the second statement in the transaction) to fail.
        store
            .interact_with_connection(|conn| {
                conn.execute("DROP TABLE partial_blockchain_nodes", []).unwrap();
                Ok(())
            })
            .await
            .unwrap();

        let result = Store::insert_block_header(&store, &header, &nodes, true).await;
        assert!(result.is_err(), "node insert must fail against the dropped table");

        // The header must not survive: a non-atomic two-transaction insert would leave it behind.
        let stored = Store::get_block_headers(&store, &[5.into()].into_iter().collect())
            .await
            .unwrap();
        assert!(stored.is_empty(), "header must roll back when the node insert fails");
    }

    /// Tests that large stored MMRs are built consistently throughout multiple prunes
    #[tokio::test]
    async fn partial_mmr_reconstructs_after_multiple_prune() {
        // Setup (mock a large MMR to work with, with a partial tracked set)
        // ----------------------------------------------------------------------------------------

        let store = create_test_store().await;
        const TOTAL_BLOCKS: usize = 7300;

        let block_headers: Vec<BlockHeader> = (0..TOTAL_BLOCKS)
            .map(|block_num| BlockHeader::mock(u32::try_from(block_num).unwrap(), None, None, &[]))
            .collect();

        let mut mmr = Mmr::default();
        for header in &block_headers {
            mmr.add(header.commitment()).expect("valid MMR append");
        }

        let mut tracked_set: BTreeSet<usize> = (0..(TOTAL_BLOCKS - 1)).step_by(97).collect();
        tracked_set.insert(TOTAL_BLOCKS - 2);
        let tracked_blocks: Vec<usize> = tracked_set.iter().copied().collect();

        let mut tracked_nodes: BTreeMap<InOrderIndex, Word> = BTreeMap::new();
        for &block_num in &tracked_blocks {
            let header = &block_headers[block_num];
            tracked_nodes.insert(InOrderIndex::from_leaf_pos(block_num), header.commitment());

            let proof = mmr.open(block_num).expect("valid proof");
            let mut idx = InOrderIndex::from_leaf_pos(block_num);
            for node in proof.merkle_path().nodes() {
                tracked_nodes.insert(idx.sibling(), *node);
                idx = idx.parent();
            }
        }
        let tracked_nodes: Vec<(InOrderIndex, Word)> = tracked_nodes.into_iter().collect();

        let peaks_by_block: Vec<MmrPeaks> = (0..TOTAL_BLOCKS)
            .map(|block_num| {
                mmr.peaks_at(Forest::new(block_num).expect("valid forest"))
                    .expect("valid peaks")
            })
            .collect();

        // Save blocks and nodes
        store
            .interact_with_connection(move |conn| {
                let tx = conn.transaction().unwrap();
                for (block_num, block_header) in block_headers.iter().enumerate() {
                    let has_notes = tracked_set.contains(&block_num);
                    SqliteStore::insert_block_header_tx(&tx, block_header, has_notes).unwrap();
                }

                SqliteStore::insert_partial_blockchain_nodes_tx(&tx, &tracked_nodes).unwrap();
                tx.commit().unwrap();
                Ok(())
            })
            .await
            .unwrap();

        let prune_heights = [
            TOTAL_BLOCKS / 5,
            (TOTAL_BLOCKS * 2) / 5,
            (TOTAL_BLOCKS * 3) / 5,
            TOTAL_BLOCKS - 1,
        ];

        // Tests/assertions
        // ----------------------------------------------------------------------------------------

        let mut previous_remaining: Option<i64> = None;
        for height in prune_heights {
            let height_i64 = i64::try_from(height).expect("fits in i64");
            let peaks_bytes = peaks_by_block[height].peaks().to_vec().to_bytes();

            // Update sync height (and the matching MMR peaks) to simulate having synced further
            store
                .interact_with_connection(move |conn| {
                    conn.execute(
                        "UPDATE blockchain_checkpoint SET block_num = ?, partial_blockchain_peaks = ?",
                        params![height_i64, peaks_bytes],
                    )
                    .unwrap();
                    Ok(())
                })
                .await
                .unwrap();

            // Prune
            store.untrack_and_prune_irrelevant_blocks(&[], &[]).await.unwrap();

            // Assert blocks
            let remaining_headers: i64 = store
                .interact_with_connection(|conn| {
                    let count = conn
                        .query_row("SELECT COUNT(*) FROM block_headers", [], |row| row.get(0))
                        .unwrap();
                    Ok(count)
                })
                .await
                .unwrap();
            if let Some(previous) = previous_remaining {
                assert!(remaining_headers < previous);
            } else {
                assert!(remaining_headers < i64::try_from(TOTAL_BLOCKS).unwrap());
            }
            previous_remaining = Some(remaining_headers);
        }

        // Try build MMR
        let partial_mmr = Store::get_current_partial_mmr(&store).await.unwrap();
        assert_eq!(partial_mmr.peaks().hash_peaks(), mmr.peaks().hash_peaks());

        for block_num in tracked_blocks {
            let partial_proof = partial_mmr.open(block_num).expect("partial mmr query succeeds");
            assert!(partial_proof.is_some());
            assert_eq!(
                partial_proof.unwrap().merkle_path(),
                mmr.open(block_num).unwrap().merkle_path()
            );
        }
    }

    /// Collects authentication nodes for a set of tracked leaves in an MMR.
    fn collect_auth_nodes(
        mmr: &Mmr,
        block_headers: &[BlockHeader],
        tracked: &BTreeSet<usize>,
    ) -> Vec<(InOrderIndex, Word)> {
        let mut nodes: BTreeMap<InOrderIndex, Word> = BTreeMap::new();
        for &block_num in tracked {
            nodes.insert(
                InOrderIndex::from_leaf_pos(block_num),
                block_headers[block_num].commitment(),
            );
            let proof = mmr.open(block_num).expect("valid proof");
            let mut idx = InOrderIndex::from_leaf_pos(block_num);
            for node in proof.merkle_path().nodes() {
                nodes.insert(idx.sibling(), *node);
                idx = idx.parent();
            }
        }
        nodes.into_iter().collect()
    }

    /// Tests that `untrack_and_prune_irrelevant_blocks` removes redundant authentication nodes for
    /// untracked blocks while preserving nodes needed by blocks that remain tracked.
    #[tokio::test]
    async fn prune_irrelevant_blocks_removes_redundant_auth_nodes() {
        let store = create_test_store().await;
        const TOTAL_BLOCKS: usize = 16;

        let headers: Vec<BlockHeader> = (0..TOTAL_BLOCKS)
            .map(|n| BlockHeader::mock(u32::try_from(n).unwrap(), None, None, &[]))
            .collect();
        let mut mmr = Mmr::default();
        for h in &headers {
            mmr.add(h.commitment()).expect("valid MMR append");
        }

        // Track blocks 3 and 10; we will untrack 3 later.
        let tracked: BTreeSet<usize> = [3, 10].into();
        let auth_nodes = collect_auth_nodes(&mmr, &headers, &tracked);
        let tip_peaks_bytes = mmr
            .peaks_at(Forest::new(TOTAL_BLOCKS - 1).expect("valid forest"))
            .unwrap()
            .peaks()
            .to_vec()
            .to_bytes();

        // Persist everything.
        let headers_clone = headers.clone();
        store
            .interact_with_connection(move |conn| {
                let tx = conn.transaction().unwrap();
                for (i, header) in headers_clone.iter().enumerate().take(TOTAL_BLOCKS) {
                    SqliteStore::insert_block_header_tx(&tx, header, tracked.contains(&i)).unwrap();
                }
                SqliteStore::insert_partial_blockchain_nodes_tx(&tx, &auth_nodes).unwrap();
                tx.execute(
                    "UPDATE blockchain_checkpoint SET block_num = ?, partial_blockchain_peaks = ?",
                    params![i64::try_from(TOTAL_BLOCKS - 1).unwrap(), tip_peaks_bytes],
                )
                .unwrap();
                tx.commit().unwrap();
                Ok(())
            })
            .await
            .unwrap();

        // Untrack block 3 via the PartialMmr, then prune.
        let mut partial_mmr = Store::get_current_partial_mmr(&store).await.unwrap();
        let removed: Vec<InOrderIndex> =
            partial_mmr.untrack(3).into_iter().map(|(idx, _)| idx).collect();
        assert!(!removed.is_empty(), "untracking should remove at least one node");

        store
            .untrack_and_prune_irrelevant_blocks(&[BlockNumber::from(3u32)], &removed)
            .await
            .unwrap();

        // Block 3 header should be deleted, block 10 should still be provable.
        let rebuilt = Store::get_current_partial_mmr(&store).await.unwrap();
        assert_eq!(rebuilt.peaks().hash_peaks(), mmr.peaks().hash_peaks());

        let proof_10 = rebuilt.open(10).expect("open succeeds");
        assert!(proof_10.is_some(), "block 10 should still be provable");

        let proof_3 = rebuilt.open(3).expect("open succeeds");
        assert!(proof_3.is_none(), "block 3 should no longer be provable");
    }
}