From 7aeec0c732d8ec0e1b9b90d6f1bb2ab25dc15505 Mon Sep 17 00:00:00 2001 From: Michael Zaikin Date: Wed, 11 Feb 2026 15:37:38 +0000 Subject: [PATCH] feat(discovery-core): add paginated last-note-index search with bisection --- .../src/discovery/last_note_index.rs | 332 ++++++++++++++++-- crates/discovery-core/src/discovery/notes.rs | 8 +- 2 files changed, 309 insertions(+), 31 deletions(-) diff --git a/crates/discovery-core/src/discovery/last_note_index.rs b/crates/discovery-core/src/discovery/last_note_index.rs index 04d5bf70d..a209964c0 100644 --- a/crates/discovery-core/src/discovery/last_note_index.rs +++ b/crates/discovery-core/src/discovery/last_note_index.rs @@ -1,23 +1,26 @@ //! Note index boundary probing. //! -//! Probes note existence at exponentially increasing indices — no decryption, -//! no nullifier checks. Used by notes discovery to find `max_note_index` -//! for a linear scan. +//! Finds the last note index in a subchannel via exponential search + bisection. +//! No decryption, no nullifier checks — just probes note existence. +//! +//! [`exponential_probe`] is shared between notes discovery (finds `max_note_index` +//! for a linear scan) and [`find_last_note_index_paginated`] (outgoing channels). use std::collections::HashMap; +use std::future::Future; use starknet_types_core::felt::Felt; +use crate::discovery::cursor::SubchannelCursor; use crate::discovery::{DiscoveryError, COST_NOTE_PROBING}; use crate::io_budget::IoBudget; use crate::privacy_pool::hashes::compute_note_id; use crate::privacy_pool::views::IViews; -/// Maximum offset for exponential probing. Caps the jump distance to avoid -/// overshooting into sparse index space (offsets: 0, 1, 3, 7, …, 1023). -/// TODO: make configurable -const MAX_PROBE_OFFSET: u64 = 1024; +/// Default max probe offset for notes discovery. Caps the jump distance to +/// avoid overshooting into sparse index space (offsets: 0, 1, 3, 7, …, 1023). +pub const DEFAULT_MAX_PROBE_OFFSET: u64 = 1024; /// Result of a batched exponential probe. pub struct ExponentialProbeResult { @@ -27,46 +30,148 @@ pub struct ExponentialProbeResult { /// Last index confirmed to exist. /// `None` if the first probe missed (empty subchannel) or budget exhausted. pub last_found_index: Option, + /// First index where no note exists (`None` if all probes hit = need more). + /// Used by [`find_last_note_index_paginated`] to set the bisection upper bound. + pub first_empty_index: Option, /// Whether the probe conclusively found the boundary (sentinel or start-of-empty). /// `false` when the batch was truncated by budget before reaching a boundary. pub probe_complete: bool, } +/// Finds the last note index via [`exponential_probe`] + [`bisect_boundary`]. +/// +/// Runs in two phases, resumable via cursor: +/// 1. **Ascending**: batched exponential probing to find bounds (`lo`, `hi`). +/// 2. **Bisection**: sequential binary search for exact boundary. +/// +/// Returns `(last_index, has_more)`. When `has_more` is `false`, the search is complete. +pub async fn find_last_note_index_paginated( + pool: &S, + channel_key: Felt, + token: Felt, + cursor: &mut SubchannelCursor, + budget: &IoBudget, +) -> Result<(Option, bool), DiscoveryError> { + // Phase 1: Ascending (find bounds) + if cursor.max_note_index.is_none() { + let start = cursor.start_index(); + let result = exponential_probe(pool, channel_key, token, start, None, None, budget).await?; + + if let Some(last_found_index) = result.last_found_index { + cursor.last_note_index = Some(last_found_index); + } + if let Some(index) = result.first_empty_index { + cursor.max_note_index = Some(index); + } + + // Budget exhausted, all probes hit, or empty subchannel — nothing to bisect yet. + if cursor.max_note_index.is_none() || cursor.last_note_index.is_none() { + let has_more = !result.probe_complete; + return Ok((cursor.last_note_index, has_more)); + } + } + + // Phase 2: Bisection (narrow down to exact boundary, sequential) + // Probe returns `(Option, bool)`: + // `Some(idx)` = note exists at idx, `None` = absent. + // `bool` = budget exhausted (true = stop early). + let probe = |idx: u64| { + let budget = budget.clone(); + async move { + if !budget.consume(COST_NOTE_PROBING) { + return Ok((None, true)); + } + let note_id = compute_note_id(channel_key, token, idx); + if pool.get_note(note_id).await? != Felt::ZERO { + Ok((Some(idx), false)) + } else { + Ok((None, false)) + } + } + }; + let complete = bisect_boundary(&probe, cursor).await?; + Ok((cursor.last_note_index, !complete)) +} + +/// Binary search between `lo` (exists) and `hi` (absent) to find +/// the exact last occupied index. +/// +/// Requires both `cursor.last_note_index` (lo) and `cursor.max_note_index` (hi) +/// to be `Some`. The probe closure returns `(Option, bool)` — see +/// [`find_last_note_index_paginated`] for the contract. +/// +/// Updates cursor in place. Returns `true` if complete, `false` if budget exhausted. +async fn bisect_boundary( + probe: F, + cursor: &mut SubchannelCursor, +) -> Result +where + F: Fn(u64) -> Fut, + Fut: Future, bool), DiscoveryError>>, +{ + let (mut lo, mut hi) = match (cursor.last_note_index, cursor.max_note_index) { + (Some(lo), Some(hi)) => (lo, hi), + _ => unreachable!("bisect_boundary called without both lo and hi"), + }; + + while lo + 1 < hi { + let mid = lo + (hi - lo) / 2; + match probe(mid).await? { + (Some(_), _) => lo = mid, + (None, false) => hi = mid, + (None, true) => { + cursor.last_note_index = Some(lo); + cursor.max_note_index = Some(hi); + return Ok(false); + } + } + } + + cursor.last_note_index = Some(lo); + cursor.max_note_index = Some(hi); + Ok(true) +} + /// Probes note existence at exponentially increasing indices in a single batch. /// /// From `start`, probes at offsets `0, 1, 3, 7, 15, ..., 2^k - 1` capped at -/// [`MAX_PROBE_OFFSET`]. The `2^k - 1` pattern is denser at the start than +/// `max_probe_offset`. The `2^k - 1` pattern is denser at the start than /// powers of 2, yielding more cache hits for small subchannels. /// -/// `prior_max` narrows the search when a previous probe already found a bound: -/// - `None` (first discovery): offsets up to `MAX_PROBE_OFFSET` +/// `prior_max_index` narrows the search when a previous probe already found a bound: +/// - `None` (first discovery): offsets up to `max_probe_offset` /// - `Some(m)` (re-probe after scan): range capped at `m * 2` /// +/// `max_probe_offset` caps the maximum offset per batch. Use +/// `Some(`[`DEFAULT_MAX_PROBE_OFFSET`]`)` for notes discovery or `None` for +/// unbounded probing (e.g. last-note-index search). +/// /// All probed notes that exist are cached in the result for use by the /// linear scan phase. /// /// Issues exactly one batch of probes via `pool.get_notes_batch()`. Callers /// handle iteration via cursor-based pagination. -pub async fn exponential_ascend( +pub async fn exponential_probe( pool: &S, channel_key: Felt, token: Felt, start_index: u64, prior_max_index: Option, + max_probe_offset: Option, budget: &IoBudget, ) -> Result { let max_index = match prior_max_index { - None => start_index.saturating_add(MAX_PROBE_OFFSET), + None => start_index.saturating_add(max_probe_offset.unwrap_or(u64::MAX)), Some(prior_max) => prior_max.saturating_mul(2).max(start_index), }; let max_probe_offset = max_index.saturating_sub(start_index); - // Offsets: [0, 1, 3, 7, 15, ..., 2^k - 1] where 2^k - 1 <= min(max_probe_offset, MAX_PROBE_OFFSET). + // Offsets: [0, 1, 3, 7, 15, ..., 2^k - 1] where 2^k - 1 <= max_probe_offset. let offsets: Vec = std::iter::once(0) .chain( (1..64) .map(|exp| (1u64 << exp) - 1) - .take_while(|&offset| offset <= max_probe_offset.min(MAX_PROBE_OFFSET)), + .take_while(|&offset| offset <= max_probe_offset), ) .collect(); @@ -76,6 +181,7 @@ pub async fn exponential_ascend( return Ok(ExponentialProbeResult { cache: HashMap::new(), last_found_index: None, + first_empty_index: None, probe_complete: !budget_exhausted, }); } @@ -108,6 +214,7 @@ pub async fn exponential_ascend( Ok(ExponentialProbeResult { cache, last_found_index, + first_empty_index, probe_complete, }) } @@ -118,6 +225,7 @@ mod tests { use crate::privacy_pool::hashes::compute_note_id; use crate::privacy_pool::storage_slots; use crate::storage_backend::MockBackend; + use crate::test_fixtures::{get_channel_key, get_subchannel_token, load_devnet_fixture}; const CK: Felt = Felt::from_hex_unchecked("0x12345"); const TK: Felt = Felt::from_hex_unchecked("0x67890"); @@ -136,12 +244,20 @@ mod tests { } #[tokio::test] - async fn test_exponential_ascend_empty() { + async fn test_exponential_probe_empty() { let backend = mock_with_notes(None); let budget = IoBudget::new(100); - let result = exponential_ascend(&backend, CK, TK, 0, None, &budget) - .await - .unwrap(); + let result = exponential_probe( + &backend, + CK, + TK, + 0, + None, + Some(DEFAULT_MAX_PROBE_OFFSET), + &budget, + ) + .await + .unwrap(); assert_eq!(result.last_found_index, None); assert!(result.cache.is_empty()); @@ -149,13 +265,21 @@ mod tests { } #[tokio::test] - async fn test_exponential_ascend_one_element() { + async fn test_exponential_probe_one_element() { // Elements at 0 only. Probes: offset 0→0 (hit), 1→1 (miss). let backend = mock_with_notes(Some(0)); let budget = IoBudget::new(100); - let result = exponential_ascend(&backend, CK, TK, 0, None, &budget) - .await - .unwrap(); + let result = exponential_probe( + &backend, + CK, + TK, + 0, + None, + Some(DEFAULT_MAX_PROBE_OFFSET), + &budget, + ) + .await + .unwrap(); assert_eq!(result.last_found_index, Some(0)); assert_eq!(result.cache.len(), 1); @@ -163,14 +287,22 @@ mod tests { } #[tokio::test] - async fn test_exponential_ascend_multiple_elements() { + async fn test_exponential_probe_multiple_elements() { // Elements at 0..=4. Probes: offset 0→0 (hit), 1→1 (hit), // 3→3 (hit), 7→7 (miss) let backend = mock_with_notes(Some(4)); let budget = IoBudget::new(100); - let result = exponential_ascend(&backend, CK, TK, 0, None, &budget) - .await - .unwrap(); + let result = exponential_probe( + &backend, + CK, + TK, + 0, + None, + Some(DEFAULT_MAX_PROBE_OFFSET), + &budget, + ) + .await + .unwrap(); assert_eq!(result.last_found_index, Some(3)); assert_eq!(result.cache.len(), 3); @@ -182,12 +314,154 @@ mod tests { let backend = mock_with_notes(Some(100)); // Budget for only 2 probes: offsets 0, 1 let budget = IoBudget::new(2 * COST_NOTE_PROBING); - let result = exponential_ascend(&backend, CK, TK, 0, None, &budget) - .await - .unwrap(); + let result = exponential_probe( + &backend, + CK, + TK, + 0, + None, + Some(DEFAULT_MAX_PROBE_OFFSET), + &budget, + ) + .await + .unwrap(); assert_eq!(result.last_found_index, Some(1)); assert_eq!(result.cache.len(), 2); assert!(!result.probe_complete); } + + /// Helper: mock bisection probe. + fn mock_probe( + last_index: Option, + ) -> impl Fn(u64) -> std::future::Ready, bool), DiscoveryError>> { + move |idx| { + let exists = last_index.is_some_and(|last| idx <= last); + std::future::ready(Ok(if exists { + (Some(idx), false) + } else { + (None, false) + })) + } + } + + #[tokio::test] + async fn test_bisect_boundary_adjacent() { + let mut cursor = SubchannelCursor { + note_discovery_complete: false, + last_note_index: Some(0), + max_note_index: Some(1), + }; + let complete = bisect_boundary(mock_probe(Some(0)), &mut cursor) + .await + .unwrap(); + assert!(complete); + assert_eq!(cursor.last_note_index, Some(0)); + } + + #[tokio::test] + async fn test_bisect_boundary_gap() { + let mut cursor = SubchannelCursor { + note_discovery_complete: false, + last_note_index: Some(3), + max_note_index: Some(7), + }; + let complete = bisect_boundary(mock_probe(Some(4)), &mut cursor) + .await + .unwrap(); + assert!(complete); + assert_eq!(cursor.last_note_index, Some(4)); + assert_eq!(cursor.max_note_index, Some(5)); + } + + #[tokio::test] + async fn test_find_last_note_index_with_fixture() { + let fixture = load_devnet_fixture(); + let backend = MockBackend::new(fixture.slots); + + let channel_key = get_channel_key( + &backend, + fixture.constants.alice_address, + &fixture.constants.alice_viewing_key, + ) + .await + .expect("Alice should have at least one channel"); + + let token = get_subchannel_token(&backend, channel_key) + .await + .expect("Alice's channel should have at least one subchannel"); + + let mut cursor = SubchannelCursor::default(); + let budget = IoBudget::new(100); + + let (last_index, has_more) = + find_last_note_index_paginated(&backend, channel_key, token, &mut cursor, &budget) + .await + .unwrap(); + + // Alice has 1 note at index 0 + assert_eq!(last_index, Some(0)); + assert!(!has_more); + } + + #[tokio::test] + async fn test_find_last_note_index_empty() { + let backend = MockBackend::empty(); + let channel_key = Felt::from_hex_unchecked("0x12345"); + let token = Felt::from_hex_unchecked("0x67890"); + + let mut cursor = SubchannelCursor::default(); + let budget = IoBudget::new(100); + + let (last_index, has_more) = + find_last_note_index_paginated(&backend, channel_key, token, &mut cursor, &budget) + .await + .unwrap(); + + assert_eq!(last_index, None); + assert!(!has_more); + } + + #[tokio::test] + async fn test_find_last_note_index_budget_exhausted_ascending() { + let fixture = load_devnet_fixture(); + let backend = MockBackend::new(fixture.slots); + + let channel_key = get_channel_key( + &backend, + fixture.constants.alice_address, + &fixture.constants.alice_viewing_key, + ) + .await + .expect("Alice should have at least one channel"); + + let token = get_subchannel_token(&backend, channel_key) + .await + .expect("Alice's channel should have at least one subchannel"); + + let mut cursor = SubchannelCursor::default(); + + // Budget for 1 probe: batch gets offset 0 only (note at 0 exists), + // but no first_empty found — budget_exhausted (all probes hit). + let budget = IoBudget::new(COST_NOTE_PROBING); + let (last_index, has_more) = + find_last_note_index_paginated(&backend, channel_key, token, &mut cursor, &budget) + .await + .unwrap(); + + assert_eq!(last_index, Some(0)); + assert!(has_more, "budget exhausted during ascending"); + assert_eq!(cursor.last_note_index, Some(0)); + assert!(cursor.max_note_index.is_none(), "sentinel not found yet"); + + // Resume with more budget + let budget = IoBudget::new(100); + let (last_index, has_more) = + find_last_note_index_paginated(&backend, channel_key, token, &mut cursor, &budget) + .await + .unwrap(); + + assert_eq!(last_index, Some(0)); + assert!(!has_more); + } } diff --git a/crates/discovery-core/src/discovery/notes.rs b/crates/discovery-core/src/discovery/notes.rs index c25736607..cb54d197b 100644 --- a/crates/discovery-core/src/discovery/notes.rs +++ b/crates/discovery-core/src/discovery/notes.rs @@ -19,7 +19,9 @@ use starknet_types_core::felt::Felt; use tracing::{debug, trace}; use crate::discovery::cursor::SubchannelCursor; -use crate::discovery::last_note_index::{exponential_ascend, ExponentialProbeResult}; +use crate::discovery::last_note_index::{ + exponential_probe, ExponentialProbeResult, DEFAULT_MAX_PROBE_OFFSET, +}; use crate::discovery::{DiscoveryError, COST_NOTE, COST_NOTE_PROBING}; use crate::io_budget::IoBudget; use crate::privacy_pool::decryption::{decrypt_note_amount, unpack_note_amount}; @@ -164,12 +166,14 @@ async fn probe_note_boundary( cache, last_found_index, probe_complete, - } = exponential_ascend( + .. + } = exponential_probe( pool, channel_key, token, start_index, cursor.max_note_index, + Some(DEFAULT_MAX_PROBE_OFFSET), budget, ) .await?;