miden-node-store 0.17.0-rc.4

use std::collections::{BTreeMap, BTreeSet, HashSet};
use std::ops::RangeInclusive;

use diesel::dsl::sql;
use diesel::prelude::{ExpressionMethods, QueryDsl, Queryable, QueryableByName, Selectable};
use diesel::query_dsl::methods::SelectDsl;
use diesel::sql_types::BigInt;
use diesel::sqlite::Sqlite;
use diesel::{
    JoinOnDsl,
    NullableExpressionMethods,
    OptionalExtension,
    RunQueryDsl,
    SelectableHelper,
    SqliteConnection,
};
use miden_node_utils::limiter::{
    QueryParamLimiter,
    QueryParamNoteCommitmentLimit,
    QueryParamNoteTagLimit,
};
use miden_protocol::Word;
use miden_protocol::account::AccountId;
use miden_protocol::block::{BlockNoteIndex, BlockNumber};
use miden_protocol::crypto::merkle::SparseMerklePath;
use miden_protocol::note::{
    NoteAssets,
    NoteAttachments,
    NoteDetails,
    NoteId,
    NoteInclusionProof,
    NoteMetadata,
    NoteRecipient,
    NoteScript,
    NoteStorage,
    NoteTag,
    NoteType,
    Nullifier,
    PartialNoteMetadata,
};
use miden_protocol::utils::serde::{Deserializable, Serializable};

use crate::db::models::conv::{SqlTypeConvert, raw_sql_to_idx};
use crate::db::models::queries::select_block_header_by_block_num;
use crate::db::models::{serialize_vec, vec_raw_try_into};
use crate::db::{DatabaseError, NoteRecord, NoteSyncRecord, NoteSyncUpdate, schema};
use crate::errors::NoteSyncError;

/// Estimated byte size of a [`NoteSyncUpdate`] excluding its notes.
///
/// Includes a canonical header with validator keys, a scheduled protocol configuration, and an
/// MMR proof with 32 siblings.
pub(crate) const NOTE_SYNC_BLOCK_OVERHEAD_BYTES: usize = 1800;

/// Estimated byte size of a single [`NoteSyncRecord`].
///
/// Includes a note ID, an index, compact metadata with four attachment entries, and a sparse
/// Merkle path with 16 siblings.
pub(crate) const NOTE_SYNC_RECORD_BYTES: usize = 900;

/// Select notes matching the given tags within a block range.
///
/// # Parameters
/// * `note_tags`: List of note tags to filter by
///     - Limit: 0 <= count <= 1000
/// * `block_range`: Range of blocks to search (inclusive)
///
/// # Returns
///
/// All matching notes from the first block within the range containing a matching note. If no
/// matching notes are found at all, then an empty vector is returned.
///
/// # Raw SQL
///
/// ```sql
/// SELECT
///     committed_at,
///     batch_index,
///     note_index,
///     note_id,
///     note_type,
///     sender,
///     tag,
///     attachment,
///     inclusion_path
/// FROM
///     notes
/// WHERE
///     committed_at = (
///         SELECT
///             committed_at
///         FROM
///             notes
///         WHERE
///             tag IN (?1) AND
///             committed_at >= ?2 AND
///             committed_at <= ?3
///         ORDER BY
///             committed_at ASC
///         LIMIT 1
///     ) AND
///     tag IN (?1)
/// ORDER BY
///     committed_at ASC, batch_index ASC, note_index ASC
/// ```
pub(crate) fn select_notes_since_block_by_tag(
    conn: &mut SqliteConnection,
    note_tags: &[u32],
    block_range: RangeInclusive<BlockNumber>,
) -> Result<Vec<NoteSyncRecord>, DatabaseError> {
    QueryParamNoteTagLimit::check(note_tags.len())?;
    let desired_note_tags: Vec<i64> = note_tags.iter().copied().map(i64::from).collect();
    // The column stores tags as unsigned 32-bit values. The Diesel schema declares the column as a
    // 32-bit `Integer`, which cannot bind tags at or above 2^31.
    let tag_column = sql::<BigInt>("notes.tag");
    let start_block_num = block_range.start().to_raw_sql();
    let end_block_num = block_range.end().to_raw_sql();

    let Some(desired_block_num): Option<i64> =
        SelectDsl::select(schema::notes::table, schema::notes::committed_at)
            .filter(tag_column.clone().eq_any(&desired_note_tags))
            .filter(schema::notes::committed_at.ge(start_block_num))
            .filter(schema::notes::committed_at.le(end_block_num))
            .order_by(schema::notes::committed_at.asc())
            .limit(1)
            .get_result(conn)
            .optional()?
    else {
        return Ok(Vec::new());
    };

    let notes = SelectDsl::select(schema::notes::table, NoteSyncRecordRawRow::as_select())
        .filter(schema::notes::committed_at.eq(desired_block_num))
        .filter(tag_column.eq_any(&desired_note_tags))
        .order_by((
            schema::notes::committed_at.asc(),
            schema::notes::batch_index.asc(),
            schema::notes::note_index.asc(),
        ))
        .get_results::<NoteSyncRecordRawRow>(conn)
        .map_err(DatabaseError::from)?;

    vec_raw_try_into(notes)
}

/// Select all notes matching the given set of identifiers
///
/// # Raw SQL
///
/// ```sql
/// SELECT
///     notes.committed_at,
///     notes.batch_index,
///     notes.note_index,
///     notes.note_id,
///     notes.note_type,
///     notes.sender,
///     notes.tag,
///     notes.attachment,
///     notes.assets,
///     notes.storage,
///     notes.serial_num,
///     notes.inclusion_path,
///     note_scripts.script
/// FROM notes
/// LEFT JOIN note_scripts ON notes.script_root = note_scripts.script_root
/// WHERE note_id IN (?1)
/// ```
pub(crate) fn select_notes_by_id(
    conn: &mut SqliteConnection,
    note_ids: &[NoteId],
) -> Result<Vec<NoteRecord>, DatabaseError> {
    let note_ids = serialize_vec(note_ids);
    let q = schema::notes::table
        .left_join(
            schema::note_scripts::table
                .on(schema::notes::script_root.eq(schema::note_scripts::script_root.nullable())),
        )
        .filter(schema::notes::note_id.eq_any(&note_ids));
    let raw: Vec<_> = SelectDsl::select(
        q,
        (NoteRecordRawRow::as_select(), schema::note_scripts::script.nullable()),
    )
    .load::<(NoteRecordRawRow, Option<Vec<u8>>)>(conn)?;
    let records = vec_raw_try_into::<NoteRecord, NoteRecordWithScriptRawJoined>(
        raw.into_iter().map(NoteRecordWithScriptRawJoined::from),
    )?;
    Ok(records)
}

/// Select the requested note IDs that the notes table contains at or before `up_to_block`.
///
/// # Raw SQL
///
/// ```sql
/// SELECT
///     notes.note_id
/// FROM notes
/// WHERE note_id IN (?1) AND committed_at <= ?2
/// ```
pub(crate) fn select_existing_note_ids(
    conn: &mut SqliteConnection,
    note_ids: &[NoteId],
    up_to_block: BlockNumber,
) -> Result<HashSet<NoteId>, DatabaseError> {
    QueryParamNoteCommitmentLimit::check(note_ids.len())?;

    let note_ids = serialize_vec(note_ids);

    let raw_note_ids = SelectDsl::select(schema::notes::table, schema::notes::note_id)
        .filter(schema::notes::note_id.eq_any(&note_ids))
        .filter(schema::notes::committed_at.le(up_to_block.to_raw_sql()))
        .load::<Vec<u8>>(conn)?;

    let note_ids = raw_note_ids
        .into_iter()
        .map(|note_id| NoteId::read_from_bytes(&note_id))
        .collect::<Result<HashSet<_>, _>>()?;

    Ok(note_ids)
}

/// Select note inclusion proofs matching the note commitments, restricted to notes committed at
/// or before `up_to_block`.
///
/// # Parameters
/// * `note_ids`: Set of note IDs to query
///     - Limit: 0 <= count <= 1000
/// * `up_to_block`: Only notes committed at or before this block are returned
///
/// # Returns
///
/// - Empty map if no matching `note`.
/// - Otherwise, note inclusion proofs, which `note_id` matches the `NoteId` as bytes.
///
/// # Raw SQL
///
/// ```sql
/// SELECT
///     committed_at,
///     note_id,
///     batch_index,
///     note_index,
///     inclusion_path
/// FROM
///     notes
/// WHERE
///     note_id IN (?1) AND
///     committed_at <= ?2
/// ORDER BY
///     committed_at ASC
/// ```
pub(crate) fn select_note_inclusion_proofs(
    conn: &mut SqliteConnection,
    note_commitments: &BTreeSet<Word>,
    up_to_block: BlockNumber,
) -> Result<BTreeMap<NoteId, NoteInclusionProof>, DatabaseError> {
    QueryParamNoteCommitmentLimit::check(note_commitments.len())?;

    let note_commitments = serialize_vec(note_commitments.iter());

    let raw_notes = SelectDsl::select(
        schema::notes::table,
        (
            schema::notes::committed_at,
            schema::notes::note_id,
            schema::notes::batch_index,
            schema::notes::note_index,
            schema::notes::inclusion_path,
        ),
    )
    .filter(schema::notes::note_id.eq_any(note_commitments))
    .filter(schema::notes::committed_at.le(up_to_block.to_raw_sql()))
    .order_by(schema::notes::committed_at.asc())
    .load::<(i64, Vec<u8>, i32, i32, Vec<u8>)>(conn)?;

    raw_notes
        .iter()
        .map(|(block_num, note_id, batch_index, note_index, merkle_path)| {
            let note_id = NoteId::read_from_bytes(&note_id[..])?;
            let block_num = BlockNumber::from_raw_sql(*block_num)?;
            let node_index_in_block =
                BlockNoteIndex::new(raw_sql_to_idx(*batch_index), raw_sql_to_idx(*note_index))
                    .expect("batch and note index from DB should be valid")
                    .leaf_index_value();
            let merkle_path = miden_node_persistence::decode::<SparseMerklePath>(&merkle_path[..])?;
            let proof = NoteInclusionProof::new(block_num, node_index_in_block, merkle_path)?;
            Ok((note_id, proof))
        })
        .collect::<Result<BTreeMap<_, _>, _>>()
}

/// Select note sync records matching the given note commitments.
///
/// # Parameters
/// * `note_commitments`: Slice of note commitments to query
///     - Limit: 0 <= count <= 1000
///
/// # Returns
///
/// - Empty map if no matching `note`.
/// - Otherwise, note sync records keyed by `NoteId`.
///
/// # Raw SQL
///
/// ```sql
/// SELECT
///     committed_at,
///     batch_index,
///     note_index,
///     note_id,
///     note_commitment,
///     note_type,
///     sender,
///     tag,
///     attachment,
///     inclusion_path
/// FROM
///     notes
/// WHERE
///     note_commitment IN (?1)
/// ORDER BY
///     committed_at ASC
/// ```
pub(crate) fn select_note_sync_records(
    conn: &mut SqliteConnection,
    note_ids: &[NoteId],
) -> Result<BTreeMap<NoteId, NoteSyncRecord>, DatabaseError> {
    QueryParamNoteCommitmentLimit::check(note_ids.len())?;

    let note_id_bytes: Vec<Vec<u8>> = note_ids.iter().map(|id| id.as_word().to_bytes()).collect();

    let raw_notes = SelectDsl::select(schema::notes::table, NoteSyncRecordRawRow::as_select())
        .filter(schema::notes::note_id.eq_any(note_id_bytes))
        .order_by(schema::notes::committed_at.asc())
        .load::<NoteSyncRecordRawRow>(conn)?;

    raw_notes
        .into_iter()
        .map(|raw_note| {
            let note: NoteSyncRecord = raw_note.try_into()?;
            Ok((note.note_id, note))
        })
        .collect()
}

/// Maps each given nullifier to its note ID.
///
/// Only public notes have a nullifier stored (`notes.nullifier` is NULL for private notes), so
/// private notes never match and are absent from the result.
///
/// ```sql
/// SELECT
///     nullifier,
///     note_id
/// FROM
///     notes
/// WHERE
///     nullifier IN (?1)
/// ```
pub(crate) fn select_note_ids_by_nullifier(
    conn: &mut SqliteConnection,
    nullifiers: &[Nullifier],
) -> Result<BTreeMap<Nullifier, NoteId>, DatabaseError> {
    if nullifiers.is_empty() {
        return Ok(BTreeMap::new());
    }

    let nullifier_bytes: Vec<Vec<u8>> = nullifiers.iter().map(Nullifier::to_bytes).collect();
    let pairs =
        SelectDsl::select(schema::notes::table, (schema::notes::nullifier, schema::notes::note_id))
            .filter(schema::notes::nullifier.eq_any(nullifier_bytes))
            .load::<(Option<Vec<u8>>, Vec<u8>)>(conn)?;

    let mut note_ids_by_nullifier = BTreeMap::new();
    for (nullifier, note_id) in pairs {
        let Some(nullifier) = nullifier else { continue };
        let nullifier = Nullifier::read_from_bytes(&nullifier)?;
        let note_id = NoteId::read_from_bytes(&note_id)?;
        note_ids_by_nullifier.insert(nullifier, note_id);
    }
    Ok(note_ids_by_nullifier)
}

/// Returns the script for a note by its root.
///
/// ```sql
/// SELECT
///     script_root,
///     script
/// FROM
///     note_scripts
/// WHERE
///     script_root = ?1
/// ```
pub(crate) fn select_note_script_by_root(
    conn: &mut SqliteConnection,
    root: Word,
) -> Result<Option<NoteScript>, DatabaseError> {
    let raw = SelectDsl::select(schema::note_scripts::table, schema::note_scripts::script)
        .filter(schema::note_scripts::script_root.eq(root.to_bytes()))
        .get_result::<Vec<u8>>(conn)
        .optional()?;

    raw.as_ref()
        .map(|bytes| miden_node_persistence::decode::<NoteScript>(bytes))
        .transpose()
        .map_err(Into::into)
}

/// Loads the data necessary for a note sync across all matching blocks in the given range.
///
/// Returns one [`NoteSyncUpdate`] per block that contains at least one note matching the
/// requested tags, ordered by block number ascending.
pub(crate) fn get_note_sync_multi(
    conn: &mut SqliteConnection,
    note_tags: &[u32],
    block_range: RangeInclusive<BlockNumber>,
    max_response_payload_bytes: usize,
) -> Result<Vec<NoteSyncUpdate>, NoteSyncError> {
    let mut current_from = *block_range.start();
    let block_end = *block_range.end();
    let mut updates = Vec::new();
    let mut accumulated_size = 0usize;

    loop {
        let notes = select_notes_since_block_by_tag(conn, note_tags, current_from..=block_end)?;

        let Some(block_num) = notes.first().map(|note| note.block_num) else {
            break;
        };

        accumulated_size += NOTE_SYNC_BLOCK_OVERHEAD_BYTES + notes.len() * NOTE_SYNC_RECORD_BYTES;

        if !updates.is_empty() && accumulated_size > max_response_payload_bytes {
            break;
        }

        let block_header = select_block_header_by_block_num(conn, Some(block_num))?
            .ok_or(NoteSyncError::EmptyBlockHeadersTable)?;
        updates.push(NoteSyncUpdate { notes, block_header });
        current_from = block_num + 1;
    }

    Ok(updates)
}

#[derive(Debug, Clone, PartialEq, Selectable, Queryable, QueryableByName)]
#[diesel(table_name = schema::notes)]
#[diesel(check_for_backend(Sqlite))]
pub struct NoteSyncRecordRawRow {
    pub committed_at: i64, // BlockNumber
    #[diesel(embed)]
    pub block_note_index: BlockNoteIndexRawRow,
    pub note_id: Vec<u8>, // BlobDigest
    #[diesel(embed)]
    pub metadata: NoteMetadataRawRow,
    pub inclusion_path: Vec<u8>, // SparseMerklePath
}

impl TryInto<NoteSyncRecord> for NoteSyncRecordRawRow {
    type Error = DatabaseError;
    fn try_into(self) -> Result<NoteSyncRecord, Self::Error> {
        let block_num = BlockNumber::from_raw_sql(self.committed_at)?;
        let note_index = self.block_note_index.try_into()?;

        let note_id = NoteId::from_raw(Word::read_from_bytes(&self.note_id[..])?);
        let inclusion_path =
            miden_node_persistence::decode::<SparseMerklePath>(&self.inclusion_path[..])?;
        let (metadata, attachments) = self.metadata.try_into()?;
        Ok(NoteSyncRecord {
            block_num,
            note_index,
            note_id,
            metadata,
            attachments,
            inclusion_path,
        })
    }
}

#[derive(Debug, Clone, PartialEq, Selectable, Queryable, QueryableByName)]
#[diesel(table_name = schema::notes)]
#[diesel(check_for_backend(Sqlite))]
pub struct NoteDetailsRawRow {
    pub assets: Option<Vec<u8>>,
    pub storage: Option<Vec<u8>>,
    pub serial_num: Option<Vec<u8>>,
}

// Note: One cannot use `#[diesel(embed)]` to structure this, it will yield a significant amount of
// errors when used with join and debugging is painful to put it mildly.
#[derive(Debug, Clone, PartialEq, Queryable)]
pub struct NoteRecordWithScriptRawJoined {
    pub committed_at: i64,

    pub batch_index: i32,
    pub note_index: i32, // index within batch
    // #[diesel(embed)]
    // pub note_index: BlockNoteIndexRaw,
    pub note_id: Vec<u8>,

    pub note_type: i32,
    pub sender: Vec<u8>, // AccountId
    pub tag: i32,
    pub attachment: Vec<u8>,
    // #[diesel(embed)]
    // pub metadata: NoteMetadataRaw,
    pub assets: Option<Vec<u8>>,
    pub storage: Option<Vec<u8>>,
    pub serial_num: Option<Vec<u8>>,

    // #[diesel(embed)]
    // pub details: NoteDetailsRaw,
    pub inclusion_path: Vec<u8>,
    pub script: Option<Vec<u8>>, // not part of notes::table!
}

impl From<(NoteRecordRawRow, Option<Vec<u8>>)> for NoteRecordWithScriptRawJoined {
    fn from((note, script): (NoteRecordRawRow, Option<Vec<u8>>)) -> Self {
        let NoteRecordRawRow {
            committed_at,
            batch_index,
            note_index,
            note_id,
            note_type,
            sender,
            tag,
            attachment,
            assets,
            storage,
            serial_num,
            inclusion_path,
        } = note;
        Self {
            committed_at,
            batch_index,
            note_index,
            note_id,
            note_type,
            sender,
            tag,
            attachment,
            assets,
            storage,
            serial_num,
            inclusion_path,
            script,
        }
    }
}

impl TryInto<NoteRecord> for NoteRecordWithScriptRawJoined {
    type Error = DatabaseError;
    fn try_into(self) -> Result<NoteRecord, Self::Error> {
        // let (raw, script) = self;
        let raw = self;
        let NoteRecordWithScriptRawJoined {
            committed_at,

            batch_index,
            note_index,
            // block note index ^^^
            note_id,

            note_type,
            sender,
            tag,
            attachment,
            // metadata ^^^,
            assets,
            storage,
            serial_num,
            // details ^^^,
            inclusion_path,
            script,
            ..
        } = raw;
        let index = BlockNoteIndexRawRow { batch_index, note_index };
        let metadata = NoteMetadataRawRow { note_type, sender, tag, attachment };
        let details = NoteDetailsRawRow { assets, storage, serial_num };

        let (metadata, attachments) = metadata.try_into()?;
        let committed_at = BlockNumber::from_raw_sql(committed_at)?;
        let note_id = Word::read_from_bytes(&note_id[..])?;
        let script = script
            .map(|script| miden_node_persistence::decode::<NoteScript>(&script[..]))
            .transpose()?;
        let details = if let NoteDetailsRawRow {
            assets: Some(assets),
            storage: Some(storage),
            serial_num: Some(serial_num),
        } = details
        {
            let storage = miden_node_persistence::decode::<NoteStorage>(&storage[..])?;
            let serial_num = Word::read_from_bytes(&serial_num[..])?;
            let script =
                script.ok_or_else(|| {
                    miden_node_db::DatabaseError::conversiont_from_sql::<
                        NoteRecipient,
                        DatabaseError,
                        _,
                    >(None)
                })?;
            let recipient = NoteRecipient::new(serial_num, script, storage);
            let assets = miden_node_persistence::decode::<NoteAssets>(&assets[..])?;
            Some(NoteDetails::new(assets, recipient))
        } else {
            None
        };
        let inclusion_path =
            miden_node_persistence::decode::<SparseMerklePath>(&inclusion_path[..])?;
        let note_index = index.try_into()?;
        Ok(NoteRecord {
            block_num: committed_at,
            note_index,
            note_id,
            metadata,
            details,
            attachments,
            inclusion_path,
        })
    }
}

#[derive(Debug, Clone, PartialEq, Selectable, Queryable, QueryableByName)]
#[diesel(table_name = schema::notes)]
#[diesel(check_for_backend(Sqlite))]
pub struct NoteRecordRawRow {
    pub committed_at: i64,

    pub batch_index: i32,
    pub note_index: i32, // index within batch
    pub note_id: Vec<u8>,

    pub note_type: i32,
    pub sender: Vec<u8>, // AccountId
    pub tag: i32,
    pub attachment: Vec<u8>,

    pub assets: Option<Vec<u8>>,
    pub storage: Option<Vec<u8>>,
    pub serial_num: Option<Vec<u8>>,

    pub inclusion_path: Vec<u8>,
}

#[derive(Debug, Clone, PartialEq, Selectable, Queryable, QueryableByName)]
#[diesel(table_name = schema::notes)]
#[diesel(check_for_backend(Sqlite))]
pub struct NoteMetadataRawRow {
    note_type: i32,
    sender: Vec<u8>, // AccountId
    tag: i32,
    attachment: Vec<u8>,
}

#[expect(clippy::cast_sign_loss, clippy::cast_possible_truncation)]
impl TryInto<(NoteMetadata, NoteAttachments)> for NoteMetadataRawRow {
    type Error = DatabaseError;
    fn try_into(self) -> Result<(NoteMetadata, NoteAttachments), Self::Error> {
        let sender = AccountId::read_from_bytes(&self.sender[..])?;
        let note_type = NoteType::try_from(self.note_type as u8)
            .map_err(miden_node_db::DatabaseError::conversiont_from_sql::<NoteType, _, _>)?;
        let tag = NoteTag::new(self.tag as u32);
        let attachments = miden_node_persistence::decode::<NoteAttachments>(&self.attachment)?;
        let partial = PartialNoteMetadata::new(sender, note_type).with_tag(tag);
        let metadata = NoteMetadata::new(partial, &attachments);
        Ok((metadata, attachments))
    }
}

#[derive(Debug, Clone, PartialEq, Selectable, Queryable, QueryableByName)]
#[diesel(table_name = schema::notes)]
#[diesel(check_for_backend(Sqlite))]
pub struct BlockNoteIndexRawRow {
    pub batch_index: i32,
    pub note_index: i32, // index within batch
}

#[expect(clippy::cast_sign_loss, reason = "Indices are cast to usize for ease of use")]
impl TryInto<BlockNoteIndex> for BlockNoteIndexRawRow {
    type Error = DatabaseError;
    fn try_into(self) -> Result<BlockNoteIndex, Self::Error> {
        let batch_index = self.batch_index as usize;
        let note_index = self.note_index as usize;
        let index = BlockNoteIndex::new(batch_index, note_index).ok_or_else(|| {
            miden_node_db::DatabaseError::conversiont_from_sql::<BlockNoteIndex, DatabaseError, _>(
                None,
            )
        })?;
        Ok(index)
    }
}