//! Pure recursive composite executor tests.

#![allow(clippy::expect_used)]

use context_core::{OccurrenceId, PointId, ProfileId, SourceAuthority, SourceKey, SourceVersion};
use context_query::{
    Cancellation, Candidate, CandidateBranch, CandidatePage, CandidateProvenance, CandidateSource,
    CandidateSourceKind, Completion, ExecutionBudget, ExecutionState, ExternalRerankPage,
    ExternalReranker, FilterCandidateBatch, FilterCandidateSource, Formula, Fusion,
    HydratedCandidate, MultiProfileBranch, PortBudget, ProfileName, QueryClock, QueryError,
    QueryExecutor, QueryIr, QueryKind, RecheckPage, ScoreOrder, SourceReadiness, SourceRechecker,
    StageDiagnostic, StageKind, TelemetrySink, TopologyExpander, build_multi_profile_query,
};
use std::cell::Cell;

#[derive(Default)]
struct RoutingSource {
    calls: usize,
    readiness_calls: usize,
    unavailable_second_branch: bool,
    partial_pages: bool,
    retained_memory_bytes: usize,
    observed_memory_budgets: Vec<usize>,
}

impl CandidateSource for RoutingSource {
    fn readiness(
        &mut self,
        query: &QueryIr,
        _budget: PortBudget,
    ) -> Result<SourceReadiness, QueryError> {
        self.readiness_calls += 1;
        if self.unavailable_second_branch && is_second_branch(query) {
            Ok(SourceReadiness::NotReady {
                reason: context_query::ReadinessReason::GenerationMissing,
            })
        } else {
            Ok(SourceReadiness::Ready)
        }
    }

    fn candidates(
        &mut self,
        query: &QueryIr,
        _filter: Option<&FilterCandidateBatch>,
        limit: usize,
        budget: PortBudget,
    ) -> Result<CandidatePage, QueryError> {
        self.calls += 1;
        self.observed_memory_budgets.push(budget.max_memory_bytes());
        let rows = if is_second_branch(query) {
            vec![candidate(2, 0.9), candidate(3, 0.1)]
        } else {
            vec![candidate(1, 0.9), candidate(2, 0.1)]
        };
        Ok(
            CandidatePage::new(rows.into_iter().take(limit).collect(), !self.partial_pages)
                .with_retained_memory_bytes(self.retained_memory_bytes),
        )
    }
}

#[derive(Default)]
struct ExactRechecker;

impl SourceRechecker for ExactRechecker {
    fn recheck(
        &mut self,
        _query: &QueryIr,
        candidates: &[Candidate],
        limit: usize,
        _budget: PortBudget,
    ) -> Result<RecheckPage, QueryError> {
        let rows = candidates
            .iter()
            .take(limit)
            .map(|candidate| {
                HydratedCandidate::new(
                    candidate.point_id(),
                    SourceKey::new(candidate.point_id().get().to_string())?,
                    candidate.approximate_score(),
                )
            })
            .collect::<Result<Vec<_>, _>>()?;
        Ok(RecheckPage::new(rows, candidates.len()))
    }
}

#[derive(Default)]
struct Diagnostics(Vec<StageDiagnostic>);

impl TelemetrySink for Diagnostics {
    fn record(&mut self, diagnostic: &StageDiagnostic) -> Result<(), QueryError> {
        self.0.push(diagnostic.clone());
        Ok(())
    }
}

struct NeverCancelled;

impl Cancellation for NeverCancelled {
    fn is_cancelled(&self) -> bool {
        false
    }
}

#[derive(Default)]
struct CountingFilter {
    calls: usize,
}

impl FilterCandidateSource for CountingFilter {
    fn filter_candidates(
        &mut self,
        _query: &QueryIr,
        limit: usize,
        _budget: PortBudget,
    ) -> Result<FilterCandidateBatch, QueryError> {
        self.calls += 1;
        Ok(FilterCandidateBatch::new(
            (1..=limit as u64).map(PointId::new).collect(),
            limit,
            true,
        ))
    }
}

struct CancelOnCall {
    calls: Cell<usize>,
    call: usize,
}

impl Cancellation for CancelOnCall {
    fn is_cancelled(&self) -> bool {
        let calls = self.calls.get() + 1;
        self.calls.set(calls);
        calls >= self.call
    }
}

fn candidate(point_id: u64, score: f64) -> Candidate {
    Candidate::new(
        PointId::new(point_id),
        score,
        CandidateProvenance::new(
            OccurrenceId::new(point_id.saturating_add(1))
                .expect("saturating increment is non-zero"),
            CandidateBranch::DenseAnn,
            CandidateSourceKind::Hnsw,
            ScoreOrder::LowerIsBetter,
            SourceAuthority::DerivedArtifact,
        ),
    )
    .expect("candidate fixture should be finite")
}

fn branch(first_dimension: f32) -> QueryIr {
    QueryIr::nearest(
        None,
        vec![first_dimension, 1.0],
        ScoreOrder::HigherIsBetter,
        None,
        3,
    )
    .expect("branch query should be valid")
}

fn filtered_branch(first_dimension: f32) -> QueryIr {
    QueryIr::nearest(
        None,
        vec![first_dimension, 1.0],
        ScoreOrder::HigherIsBetter,
        Some(serde_json::json!({
            "must": [{"key": "tenant", "match": {"value": "acme"}}]
        })),
        3,
    )
    .expect("filtered branch query should be valid")
}

fn is_second_branch(query: &QueryIr) -> bool {
    matches!(
        query.kind(),
        QueryKind::Nearest { vector, .. } if vector.as_slice()[0] < 0.0
    )
}

fn budget(stages: usize) -> ExecutionBudget {
    ExecutionBudget::new(8, 8, 8, stages, 2, 3).expect("test budget should be valid")
}

fn execute(
    query: &QueryIr,
    source: &mut RoutingSource,
    stages: usize,
) -> context_query::ExecutionOutcome {
    QueryExecutor::new(
        source,
        None,
        &mut ExactRechecker,
        &mut Diagnostics::default(),
        &NeverCancelled,
    )
    .execute(query, budget(stages))
    .expect("composite execution should succeed")
}

struct ProfileSource {
    score_scale: f64,
}

impl CandidateSource for ProfileSource {
    fn readiness(
        &mut self,
        _query: &QueryIr,
        _budget: PortBudget,
    ) -> Result<SourceReadiness, QueryError> {
        Ok(SourceReadiness::Ready)
    }

    fn candidates(
        &mut self,
        query: &QueryIr,
        _filter: Option<&FilterCandidateBatch>,
        limit: usize,
        _budget: PortBudget,
    ) -> Result<CandidatePage, QueryError> {
        let QueryKind::ProfileNearest { profile, .. } = query.kind() else {
            return Err(QueryError::PortFailure {
                stage: "profile_fixture",
                message: "expected profile-nearest leaf".to_owned(),
            });
        };
        let (profile_id, source_version, point_ids) = if profile.as_str() == "legacy" {
            (10, 100, [1, 2])
        } else {
            (20, 200, [3, 2])
        };
        let rows = point_ids
            .into_iter()
            .enumerate()
            .take(limit)
            .map(|(rank, point_id)| {
                Candidate::new(
                    PointId::new(point_id),
                    self.score_scale
                        * f64::from(
                            u32::try_from(rank.saturating_add(1)).expect("bounded fixture rank"),
                        ),
                    CandidateProvenance::new(
                        OccurrenceId::new(profile_id * 1_000 + point_id)
                            .expect("profile occurrence"),
                        CandidateBranch::MultiProfile,
                        CandidateSourceKind::Hnsw,
                        ScoreOrder::LowerIsBetter,
                        SourceAuthority::DerivedArtifact,
                    )
                    .with_profile(ProfileId::new(profile_id).expect("profile identity"))
                    .with_source_version(
                        SourceVersion::new(source_version).expect("source version"),
                    ),
                )
            })
            .collect::<Result<Vec<_>, _>>()?;
        Ok(CandidatePage::new(rows, true))
    }
}

fn profile_branch(name: &str, hash: u64, weight: f64) -> MultiProfileBranch {
    MultiProfileBranch::new(
        ProfileName::new(name).expect("profile name"),
        hash,
        "[1,0,0]".to_owned(),
        2,
        weight,
    )
    .expect("profile branch")
}

fn execute_profile_query(
    score_scale: f64,
    budget: ExecutionBudget,
) -> context_query::ExecutionOutcome {
    let query = build_multi_profile_query(
        vec![
            profile_branch("legacy", 11, 2.0),
            profile_branch("modern", 22, 1.0),
        ],
        None,
        60,
        3,
    )
    .expect("multi-profile query");
    QueryExecutor::new(
        &mut ProfileSource { score_scale },
        None,
        &mut ExactRechecker,
        &mut Diagnostics::default(),
        &NeverCancelled,
    )
    .execute(&query, budget)
    .expect("multi-profile execution")
}

#[test]
fn profile_nearest_executor_uses_rank_only_fusion_and_retains_provenance() {
    let budget = ExecutionBudget::new(8, 8, 8, 8, 2, 3).expect("profile budget");
    let small_scores = execute_profile_query(0.001, budget);
    let large_scores = execute_profile_query(1_000_000.0, budget);

    let point_ids = |outcome: &context_query::ExecutionOutcome| {
        outcome
            .points()
            .iter()
            .map(|point| point.point_id().get())
            .collect::<Vec<_>>()
    };
    assert_eq!(point_ids(&small_scores), vec![2, 1, 3]);
    assert_eq!(point_ids(&small_scores), point_ids(&large_scores));
    assert_eq!(small_scores.points().len(), 3);

    let shared = &small_scores.points()[0];
    assert_eq!(shared.point_id(), PointId::new(2));
    assert_eq!(shared.contributions().len(), 2);
    assert_eq!(
        shared
            .contributions()
            .iter()
            .map(|contribution| {
                let provenance = contribution.provenance();
                (
                    provenance.branch(),
                    provenance.source(),
                    provenance.profile().map(ProfileId::get),
                    provenance.source_version().map(SourceVersion::get),
                )
            })
            .collect::<Vec<_>>(),
        vec![
            (
                CandidateBranch::MultiProfile,
                CandidateSourceKind::Hnsw,
                Some(10),
                Some(100),
            ),
            (
                CandidateBranch::MultiProfile,
                CandidateSourceKind::Hnsw,
                Some(20),
                Some(200),
            ),
        ]
    );
}

#[test]
fn prefetch_carries_retained_candidate_memory_between_branch_budgets() {
    let query = QueryIr::new(
        QueryKind::Prefetch {
            branches: vec![branch(1.0), branch(-1.0)],
            fusion: Fusion::STANDARD_RRF,
        },
        ScoreOrder::HigherIsBetter,
        None,
        3,
    )
    .expect("prefetch should be valid");
    let mut source = RoutingSource {
        retained_memory_bytes: 1_024,
        ..Default::default()
    };
    let budget = ExecutionBudget::new(8, 8, 8, 8, 2, 3)
        .and_then(|budget| budget.with_resource_limits(100, 1024 * 1024, 1024 * 1024, 10_000))
        .expect("prefetch memory budget");
    let outcome = QueryExecutor::new(
        &mut source,
        None,
        &mut ExactRechecker,
        &mut Diagnostics::default(),
        &NeverCancelled,
    )
    .execute(&query, budget)
    .expect("prefetch should account retained adapter state");

    assert_eq!(source.observed_memory_budgets.len(), 2);
    assert!(source.observed_memory_budgets[1] < source.observed_memory_budgets[0]);
    assert!(outcome.usage().memory_bytes() >= 2 * source.retained_memory_bytes);
}

#[test]
fn profile_nearest_executor_breaks_equal_rank_fusion_scores_by_point_id() {
    let query = build_multi_profile_query(
        vec![
            profile_branch("legacy", 11, 1.0),
            profile_branch("modern", 22, 1.0),
        ],
        None,
        60,
        3,
    )
    .expect("equal-weight multi-profile query");
    let outcome = QueryExecutor::new(
        &mut ProfileSource { score_scale: 1.0 },
        None,
        &mut ExactRechecker,
        &mut Diagnostics::default(),
        &NeverCancelled,
    )
    .execute(
        &query,
        ExecutionBudget::new(8, 8, 8, 8, 2, 3).expect("tie budget"),
    )
    .expect("equal-weight execution");

    assert_eq!(
        outcome
            .points()
            .iter()
            .map(|point| point.point_id().get())
            .collect::<Vec<_>>(),
        vec![2, 1, 3]
    );
    assert_eq!(outcome.points()[1].score(), outcome.points()[2].score());
}

#[test]
fn profile_nearest_executor_fails_closed_at_candidate_budget() {
    let budget = ExecutionBudget::new(2, 8, 8, 8, 2, 3).expect("bounded profile budget");
    let outcome = execute_profile_query(1.0, budget);
    assert_eq!(outcome.completion(), Completion::BudgetExhausted);
    assert!(outcome.points().is_empty());
    assert!(outcome.usage().candidates() <= 2);
}

#[test]
fn profile_nearest_executor_cancels_before_profile_port_work() {
    let query = build_multi_profile_query(vec![profile_branch("legacy", 11, 1.0)], None, 60, 2)
        .expect("multi-profile query");
    let mut source = ProfileSource { score_scale: 1.0 };
    let outcome = QueryExecutor::new(
        &mut source,
        None,
        &mut ExactRechecker,
        &mut Diagnostics::default(),
        &CancelOnCall {
            calls: Cell::new(0),
            call: 1,
        },
    )
    .execute(
        &query,
        ExecutionBudget::new(4, 4, 4, 4, 1, 2).expect("cancellation budget"),
    )
    .expect("cancelled execution");

    assert_eq!(outcome.completion(), Completion::Cancelled);
    assert!(outcome.points().is_empty());
}

#[test]
fn prefetch_uses_rrf_with_deduplication_and_deterministic_ties() {
    let query = QueryIr::new(
        QueryKind::Prefetch {
            branches: vec![branch(1.0), branch(-1.0)],
            fusion: Fusion::STANDARD_RRF,
        },
        ScoreOrder::HigherIsBetter,
        None,
        3,
    )
    .expect("prefetch should be valid");
    let outcome = execute(&query, &mut RoutingSource::default(), 8);

    assert_eq!(outcome.completion(), Completion::Complete);
    assert_eq!(
        outcome
            .points()
            .iter()
            .map(|point| point.point_id().get())
            .collect::<Vec<_>>(),
        vec![2, 1, 3]
    );
    assert_eq!(outcome.usage().candidates(), 4);
    assert_eq!(outcome.usage().rechecks(), 4);
    assert!(
        outcome
            .diagnostics()
            .iter()
            .any(|diagnostic| diagnostic.stage() == StageKind::Fusion
                && diagnostic.strategy() == "reciprocal_rank_fusion")
    );
}

#[test]
fn weighted_prefetch_uses_rank_only_weighted_rrf() {
    let weighted = |query, weight| {
        QueryIr::new(
            QueryKind::Weighted {
                query: Box::new(query),
                weight,
            },
            ScoreOrder::HigherIsBetter,
            None,
            3,
        )
        .expect("weighted branch should be valid")
    };
    let query = QueryIr::new(
        QueryKind::Prefetch {
            branches: vec![weighted(branch(1.0), 3.0), weighted(branch(-1.0), 1.0)],
            fusion: Fusion::WeightedRrf { rank_constant: 60 },
        },
        ScoreOrder::HigherIsBetter,
        None,
        3,
    )
    .expect("weighted prefetch should be valid");
    let outcome = execute(&query, &mut RoutingSource::default(), 8);

    assert_eq!(
        outcome
            .points()
            .iter()
            .map(|point| point.point_id().get())
            .collect::<Vec<_>>(),
        vec![2, 1, 3]
    );
    assert_eq!(
        outcome
            .diagnostics()
            .iter()
            .filter(|diagnostic| diagnostic.strategy() == "weighted_reciprocal_rank_fusion")
            .count(),
        1
    );
}

#[test]
fn prefetch_executes_direct_weighted_branch_limits() {
    let weighted = QueryIr::new(
        QueryKind::Weighted {
            query: Box::new(branch(1.0)),
            weight: 2.0,
        },
        ScoreOrder::HigherIsBetter,
        None,
        1,
    )
    .expect("weighted branch should be valid");
    let query = QueryIr::new(
        QueryKind::Prefetch {
            branches: vec![weighted],
            fusion: Fusion::WeightedRrf { rank_constant: 60 },
        },
        ScoreOrder::HigherIsBetter,
        None,
        3,
    )
    .expect("prefetch should be valid");
    let outcome = execute(&query, &mut RoutingSource::default(), 8);

    assert_eq!(outcome.points().len(), 1);
    assert_eq!(outcome.points()[0].point_id().get(), 1);
    assert!(
        outcome
            .diagnostics()
            .iter()
            .all(|diagnostic| diagnostic.strategy() != "weighted_score")
    );
}

#[derive(Default)]
struct ExtremeScoreSource;

impl CandidateSource for ExtremeScoreSource {
    fn readiness(
        &mut self,
        _query: &QueryIr,
        _budget: PortBudget,
    ) -> Result<SourceReadiness, QueryError> {
        Ok(SourceReadiness::Ready)
    }

    fn candidates(
        &mut self,
        _query: &QueryIr,
        _filter: Option<&FilterCandidateBatch>,
        _limit: usize,
        _budget: PortBudget,
    ) -> Result<CandidatePage, QueryError> {
        Ok(CandidatePage::new(
            vec![candidate(1, f64::MAX), candidate(2, f64::MAX / 2.0)],
            true,
        ))
    }
}

#[test]
fn weighted_rrf_treats_extreme_weights_as_rank_metadata() {
    let weighted = QueryIr::new(
        QueryKind::Weighted {
            query: Box::new(branch(1.0)),
            weight: 2.0,
        },
        ScoreOrder::HigherIsBetter,
        None,
        2,
    )
    .expect("weighted branch");
    let query = QueryIr::new(
        QueryKind::Prefetch {
            branches: vec![weighted],
            fusion: Fusion::WeightedRrf { rank_constant: 60 },
        },
        ScoreOrder::HigherIsBetter,
        None,
        2,
    )
    .expect("weighted prefetch");
    let outcome = QueryExecutor::new(
        &mut ExtremeScoreSource,
        None,
        &mut ExactRechecker,
        &mut Diagnostics::default(),
        &NeverCancelled,
    )
    .execute(&query, budget(8))
    .expect("rank-only weighted RRF must not multiply source scores");

    assert_eq!(outcome.completion(), Completion::Complete);
    assert_eq!(outcome.points()[0].point_id(), PointId::new(1));
    assert!(
        outcome
            .diagnostics()
            .iter()
            .all(|diagnostic| diagnostic.strategy() != "weighted_score")
    );
}

#[test]
fn fusion_handles_maximal_cross_branch_overlap_with_linear_metadata_work() {
    let branches = vec![branch(1.0); 32];
    let query = QueryIr::new(
        QueryKind::Prefetch {
            branches,
            fusion: Fusion::STANDARD_RRF,
        },
        ScoreOrder::HigherIsBetter,
        None,
        3,
    )
    .expect("maximal-overlap prefetch");
    let budget =
        ExecutionBudget::new(128, 128, 128, 128, 8, 3).expect("overlap budget should be valid");
    let outcome = QueryExecutor::new(
        &mut RoutingSource::default(),
        None,
        &mut ExactRechecker,
        &mut Diagnostics::default(),
        &NeverCancelled,
    )
    .execute(&query, budget)
    .expect("maximal overlap should remain bounded");

    assert_eq!(outcome.completion(), Completion::Complete);
    assert_eq!(outcome.points().len(), 2);
    assert!(
        outcome
            .points()
            .iter()
            .all(|point| point.contributions().len() == 32)
    );
    assert!(outcome.usage().comparisons() < 1_000);

    let exact_memory = outcome.usage().memory_bytes();
    let exact_budget = ExecutionBudget::new(128, 128, 128, 128, 8, 3)
        .expect("overlap budget")
        .with_resource_limits(1_000, exact_memory, 1024 * 1024, 10_000)
        .expect("exact fusion memory budget");
    let exact = QueryExecutor::new(
        &mut RoutingSource::default(),
        None,
        &mut ExactRechecker,
        &mut Diagnostics::default(),
        &NeverCancelled,
    )
    .execute(&query, exact_budget)
    .expect("exact projected fusion memory should be accepted");
    assert_eq!(exact.completion(), Completion::Complete);

    let below_budget = ExecutionBudget::new(128, 128, 128, 128, 8, 3)
        .expect("overlap budget")
        .with_resource_limits(1_000, exact_memory.saturating_sub(1), 1024 * 1024, 10_000)
        .expect("one-byte-short fusion memory budget");
    let below = QueryExecutor::new(
        &mut RoutingSource::default(),
        None,
        &mut ExactRechecker,
        &mut Diagnostics::default(),
        &NeverCancelled,
    )
    .execute(&query, below_budget)
    .expect("one-byte-short fusion memory is a typed outcome");
    assert_eq!(below.completion(), Completion::BudgetExhausted);
    assert!(below.points().is_empty());
}

#[test]
fn fusion_many_unique_single_contributions_obey_the_exact_memory_boundary() {
    let query = QueryIr::new(
        QueryKind::Prefetch {
            branches: vec![branch(1.0), branch(-1.0)],
            fusion: Fusion::STANDARD_RRF,
        },
        ScoreOrder::HigherIsBetter,
        None,
        3,
    )
    .expect("partially disjoint prefetch");
    let base = budget(8);
    let outcome = QueryExecutor::new(
        &mut RoutingSource::default(),
        None,
        &mut ExactRechecker,
        &mut Diagnostics::default(),
        &NeverCancelled,
    )
    .execute(&query, base)
    .expect("partially disjoint fusion");
    assert_eq!(outcome.completion(), Completion::Complete);
    assert_eq!(outcome.points().len(), 3);
    assert_eq!(outcome.points()[0].contributions().len(), 2);
    assert_eq!(outcome.points()[1].contributions().len(), 1);
    assert_eq!(outcome.points()[2].contributions().len(), 1);

    let exact_memory = outcome.usage().memory_bytes();
    let exact = budget(8)
        .with_resource_limits(1_000, exact_memory, 1024 * 1024, 10_000)
        .expect("exact unique-fusion memory budget");
    let exact = QueryExecutor::new(
        &mut RoutingSource::default(),
        None,
        &mut ExactRechecker,
        &mut Diagnostics::default(),
        &NeverCancelled,
    )
    .execute(&query, exact)
    .expect("exact unique-fusion memory boundary");
    assert_eq!(exact.completion(), Completion::Complete);

    let below = budget(8)
        .with_resource_limits(1_000, exact_memory.saturating_sub(1), 1024 * 1024, 10_000)
        .expect("one-byte-short unique-fusion memory budget");
    let below = QueryExecutor::new(
        &mut RoutingSource::default(),
        None,
        &mut ExactRechecker,
        &mut Diagnostics::default(),
        &NeverCancelled,
    )
    .execute(&query, below)
    .expect("one-byte-short unique-fusion memory outcome");
    assert_eq!(below.completion(), Completion::BudgetExhausted);
    assert!(below.points().is_empty());
}

#[test]
fn formula_threshold_and_rerank_execute_in_tree_order() {
    let formula = QueryIr::new(
        QueryKind::Formula {
            query: Box::new(branch(1.0)),
            formula: Formula::new("$score * 2").expect("formula text should be bounded"),
        },
        ScoreOrder::HigherIsBetter,
        None,
        3,
    )
    .expect("formula node should be valid");
    let threshold = QueryIr::new(
        QueryKind::ScoreThreshold {
            query: Box::new(formula),
            minimum: Some(0.5),
            maximum: None,
        },
        ScoreOrder::HigherIsBetter,
        None,
        3,
    )
    .expect("threshold node should be valid");
    let query = QueryIr::new(
        QueryKind::Rerank {
            query: Box::new(threshold),
        },
        ScoreOrder::HigherIsBetter,
        None,
        1,
    )
    .expect("rerank node should be valid");
    let outcome = execute(&query, &mut RoutingSource::default(), 8);

    assert_eq!(outcome.points().len(), 1);
    assert_eq!(outcome.points()[0].point_id().get(), 1);
    assert_eq!(outcome.points()[0].score(), 1.8);
    assert_eq!(
        outcome
            .diagnostics()
            .iter()
            .map(StageDiagnostic::stage)
            .collect::<Vec<_>>()
            .last(),
        Some(&StageKind::Rerank)
    );
}

#[test]
fn prefetch_propagates_unavailable_sources_and_global_budget_exhaustion() {
    let query = QueryIr::new(
        QueryKind::Prefetch {
            branches: vec![branch(1.0), branch(-1.0)],
            fusion: Fusion::STANDARD_RRF,
        },
        ScoreOrder::HigherIsBetter,
        None,
        3,
    )
    .expect("prefetch should be valid");
    let unavailable = execute(
        &query,
        &mut RoutingSource {
            unavailable_second_branch: true,
            ..Default::default()
        },
        8,
    );
    assert!(matches!(
        unavailable.state(),
        ExecutionState::NotReady { .. }
    ));
    assert!(unavailable.points().is_empty());

    let exhausted = execute(&query, &mut RoutingSource::default(), 3);
    assert_eq!(exhausted.completion(), Completion::BudgetExhausted);
    assert!(exhausted.points().is_empty());
    assert!(exhausted.usage().stages() <= 3);
}

#[test]
fn wrapped_filtered_branches_cannot_exceed_the_global_filter_budget() {
    let wrapped = QueryIr::new(
        QueryKind::ScoreThreshold {
            query: Box::new(filtered_branch(-1.0)),
            minimum: None,
            maximum: None,
        },
        ScoreOrder::HigherIsBetter,
        None,
        3,
    )
    .expect("wrapped filtered branch should be valid");
    let query = QueryIr::new(
        QueryKind::Prefetch {
            branches: vec![filtered_branch(1.0), wrapped],
            fusion: Fusion::STANDARD_RRF,
        },
        ScoreOrder::HigherIsBetter,
        None,
        3,
    )
    .expect("prefetch should be valid");
    let mut source = RoutingSource::default();
    let mut filter = CountingFilter::default();
    let outcome = QueryExecutor::new(
        &mut source,
        Some(&mut filter),
        &mut ExactRechecker,
        &mut Diagnostics::default(),
        &NeverCancelled,
    )
    .execute(
        &query,
        ExecutionBudget::new(8, 2, 8, 8, 2, 3).expect("budget should be valid"),
    )
    .expect("execution should remain bounded");

    assert_eq!(outcome.completion(), Completion::BudgetExhausted);
    assert_eq!(outcome.usage().filter_candidates(), 2);
    assert_eq!(filter.calls, 1);
}

#[test]
fn invalid_formula_fails_before_any_candidate_work() {
    let query = QueryIr::new(
        QueryKind::Formula {
            query: Box::new(branch(1.0)),
            formula: Formula::new("system($score)").expect("opaque text remains constructible"),
        },
        ScoreOrder::HigherIsBetter,
        None,
        3,
    )
    .expect("opaque formula plan should remain constructible");
    let mut source = RoutingSource::default();
    let error = QueryExecutor::new(
        &mut source,
        None,
        &mut ExactRechecker,
        &mut Diagnostics::default(),
        &NeverCancelled,
    )
    .execute(&query, budget(8))
    .expect_err("invalid executable formula should fail");

    assert!(matches!(
        error,
        QueryError::InvalidInput {
            field: "formula",
            ..
        }
    ));
    assert_eq!(source.calls, 0);
    assert_eq!(source.readiness_calls, 0);
}

#[test]
fn post_processing_applies_to_authoritative_partial_results() {
    let partial = || RoutingSource {
        partial_pages: true,
        ..Default::default()
    };
    let wrap = |kind| {
        QueryIr::new(kind, ScoreOrder::HigherIsBetter, None, 1).expect("wrapper should be valid")
    };

    let weighted = wrap(QueryKind::Weighted {
        query: Box::new(branch(1.0)),
        weight: 2.0,
    });
    let mut source = partial();
    let outcome = execute(&weighted, &mut source, 8);
    assert_eq!(outcome.completion(), Completion::BudgetExhausted);
    assert_eq!(outcome.points()[0].score(), 1.8);

    let threshold = wrap(QueryKind::ScoreThreshold {
        query: Box::new(branch(1.0)),
        minimum: Some(0.5),
        maximum: None,
    });
    let mut source = partial();
    let outcome = execute(&threshold, &mut source, 8);
    assert_eq!(outcome.completion(), Completion::BudgetExhausted);
    assert_eq!(outcome.points().len(), 1);

    let formula = wrap(QueryKind::Formula {
        query: Box::new(branch(1.0)),
        formula: Formula::new("$score + 1").expect("formula should be valid"),
    });
    let mut source = partial();
    let outcome = execute(&formula, &mut source, 8);
    assert_eq!(outcome.completion(), Completion::BudgetExhausted);
    assert_eq!(outcome.points()[0].score(), 1.9);

    let rerank = wrap(QueryKind::Rerank {
        query: Box::new(branch(1.0)),
    });
    let mut source = partial();
    let outcome = execute(&rerank, &mut source, 8);
    assert_eq!(outcome.completion(), Completion::BudgetExhausted);
    assert_eq!(outcome.points().len(), 1);
}

#[test]
fn transform_cancellation_never_returns_points() {
    let query = QueryIr::new(
        QueryKind::Weighted {
            query: Box::new(branch(1.0)),
            weight: 2.0,
        },
        ScoreOrder::HigherIsBetter,
        None,
        3,
    )
    .expect("weighted query should be valid");
    let cancellation = CancelOnCall {
        calls: Cell::new(0),
        call: 8,
    };
    let outcome = QueryExecutor::new(
        &mut RoutingSource::default(),
        None,
        &mut ExactRechecker,
        &mut Diagnostics::default(),
        &cancellation,
    )
    .execute(&query, budget(8))
    .expect("cancellation should be an outcome");

    assert_eq!(outcome.completion(), Completion::Cancelled);
    assert!(outcome.points().is_empty());
}

#[derive(Default)]
struct PerLeafBudgetSource {
    requested: Vec<usize>,
}

impl CandidateSource for PerLeafBudgetSource {
    fn readiness(
        &mut self,
        _query: &QueryIr,
        _budget: PortBudget,
    ) -> Result<SourceReadiness, QueryError> {
        Ok(SourceReadiness::Ready)
    }

    fn candidate_limit(
        &mut self,
        query: &QueryIr,
        remaining: usize,
        _budget: PortBudget,
    ) -> Result<usize, QueryError> {
        Ok(query.limit().min(remaining))
    }

    fn candidates(
        &mut self,
        query: &QueryIr,
        _filter: Option<&FilterCandidateBatch>,
        limit: usize,
        _budget: PortBudget,
    ) -> Result<CandidatePage, QueryError> {
        self.requested.push(limit);
        let offset = if is_second_branch(query) { 100 } else { 0 };
        Ok(CandidatePage::new(
            (1..=limit)
                .map(|point_id| candidate(offset + point_id as u64, 1.0))
                .collect(),
            true,
        ))
    }
}

#[test]
fn prefetch_reserves_candidate_work_per_leaf() {
    let query = QueryIr::new(
        QueryKind::Prefetch {
            branches: vec![branch(1.0), branch(-1.0)],
            fusion: Fusion::STANDARD_RRF,
        },
        ScoreOrder::HigherIsBetter,
        None,
        3,
    )
    .expect("prefetch should be valid");
    let mut source = PerLeafBudgetSource::default();
    let outcome = QueryExecutor::new(
        &mut source,
        None,
        &mut ExactRechecker,
        &mut Diagnostics::default(),
        &NeverCancelled,
    )
    .execute(
        &query,
        ExecutionBudget::new(6, 1, 6, 16, 2, 3).expect("budget should be valid"),
    )
    .expect("both branches should execute");

    assert_eq!(outcome.completion(), Completion::Complete);
    assert_eq!(source.requested, vec![3, 3]);
    assert_eq!(outcome.usage().candidates(), 6);
}

#[test]
fn fusion_retains_every_branch_occurrence_and_rank_contribution() {
    let query = QueryIr::new(
        QueryKind::Prefetch {
            branches: vec![branch(1.0), branch(-1.0)],
            fusion: Fusion::Rrf { rank_constant: 10 },
        },
        ScoreOrder::HigherIsBetter,
        None,
        3,
    )
    .expect("prefetch should be valid");
    let outcome = execute(&query, &mut RoutingSource::default(), 8);
    let shared = outcome
        .points()
        .iter()
        .find(|point| point.point_id() == PointId::new(2))
        .expect("shared point should survive fusion");

    assert_eq!(shared.contributions().len(), 2);
    assert!(shared.contributions().iter().all(|contribution| {
        contribution.provenance().branch() == CandidateBranch::DenseAnn
            && contribution.fusion_contribution().is_some()
    }));
    let contribution_sum = shared
        .contributions()
        .iter()
        .filter_map(|contribution| contribution.fusion_contribution())
        .sum::<f64>();
    assert!((contribution_sum - shared.score()).abs() < f64::EPSILON);
}

struct FixedClock(Cell<u64>);

impl QueryClock for FixedClock {
    fn now_micros(&self) -> u64 {
        let now = self.0.get();
        self.0.set(now.saturating_add(100));
        now
    }
}

struct ManualClock<'a>(&'a Cell<u64>);

impl QueryClock for ManualClock<'_> {
    fn now_micros(&self) -> u64 {
        self.0.get()
    }
}

struct DeadlineSource<'a> {
    clock: &'a Cell<u64>,
    candidate_calls: &'a Cell<usize>,
}

impl CandidateSource for DeadlineSource<'_> {
    fn readiness(
        &mut self,
        _query: &QueryIr,
        _budget: PortBudget,
    ) -> Result<SourceReadiness, QueryError> {
        self.clock.set(100);
        Ok(SourceReadiness::Ready)
    }

    fn candidates(
        &mut self,
        _query: &QueryIr,
        _filter: Option<&FilterCandidateBatch>,
        _limit: usize,
        _budget: PortBudget,
    ) -> Result<CandidatePage, QueryError> {
        self.candidate_calls
            .set(self.candidate_calls.get().saturating_add(1));
        Ok(CandidatePage::new(vec![candidate(1, 0.0)], true))
    }
}

struct CountingRechecker<'a>(&'a Cell<usize>);

impl SourceRechecker for CountingRechecker<'_> {
    fn recheck(
        &mut self,
        _query: &QueryIr,
        _candidates: &[Candidate],
        _limit: usize,
        _budget: PortBudget,
    ) -> Result<RecheckPage, QueryError> {
        self.0.set(self.0.get().saturating_add(1));
        Ok(RecheckPage::new(Vec::new(), 0))
    }
}

#[test]
fn elapsed_deadline_stops_later_ports_after_the_first_boundary_overrun() {
    let clock = Cell::new(0);
    let candidate_calls = Cell::new(0);
    let recheck_calls = Cell::new(0);
    let budget = ExecutionBudget::new(8, 8, 8, 8, 4, 4)
        .expect("base budget")
        .with_resource_limits(100, 1024 * 1024, 1024 * 1024, 50)
        .expect("resource limits");
    let outcome = QueryExecutor::new(
        &mut DeadlineSource {
            clock: &clock,
            candidate_calls: &candidate_calls,
        },
        None,
        &mut CountingRechecker(&recheck_calls),
        &mut Diagnostics::default(),
        &NeverCancelled,
    )
    .with_clock(&ManualClock(&clock))
    .execute(&branch(1.0), budget)
    .expect("deadline exhaustion should be typed");

    assert_eq!(outcome.completion(), Completion::BudgetExhausted);
    assert!(outcome.points().is_empty());
    assert_eq!(outcome.usage().elapsed_micros(), 100);
    assert_eq!(candidate_calls.get(), 0);
    assert_eq!(recheck_calls.get(), 0);
}

#[test]
fn elapsed_and_comparison_budgets_fail_closed_without_partial_points() {
    let query = branch(1.0);
    let budget = ExecutionBudget::new(8, 8, 8, 8, 4, 4)
        .expect("base budget")
        .with_resource_limits(1, 1024 * 1024, 1024 * 1024, 50)
        .expect("resource limits");
    let mut source = RoutingSource::default();
    let outcome = QueryExecutor::new(
        &mut source,
        None,
        &mut ExactRechecker,
        &mut Diagnostics::default(),
        &NeverCancelled,
    )
    .with_clock(&FixedClock(Cell::new(0)))
    .execute(&query, budget)
    .expect("budget exhaustion is typed");

    assert_eq!(outcome.completion(), Completion::BudgetExhausted);
    assert!(outcome.points().is_empty());
    assert!(outcome.usage().comparisons() > 1 || outcome.usage().elapsed_micros() > 50);
}

#[test]
fn exact_comparison_boundary_completes() {
    let query = branch(1.0);
    let budget = ExecutionBudget::new(8, 8, 8, 8, 4, 4)
        .expect("base budget")
        .with_resource_limits(4, 1024 * 1024, 1024 * 1024, 10_000)
        .expect("resource limits");
    let mut source = RoutingSource::default();
    let outcome = QueryExecutor::new(
        &mut source,
        None,
        &mut ExactRechecker,
        &mut Diagnostics::default(),
        &NeverCancelled,
    )
    .with_clock(&FixedClock(Cell::new(0)))
    .execute(&query, budget)
    .expect("the inclusive comparison maximum should be valid");

    assert_eq!(outcome.completion(), Completion::Complete);
    assert_eq!(outcome.usage().comparisons(), 4);
    assert_eq!(outcome.points().len(), 2);
}

struct FakeExternalReranker {
    revision: u64,
    exhausted: bool,
    calls: usize,
}

impl ExternalReranker for FakeExternalReranker {
    fn rerank(
        &mut self,
        _query: &QueryIr,
        rows: &[HydratedCandidate],
        limit: usize,
        _budget: PortBudget,
    ) -> Result<ExternalRerankPage, QueryError> {
        self.calls = self.calls.saturating_add(1);
        let reranked = rows
            .iter()
            .rev()
            .take(limit)
            .enumerate()
            .map(|(rank, row)| {
                let rank = u32::try_from(rank).expect("bounded test rank fits u32");
                HydratedCandidate::new(
                    row.point_id(),
                    row.source_key().clone(),
                    100.0 - f64::from(rank),
                )
            })
            .collect::<Result<Vec<_>, _>>()?;
        Ok(ExternalRerankPage::new(
            reranked,
            rows.len(),
            self.exhausted,
            self.revision,
        ))
    }
}

#[test]
fn external_rerank_port_preserves_authoritative_provenance_and_revision() {
    let inner = branch(1.0);
    let mut baseline_source = RoutingSource::default();
    let baseline = QueryExecutor::new(
        &mut baseline_source,
        None,
        &mut ExactRechecker,
        &mut Diagnostics::default(),
        &NeverCancelled,
    )
    .execute(&inner, budget(8))
    .expect("inner query should execute");
    let query = QueryIr::new(
        QueryKind::ExternalRerank {
            query: Box::new(inner),
            model_revision: 7,
        },
        ScoreOrder::HigherIsBetter,
        None,
        2,
    )
    .expect("external rerank query");
    let mut source = RoutingSource::default();
    let mut reranker = FakeExternalReranker {
        revision: 7,
        exhausted: true,
        calls: 0,
    };
    let outcome = QueryExecutor::new(
        &mut source,
        None,
        &mut ExactRechecker,
        &mut Diagnostics::default(),
        &NeverCancelled,
    )
    .with_external_reranker(&mut reranker)
    .execute(
        &query,
        budget(8)
            .with_resource_limits(1_000, 1024 * 1024, 6, 10_000)
            .expect("short-key rerank budget"),
    )
    .expect("external rerank should execute");

    assert_eq!(outcome.completion(), Completion::Complete);
    let authority_key_bytes = outcome
        .points()
        .iter()
        .map(|row| row.source_key().as_str().len())
        .sum::<usize>();
    assert_eq!(
        outcome.usage().hydration_bytes(),
        baseline
            .usage()
            .hydration_bytes()
            .saturating_add(
                baseline
                    .points()
                    .iter()
                    .map(|row| row.source_key().as_str().len())
                    .sum::<usize>()
            )
            .saturating_add(authority_key_bytes),
        "usage must charge the all-candidate authority pass and winner recheck"
    );
    assert!(
        outcome
            .points()
            .iter()
            .all(|row| !row.contributions().is_empty())
    );
    assert_eq!(
        outcome.diagnostics().last().map(StageDiagnostic::stage),
        Some(StageKind::ExternalRerank)
    );
}

struct FakeTopology {
    calls: usize,
}

struct EmptyTopology;

impl TopologyExpander for EmptyTopology {
    fn expand(
        &mut self,
        _query: &QueryIr,
        _seeds: &[HydratedCandidate],
        _max_depth: usize,
        _limit: usize,
        _budget: PortBudget,
    ) -> Result<CandidatePage, QueryError> {
        Ok(CandidatePage::with_scored_count(Vec::new(), 2, true).with_strategy("empty_topology"))
    }
}

impl TopologyExpander for FakeTopology {
    fn expand(
        &mut self,
        _query: &QueryIr,
        _seeds: &[HydratedCandidate],
        _max_depth: usize,
        _limit: usize,
        _budget: PortBudget,
    ) -> Result<CandidatePage, QueryError> {
        self.calls = self.calls.saturating_add(1);
        let provenance = CandidateProvenance::new(
            OccurrenceId::new(99).expect("nonzero occurrence"),
            CandidateBranch::Topology,
            CandidateSourceKind::Topology,
            ScoreOrder::HigherIsBetter,
            SourceAuthority::DerivedArtifact,
        );
        Ok(CandidatePage::with_scored_count(
            vec![Candidate::new(PointId::new(3), 0.75, provenance)?],
            2,
            true,
        )
        .with_expansion_count(1)
        .with_strategy("fake_topology"))
    }
}

#[test]
fn topology_port_rechecks_expanded_candidates_and_retains_provenance() {
    let query = QueryIr::new(
        QueryKind::TopologyExpand {
            query: Box::new(branch(1.0)),
            max_depth: 2,
        },
        ScoreOrder::HigherIsBetter,
        None,
        3,
    )
    .expect("topology query");
    let mut source = RoutingSource::default();
    let mut topology = FakeTopology { calls: 0 };
    let outcome = QueryExecutor::new(
        &mut source,
        None,
        &mut ExactRechecker,
        &mut Diagnostics::default(),
        &NeverCancelled,
    )
    .with_topology_expander(&mut topology)
    .execute(&query, budget(8))
    .expect("topology expansion should execute");

    assert_eq!(outcome.completion(), Completion::Complete);
    assert_eq!(outcome.points()[0].point_id(), PointId::new(3));
    assert_eq!(
        outcome.points()[0].contributions()[0].provenance().source(),
        CandidateSourceKind::Topology
    );
}

#[test]
fn topology_reserves_expansion_and_recheck_stages_before_calling_the_expander() {
    let query = QueryIr::new(
        QueryKind::TopologyExpand {
            query: Box::new(branch(1.0)),
            max_depth: 2,
        },
        ScoreOrder::HigherIsBetter,
        None,
        3,
    )
    .expect("topology query");
    let mut source = RoutingSource::default();
    let mut topology = FakeTopology { calls: 0 };
    let outcome = QueryExecutor::new(
        &mut source,
        None,
        &mut ExactRechecker,
        &mut Diagnostics::default(),
        &NeverCancelled,
    )
    .with_topology_expander(&mut topology)
    .execute(&query, budget(3))
    .expect("stage exhaustion should be typed");

    assert_eq!(topology.calls, 0);
    assert_eq!(outcome.completion(), Completion::BudgetExhausted);
    assert!(outcome.points().is_empty());
    assert_eq!(outcome.usage().stages(), 2);
}

#[test]
fn topology_accepts_the_exact_two_stage_boundary() {
    let query = QueryIr::new(
        QueryKind::TopologyExpand {
            query: Box::new(branch(1.0)),
            max_depth: 2,
        },
        ScoreOrder::HigherIsBetter,
        None,
        3,
    )
    .expect("topology query");
    let mut source = RoutingSource::default();
    let mut topology = FakeTopology { calls: 0 };
    let outcome = QueryExecutor::new(
        &mut source,
        None,
        &mut ExactRechecker,
        &mut Diagnostics::default(),
        &NeverCancelled,
    )
    .with_topology_expander(&mut topology)
    .execute(&query, budget(4))
    .expect("exact stage boundary should execute");

    assert_eq!(topology.calls, 1);
    assert_eq!(outcome.completion(), Completion::Complete);
    assert_eq!(outcome.usage().stages(), 4);
}

#[test]
fn empty_topology_completes_at_the_exact_comparison_boundary() {
    let query = QueryIr::new(
        QueryKind::TopologyExpand {
            query: Box::new(branch(1.0)),
            max_depth: 2,
        },
        ScoreOrder::HigherIsBetter,
        None,
        3,
    )
    .expect("topology query");
    let exact = budget(4)
        .with_resource_limits(6, 1024 * 1024, 1024 * 1024, 10_000)
        .expect("exact comparison budget");
    let mut source = RoutingSource::default();
    let outcome = QueryExecutor::new(
        &mut source,
        None,
        &mut ExactRechecker,
        &mut Diagnostics::default(),
        &NeverCancelled,
    )
    .with_topology_expander(&mut EmptyTopology)
    .execute(&query, exact)
    .expect("empty exhausted topology page should complete");

    assert_eq!(outcome.completion(), Completion::Complete);
    assert_eq!(outcome.usage().comparisons(), 6);
    assert_eq!(outcome.usage().stages(), 3);
    assert!(outcome.points().is_empty());
    assert_eq!(
        outcome.diagnostics().last().map(StageDiagnostic::stage),
        Some(StageKind::TopologyExpansion)
    );
}

#[test]
fn composite_allocations_exhaust_tiny_memory_and_hydration_budgets_without_points() {
    let leaf = execute(&branch(1.0), &mut RoutingSource::default(), 8);

    let prefetch = QueryIr::new(
        QueryKind::Prefetch {
            branches: vec![branch(1.0), branch(-1.0)],
            fusion: Fusion::STANDARD_RRF,
        },
        ScoreOrder::HigherIsBetter,
        None,
        3,
    )
    .expect("prefetch query");
    let mut prefetch_source = RoutingSource::default();
    let prefetch_budget = ExecutionBudget::new(8, 8, 8, 16, 4, 3)
        .expect("prefetch budget")
        .with_resource_limits(
            1_000,
            leaf.usage().memory_bytes().saturating_add(64),
            1024 * 1024,
            10_000,
        )
        .expect("tiny prefetch memory budget");
    let prefetch_outcome = QueryExecutor::new(
        &mut prefetch_source,
        None,
        &mut ExactRechecker,
        &mut Diagnostics::default(),
        &NeverCancelled,
    )
    .execute(&prefetch, prefetch_budget)
    .expect("prefetch memory exhaustion should be typed");
    assert_eq!(prefetch_outcome.completion(), Completion::BudgetExhausted);
    assert!(prefetch_outcome.points().is_empty());

    let external = QueryIr::new(
        QueryKind::ExternalRerank {
            query: Box::new(branch(1.0)),
            model_revision: 7,
        },
        ScoreOrder::HigherIsBetter,
        None,
        2,
    )
    .expect("external rerank query");
    let mut external_source = RoutingSource::default();
    let mut reranker = FakeExternalReranker {
        revision: 7,
        exhausted: true,
        calls: 0,
    };
    let external_budget = ExecutionBudget::new(8, 8, 8, 16, 4, 3)
        .expect("external budget")
        .with_resource_limits(
            1_000,
            1024 * 1024,
            leaf.usage().hydration_bytes().saturating_add(1),
            10_000,
        )
        .expect("tiny external hydration budget");
    let external_outcome = QueryExecutor::new(
        &mut external_source,
        None,
        &mut ExactRechecker,
        &mut Diagnostics::default(),
        &NeverCancelled,
    )
    .with_external_reranker(&mut reranker)
    .execute(&external, external_budget)
    .expect("external hydration exhaustion should be typed");
    assert_eq!(reranker.calls, 0);
    assert_eq!(external_outcome.completion(), Completion::BudgetExhausted);
    assert!(external_outcome.points().is_empty());

    let topology_query = QueryIr::new(
        QueryKind::TopologyExpand {
            query: Box::new(branch(1.0)),
            max_depth: 1,
        },
        ScoreOrder::HigherIsBetter,
        None,
        2,
    )
    .expect("topology query");
    let mut topology_source = RoutingSource::default();
    let mut topology = FakeTopology { calls: 0 };
    let topology_budget = ExecutionBudget::new(8, 8, 8, 16, 4, 3)
        .expect("topology budget")
        .with_resource_limits(
            1_000,
            leaf.usage().memory_bytes().saturating_add(1),
            1024 * 1024,
            10_000,
        )
        .expect("tiny topology memory budget");
    let topology_outcome = QueryExecutor::new(
        &mut topology_source,
        None,
        &mut ExactRechecker,
        &mut Diagnostics::default(),
        &NeverCancelled,
    )
    .with_topology_expander(&mut topology)
    .execute(&topology_query, topology_budget)
    .expect("topology memory exhaustion should be typed");
    assert_eq!(topology.calls, 0);
    assert_eq!(topology_outcome.completion(), Completion::BudgetExhausted);
    assert!(topology_outcome.points().is_empty());
}

struct FixedWorkRechecker {
    comparisons: usize,
    calls: usize,
}

impl SourceRechecker for FixedWorkRechecker {
    fn recheck(
        &mut self,
        _query: &QueryIr,
        candidates: &[Candidate],
        limit: usize,
        _budget: PortBudget,
    ) -> Result<RecheckPage, QueryError> {
        self.calls = self.calls.saturating_add(1);
        let rows = candidates
            .iter()
            .take(limit)
            .map(|candidate| {
                HydratedCandidate::new(
                    candidate.point_id(),
                    SourceKey::new(candidate.point_id().get().to_string())?,
                    candidate.approximate_score(),
                )
            })
            .collect::<Result<Vec<_>, QueryError>>()?;
        Ok(RecheckPage::new(rows, self.comparisons))
    }
}

#[test]
fn two_branch_recheck_work_uses_one_global_comparison_budget() {
    let query = QueryIr::new(
        QueryKind::Prefetch {
            branches: vec![branch(1.0), branch(-1.0)],
            fusion: Fusion::STANDARD_RRF,
        },
        ScoreOrder::HigherIsBetter,
        None,
        3,
    )
    .expect("prefetch query");
    let mut source = RoutingSource::default();
    let mut rechecker = FixedWorkRechecker {
        comparisons: 3,
        calls: 0,
    };
    let outcome = QueryExecutor::new(
        &mut source,
        None,
        &mut rechecker,
        &mut Diagnostics::default(),
        &NeverCancelled,
    )
    .execute(
        &query,
        ExecutionBudget::new(8, 8, 8, 16, 4, 3)
            .expect("base budget")
            .with_resource_limits(10, 1024 * 1024, 1024 * 1024, 10_000)
            .expect("comparison budget"),
    )
    .expect("global recheck budget exhaustion should be typed");

    assert_eq!(rechecker.calls, 2);
    assert_eq!(outcome.completion(), Completion::BudgetExhausted);
    assert!(outcome.points().is_empty());
    assert_eq!(outcome.usage().comparisons(), 10);
}

#[test]
fn candidate_expansion_work_is_globally_enforced() {
    struct ExpandingSource;
    impl CandidateSource for ExpandingSource {
        fn readiness(
            &mut self,
            _query: &QueryIr,
            _budget: PortBudget,
        ) -> Result<SourceReadiness, QueryError> {
            Ok(SourceReadiness::Ready)
        }

        fn candidates(
            &mut self,
            _query: &QueryIr,
            _filter: Option<&FilterCandidateBatch>,
            _limit: usize,
            _budget: PortBudget,
        ) -> Result<CandidatePage, QueryError> {
            Ok(CandidatePage::new(vec![candidate(1, 1.0)], true).with_expansion_count(3))
        }
    }

    let error = QueryExecutor::new(
        &mut ExpandingSource,
        None,
        &mut ExactRechecker,
        &mut Diagnostics::default(),
        &NeverCancelled,
    )
    .execute(&branch(1.0), budget(8))
    .expect_err("expansion overrun must fail closed");
    assert!(matches!(
        error,
        QueryError::PortContractViolation {
            stage: "candidate_expansions",
            requested: 2,
            returned: 3,
        }
    ));
}