Skip to main content

fallow_engine/
similar_code.rs

1//! Pure, bounded evaluation of provider-supplied code vectors.
2//!
3//! This prototype deliberately owns no extraction, model, network, subprocess,
4//! or public output behavior. Orchestration validates provider consent and then
5//! passes vectors into this deterministic layer.
6
7use std::collections::VecDeque;
8use std::fmt;
9use std::mem::size_of;
10
11use rustc_hash::FxHashMap;
12use sha2::{Digest, Sha256};
13
14pub use fallow_types::similar_code::{
15    SIMILAR_CODE_EXTRACTION_SEMANTICS_VERSION as EXTRACTION_SEMANTICS_VERSION,
16    SimilarCodeFunctionLocation as FunctionLocation, SimilarCodeSourceDigest,
17};
18
19/// One extracted function and its provider-supplied vector.
20#[derive(Debug, Clone)]
21pub struct FunctionVector {
22    /// Stable source location for this run.
23    pub location: FunctionLocation,
24    /// Full SHA-256 digest of the exact extracted function source.
25    pub source_sha256: SimilarCodeSourceDigest,
26    /// Version of the extraction semantics that produced the function.
27    pub extraction_semantics_version: u32,
28    /// Dense vector values returned by the provider.
29    pub values: Vec<f32>,
30}
31
32/// Source identity available before provider inference.
33#[derive(Debug, Clone, Copy)]
34pub struct SimilarCodeSelectionInput<'a> {
35    /// Stable source occurrence for this run.
36    pub location: &'a FunctionLocation,
37    /// Full SHA-256 digest of the exact source fragment.
38    pub source_sha256: SimilarCodeSourceDigest,
39    /// Whether this function satisfies every active reporting-scope predicate.
40    pub in_scope: bool,
41}
42
43/// Hard limits for one candidate-evaluation run.
44#[derive(Debug, Clone, Copy, PartialEq, Eq)]
45pub struct SimilarCodeLimits {
46    /// Required vector width.
47    pub dimensions: usize,
48    /// Maximum functions considered after deterministic sorting.
49    pub max_functions: usize,
50    /// Maximum pairwise cosine comparisons.
51    pub max_comparisons: usize,
52    /// Maximum candidates retained after per-function neighbor filtering.
53    pub max_candidates: usize,
54    /// Maximum retained candidates involving any one function.
55    pub max_neighbors_per_function: usize,
56    /// Maximum bytes represented by vectors considered in this run.
57    pub max_vector_bytes: usize,
58}
59
60impl SimilarCodeLimits {
61    /// Construct default work limits for a provider-declared vector width.
62    #[must_use]
63    pub const fn for_dimensions(dimensions: usize) -> Self {
64        Self {
65            dimensions,
66            max_functions: 10_000,
67            max_comparisons: 1_000_000,
68            max_candidates: 4_096,
69            max_neighbors_per_function: 20,
70            max_vector_bytes: 256 * 1024 * 1024,
71        }
72    }
73}
74
75/// Why an evaluation omitted otherwise eligible work.
76#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
77#[non_exhaustive]
78pub enum SimilarCodeSkipReason {
79    /// Functions exceeded the configured function limit.
80    FunctionLimit,
81    /// Functions exceeded the configured vector-memory limit.
82    VectorMemoryLimit,
83    /// Pairwise checks exceeded the comparison limit.
84    ComparisonLimit,
85    /// Threshold-passing pairs exceeded the candidate limit.
86    CandidateLimit,
87    /// Candidate pairs exceeded a per-function neighbor limit.
88    NeighborLimit,
89}
90
91/// Omitted-work count for one stable reason.
92#[derive(Debug, Clone, Copy, PartialEq, Eq)]
93pub struct SimilarCodeSkip {
94    /// Stable skip reason.
95    pub reason: SimilarCodeSkipReason,
96    /// Number of functions, comparisons, or candidates omitted.
97    pub count: usize,
98}
99
100/// Deterministic bounded corpus chosen before provider inference.
101#[derive(Debug, Clone, PartialEq, Eq)]
102pub struct SimilarCodeCorpusSelection {
103    /// Indices into the caller's input slice, in stable occurrence order.
104    pub selected_indices: Vec<usize>,
105    /// Scope membership aligned with `selected_indices`.
106    pub selected_in_scope: Vec<bool>,
107    /// Typed omissions caused by function, vector-memory, or comparison limits.
108    pub skipped: Vec<SimilarCodeSkip>,
109}
110
111/// Whether all eligible work within the supplied corpus was evaluated.
112#[derive(Debug, Clone, Copy, PartialEq, Eq)]
113pub enum SimilarCodeCompletionStatus {
114    /// No configured limit omitted eligible work.
115    Complete,
116    /// One or more configured limits omitted eligible work.
117    Partial,
118}
119
120/// Machine-readable completion evidence for one run.
121#[derive(Debug, Clone, PartialEq, Eq)]
122pub struct SimilarCodeCompletion {
123    /// Complete or partial state.
124    pub status: SimilarCodeCompletionStatus,
125    /// Limits applied to the run.
126    pub limits: SimilarCodeLimits,
127    /// Functions considered after deterministic limiting.
128    pub functions_considered: usize,
129    /// Pairwise comparisons actually performed.
130    pub comparisons_performed: usize,
131    /// Omitted work in stable reason order.
132    pub skipped: Vec<SimilarCodeSkip>,
133}
134
135/// Candidate verification state.
136#[derive(Debug, Clone, Copy, PartialEq, Eq)]
137#[non_exhaustive]
138pub enum SimilarCodeVerificationStatus {
139    /// Provider similarity has not been adjudicated by a human or verifier.
140    Unverified,
141}
142
143/// One advisory similar-code pair.
144#[derive(Debug, Clone, PartialEq)]
145pub struct SimilarCodeCandidate {
146    /// Snapshot identity over content, occurrences, and extraction semantics.
147    pub candidate_id: String,
148    /// Content-only identity stable when both functions move without changing.
149    pub review_key: String,
150    /// Canonically ordered first source location.
151    pub left: FunctionLocation,
152    /// Canonically ordered second source location.
153    pub right: FunctionLocation,
154    /// Cosine similarity in the closed range from -1 to 1.
155    pub similarity: f64,
156    /// Explicit adjudication state. Prototype results always start unverified.
157    pub verification_status: SimilarCodeVerificationStatus,
158}
159
160/// Results and boundedness evidence for one candidate-evaluation run.
161#[derive(Debug, Clone, PartialEq)]
162pub struct SimilarCodeEvaluation {
163    /// Candidates ranked by similarity and then stable source identity.
164    pub candidates: Vec<SimilarCodeCandidate>,
165    /// Machine-readable completion evidence.
166    pub completion: SimilarCodeCompletion,
167}
168
169/// Invalid input rejected before candidate output is produced.
170#[derive(Debug, Clone, PartialEq, Eq)]
171#[non_exhaustive]
172pub enum SimilarCodeError {
173    /// The configured vector width was zero.
174    ZeroDimensions,
175    /// The similarity threshold was non-finite or outside -1 through 1.
176    InvalidThreshold,
177    /// A vector used an unexpected extraction version.
178    ExtractionVersion {
179        /// Source location of the invalid vector.
180        location: FunctionLocation,
181        /// Expected extraction version.
182        expected: u32,
183        /// Actual extraction version.
184        actual: u32,
185    },
186    /// A vector width differed from the configured width.
187    DimensionMismatch {
188        /// Source location of the invalid vector.
189        location: FunctionLocation,
190        /// Expected vector width.
191        expected: usize,
192        /// Actual vector width.
193        actual: usize,
194    },
195    /// A vector contained NaN or infinity.
196    NonFiniteVector {
197        /// Source location of the invalid vector.
198        location: FunctionLocation,
199    },
200    /// A vector had zero magnitude and cannot be compared by cosine similarity.
201    ZeroMagnitudeVector {
202        /// Source location of the invalid vector.
203        location: FunctionLocation,
204    },
205    /// More than one vector claimed the same stable function identity.
206    DuplicateFunctionIdentity {
207        /// Source location shared by the duplicate inputs.
208        location: FunctionLocation,
209    },
210    /// Preselected vector or scope count did not match the source selection contract.
211    SelectionLengthMismatch {
212        /// Number of aligned entries expected from selected source indices.
213        expected: usize,
214        /// Number of aligned entries supplied.
215        actual: usize,
216    },
217    /// A caller-supplied preselection exceeded one current hard limit.
218    SelectionLimitExceeded {
219        /// Limit that the preselection exceeded.
220        reason: SimilarCodeSkipReason,
221        /// Observed functions, bytes, or comparisons.
222        observed: usize,
223        /// Maximum admitted functions, bytes, or comparisons.
224        limit: usize,
225    },
226}
227
228impl fmt::Display for SimilarCodeError {
229    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
230        match self {
231            Self::ZeroDimensions => formatter.write_str("similar-code dimensions must be positive"),
232            Self::InvalidThreshold => {
233                formatter.write_str("similar-code threshold must be finite and between -1 and 1")
234            }
235            Self::ExtractionVersion {
236                location,
237                expected,
238                actual,
239            } => write!(
240                formatter,
241                "{}:{} uses extraction version {actual}, expected {expected}",
242                location.file, location.start_line
243            ),
244            Self::DimensionMismatch {
245                location,
246                expected,
247                actual,
248            } => write!(
249                formatter,
250                "{}:{} has {actual} vector dimensions, expected {expected}",
251                location.file, location.start_line
252            ),
253            Self::NonFiniteVector { location } => write!(
254                formatter,
255                "{}:{} contains a non-finite vector value",
256                location.file, location.start_line
257            ),
258            Self::ZeroMagnitudeVector { location } => write!(
259                formatter,
260                "{}:{} has a zero-magnitude vector",
261                location.file, location.start_line
262            ),
263            Self::DuplicateFunctionIdentity { location } => write!(
264                formatter,
265                "{}:{}:{} has duplicate similar-code vector input",
266                location.file, location.start_line, location.start_column_utf8
267            ),
268            Self::SelectionLengthMismatch { expected, actual } => write!(
269                formatter,
270                "similar-code selection expected {expected} aligned entries, received {actual}"
271            ),
272            Self::SelectionLimitExceeded {
273                reason,
274                observed,
275                limit,
276            } => write!(
277                formatter,
278                "similar-code preselection exceeded {reason:?}: observed {observed}, limit {limit}"
279            ),
280        }
281    }
282}
283
284impl std::error::Error for SimilarCodeError {}
285
286/// Identity for one vector-cache entry.
287#[derive(Debug, Clone, PartialEq, Eq, Hash)]
288pub struct VectorCacheKey {
289    /// Full SHA-256 digest of the exact function source.
290    pub function_source_sha256: SimilarCodeSourceDigest,
291    /// Extraction-semantics version.
292    pub extraction_semantics_version: u32,
293    /// Stable model identifier.
294    pub model_id: String,
295    /// Immutable model revision or artifact digest.
296    pub model_revision: String,
297    /// Vector width.
298    pub dimensions: usize,
299    /// Digest of provider parameters that influence vector output.
300    pub provider_parameter_digest: u64,
301}
302
303/// Separate FIFO cache with a bounded vector-payload budget.
304///
305/// Source fragments are never stored. The byte budget covers vector values,
306/// while map, queue, key, and allocator overhead remain normal process memory.
307#[derive(Debug)]
308pub struct SimilarCodeVectorCache {
309    max_payload_bytes: usize,
310    used_payload_bytes: usize,
311    entries: FxHashMap<VectorCacheKey, Box<[f32]>>,
312    insertion_order: VecDeque<VectorCacheKey>,
313}
314
315impl SimilarCodeVectorCache {
316    /// Create an empty cache with a vector-payload byte budget.
317    #[must_use]
318    pub fn new(max_bytes: usize) -> Self {
319        Self {
320            max_payload_bytes: max_bytes,
321            used_payload_bytes: 0,
322            entries: FxHashMap::default(),
323            insertion_order: VecDeque::new(),
324        }
325    }
326
327    /// Return a cached vector without changing eviction order.
328    #[must_use]
329    pub fn get(&self, key: &VectorCacheKey) -> Option<&[f32]> {
330        self.entries.get(key).map(AsRef::as_ref)
331    }
332
333    /// Insert a vector, evicting oldest entries until it fits.
334    ///
335    /// Returns `false` when the vector itself exceeds the cache budget or its
336    /// width does not match the cache key.
337    pub fn insert(&mut self, key: VectorCacheKey, values: Vec<f32>) -> bool {
338        if values.len() != key.dimensions
339            || values.iter().any(|value| !value.is_finite())
340            || values.iter().all(|value| is_zero(*value))
341        {
342            return false;
343        }
344        let bytes = vector_bytes(values.len());
345        if bytes > self.max_payload_bytes {
346            return false;
347        }
348
349        if let Some(existing) = self.entries.get_mut(&key) {
350            self.used_payload_bytes = self
351                .used_payload_bytes
352                .saturating_sub(vector_bytes(existing.len()))
353                .saturating_add(bytes);
354            *existing = values.into_boxed_slice();
355            return true;
356        }
357
358        while self.used_payload_bytes.saturating_add(bytes) > self.max_payload_bytes {
359            let Some(oldest) = self.insertion_order.pop_front() else {
360                return false;
361            };
362            if let Some(removed) = self.entries.remove(&oldest) {
363                self.used_payload_bytes = self
364                    .used_payload_bytes
365                    .saturating_sub(vector_bytes(removed.len()));
366            }
367        }
368
369        self.used_payload_bytes = self.used_payload_bytes.saturating_add(bytes);
370        self.insertion_order.push_back(key.clone());
371        self.entries.insert(key, values.into_boxed_slice());
372        true
373    }
374
375    /// Bytes represented by cached vector payloads.
376    #[must_use]
377    pub const fn used_payload_bytes(&self) -> usize {
378        self.used_payload_bytes
379    }
380
381    /// Number of cached vectors.
382    #[must_use]
383    pub fn len(&self) -> usize {
384        self.entries.len()
385    }
386
387    /// Whether the cache is empty.
388    #[must_use]
389    pub fn is_empty(&self) -> bool {
390        self.entries.is_empty()
391    }
392}
393
394#[derive(Debug, Clone)]
395struct ScoredPair {
396    left: usize,
397    right: usize,
398    similarity: f64,
399}
400
401#[derive(Debug, Clone, Copy)]
402struct SelectedVector {
403    index: usize,
404    inverse_norm: f64,
405    in_scope: bool,
406}
407
408/// Validate vectors without generating pair candidates.
409pub fn validate_function_vectors(
410    vectors: &[FunctionVector],
411    dimensions: usize,
412    extraction_semantics_version: u32,
413) -> Result<(), SimilarCodeError> {
414    if dimensions == 0 {
415        return Err(SimilarCodeError::ZeroDimensions);
416    }
417
418    for vector in vectors {
419        let _inverse_norm = validate_vector(vector, dimensions, extraction_semantics_version)?;
420    }
421    Ok(())
422}
423
424/// Evaluate provider-supplied vectors with bounded brute-force cosine search.
425pub fn evaluate_similar_code(
426    vectors: &[FunctionVector],
427    threshold: f64,
428    limits: SimilarCodeLimits,
429    extraction_semantics_version: u32,
430) -> Result<SimilarCodeEvaluation, SimilarCodeError> {
431    if limits.dimensions == 0 {
432        return Err(SimilarCodeError::ZeroDimensions);
433    }
434    if !threshold.is_finite() || !(-1.0..=1.0).contains(&threshold) {
435        return Err(SimilarCodeError::InvalidThreshold);
436    }
437    let mut skips = FxHashMap::default();
438    let selected = select_vectors(vectors, limits, extraction_semantics_version, &mut skips)?;
439    let (ranked, comparisons_performed) = score_pairs(vectors, &selected, threshold, limits);
440    let candidates = build_candidates(
441        vectors,
442        &selected,
443        ranked,
444        limits,
445        extraction_semantics_version,
446        &mut skips,
447    );
448
449    let mut skipped = skips
450        .into_iter()
451        .filter(|(_, count)| *count > 0)
452        .map(|(reason, count)| SimilarCodeSkip { reason, count })
453        .collect::<Vec<_>>();
454    skipped.sort_by_key(|skip| skip.reason);
455
456    Ok(SimilarCodeEvaluation {
457        candidates,
458        completion: SimilarCodeCompletion {
459            status: if skipped.is_empty() {
460                SimilarCodeCompletionStatus::Complete
461            } else {
462                SimilarCodeCompletionStatus::Partial
463            },
464            limits,
465            functions_considered: selected.len(),
466            comparisons_performed,
467            skipped,
468        },
469    })
470}
471
472/// Evaluate only the vectors produced for a prior source-corpus selection.
473///
474/// `vectors` must follow `selection.selected_indices` order. This preserves
475/// both the pre-inference omissions and scope membership selected by the caller.
476pub fn evaluate_selected_similar_code(
477    vectors: &[FunctionVector],
478    selection: &SimilarCodeCorpusSelection,
479    threshold: f64,
480    limits: SimilarCodeLimits,
481    extraction_semantics_version: u32,
482) -> Result<SimilarCodeEvaluation, SimilarCodeError> {
483    if selection.selected_in_scope.len() != selection.selected_indices.len() {
484        return Err(SimilarCodeError::SelectionLengthMismatch {
485            expected: selection.selected_indices.len(),
486            actual: selection.selected_in_scope.len(),
487        });
488    }
489    if vectors.len() != selection.selected_indices.len() {
490        return Err(SimilarCodeError::SelectionLengthMismatch {
491            expected: selection.selected_indices.len(),
492            actual: vectors.len(),
493        });
494    }
495
496    if limits.dimensions == 0 {
497        return Err(SimilarCodeError::ZeroDimensions);
498    }
499    if !threshold.is_finite() || !(-1.0..=1.0).contains(&threshold) {
500        return Err(SimilarCodeError::InvalidThreshold);
501    }
502    validate_preselection_limits(&selection.selected_in_scope, limits)?;
503    let selection_inputs = vectors
504        .iter()
505        .zip(&selection.selected_in_scope)
506        .map(|(vector, &in_scope)| SimilarCodeSelectionInput {
507            location: &vector.location,
508            source_sha256: vector.source_sha256,
509            in_scope,
510        })
511        .collect::<Vec<_>>();
512    validated_occurrence_identities(&selection_inputs)?;
513    let selected = vectors
514        .iter()
515        .zip(&selection.selected_in_scope)
516        .enumerate()
517        .map(|(index, (vector, &in_scope))| {
518            let inverse_norm =
519                validate_vector(vector, limits.dimensions, extraction_semantics_version)?;
520            Ok(SelectedVector {
521                index,
522                inverse_norm,
523                in_scope,
524            })
525        })
526        .collect::<Result<Vec<_>, SimilarCodeError>>()?;
527    let mut skips = FxHashMap::default();
528    let (ranked, comparisons_performed) = score_pairs(vectors, &selected, threshold, limits);
529    let candidates = build_candidates(
530        vectors,
531        &selected,
532        ranked,
533        limits,
534        extraction_semantics_version,
535        &mut skips,
536    );
537    let mut evaluation = SimilarCodeEvaluation {
538        candidates,
539        completion: SimilarCodeCompletion {
540            status: SimilarCodeCompletionStatus::Complete,
541            limits,
542            functions_considered: selected.len(),
543            comparisons_performed,
544            skipped: skips
545                .into_iter()
546                .map(|(reason, count)| SimilarCodeSkip { reason, count })
547                .collect(),
548        },
549    };
550    let mut skipped = evaluation
551        .completion
552        .skipped
553        .drain(..)
554        .map(|skip| (skip.reason, skip.count))
555        .collect::<FxHashMap<_, _>>();
556    for skip in &selection.skipped {
557        record_skip(&mut skipped, skip.reason, skip.count);
558    }
559    evaluation.completion.skipped = skipped
560        .into_iter()
561        .map(|(reason, count)| SimilarCodeSkip { reason, count })
562        .collect();
563    evaluation
564        .completion
565        .skipped
566        .sort_by_key(|skip| skip.reason);
567    if !evaluation.completion.skipped.is_empty() {
568        evaluation.completion.status = SimilarCodeCompletionStatus::Partial;
569    }
570    Ok(evaluation)
571}
572
573fn validate_preselection_limits(
574    selected_in_scope: &[bool],
575    limits: SimilarCodeLimits,
576) -> Result<(), SimilarCodeError> {
577    let functions = selected_in_scope.len();
578    if functions > limits.max_functions {
579        return Err(SimilarCodeError::SelectionLimitExceeded {
580            reason: SimilarCodeSkipReason::FunctionLimit,
581            observed: functions,
582            limit: limits.max_functions,
583        });
584    }
585    let vector_bytes = functions.saturating_mul(vector_bytes(limits.dimensions));
586    if vector_bytes > limits.max_vector_bytes {
587        return Err(SimilarCodeError::SelectionLimitExceeded {
588            reason: SimilarCodeSkipReason::VectorMemoryLimit,
589            observed: vector_bytes,
590            limit: limits.max_vector_bytes,
591        });
592    }
593    let scoped_functions = selected_in_scope
594        .iter()
595        .filter(|&&in_scope| in_scope)
596        .count();
597    let comparisons = scoped_pair_count(functions, scoped_functions);
598    if comparisons > limits.max_comparisons {
599        return Err(SimilarCodeError::SelectionLimitExceeded {
600            reason: SimilarCodeSkipReason::ComparisonLimit,
601            observed: comparisons,
602            limit: limits.max_comparisons,
603        });
604    }
605    Ok(())
606}
607
608/// Select a deterministic fair corpus before provider inference.
609///
610/// In-scope functions receive deterministic priority, while both the scoped and
611/// background partitions use normalized source occurrence plus full source
612/// digest for fair selection. The subset is returned in stable occurrence order.
613pub fn select_similar_code_corpus(
614    functions: &[SimilarCodeSelectionInput<'_>],
615    limits: SimilarCodeLimits,
616) -> Result<SimilarCodeCorpusSelection, SimilarCodeError> {
617    if limits.dimensions == 0 {
618        return Err(SimilarCodeError::ZeroDimensions);
619    }
620
621    let scoped_functions = functions
622        .iter()
623        .filter(|function| function.in_scope)
624        .count();
625    if scoped_functions == 0 {
626        return Ok(SimilarCodeCorpusSelection {
627            selected_indices: Vec::new(),
628            selected_in_scope: Vec::new(),
629            skipped: Vec::new(),
630        });
631    }
632
633    let occurrence_identities = validated_occurrence_identities(functions)?;
634
635    let memory_function_limit = limits
636        .max_vector_bytes
637        .checked_div(vector_bytes(limits.dimensions))
638        .unwrap_or(0);
639    let function_limit = functions.len().min(limits.max_functions);
640    let memory_considered = function_limit.min(memory_function_limit);
641    let considered = functions_within_scoped_comparison_budget(
642        memory_considered,
643        scoped_functions,
644        limits.max_comparisons,
645    );
646
647    let order = selected_corpus_order(functions, &occurrence_identities, considered);
648
649    let mut skips = FxHashMap::default();
650    record_skip(
651        &mut skips,
652        SimilarCodeSkipReason::FunctionLimit,
653        functions.len().saturating_sub(function_limit),
654    );
655    record_skip(
656        &mut skips,
657        SimilarCodeSkipReason::VectorMemoryLimit,
658        function_limit.saturating_sub(memory_considered),
659    );
660    record_skip(
661        &mut skips,
662        SimilarCodeSkipReason::ComparisonLimit,
663        scoped_pair_count(memory_considered, scoped_functions.min(memory_considered))
664            .saturating_sub(scoped_pair_count(
665                considered,
666                scoped_functions.min(considered),
667            )),
668    );
669    let mut skipped = skips
670        .into_iter()
671        .map(|(reason, count)| SimilarCodeSkip { reason, count })
672        .collect::<Vec<_>>();
673    skipped.sort_by_key(|skip| skip.reason);
674
675    let selected_in_scope = order
676        .iter()
677        .map(|&index| functions[index].in_scope)
678        .collect();
679    Ok(SimilarCodeCorpusSelection {
680        selected_indices: order,
681        selected_in_scope,
682        skipped,
683    })
684}
685
686fn validated_occurrence_identities(
687    functions: &[SimilarCodeSelectionInput<'_>],
688) -> Result<Vec<String>, SimilarCodeError> {
689    let occurrence_identities = functions
690        .iter()
691        .map(|function| occurrence_identity_for_location(function.location))
692        .collect::<Vec<_>>();
693    let mut identity_order = (0..functions.len()).collect::<Vec<_>>();
694    identity_order.sort_by(|&left, &right| {
695        occurrence_identities[left]
696            .cmp(&occurrence_identities[right])
697            .then_with(|| {
698                functions[left]
699                    .source_sha256
700                    .cmp(&functions[right].source_sha256)
701            })
702    });
703    for duplicate in identity_order.windows(2) {
704        if occurrence_identities[duplicate[0]] == occurrence_identities[duplicate[1]] {
705            return Err(SimilarCodeError::DuplicateFunctionIdentity {
706                location: functions[duplicate[0]].location.clone(),
707            });
708        }
709    }
710    Ok(occurrence_identities)
711}
712
713fn selected_corpus_order(
714    functions: &[SimilarCodeSelectionInput<'_>],
715    occurrence_identities: &[String],
716    considered: usize,
717) -> Vec<usize> {
718    let selection_keys = functions.iter().map(selection_key).collect::<Vec<_>>();
719    let compare_selection = |&left: &usize, &right: &usize| {
720        selection_keys[left]
721            .cmp(&selection_keys[right])
722            .then_with(|| {
723                functions[left]
724                    .source_sha256
725                    .cmp(&functions[right].source_sha256)
726            })
727            .then_with(|| occurrence_identities[left].cmp(&occurrence_identities[right]))
728    };
729    let mut scoped = (0..functions.len())
730        .filter(|&index| functions[index].in_scope)
731        .collect::<Vec<_>>();
732    let mut background = (0..functions.len())
733        .filter(|&index| !functions[index].in_scope)
734        .collect::<Vec<_>>();
735    scoped.sort_by(compare_selection);
736    background.sort_by(compare_selection);
737    scoped.truncate(considered);
738    background.truncate(considered.saturating_sub(scoped.len()));
739    scoped.extend(background);
740    scoped.sort_by(|&left, &right| {
741        occurrence_identities[left]
742            .cmp(&occurrence_identities[right])
743            .then_with(|| {
744                functions[left]
745                    .source_sha256
746                    .cmp(&functions[right].source_sha256)
747            })
748    });
749    scoped
750}
751
752fn select_vectors(
753    vectors: &[FunctionVector],
754    limits: SimilarCodeLimits,
755    extraction_semantics_version: u32,
756    skips: &mut FxHashMap<SimilarCodeSkipReason, usize>,
757) -> Result<Vec<SelectedVector>, SimilarCodeError> {
758    let functions = vectors
759        .iter()
760        .map(|vector| SimilarCodeSelectionInput {
761            location: &vector.location,
762            source_sha256: vector.source_sha256,
763            in_scope: true,
764        })
765        .collect::<Vec<_>>();
766    let selection = select_similar_code_corpus(&functions, limits)?;
767    for skip in selection.skipped {
768        record_skip(skips, skip.reason, skip.count);
769    }
770    selection
771        .selected_indices
772        .into_iter()
773        .zip(selection.selected_in_scope)
774        .map(|(index, in_scope)| {
775            let inverse_norm = validate_vector(
776                &vectors[index],
777                limits.dimensions,
778                extraction_semantics_version,
779            )?;
780            Ok(SelectedVector {
781                index,
782                inverse_norm,
783                in_scope,
784            })
785        })
786        .collect::<Result<Vec<_>, SimilarCodeError>>()
787}
788
789fn score_pairs(
790    vectors: &[FunctionVector],
791    selected: &[SelectedVector],
792    threshold: f64,
793    limits: SimilarCodeLimits,
794) -> (Vec<ScoredPair>, usize) {
795    let scoped_functions = selected.iter().filter(|vector| vector.in_scope).count();
796    let possible_comparisons = scoped_pair_count(selected.len(), scoped_functions);
797    debug_assert!(possible_comparisons <= limits.max_comparisons);
798    let mut comparisons_performed = 0usize;
799    let mut ranked = Vec::with_capacity(possible_comparisons);
800
801    for left in 0..selected.len() {
802        for right in left + 1..selected.len() {
803            if !selected[left].in_scope && !selected[right].in_scope {
804                continue;
805            }
806            comparisons_performed += 1;
807            let similarity = cosine_similarity(
808                &vectors[selected[left].index],
809                &vectors[selected[right].index],
810                selected[left].inverse_norm,
811                selected[right].inverse_norm,
812            );
813            if similarity < threshold {
814                continue;
815            }
816            ranked.push(ScoredPair {
817                left,
818                right,
819                similarity,
820            });
821        }
822    }
823
824    ranked.sort_by(|left, right| {
825        right
826            .similarity
827            .total_cmp(&left.similarity)
828            .then_with(|| left.left.cmp(&right.left))
829            .then_with(|| left.right.cmp(&right.right))
830    });
831    (ranked, comparisons_performed)
832}
833
834fn build_candidates(
835    vectors: &[FunctionVector],
836    selected: &[SelectedVector],
837    ranked: Vec<ScoredPair>,
838    limits: SimilarCodeLimits,
839    extraction_semantics_version: u32,
840    skips: &mut FxHashMap<SimilarCodeSkipReason, usize>,
841) -> Vec<SimilarCodeCandidate> {
842    let mut neighbors = vec![0usize; selected.len()];
843    let mut candidates = Vec::with_capacity(ranked.len().min(limits.max_candidates));
844    for pair in ranked {
845        if candidates.len() >= limits.max_candidates {
846            record_skip(skips, SimilarCodeSkipReason::CandidateLimit, 1);
847            continue;
848        }
849        if neighbors[pair.left] >= limits.max_neighbors_per_function
850            || neighbors[pair.right] >= limits.max_neighbors_per_function
851        {
852            record_skip(skips, SimilarCodeSkipReason::NeighborLimit, 1);
853            continue;
854        }
855        neighbors[pair.left] += 1;
856        neighbors[pair.right] += 1;
857
858        let left = &vectors[selected[pair.left].index];
859        let right = &vectors[selected[pair.right].index];
860        let (left, right) = canonical_pair(left, right);
861        candidates.push(SimilarCodeCandidate {
862            candidate_id: candidate_id(left, right, extraction_semantics_version),
863            review_key: review_key(left, right, extraction_semantics_version),
864            left: normalized_location(&left.location),
865            right: normalized_location(&right.location),
866            similarity: pair.similarity,
867            verification_status: SimilarCodeVerificationStatus::Unverified,
868        });
869    }
870    candidates
871}
872
873fn validate_vector(
874    vector: &FunctionVector,
875    dimensions: usize,
876    extraction_semantics_version: u32,
877) -> Result<f64, SimilarCodeError> {
878    if vector.extraction_semantics_version != extraction_semantics_version {
879        return Err(SimilarCodeError::ExtractionVersion {
880            location: vector.location.clone(),
881            expected: extraction_semantics_version,
882            actual: vector.extraction_semantics_version,
883        });
884    }
885    if vector.values.len() != dimensions {
886        return Err(SimilarCodeError::DimensionMismatch {
887            location: vector.location.clone(),
888            expected: dimensions,
889            actual: vector.values.len(),
890        });
891    }
892    if vector.values.iter().any(|value| !value.is_finite()) {
893        return Err(SimilarCodeError::NonFiniteVector {
894            location: vector.location.clone(),
895        });
896    }
897    if vector.values.iter().all(|value| is_zero(*value)) {
898        return Err(SimilarCodeError::ZeroMagnitudeVector {
899            location: vector.location.clone(),
900        });
901    }
902    let squared_norm = vector.values.iter().fold(0.0, |norm, &value| {
903        let value = f64::from(value);
904        value.mul_add(value, norm)
905    });
906    Ok(squared_norm.sqrt().recip())
907}
908
909fn cosine_similarity(
910    left: &FunctionVector,
911    right: &FunctionVector,
912    left_inverse_norm: f64,
913    right_inverse_norm: f64,
914) -> f64 {
915    let mut dot = 0.0;
916    for (&left_value, &right_value) in left.values.iter().zip(&right.values) {
917        let left_value = f64::from(left_value);
918        let right_value = f64::from(right_value);
919        dot = left_value.mul_add(right_value, dot);
920    }
921    (dot * left_inverse_norm * right_inverse_norm).clamp(-1.0, 1.0)
922}
923
924const fn is_zero(value: f32) -> bool {
925    value.to_bits().trailing_zeros() >= 31
926}
927
928fn candidate_id(
929    left: &FunctionVector,
930    right: &FunctionVector,
931    extraction_semantics_version: u32,
932) -> String {
933    let mut hasher = Sha256::new();
934    hasher.update(b"fallow:similar-code:candidate-snapshot:v1\0");
935    hasher.update(extraction_semantics_version.to_be_bytes());
936    update_snapshot_identity(&mut hasher, left);
937    update_snapshot_identity(&mut hasher, right);
938    format_digest_id("similar-code:candidate:v1:", hasher.finalize().as_ref())
939}
940
941fn review_key(
942    left: &FunctionVector,
943    right: &FunctionVector,
944    extraction_semantics_version: u32,
945) -> String {
946    let (first, second) = if left.source_sha256 <= right.source_sha256 {
947        (left.source_sha256, right.source_sha256)
948    } else {
949        (right.source_sha256, left.source_sha256)
950    };
951    let mut hasher = Sha256::new();
952    hasher.update(b"fallow:similar-code:review-key:v1\0");
953    hasher.update(extraction_semantics_version.to_be_bytes());
954    hasher.update(first.as_bytes());
955    hasher.update(second.as_bytes());
956    format_digest_id("similar-code:review:v1:", hasher.finalize().as_ref())
957}
958
959fn canonical_pair<'a>(
960    left: &'a FunctionVector,
961    right: &'a FunctionVector,
962) -> (&'a FunctionVector, &'a FunctionVector) {
963    let left_occurrence = occurrence_identity(left);
964    let right_occurrence = occurrence_identity(right);
965    if (left.source_sha256, left_occurrence) <= (right.source_sha256, right_occurrence) {
966        (left, right)
967    } else {
968        (right, left)
969    }
970}
971
972fn update_snapshot_identity(hasher: &mut Sha256, vector: &FunctionVector) {
973    hasher.update(vector.source_sha256.as_bytes());
974    let occurrence = occurrence_identity(vector);
975    hasher.update(
976        u64::try_from(occurrence.len())
977            .unwrap_or(u64::MAX)
978            .to_be_bytes(),
979    );
980    hasher.update(occurrence.as_bytes());
981}
982
983fn selection_key(function: &SimilarCodeSelectionInput<'_>) -> [u8; 32] {
984    let mut hasher = Sha256::new();
985    hasher.update(b"fallow:similar-code:corpus-selection:v1\0");
986    hasher.update(function.source_sha256.as_bytes());
987    let occurrence = occurrence_identity_for_location(function.location);
988    hasher.update(
989        u64::try_from(occurrence.len())
990            .unwrap_or(u64::MAX)
991            .to_be_bytes(),
992    );
993    hasher.update(occurrence.as_bytes());
994    hasher.finalize().into()
995}
996
997fn format_digest_id(prefix: &str, digest: &[u8]) -> String {
998    const HEX: &[u8; 16] = b"0123456789abcdef";
999    let mut output = String::with_capacity(prefix.len().saturating_add(digest.len() * 2));
1000    output.push_str(prefix);
1001    for byte in digest {
1002        output.push(char::from(HEX[usize::from(byte >> 4)]));
1003        output.push(char::from(HEX[usize::from(byte & 0x0f)]));
1004    }
1005    output
1006}
1007
1008fn occurrence_identity(vector: &FunctionVector) -> String {
1009    occurrence_identity_for_location(&vector.location)
1010}
1011
1012fn occurrence_identity_for_location(location: &FunctionLocation) -> String {
1013    let location = normalized_location(location);
1014    let path = location.file;
1015    format!(
1016        "{}:{path}:{}:{}:{}:{}:{}:{}",
1017        path.len(),
1018        location.start_byte,
1019        location.end_byte,
1020        location.start_line,
1021        location.start_column_utf8,
1022        location.end_line,
1023        location.end_column_utf8
1024    )
1025}
1026
1027fn normalized_location(location: &FunctionLocation) -> FunctionLocation {
1028    FunctionLocation {
1029        file: location.file.replace('\\', "/"),
1030        start_byte: location.start_byte,
1031        end_byte: location.end_byte,
1032        start_line: location.start_line,
1033        start_column_utf8: location.start_column_utf8,
1034        end_line: location.end_line,
1035        end_column_utf8: location.end_column_utf8,
1036    }
1037}
1038
1039fn record_skip(
1040    skips: &mut FxHashMap<SimilarCodeSkipReason, usize>,
1041    reason: SimilarCodeSkipReason,
1042    count: usize,
1043) {
1044    if count > 0 {
1045        let value = skips.entry(reason).or_default();
1046        *value = value.saturating_add(count);
1047    }
1048}
1049
1050const fn vector_bytes(dimensions: usize) -> usize {
1051    dimensions.saturating_mul(size_of::<f32>())
1052}
1053
1054const fn pair_count(functions: usize) -> usize {
1055    functions.saturating_mul(functions.saturating_sub(1)) / 2
1056}
1057
1058const fn scoped_pair_count(functions: usize, scoped_functions: usize) -> usize {
1059    pair_count(functions).saturating_sub(pair_count(functions.saturating_sub(scoped_functions)))
1060}
1061
1062fn functions_within_scoped_comparison_budget(
1063    max_functions: usize,
1064    scoped_functions: usize,
1065    max_comparisons: usize,
1066) -> usize {
1067    if scoped_functions == 0 {
1068        return 0;
1069    }
1070    let mut low = 0usize;
1071    let mut high = max_functions;
1072    while low < high {
1073        let middle = low + (high - low).div_ceil(2);
1074        let comparisons = scoped_pair_count(middle, scoped_functions.min(middle));
1075        if comparisons <= max_comparisons {
1076            low = middle;
1077        } else {
1078            high = middle.saturating_sub(1);
1079        }
1080    }
1081    low
1082}
1083
1084#[cfg(test)]
1085mod tests {
1086    use super::*;
1087
1088    const DIMENSIONS: usize = 256;
1089
1090    fn digest(seed: u64) -> SimilarCodeSourceDigest {
1091        let mut bytes = [0; 32];
1092        bytes[24..].copy_from_slice(&seed.to_be_bytes());
1093        SimilarCodeSourceDigest::new(bytes)
1094    }
1095
1096    fn vector(file: &str, hash: u64, first: f32, second: f32) -> FunctionVector {
1097        let mut values = vec![0.0; DIMENSIONS];
1098        values[0] = first;
1099        values[1] = second;
1100        FunctionVector {
1101            location: FunctionLocation {
1102                file: file.into(),
1103                start_byte: 0,
1104                end_byte: 100,
1105                start_line: 1,
1106                start_column_utf8: 0,
1107                end_line: 10,
1108                end_column_utf8: 1,
1109            },
1110            source_sha256: digest(hash),
1111            extraction_semantics_version: EXTRACTION_SEMANTICS_VERSION,
1112            values,
1113        }
1114    }
1115
1116    fn limits() -> SimilarCodeLimits {
1117        SimilarCodeLimits {
1118            dimensions: DIMENSIONS,
1119            max_functions: 100,
1120            max_comparisons: 100,
1121            max_candidates: 100,
1122            max_neighbors_per_function: 100,
1123            max_vector_bytes: 100 * DIMENSIONS * size_of::<f32>(),
1124        }
1125    }
1126
1127    #[test]
1128    fn recorded_vectors_find_only_the_similar_pair() {
1129        let result = evaluate_similar_code(
1130            &[
1131                vector("src/a.ts", 0x10, 1.0, 0.0),
1132                vector("src/b.ts", 0x20, 0.98, 0.02),
1133                vector("src/c.ts", 0x30, 0.0, 1.0),
1134            ],
1135            0.95,
1136            limits(),
1137            EXTRACTION_SEMANTICS_VERSION,
1138        )
1139        .unwrap();
1140
1141        assert_eq!(result.candidates.len(), 1);
1142        assert!(
1143            result.candidates[0]
1144                .candidate_id
1145                .starts_with("similar-code:candidate:v1:")
1146        );
1147        assert!(
1148            result.candidates[0]
1149                .review_key
1150                .starts_with("similar-code:review:v1:")
1151        );
1152        assert_eq!(
1153            result.candidates[0].verification_status,
1154            SimilarCodeVerificationStatus::Unverified
1155        );
1156        assert_eq!(
1157            result.completion.status,
1158            SimilarCodeCompletionStatus::Complete
1159        );
1160    }
1161
1162    #[test]
1163    fn candidate_identity_and_order_ignore_input_order() {
1164        let forward = vec![
1165            vector("src/a.ts", 0x30, 1.0, 0.0),
1166            vector("src/b.ts", 0x10, 1.0, 0.0),
1167            vector("src/c.ts", 0x20, 1.0, 0.0),
1168        ];
1169        let mut reverse = forward.clone();
1170        reverse.reverse();
1171
1172        let forward =
1173            evaluate_similar_code(&forward, 0.9, limits(), EXTRACTION_SEMANTICS_VERSION).unwrap();
1174        let reverse =
1175            evaluate_similar_code(&reverse, 0.9, limits(), EXTRACTION_SEMANTICS_VERSION).unwrap();
1176
1177        assert_eq!(forward, reverse);
1178    }
1179
1180    #[test]
1181    fn candidate_identity_and_output_normalize_path_separators() {
1182        let unix = evaluate_similar_code(
1183            &[
1184                vector("src/a.ts", 0x10, 1.0, 0.0),
1185                vector("src/b.ts", 0x20, 1.0, 0.0),
1186            ],
1187            0.9,
1188            limits(),
1189            EXTRACTION_SEMANTICS_VERSION,
1190        )
1191        .unwrap();
1192        let windows = evaluate_similar_code(
1193            &[
1194                vector("src\\a.ts", 0x10, 1.0, 0.0),
1195                vector("src\\b.ts", 0x20, 1.0, 0.0),
1196            ],
1197            0.9,
1198            limits(),
1199            EXTRACTION_SEMANTICS_VERSION,
1200        )
1201        .unwrap();
1202
1203        assert_eq!(unix, windows);
1204    }
1205
1206    #[test]
1207    fn repeated_content_at_distinct_locations_has_unique_candidate_ids() {
1208        let result = evaluate_similar_code(
1209            &[
1210                vector("src/a.ts", 1, 1.0, 0.0),
1211                vector("src/b.ts", 1, 1.0, 0.0),
1212                vector("src/c.ts", 1, 1.0, 0.0),
1213            ],
1214            0.9,
1215            limits(),
1216            EXTRACTION_SEMANTICS_VERSION,
1217        )
1218        .unwrap();
1219        let ids = result
1220            .candidates
1221            .iter()
1222            .map(|candidate| candidate.candidate_id.as_str())
1223            .collect::<rustc_hash::FxHashSet<_>>();
1224        let review_keys = result
1225            .candidates
1226            .iter()
1227            .map(|candidate| candidate.review_key.as_str())
1228            .collect::<rustc_hash::FxHashSet<_>>();
1229
1230        assert_eq!(ids.len(), result.candidates.len());
1231        assert_eq!(review_keys.len(), 1);
1232    }
1233
1234    #[test]
1235    fn review_key_survives_moves_while_candidate_id_tracks_snapshot() {
1236        let original = evaluate_similar_code(
1237            &[
1238                vector("src/a.ts", 1, 1.0, 0.0),
1239                vector("src/b.ts", 2, 1.0, 0.0),
1240            ],
1241            0.9,
1242            limits(),
1243            EXTRACTION_SEMANTICS_VERSION,
1244        )
1245        .unwrap();
1246        let moved = evaluate_similar_code(
1247            &[
1248                vector("packages/core/a.ts", 1, 1.0, 0.0),
1249                vector("packages/core/b.ts", 2, 1.0, 0.0),
1250            ],
1251            0.9,
1252            limits(),
1253            EXTRACTION_SEMANTICS_VERSION,
1254        )
1255        .unwrap();
1256
1257        assert_ne!(
1258            original.candidates[0].candidate_id,
1259            moved.candidates[0].candidate_id
1260        );
1261        assert_eq!(
1262            original.candidates[0].review_key,
1263            moved.candidates[0].review_key
1264        );
1265    }
1266
1267    #[test]
1268    fn comparison_budget_selects_a_stable_fair_corpus_and_checks_every_pair() {
1269        let vectors = (0..6)
1270            .map(|index| vector(&format!("src/{index}.ts"), index, 1.0, 0.0))
1271            .collect::<Vec<_>>();
1272        let mut bounded = limits();
1273        bounded.max_functions = vectors.len();
1274        bounded.max_comparisons = 3;
1275        bounded.max_candidates = 3;
1276
1277        let inputs = vectors
1278            .iter()
1279            .map(|vector| SimilarCodeSelectionInput {
1280                location: &vector.location,
1281                source_sha256: vector.source_sha256,
1282                in_scope: true,
1283            })
1284            .collect::<Vec<_>>();
1285        let corpus = select_similar_code_corpus(&inputs, bounded).unwrap();
1286        let expected = corpus
1287            .selected_indices
1288            .iter()
1289            .copied()
1290            .map(|index| vectors[index].location.file.as_str())
1291            .collect::<rustc_hash::FxHashSet<_>>();
1292
1293        let direct =
1294            evaluate_similar_code(&vectors, 0.9, bounded, EXTRACTION_SEMANTICS_VERSION).unwrap();
1295        let embedded = corpus
1296            .selected_indices
1297            .iter()
1298            .map(|&index| vectors[index].clone())
1299            .collect::<Vec<_>>();
1300        let result = evaluate_selected_similar_code(
1301            &embedded,
1302            &corpus,
1303            0.9,
1304            bounded,
1305            EXTRACTION_SEMANTICS_VERSION,
1306        )
1307        .unwrap();
1308        let selected = result
1309            .candidates
1310            .iter()
1311            .flat_map(|candidate| [candidate.left.file.as_str(), candidate.right.file.as_str()])
1312            .collect::<rustc_hash::FxHashSet<_>>();
1313
1314        assert_eq!(result.completion.functions_considered, 3);
1315        assert_eq!(result.completion.comparisons_performed, 3);
1316        assert_eq!(selected, expected);
1317        assert_eq!(result.completion.skipped, corpus.skipped);
1318        assert_eq!(result, direct);
1319    }
1320
1321    #[test]
1322    fn scoped_functions_receive_priority_inside_the_comparison_budget() {
1323        let vectors = (0..6)
1324            .map(|index| vector(&format!("src/{index}.ts"), index, 1.0, 0.0))
1325            .collect::<Vec<_>>();
1326        let mut bounded = limits();
1327        bounded.max_functions = vectors.len();
1328        bounded.max_comparisons = 3;
1329        let inputs = vectors
1330            .iter()
1331            .enumerate()
1332            .map(|(index, vector)| SimilarCodeSelectionInput {
1333                location: &vector.location,
1334                source_sha256: vector.source_sha256,
1335                in_scope: index == 5,
1336            })
1337            .collect::<Vec<_>>();
1338
1339        let selection = select_similar_code_corpus(&inputs, bounded).unwrap();
1340
1341        assert!(selection.selected_indices.contains(&5));
1342        assert_eq!(selection.selected_indices.len(), 4);
1343        assert_eq!(
1344            selection
1345                .selected_in_scope
1346                .iter()
1347                .filter(|&&value| value)
1348                .count(),
1349            1
1350        );
1351    }
1352
1353    #[test]
1354    fn empty_scope_does_not_admit_or_embed_background_functions() {
1355        let vectors = (0..6)
1356            .map(|index| vector(&format!("src/{index}.ts"), index, 1.0, 0.0))
1357            .collect::<Vec<_>>();
1358        let inputs = vectors
1359            .iter()
1360            .map(|vector| SimilarCodeSelectionInput {
1361                location: &vector.location,
1362                source_sha256: vector.source_sha256,
1363                in_scope: false,
1364            })
1365            .collect::<Vec<_>>();
1366
1367        let mut bounded = limits();
1368        bounded.max_functions = 1;
1369        bounded.max_vector_bytes = vector_bytes(DIMENSIONS);
1370        bounded.max_comparisons = 0;
1371        let selection = select_similar_code_corpus(&inputs, bounded).unwrap();
1372
1373        assert!(selection.selected_indices.is_empty());
1374        assert!(selection.selected_in_scope.is_empty());
1375        assert!(selection.skipped.is_empty());
1376    }
1377
1378    #[test]
1379    fn scoped_evaluation_rejects_background_only_pairs() {
1380        let vectors = vec![
1381            vector("src/scoped.ts", 1, 1.0, 0.0),
1382            vector("src/background-a.ts", 2, 1.0, 0.0),
1383            vector("src/background-b.ts", 3, 1.0, 0.0),
1384        ];
1385        let selection = SimilarCodeCorpusSelection {
1386            selected_indices: vec![0, 1, 2],
1387            selected_in_scope: vec![true, false, false],
1388            skipped: Vec::new(),
1389        };
1390
1391        let result = evaluate_selected_similar_code(
1392            &vectors,
1393            &selection,
1394            0.9,
1395            limits(),
1396            EXTRACTION_SEMANTICS_VERSION,
1397        )
1398        .unwrap();
1399
1400        assert_eq!(result.completion.comparisons_performed, 2);
1401        assert_eq!(result.candidates.len(), 2);
1402        assert!(result.candidates.iter().all(|candidate| {
1403            candidate.left.file == "src/scoped.ts" || candidate.right.file == "src/scoped.ts"
1404        }));
1405    }
1406
1407    #[test]
1408    fn preselected_evaluation_enforces_every_hard_corpus_limit() {
1409        let vectors = vec![
1410            vector("src/a.ts", 1, 1.0, 0.0),
1411            vector("src/b.ts", 2, 1.0, 0.0),
1412            vector("src/c.ts", 3, 1.0, 0.0),
1413        ];
1414        let selection = SimilarCodeCorpusSelection {
1415            selected_indices: vec![0, 1, 2],
1416            selected_in_scope: vec![true, true, true],
1417            skipped: Vec::new(),
1418        };
1419
1420        let mut bounded = limits();
1421        bounded.max_functions = 2;
1422        assert!(matches!(
1423            evaluate_selected_similar_code(
1424                &vectors,
1425                &selection,
1426                0.9,
1427                bounded,
1428                EXTRACTION_SEMANTICS_VERSION,
1429            ),
1430            Err(SimilarCodeError::SelectionLimitExceeded {
1431                reason: SimilarCodeSkipReason::FunctionLimit,
1432                ..
1433            })
1434        ));
1435
1436        bounded = limits();
1437        bounded.max_vector_bytes = vector_bytes(DIMENSIONS) * 2;
1438        assert!(matches!(
1439            evaluate_selected_similar_code(
1440                &vectors,
1441                &selection,
1442                0.9,
1443                bounded,
1444                EXTRACTION_SEMANTICS_VERSION,
1445            ),
1446            Err(SimilarCodeError::SelectionLimitExceeded {
1447                reason: SimilarCodeSkipReason::VectorMemoryLimit,
1448                ..
1449            })
1450        ));
1451
1452        bounded = limits();
1453        bounded.max_comparisons = 2;
1454        assert!(matches!(
1455            evaluate_selected_similar_code(
1456                &vectors,
1457                &selection,
1458                0.9,
1459                bounded,
1460                EXTRACTION_SEMANTICS_VERSION,
1461            ),
1462            Err(SimilarCodeError::SelectionLimitExceeded {
1463                reason: SimilarCodeSkipReason::ComparisonLimit,
1464                ..
1465            })
1466        ));
1467    }
1468
1469    #[test]
1470    fn preselected_evaluation_rejects_duplicate_physical_functions() {
1471        let duplicate = vector("src/a.ts", 1, 1.0, 0.0);
1472        let vectors = vec![duplicate.clone(), duplicate];
1473        let selection = SimilarCodeCorpusSelection {
1474            selected_indices: vec![0, 1],
1475            selected_in_scope: vec![true, true],
1476            skipped: Vec::new(),
1477        };
1478
1479        assert!(matches!(
1480            evaluate_selected_similar_code(
1481                &vectors,
1482                &selection,
1483                0.9,
1484                limits(),
1485                EXTRACTION_SEMANTICS_VERSION,
1486            ),
1487            Err(SimilarCodeError::DuplicateFunctionIdentity { .. })
1488        ));
1489    }
1490
1491    #[test]
1492    fn neighbor_filter_backfills_before_the_global_candidate_limit() {
1493        let mut bounded = limits();
1494        bounded.max_comparisons = 6;
1495        bounded.max_candidates = 2;
1496        bounded.max_neighbors_per_function = 1;
1497
1498        let result = evaluate_similar_code(
1499            &[
1500                vector("src/a.ts", 1, 1.0, 0.0),
1501                vector("src/b.ts", 2, 1.0, 0.0),
1502                vector("src/c.ts", 3, 1.0, 0.0),
1503                vector("src/d.ts", 4, 1.0, 0.0),
1504            ],
1505            0.9,
1506            bounded,
1507            EXTRACTION_SEMANTICS_VERSION,
1508        )
1509        .unwrap();
1510
1511        assert_eq!(result.candidates.len(), 2);
1512        assert!(result.completion.skipped.iter().any(|skip| {
1513            skip.reason == SimilarCodeSkipReason::NeighborLimit && skip.count == 4
1514        }));
1515        assert!(
1516            !result
1517                .completion
1518                .skipped
1519                .iter()
1520                .any(|skip| skip.reason == SimilarCodeSkipReason::CandidateLimit)
1521        );
1522    }
1523
1524    #[test]
1525    fn invalid_vectors_fail_closed() {
1526        let mut wrong_dimensions = vector("src/a.ts", 1, 1.0, 0.0);
1527        wrong_dimensions.values.pop();
1528        assert!(matches!(
1529            validate_function_vectors(
1530                &[wrong_dimensions],
1531                DIMENSIONS,
1532                EXTRACTION_SEMANTICS_VERSION
1533            ),
1534            Err(SimilarCodeError::DimensionMismatch { .. })
1535        ));
1536
1537        let mut non_finite = vector("src/a.ts", 1, 1.0, 0.0);
1538        non_finite.values[0] = f32::NAN;
1539        assert!(matches!(
1540            validate_function_vectors(&[non_finite], DIMENSIONS, EXTRACTION_SEMANTICS_VERSION),
1541            Err(SimilarCodeError::NonFiniteVector { .. })
1542        ));
1543
1544        let mut zero = vector("src/a.ts", 1, 1.0, 0.0);
1545        zero.values.fill(0.0);
1546        assert!(matches!(
1547            validate_function_vectors(&[zero], DIMENSIONS, EXTRACTION_SEMANTICS_VERSION),
1548            Err(SimilarCodeError::ZeroMagnitudeVector { .. })
1549        ));
1550    }
1551
1552    #[test]
1553    fn limits_report_partial_work_with_stable_reasons() {
1554        let mut bounded = limits();
1555        bounded.max_functions = 3;
1556        bounded.max_comparisons = 1;
1557        bounded.max_candidates = 1;
1558        bounded.max_neighbors_per_function = 0;
1559        let result = evaluate_similar_code(
1560            &[
1561                vector("src/d.ts", 4, 1.0, 0.0),
1562                vector("src/c.ts", 3, 1.0, 0.0),
1563                vector("src/b.ts", 2, 1.0, 0.0),
1564                vector("src/a.ts", 1, 1.0, 0.0),
1565            ],
1566            0.9,
1567            bounded,
1568            EXTRACTION_SEMANTICS_VERSION,
1569        )
1570        .unwrap();
1571
1572        assert!(result.candidates.is_empty());
1573        assert_eq!(
1574            result.completion.status,
1575            SimilarCodeCompletionStatus::Partial
1576        );
1577        assert_eq!(
1578            result.completion.skipped,
1579            vec![
1580                SimilarCodeSkip {
1581                    reason: SimilarCodeSkipReason::FunctionLimit,
1582                    count: 1,
1583                },
1584                SimilarCodeSkip {
1585                    reason: SimilarCodeSkipReason::ComparisonLimit,
1586                    count: 2,
1587                },
1588                SimilarCodeSkip {
1589                    reason: SimilarCodeSkipReason::NeighborLimit,
1590                    count: 1,
1591                },
1592            ]
1593        );
1594    }
1595
1596    #[test]
1597    fn pathological_candidate_limit_does_not_preallocate_without_pairs() {
1598        let mut untrusted = limits();
1599        untrusted.max_functions = usize::MAX;
1600        untrusted.max_comparisons = usize::MAX;
1601        untrusted.max_candidates = usize::MAX;
1602
1603        let result =
1604            evaluate_similar_code(&[], 0.9, untrusted, EXTRACTION_SEMANTICS_VERSION).unwrap();
1605
1606        assert!(result.candidates.is_empty());
1607        assert_eq!(
1608            result.completion.status,
1609            SimilarCodeCompletionStatus::Complete
1610        );
1611    }
1612
1613    #[test]
1614    fn duplicate_stable_function_identity_fails_independent_of_input_order() {
1615        let first = vector("src/a.ts", 1, 1.0, 0.0);
1616        let mut conflicting = first.clone();
1617        conflicting.source_sha256 = digest(2);
1618        conflicting.values[0] = 0.5;
1619        conflicting.values[1] = 0.5;
1620
1621        for vectors in [
1622            vec![first.clone(), conflicting.clone()],
1623            vec![conflicting, first],
1624        ] {
1625            assert!(matches!(
1626                evaluate_similar_code(&vectors, 0.9, limits(), EXTRACTION_SEMANTICS_VERSION,),
1627                Err(SimilarCodeError::DuplicateFunctionIdentity { .. })
1628            ));
1629        }
1630    }
1631
1632    #[test]
1633    fn vector_cache_is_separate_bounded_and_revision_keyed() {
1634        let key = |hash, revision: &str| VectorCacheKey {
1635            function_source_sha256: digest(hash),
1636            extraction_semantics_version: EXTRACTION_SEMANTICS_VERSION,
1637            model_id: "fixture-model".to_string(),
1638            model_revision: revision.to_string(),
1639            dimensions: 2,
1640            provider_parameter_digest: 7,
1641        };
1642        let mut cache = SimilarCodeVectorCache::new(2 * 2 * size_of::<f32>());
1643
1644        assert!(cache.insert(key(1, "model@a"), vec![1.0, 0.0]));
1645        assert!(!cache.insert(key(9, "model@a"), vec![f32::NAN, 0.0]));
1646        assert!(!cache.insert(key(9, "model@a"), vec![0.0, -0.0]));
1647        assert!(cache.insert(key(2, "model@a"), vec![0.0, 1.0]));
1648        assert!(cache.insert(key(3, "model@b"), vec![0.5, 0.5]));
1649
1650        assert!(cache.get(&key(1, "model@a")).is_none());
1651        assert_eq!(cache.get(&key(2, "model@a")), Some([0.0, 1.0].as_slice()));
1652        assert_eq!(cache.get(&key(3, "model@b")), Some([0.5, 0.5].as_slice()));
1653        assert_eq!(cache.len(), 2);
1654        assert_eq!(cache.used_payload_bytes(), 4 * size_of::<f32>());
1655    }
1656}