aft/commands/semantic_search/
blocks.rs

1use std::collections::{BTreeMap, HashMap, HashSet};
2use std::fmt;
3use std::path::PathBuf;
4
5use serde::{Deserialize, Serialize};
6
7use super::comparator::{sort_r3, CandidateResult, RankedTuple, SymbolOffsetRange};
8use super::evidence_descriptor::{EvidenceDescriptor, EvidenceTier};
9use super::plan_table::SearchLaneKind;
10use super::scoring::{
11    freeze_non_exact_scores, AdmittedContribution, LaneContribution, ScoringError, ScoringPolicy,
12};
13
14pub const BLOCK_DEPTHS: [usize; 5] = [200, 400, 800, 1_600, 3_200];
15pub const MAX_BLOCK_DEPTH: usize = 3_200;
16pub const MAX_PUBLIC_TOP_K: usize = 100;
17pub const MAX_OFFSET: usize = 100_000;
18
19/// The complete identity of a canonical result list. Request paging fields are
20/// intentionally absent because they select an interval without defining the list.
21#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
22pub struct CanonicalListKey {
23    pub project_root: PathBuf,
24    pub snapshot_generation: String,
25    pub normalized_query: String,
26    pub include_tests: bool,
27}
28
29#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
30pub struct LaneCandidate {
31    pub path: PathBuf,
32    pub symbol_range: Option<SymbolOffsetRange>,
33    pub evidence: EvidenceDescriptor,
34    #[serde(skip_serializing_if = "Option::is_none")]
35    pub raw_score: Option<f32>,
36    pub is_test: bool,
37}
38
39impl LaneCandidate {
40    pub fn exact(
41        path: impl Into<PathBuf>,
42        symbol_range: Option<SymbolOffsetRange>,
43        evidence: EvidenceDescriptor,
44        is_test: bool,
45    ) -> Self {
46        Self {
47            path: path.into(),
48            symbol_range,
49            evidence,
50            raw_score: None,
51            is_test,
52        }
53    }
54
55    pub fn non_exact(
56        path: impl Into<PathBuf>,
57        symbol_range: Option<SymbolOffsetRange>,
58        evidence: EvidenceDescriptor,
59        raw_score: f32,
60        is_test: bool,
61    ) -> Self {
62        Self {
63            path: path.into(),
64            symbol_range,
65            evidence,
66            raw_score: Some(raw_score),
67            is_test,
68        }
69    }
70
71    fn identity(&self) -> CandidateIdentity {
72        CandidateIdentity {
73            path: self.path.clone(),
74            symbol_range: self.symbol_range,
75        }
76    }
77}
78
79/// A lane's complete canonical order. The outer lane vector accepted by the
80/// builder may be in completion order; no ranking decision relies on that order.
81#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
82pub struct CanonicalLane {
83    pub kind: SearchLaneKind,
84    pub candidates: Vec<LaneCandidate>,
85}
86
87impl CanonicalLane {
88    pub fn new(
89        kind: SearchLaneKind,
90        candidates: Vec<LaneCandidate>,
91    ) -> Result<Self, BlockBuildError> {
92        let mut identities = HashSet::new();
93        for candidate in &candidates {
94            if !candidate.evidence.is_valid_shape() {
95                return Err(BlockBuildError::InvalidEvidence(candidate.path.clone()));
96            }
97            if !identities.insert(candidate.identity()) {
98                return Err(BlockBuildError::DuplicateLaneCandidate {
99                    lane: kind,
100                    path: candidate.path.clone(),
101                });
102            }
103            match candidate.evidence.tier {
104                EvidenceTier::Exact if candidate.raw_score.is_some() => {
105                    return Err(BlockBuildError::InvalidExactCandidate(
106                        candidate.path.clone(),
107                    ));
108                }
109                EvidenceTier::NonExact
110                    if !kind.is_scored()
111                        || candidate.raw_score.is_none_or(|score| !score.is_finite()) =>
112                {
113                    return Err(BlockBuildError::InvalidDepthLimitedCandidate {
114                        lane: kind,
115                        path: candidate.path.clone(),
116                    });
117                }
118                EvidenceTier::Exact | EvidenceTier::NonExact => {}
119            }
120        }
121        Ok(Self { kind, candidates })
122    }
123}
124
125#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
126#[serde(rename_all = "snake_case")]
127pub enum ContributionDisposition {
128    Admitted,
129    NotAdmitted,
130    DepthExempt,
131    ProvenanceOnly,
132}
133
134/// Run-scoped observation. This is attached only after ranking and scoring are
135/// frozen, keeping provenance unavailable to all ordering decisions.
136#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
137pub struct LaneAttribution {
138    pub lane: SearchLaneKind,
139    pub position: usize,
140    pub disposition: ContributionDisposition,
141}
142
143#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
144pub struct StabilityUnit {
145    pub ranked_tuple: RankedTuple,
146    pub evidence_descriptor: EvidenceDescriptor,
147}
148
149#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
150pub struct BlockEntry {
151    pub result: CandidateResult,
152    pub tier_index: usize,
153    pub admitted_contributions: Vec<AdmittedContribution>,
154    pub lane_attribution: Vec<LaneAttribution>,
155    pub r3_order_index: usize,
156}
157
158impl BlockEntry {
159    pub fn stability_unit(&self) -> StabilityUnit {
160        StabilityUnit {
161            ranked_tuple: RankedTuple {
162                file: self.result.path.clone(),
163                symbol_range: self.result.symbol_range,
164                r3_order_index: self.r3_order_index,
165                fusion_score: self.result.fusion_score,
166                lane_score: self.result.lane_score,
167            },
168            evidence_descriptor: self.result.evidence.clone(),
169        }
170    }
171}
172
173#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
174pub struct FrozenBlock {
175    pub tier_index: usize,
176    pub entries: Vec<BlockEntry>,
177}
178
179#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
180pub struct CanonicalList {
181    pub key: CanonicalListKey,
182    pub blocks: Vec<FrozenBlock>,
183}
184
185impl CanonicalList {
186    pub fn len(&self) -> usize {
187        self.blocks.iter().map(|block| block.entries.len()).sum()
188    }
189
190    pub fn is_empty(&self) -> bool {
191        self.blocks.iter().all(|block| block.entries.is_empty())
192    }
193
194    pub fn entries(&self) -> impl Iterator<Item = &BlockEntry> {
195        self.blocks.iter().flat_map(|block| block.entries.iter())
196    }
197
198    pub fn block(&self, tier_index: usize) -> Option<&FrozenBlock> {
199        self.blocks
200            .iter()
201            .find(|block| block.tier_index == tier_index)
202    }
203
204    pub fn stability_units(&self) -> Vec<StabilityUnit> {
205        self.entries().map(BlockEntry::stability_unit).collect()
206    }
207}
208
209#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
210pub struct PageRequest {
211    pub offset: usize,
212    pub top_k: usize,
213}
214
215impl PageRequest {
216    pub fn validate(self) -> Result<Self, BlockBuildError> {
217        if self.offset > MAX_OFFSET {
218            return Err(BlockBuildError::InvalidOffset(self.offset));
219        }
220        if !(1..=MAX_PUBLIC_TOP_K).contains(&self.top_k) {
221            return Err(BlockBuildError::InvalidTopK(self.top_k));
222        }
223        Ok(self)
224    }
225
226    pub fn interval_end(self) -> u64 {
227        (self.offset as u64).saturating_add(self.top_k as u64)
228    }
229
230    pub fn effective_target(self) -> u64 {
231        self.interval_end().max(2)
232    }
233}
234
235#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
236pub struct BlockReply {
237    pub canonical_list: CanonicalList,
238    pub page: Vec<BlockEntry>,
239    pub retrieval_depth: usize,
240    pub depth_tier: usize,
241    pub lanes_exhausted: bool,
242    pub lane_enumeration_counts: BTreeMap<SearchLaneKind, usize>,
243    pub effective_target: u64,
244}
245
246impl BlockReply {
247    pub fn page_stability_units(&self) -> Vec<StabilityUnit> {
248        self.page.iter().map(BlockEntry::stability_unit).collect()
249    }
250}
251
252#[derive(Debug, Clone, PartialEq, Eq)]
253pub enum BlockBuildError {
254    DuplicateLane(SearchLaneKind),
255    DuplicateLaneCandidate { lane: SearchLaneKind, path: PathBuf },
256    InvalidEvidence(PathBuf),
257    InvalidExactCandidate(PathBuf),
258    InvalidDepthLimitedCandidate { lane: SearchLaneKind, path: PathBuf },
259    ConflictingEvidence(PathBuf),
260    MissingExactEvidence(PathBuf),
261    NonExactEmptyAdmittedSet(PathBuf),
262    InvalidDepth(usize),
263    InvalidOffset(usize),
264    InvalidTopK(usize),
265    Scoring(ScoringError),
266}
267
268impl fmt::Display for BlockBuildError {
269    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
270        match self {
271            Self::DuplicateLane(lane) => {
272                write!(formatter, "lane {lane} was supplied more than once")
273            }
274            Self::DuplicateLaneCandidate { lane, path } => write!(
275                formatter,
276                "lane {lane} produced {} more than once",
277                path.display()
278            ),
279            Self::InvalidEvidence(path) => {
280                write!(
281                    formatter,
282                    "candidate {} has invalid evidence",
283                    path.display()
284                )
285            }
286            Self::InvalidExactCandidate(path) => write!(
287                formatter,
288                "exact candidate {} must carry exact evidence and no score",
289                path.display()
290            ),
291            Self::InvalidDepthLimitedCandidate { lane, path } => write!(
292                formatter,
293                "{lane} candidate {} must carry a finite raw score",
294                path.display()
295            ),
296            Self::ConflictingEvidence(path) => write!(
297                formatter,
298                "depth-limited lanes disagree on evidence for {}",
299                path.display()
300            ),
301            Self::MissingExactEvidence(path) => write!(
302                formatter,
303                "exact-attributed candidate {} has no exact descriptor",
304                path.display()
305            ),
306            Self::NonExactEmptyAdmittedSet(path) => write!(
307                formatter,
308                "non-exact candidate {} has an empty admitted set",
309                path.display()
310            ),
311            Self::InvalidDepth(depth) => {
312                write!(
313                    formatter,
314                    "block depth {depth} is not one of {BLOCK_DEPTHS:?}"
315                )
316            }
317            Self::InvalidOffset(offset) => {
318                write!(formatter, "offset {offset} exceeds maximum {MAX_OFFSET}")
319            }
320            Self::InvalidTopK(top_k) => write!(
321                formatter,
322                "topK {top_k} is outside the public range 1..={MAX_PUBLIC_TOP_K}"
323            ),
324            Self::Scoring(error) => error.fmt(formatter),
325        }
326    }
327}
328
329impl std::error::Error for BlockBuildError {}
330
331impl From<ScoringError> for BlockBuildError {
332    fn from(error: ScoringError) -> Self {
333        Self::Scoring(error)
334    }
335}
336
337#[derive(Debug, Clone, PartialEq, Eq, Hash)]
338struct CandidateIdentity {
339    path: PathBuf,
340    symbol_range: Option<SymbolOffsetRange>,
341}
342
343#[derive(Debug, Clone)]
344struct ObservedCandidate<'a> {
345    exact: Vec<(SearchLaneKind, usize, &'a LaneCandidate)>,
346    depth_limited: Vec<LaneContribution>,
347    depth_evidence: Vec<&'a EvidenceDescriptor>,
348}
349
350impl<'a> ObservedCandidate<'a> {
351    fn new() -> Self {
352        Self {
353            exact: Vec::new(),
354            depth_limited: Vec::new(),
355            depth_evidence: Vec::new(),
356        }
357    }
358
359    fn tier_index(&self) -> Option<usize> {
360        if !self.exact.is_empty() {
361            return Some(0);
362        }
363        self.depth_limited
364            .iter()
365            .filter_map(|contribution| tier_for_position(contribution.position))
366            .min()
367    }
368}
369
370/// Builds immutable tier blocks sorted by the canonical ranking rules.
371#[derive(Debug, Clone)]
372pub struct BlockBuilder {
373    key: CanonicalListKey,
374    policy: ScoringPolicy,
375    lanes: Vec<CanonicalLane>,
376}
377
378impl BlockBuilder {
379    pub fn new(
380        key: CanonicalListKey,
381        policy: ScoringPolicy,
382        lanes_in_completion_order: Vec<CanonicalLane>,
383    ) -> Result<Self, BlockBuildError> {
384        let mut kinds = HashSet::new();
385        for lane in &lanes_in_completion_order {
386            if !kinds.insert(lane.kind) {
387                return Err(BlockBuildError::DuplicateLane(lane.kind));
388            }
389        }
390        Ok(Self {
391            key,
392            policy,
393            lanes: lanes_in_completion_order,
394        })
395    }
396
397    /// Builds a public page, increasing retrieval depth only until it can
398    /// determine the requested interval and the canonical list's first two results.
399    pub fn build_for_request(&self, request: PageRequest) -> Result<BlockReply, BlockBuildError> {
400        let request = request.validate()?;
401        let initial_tier = starting_tier(request.interval_end());
402        self.build(initial_tier, Some(request))
403    }
404
405    /// Harness/reference entry point that reconstructs all blocks through one
406    /// explicit retrieval depth without changing the public topK domain.
407    pub fn build_at_depth(&self, depth: usize) -> Result<BlockReply, BlockBuildError> {
408        let tier = BLOCK_DEPTHS
409            .iter()
410            .position(|candidate| *candidate == depth)
411            .ok_or(BlockBuildError::InvalidDepth(depth))?;
412        self.build(tier, None)
413    }
414
415    fn build(
416        &self,
417        initial_tier: usize,
418        request: Option<PageRequest>,
419    ) -> Result<BlockReply, BlockBuildError> {
420        let effective_target = request.map_or(0, PageRequest::effective_target);
421        let mut reached_tier = initial_tier;
422        let mut observed = self.observe_through_depth(BLOCK_DEPTHS[reached_tier]);
423        let mut frozen_identities = HashSet::new();
424        let mut blocks = Vec::new();
425        let mut list_len = 0usize;
426
427        for tier_index in 0..=reached_tier {
428            let block =
429                self.freeze_block(tier_index, &observed, &mut frozen_identities, list_len)?;
430            list_len += block.entries.len();
431            blocks.push(block);
432        }
433
434        while request.is_some()
435            && (list_len as u64) < effective_target
436            && self.has_unconsumed_candidates(BLOCK_DEPTHS[reached_tier])
437            && reached_tier + 1 < BLOCK_DEPTHS.len()
438        {
439            reached_tier += 1;
440            observed = self.observe_through_depth(BLOCK_DEPTHS[reached_tier]);
441            let block =
442                self.freeze_block(reached_tier, &observed, &mut frozen_identities, list_len)?;
443            list_len += block.entries.len();
444            blocks.push(block);
445        }
446
447        let reached_depth = BLOCK_DEPTHS[reached_tier];
448        let lanes_exhausted = !self.has_unconsumed_candidates(reached_depth);
449        self.attach_run_scoped_attribution(&mut blocks, &observed, reached_depth);
450
451        let canonical_list = CanonicalList {
452            key: self.key.clone(),
453            blocks,
454        };
455        let page = request.map_or_else(Vec::new, |page_request| {
456            canonical_list
457                .entries()
458                .skip(page_request.offset)
459                .take(page_request.top_k)
460                .cloned()
461                .collect()
462        });
463
464        Ok(BlockReply {
465            canonical_list,
466            page,
467            retrieval_depth: reached_depth,
468            depth_tier: reached_tier,
469            lanes_exhausted,
470            lane_enumeration_counts: self.lane_enumeration_counts(reached_depth),
471            effective_target,
472        })
473    }
474
475    fn observe_through_depth(
476        &self,
477        depth: usize,
478    ) -> HashMap<CandidateIdentity, ObservedCandidate<'_>> {
479        let mut observed = HashMap::new();
480        for lane in &self.lanes {
481            for (position, candidate) in lane.candidates.iter().enumerate() {
482                if !self.key.include_tests && candidate.is_test {
483                    continue;
484                }
485                if candidate.evidence.tier == EvidenceTier::NonExact && position >= depth {
486                    continue;
487                }
488                let entry = observed
489                    .entry(candidate.identity())
490                    .or_insert_with(ObservedCandidate::new);
491                if candidate.evidence.tier == EvidenceTier::Exact {
492                    entry.exact.push((lane.kind, position, candidate));
493                } else {
494                    entry.depth_limited.push(LaneContribution {
495                        lane: lane.kind,
496                        position,
497                        raw_score: candidate.raw_score.expect("lane validated at construction"),
498                    });
499                    entry.depth_evidence.push(&candidate.evidence);
500                }
501            }
502        }
503        observed
504    }
505
506    fn freeze_block(
507        &self,
508        tier_index: usize,
509        observed: &HashMap<CandidateIdentity, ObservedCandidate<'_>>,
510        frozen_identities: &mut HashSet<CandidateIdentity>,
511        block_start: usize,
512    ) -> Result<FrozenBlock, BlockBuildError> {
513        let tier_depth = BLOCK_DEPTHS[tier_index];
514        let mut pending = Vec::new();
515
516        for (identity, candidate) in observed {
517            if frozen_identities.contains(identity) || candidate.tier_index() != Some(tier_index) {
518                continue;
519            }
520
521            let (result, admitted_contributions) = if !candidate.exact.is_empty() {
522                let exact_evidence = candidate
523                    .exact
524                    .iter()
525                    .map(|(_, _, exact)| &exact.evidence)
526                    .min_by(|left, right| {
527                        exact_evidence_rank(left).cmp(&exact_evidence_rank(right))
528                    })
529                    .ok_or_else(|| BlockBuildError::MissingExactEvidence(identity.path.clone()))?;
530                (
531                    CandidateResult::new_exact(
532                        identity.path.clone(),
533                        identity.symbol_range,
534                        exact_evidence.clone(),
535                    ),
536                    Vec::new(),
537                )
538            } else {
539                let evidence = candidate.depth_evidence.first().ok_or_else(|| {
540                    BlockBuildError::NonExactEmptyAdmittedSet(identity.path.clone())
541                })?;
542                if candidate
543                    .depth_evidence
544                    .iter()
545                    .any(|other| *other != *evidence)
546                {
547                    return Err(BlockBuildError::ConflictingEvidence(identity.path.clone()));
548                }
549                let scores =
550                    freeze_non_exact_scores(&candidate.depth_limited, tier_depth, &self.policy)
551                        .map_err(|error| match error {
552                            ScoringError::EmptyAdmittedSet => {
553                                BlockBuildError::NonExactEmptyAdmittedSet(identity.path.clone())
554                            }
555                            other => BlockBuildError::Scoring(other),
556                        })?;
557                (
558                    CandidateResult::new_non_exact(
559                        identity.path.clone(),
560                        identity.symbol_range,
561                        (*evidence).clone(),
562                        scores.fusion_score,
563                        scores.lane_score,
564                        scores.best_lane,
565                    ),
566                    scores.admitted,
567                )
568            };
569
570            pending.push(BlockEntry {
571                result,
572                tier_index,
573                admitted_contributions,
574                lane_attribution: Vec::new(),
575                r3_order_index: 0,
576            });
577            frozen_identities.insert(identity.clone());
578        }
579
580        let mut results: Vec<_> = pending.iter().map(|entry| entry.result.clone()).collect();
581        sort_r3(&mut results);
582        let mut by_identity: HashMap<_, _> = pending
583            .drain(..)
584            .map(|entry| (identity_of_result(&entry.result), entry))
585            .collect();
586        let mut entries = Vec::with_capacity(results.len());
587        for (within_block, result) in results.into_iter().enumerate() {
588            let identity = identity_of_result(&result);
589            let mut entry = by_identity
590                .remove(&identity)
591                .expect("R3 sorting preserves every deduplicated candidate");
592            entry.result = result;
593            entry.r3_order_index = block_start + within_block;
594            entries.push(entry);
595        }
596
597        Ok(FrozenBlock {
598            tier_index,
599            entries,
600        })
601    }
602
603    fn attach_run_scoped_attribution(
604        &self,
605        blocks: &mut [FrozenBlock],
606        observed: &HashMap<CandidateIdentity, ObservedCandidate<'_>>,
607        reached_depth: usize,
608    ) {
609        for entry in blocks.iter_mut().flat_map(|block| block.entries.iter_mut()) {
610            let identity = identity_of_result(&entry.result);
611            let Some(candidate) = observed.get(&identity) else {
612                continue;
613            };
614            let mut attribution = Vec::new();
615            if entry.result.evidence.tier == EvidenceTier::Exact {
616                attribution.extend(candidate.exact.iter().map(|(lane, position, _)| {
617                    LaneAttribution {
618                        lane: *lane,
619                        position: *position,
620                        disposition: ContributionDisposition::DepthExempt,
621                    }
622                }));
623            }
624            for contribution in &candidate.depth_limited {
625                if contribution.position >= reached_depth {
626                    continue;
627                }
628                let disposition = if entry.result.evidence.tier == EvidenceTier::Exact {
629                    ContributionDisposition::ProvenanceOnly
630                } else if contribution.position < BLOCK_DEPTHS[entry.tier_index] {
631                    ContributionDisposition::Admitted
632                } else {
633                    ContributionDisposition::NotAdmitted
634                };
635                attribution.push(LaneAttribution {
636                    lane: contribution.lane,
637                    position: contribution.position,
638                    disposition,
639                });
640            }
641            attribution.sort_by_key(|value| value.lane.default_plan_order_index());
642            entry.lane_attribution = attribution;
643        }
644    }
645
646    fn has_unconsumed_candidates(&self, depth: usize) -> bool {
647        self.lanes.iter().any(|lane| {
648            lane.candidates
649                .iter()
650                .enumerate()
651                .skip(depth.min(lane.candidates.len()))
652                .any(|(_, candidate)| {
653                    candidate.evidence.tier == EvidenceTier::NonExact
654                        && (self.key.include_tests || !candidate.is_test)
655                })
656        })
657    }
658
659    fn lane_enumeration_counts(&self, depth: usize) -> BTreeMap<SearchLaneKind, usize> {
660        self.lanes
661            .iter()
662            .map(|lane| {
663                let count = if lane
664                    .candidates
665                    .iter()
666                    .all(|candidate| candidate.evidence.tier == EvidenceTier::Exact)
667                {
668                    lane.candidates.len()
669                } else {
670                    lane.candidates.len().min(depth)
671                };
672                (lane.kind, count)
673            })
674            .collect()
675    }
676}
677
678pub fn tier_for_position(position: usize) -> Option<usize> {
679    BLOCK_DEPTHS.iter().position(|depth| position < *depth)
680}
681
682fn starting_tier(interval_end: u64) -> usize {
683    BLOCK_DEPTHS
684        .iter()
685        .position(|depth| (*depth as u64) >= interval_end)
686        .unwrap_or(BLOCK_DEPTHS.len() - 1)
687}
688
689fn identity_of_result(result: &CandidateResult) -> CandidateIdentity {
690    CandidateIdentity {
691        path: result.path.clone(),
692        symbol_range: result.symbol_range,
693    }
694}
695
696fn exact_evidence_rank(evidence: &EvidenceDescriptor) -> (usize, usize, usize, usize) {
697    use super::evidence_descriptor::EvidenceKind;
698
699    let kind = match evidence.kind {
700        EvidenceKind::Definition => 0,
701        EvidenceKind::E1 => 1,
702        EvidenceKind::Anchored => 2,
703        EvidenceKind::E2 => 3,
704        EvidenceKind::None => usize::MAX,
705    };
706    (
707        kind,
708        usize::MAX - evidence.occurrences.unwrap_or(0),
709        usize::MAX - evidence.matched_span.unwrap_or(0),
710        evidence.window_lines.unwrap_or(usize::MAX),
711    )
712}
aft/commands/semantic_search/blocks.rs

aft/commands/semantic_search/
blocks.rs
