armdb 0.14.2

use crate::DbResult;

/// Configuration for bounded pagination over prepared row chunks.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PageSpec {
    /// Maximum number of accepted items to return.
    pub limit: usize,
    /// Maximum number of rows prepared together.
    ///
    /// Recommended: `limit`. A smaller chunk prepares more often; a larger one
    /// prepares rows the page may never reach.
    pub chunk: usize,
    /// Optional maximum number of examined rows.
    pub max_examined: Option<usize>,
}

/// Why a chunked page stopped reading rows.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StopReason {
    /// The requested number of accepted items was collected.
    ///
    /// This takes precedence over both other reasons: over
    /// [`StopReason::Exhausted`] when the filling row is the input's final row,
    /// and over [`StopReason::BudgetExhausted`] when the filling row is also the
    /// budget's final row. No look-ahead is performed, so a valid continuation
    /// can be empty.
    Filled,
    /// Pulling a batch found the input iterator empty.
    Exhausted,
    /// The configured examination budget was spent after its final row.
    BudgetExhausted,
}

/// A page of projected items and the cursor for its continuation.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ChunkedPage<K, R> {
    /// Accepted projected items in input order.
    pub items: Vec<R>,
    /// The cursor to resume from.
    ///
    /// With a nonzero limit, a [`StopReason::Filled`] page returns the
    /// accepted filling row's cursor. A zero limit returns
    /// [`StopReason::Filled`] with `None` without observing the input or
    /// invoking any closure. A [`StopReason::BudgetExhausted`] page returns
    /// the final examined row's cursor. An [`StopReason::Exhausted`] page has
    /// no cursor.
    pub next: Option<K>,
    /// The condition that stopped pagination.
    pub stop: StopReason,
}

/// Collect a bounded page while preparing rows in chunks before projection.
///
/// A prepared chunk can contain a prefetched suffix that `project` never
/// receives when the page fills. That suffix is discarded and, after resuming
/// from the returned cursor with a newly narrowed iterator, is prepared again.
/// Therefore `prepare` must be read-only or idempotent.
///
/// [`PageSpec::max_examined`] bounds the rows pulled from `rows`, which for a
/// coalesced merge is fewer than the rows its sources yield: one merged row can
/// fold an arbitrary run of order-equal source rows.
///
/// # Errors
///
/// A `prepare` or `project` error returns [`Err`] without a partial
/// [`ChunkedPage`]. Side effects that occurred before the error are not rolled
/// back.
pub fn page_chunks<It, T, K, C, R>(
    mut rows: It,
    spec: PageSpec,
    cursor: impl Fn(&T) -> K,
    mut prepare: impl FnMut(&[T]) -> DbResult<C>,
    mut project: impl FnMut(T, &C) -> DbResult<Option<R>>,
) -> DbResult<ChunkedPage<K, R>>
where
    It: Iterator<Item = T>,
{
    if spec.limit == 0 {
        return Ok(ChunkedPage {
            items: Vec::new(),
            next: None,
            stop: StopReason::Filled,
        });
    }

    let chunk = spec.chunk.max(1);
    let max = spec.max_examined.map(|budget| budget.max(1));
    let mut items = Vec::with_capacity(spec.limit.min(1024));
    // `chunk` often comes straight from an RPC parameter, and the batch buffer is
    // refilled every pass — allocate a chunk once, not one per pass.
    let mut batch: Vec<T> = Vec::with_capacity(chunk.min(1024));
    let mut examined = 0;
    let mut last = None;

    loop {
        // `examined < budget` holds here: the first pass starts at zero against a
        // budget of at least one, and a later pass is only reached when the inner
        // loop ran to completion without spending the budget. So `allowed >= 1`
        // and the branch below is a backstop, not a live path.
        let allowed = max.map_or(chunk, |budget| chunk.min(budget - examined));
        if allowed == 0 {
            return Ok(ChunkedPage {
                items,
                next: last,
                stop: StopReason::BudgetExhausted,
            });
        }

        batch.clear();
        batch.extend(rows.by_ref().take(allowed));
        if batch.is_empty() {
            return Ok(ChunkedPage {
                items,
                next: None,
                stop: StopReason::Exhausted,
            });
        }

        let context = prepare(&batch)?;
        for row in batch.drain(..) {
            last = Some(cursor(&row));
            examined += 1;
            if let Some(item) = project(row, &context)? {
                items.push(item);
                if items.len() == spec.limit {
                    return Ok(ChunkedPage {
                        items,
                        next: last,
                        stop: StopReason::Filled,
                    });
                }
            }
            if max.is_some_and(|budget| examined == budget) {
                return Ok(ChunkedPage {
                    items,
                    next: last,
                    stop: StopReason::BudgetExhausted,
                });
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::cell::RefCell;
    use std::rc::Rc;

    use crate::{DbError, DbResult};

    use super::{PageSpec, StopReason, page_chunks};

    #[test]
    fn filling_mid_chunk_returns_the_filling_cursor_and_discards_prefetch() -> DbResult<()> {
        let batches = Rc::new(RefCell::new(Vec::new()));
        let batches_for_prepare = batches.clone();
        let page = page_chunks(
            1..=7,
            PageSpec {
                limit: 3,
                chunk: 4,
                max_examined: None,
            },
            |row| *row,
            move |batch: &[i32]| {
                batches_for_prepare.borrow_mut().push(batch.to_vec());
                Ok(())
            },
            |row, &()| Ok(Some(row)),
        )?;

        assert_eq!(page.items, vec![1, 2, 3]);
        assert_eq!(page.next, Some(3));
        assert_eq!(page.stop, StopReason::Filled);
        assert_eq!(&*batches.borrow(), &[vec![1, 2, 3, 4]]);

        let next_batches = Rc::new(RefCell::new(Vec::new()));
        let next_batches_for_prepare = next_batches.clone();
        page_chunks(
            4..=7,
            PageSpec {
                limit: 3,
                chunk: 4,
                max_examined: None,
            },
            |row| *row,
            move |batch: &[i32]| {
                next_batches_for_prepare.borrow_mut().push(batch.to_vec());
                Ok(())
            },
            |row, &()| Ok(Some(row)),
        )?;
        assert_eq!(&*next_batches.borrow(), &[vec![4, 5, 6, 7]]);
        Ok(())
    }

    #[test]
    fn filling_on_the_final_row_keeps_a_cursor_for_an_empty_follow_up() -> DbResult<()> {
        let page = page_chunks(
            1..=3,
            PageSpec {
                limit: 3,
                chunk: 4,
                max_examined: None,
            },
            |row| *row,
            |_| Ok(()),
            |row, &()| Ok(Some(row)),
        )?;
        assert_eq!(page.items, vec![1, 2, 3]);
        assert_eq!(page.next, Some(3));
        assert_eq!(page.stop, StopReason::Filled);

        let resumed = page_chunks(
            [1, 2, 3].into_iter().filter(|row| *row > 3),
            PageSpec {
                limit: 3,
                chunk: 4,
                max_examined: None,
            },
            |row| *row,
            |_| Ok(()),
            |row, &()| Ok(Some(row)),
        )?;
        assert!(resumed.items.is_empty());
        assert_eq!(resumed.next, None);
        assert_eq!(resumed.stop, StopReason::Exhausted);
        Ok(())
    }

    #[test]
    fn all_rejected_rows_stop_at_the_examination_budget() -> DbResult<()> {
        let page = page_chunks(
            1..,
            PageSpec {
                limit: 3,
                chunk: 2,
                max_examined: Some(4),
            },
            |row| *row,
            |_| Ok(()),
            |_, &()| Ok(None::<i32>),
        )?;
        assert!(page.items.is_empty());
        assert_eq!(page.next, Some(4));
        assert_eq!(page.stop, StopReason::BudgetExhausted);

        let resumed = page_chunks(
            (1..=7).filter(|row| *row > page.next.expect("budget page has a cursor")),
            PageSpec {
                limit: 3,
                chunk: 2,
                max_examined: None,
            },
            |row| *row,
            |_| Ok(()),
            |row, &()| Ok(Some(row)),
        )?;
        assert_eq!(resumed.items, vec![5, 6, 7]);
        assert_eq!(resumed.next, Some(7));
        assert_eq!(resumed.stop, StopReason::Filled);
        Ok(())
    }

    #[test]
    fn rejected_first_chunk_tops_up_and_prepares_each_nonempty_batch() -> DbResult<()> {
        let prepared = Rc::new(RefCell::new(Vec::new()));
        let prepared_for_prepare = prepared.clone();
        let page = page_chunks(
            1..=5,
            PageSpec {
                limit: 2,
                chunk: 2,
                max_examined: None,
            },
            |row| *row,
            move |batch: &[i32]| {
                prepared_for_prepare.borrow_mut().push(batch.to_vec());
                Ok(())
            },
            |row, &()| Ok((row >= 3).then_some(row)),
        )?;
        assert_eq!(page.items, vec![3, 4]);
        assert_eq!(page.next, Some(4));
        assert_eq!(page.stop, StopReason::Filled);
        assert_eq!(&*prepared.borrow(), &[vec![1, 2], vec![3, 4]]);
        Ok(())
    }

    #[test]
    fn zero_chunk_and_zero_budget_are_normalized_to_one() -> DbResult<()> {
        for spec in [
            PageSpec {
                limit: 1,
                chunk: 0,
                max_examined: None,
            },
            PageSpec {
                limit: 1,
                chunk: 4,
                max_examined: Some(0),
            },
        ] {
            let page = page_chunks(1.., spec, |row| *row, |_| Ok(()), |row, &()| Ok(Some(row)))?;
            assert_eq!(page.items, vec![1]);
            assert_eq!(page.next, Some(1));
            assert_eq!(page.stop, StopReason::Filled);
        }
        Ok(())
    }

    #[test]
    fn zero_limit_does_not_touch_the_iterator_or_closures() -> DbResult<()> {
        let rows = std::iter::from_fn(|| -> Option<i32> { panic!("iterator must not be called") });
        let page = page_chunks(
            rows,
            PageSpec {
                limit: 0,
                chunk: 1,
                max_examined: None,
            },
            |_| -> i32 { panic!("cursor must not be called") },
            |_| -> DbResult<()> { panic!("prepare must not be called") },
            |_, _: &()| -> DbResult<Option<i32>> { panic!("project must not be called") },
        )?;
        assert!(page.items.is_empty());
        assert_eq!(page.next, None);
        assert_eq!(page.stop, StopReason::Filled);
        Ok(())
    }

    #[test]
    fn short_input_is_exhausted_without_a_cursor() -> DbResult<()> {
        let page = page_chunks(
            1..=2,
            PageSpec {
                limit: 3,
                chunk: 4,
                max_examined: None,
            },
            |row| *row,
            |_| Ok(()),
            |row, &()| Ok(Some(row)),
        )?;
        assert_eq!(page.items, vec![1, 2]);
        assert_eq!(page.next, None);
        assert_eq!(page.stop, StopReason::Exhausted);
        Ok(())
    }

    #[test]
    fn prepare_and_project_errors_are_propagated_without_a_page() {
        let prepare = page_chunks(
            1..=1,
            PageSpec {
                limit: 1,
                chunk: 1,
                max_examined: None,
            },
            |row| *row,
            |_| Err::<(), _>(DbError::Internal("prepare failed")),
            |row, &()| Ok(Some(row)),
        );
        assert!(matches!(prepare, Err(DbError::Internal("prepare failed"))));

        let project = page_chunks(
            1..=1,
            PageSpec {
                limit: 1,
                chunk: 1,
                max_examined: None,
            },
            |row| *row,
            |_| Ok(()),
            |_, &()| Err::<Option<i32>, _>(DbError::Internal("project failed")),
        );
        assert!(matches!(project, Err(DbError::Internal("project failed"))));
    }

    #[test]
    fn non_clone_rows_keep_their_cursor_before_projection() -> DbResult<()> {
        struct Row {
            key: u8,
            value: u8,
        }

        let page = page_chunks(
            vec![Row { key: 7, value: 11 }].into_iter(),
            PageSpec {
                limit: 1,
                chunk: 1,
                max_examined: None,
            },
            |row| row.key,
            |_| Ok(()),
            |row, &()| Ok(Some(row.value)),
        )?;
        assert_eq!(page.items, vec![11]);
        assert_eq!(page.next, Some(7));
        Ok(())
    }
}