code_system_graph_core/
batch.rs

1use std::collections::{BTreeMap, BTreeSet};
2
3use code_system_graph_model::{
4    ArtifactChangeKind, ArtifactFingerprint, CheckoutId, NativePath, RepoId
5};
6use serde::Serialize;
7use serde::de::DeserializeOwned;
8use thiserror::Error;
9
10use crate::{
11    EXTRACTION_CONTRACT_VERSION, ExtractionBudgets, ExtractionLimitExceeded, ExtractionResource, ExtractionTracker, IncrementalPlan
12};
13
14/// Stable identity of one extractor input within a concrete checkout.
15#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
16pub struct ArtifactKey {
17    /// Repository identity shared by linked worktrees.
18    pub repo_id: RepoId,
19    /// Concrete checkout identity.
20    pub checkout_id: CheckoutId,
21    /// Lossless repository-relative artifact path.
22    pub path: NativePath,
23    /// Extractor that owns this artifact.
24    pub extractor: String,
25}
26
27impl From<&ArtifactFingerprint> for ArtifactKey {
28    fn from(fingerprint: &ArtifactFingerprint) -> Self {
29        Self {
30            repo_id: fingerprint.repo_id.clone(),
31            checkout_id: fingerprint.checkout_id.clone(),
32            path: fingerprint.path.clone(),
33            extractor: fingerprint.extractor.clone(),
34        }
35    }
36}
37
38/// Complete transient output owned by one extractor input.
39#[derive(Debug, Clone, PartialEq, Eq)]
40pub struct ExtractorBatch<T> {
41    /// Fingerprint that identifies and invalidates this batch.
42    pub source: ArtifactFingerprint,
43    /// Deterministically ordered extracted outputs.
44    pub outputs: Vec<T>,
45}
46
47impl<T> ExtractorBatch<T> {
48    /// Creates one source-owned extractor batch.
49    #[must_use]
50    pub fn new(source: ArtifactFingerprint, outputs: Vec<T>) -> Self {
51        Self { source, outputs }
52    }
53
54    /// Returns the stable source key for planning and persistence.
55    #[must_use]
56    pub fn key(&self) -> ArtifactKey {
57        ArtifactKey::from(&self.source)
58    }
59}
60
61/// Required action for one source-owned extractor batch.
62#[derive(Debug, Clone, Copy, PartialEq, Eq)]
63pub enum BatchAction {
64    /// Execute the extractor for a newly discovered source.
65    Add,
66    /// Execute the extractor and replace an existing source batch.
67    Replace,
68    /// Reuse the previous batch without extractor work.
69    Reuse,
70    /// Remove the previous batch because its source disappeared.
71    Delete,
72}
73
74/// One deterministic source action derived from an incremental artifact plan.
75#[derive(Debug, Clone, PartialEq, Eq)]
76pub struct PlannedBatch {
77    /// Stable extractor input identity.
78    pub key: ArtifactKey,
79    /// Required action.
80    pub action: BatchAction,
81}
82
83/// Source-level extraction and deletion work for one scan.
84#[derive(Debug, Clone, PartialEq, Eq)]
85pub struct ExtractorBatchPlan {
86    /// Actions sorted by repository, checkout, path, and extractor.
87    pub batches: Vec<PlannedBatch>,
88}
89
90impl ExtractorBatchPlan {
91    /// Returns only source batches requiring extraction or deletion.
92    pub fn changed(&self) -> impl Iterator<Item = &PlannedBatch> {
93        self.batches
94            .iter()
95            .filter(|batch| batch.action != BatchAction::Reuse)
96    }
97}
98
99/// Converts artifact changes into source-owned extractor batch actions.
100#[must_use]
101pub fn plan_extractor_batches(plan: &IncrementalPlan) -> ExtractorBatchPlan {
102    let batches = plan
103        .changes
104        .iter()
105        .map(|change| PlannedBatch {
106            key: ArtifactKey {
107                repo_id: change.repo_id.clone(),
108                checkout_id: change.checkout_id.clone(),
109                path: change.path.clone(),
110                extractor: change.extractor.clone(),
111            },
112            action: match change.kind {
113                ArtifactChangeKind::Added => BatchAction::Add,
114                ArtifactChangeKind::Modified => BatchAction::Replace,
115                ArtifactChangeKind::Deleted => BatchAction::Delete,
116                ArtifactChangeKind::Unchanged => BatchAction::Reuse,
117            },
118        })
119        .collect();
120    ExtractorBatchPlan { batches }
121}
122
123/// Error returned when affected-neighborhood planning lacks a required batch.
124#[derive(Debug, Error, PartialEq, Eq)]
125pub enum BatchPlanError {
126    /// Multiple batches claim the same source key.
127    #[error("duplicate {side} extractor batch for `{extractor}` at `{path}`")]
128    DuplicateBatch {
129        /// Previous or current batch set.
130        side: &'static str,
131        /// Extractor owning the duplicate input.
132        extractor: String,
133        /// Diagnostic path display.
134        path: String,
135    },
136    /// A changed source has no required previous or current output batch.
137    #[error("missing {side} extractor batch for `{extractor}` at `{path}`")]
138    MissingBatch {
139        /// Previous or current batch set.
140        side: &'static str,
141        /// Extractor owning the missing input.
142        extractor: String,
143        /// Diagnostic path display.
144        path: String,
145    },
146    /// A source-owned output payload could not be serialized or decoded.
147    #[error("invalid extractor batch payload: {0}")]
148    InvalidPayload(String),
149    /// One per-invocation extraction resource exceeded its configured maximum.
150    #[error(transparent)]
151    ExtractionLimit(#[from] ExtractionLimitExceeded),
152    /// An output count cannot be represented by the persistence model.
153    #[error("extractor batch output count exceeds the supported range")]
154    OutputCountOverflow,
155    /// Persisted output count does not match the decoded payload.
156    #[error("extractor batch output count mismatch: stored {stored}, decoded {decoded}")]
157    OutputCountMismatch {
158        /// Count stored alongside the payload.
159        stored: u64,
160        /// Number of decoded outputs.
161        decoded: usize,
162    },
163}
164
165/// Encodes one typed extractor batch for atomic snapshot persistence.
166///
167/// # Errors
168///
169/// Returns [`BatchPlanError`] when outputs cannot be encoded or their count exceeds `u64`.
170pub fn store_extractor_batch<T: Serialize>(
171    batch: &ExtractorBatch<T>,
172    tracker: &mut ExtractionTracker,
173    source_was_lossy: bool,
174) -> Result<code_system_graph_model::StoredExtractorBatch, BatchPlanError> {
175    tracker.ensure_observations(u64::try_from(batch.outputs.len()).unwrap_or(u64::MAX))?;
176    let mut writer = tracker.bounded_json_writer();
177    if let Err(error) = serde_json::to_writer(&mut writer, &batch.outputs) {
178        if let Some(limit) = tracker.output_limit_error(&writer) {
179            return Err(limit.into());
180        }
181        return Err(BatchPlanError::InvalidPayload(error.to_string()));
182    }
183    Ok(code_system_graph_model::StoredExtractorBatch {
184        source: batch.source.clone(),
185        extractor_version: EXTRACTION_CONTRACT_VERSION.to_owned(),
186        budget_fingerprint: tracker.budgets().fingerprint(),
187        source_was_lossy,
188        output_count: u64::try_from(batch.outputs.len())
189            .map_err(|_| BatchPlanError::OutputCountOverflow)?,
190        payload: writer.into_inner(),
191    })
192}
193
194/// Decodes one persisted source-owned output batch.
195///
196/// # Errors
197///
198/// Returns [`BatchPlanError`] when the payload schema is invalid or its count is inconsistent.
199pub fn load_extractor_batch<T: DeserializeOwned>(
200    stored: &code_system_graph_model::StoredExtractorBatch,
201) -> Result<ExtractorBatch<T>, BatchPlanError> {
202    load_extractor_batch_with_budgets(stored, &ExtractionBudgets::default())
203}
204
205/// Decodes one persisted source-owned output batch after enforcing effective budgets.
206///
207/// The count and byte checks deliberately precede deserialization so a corrupt or untrusted
208/// persisted batch cannot force allocations beyond the active extraction policy.
209///
210/// # Errors
211///
212/// Returns [`BatchPlanError`] when the payload exceeds `budgets`, its schema is invalid, or its
213/// count is inconsistent.
214pub fn load_extractor_batch_with_budgets<T: DeserializeOwned>(
215    stored: &code_system_graph_model::StoredExtractorBatch,
216    budgets: &ExtractionBudgets,
217) -> Result<ExtractorBatch<T>, BatchPlanError> {
218    if stored.output_count > budgets.max_observations_per_artifact {
219        return Err(ExtractionLimitExceeded {
220            artifact: stored.source.path.display.clone(),
221            extractor: stored.source.extractor.clone(),
222            resource: ExtractionResource::Observations,
223            observed: stored.output_count,
224            maximum: budgets.max_observations_per_artifact,
225        }
226        .into());
227    }
228    let observed = u64::try_from(stored.payload.len()).unwrap_or(u64::MAX);
229    if observed > budgets.max_serialized_output_bytes_per_artifact {
230        return Err(ExtractionLimitExceeded {
231            artifact: stored.source.path.display.clone(),
232            extractor: stored.source.extractor.clone(),
233            resource: ExtractionResource::SerializedOutputBytes,
234            observed,
235            maximum: budgets.max_serialized_output_bytes_per_artifact,
236        }
237        .into());
238    }
239    let outputs: Vec<T> = serde_json::from_slice(&stored.payload)
240        .map_err(|error| BatchPlanError::InvalidPayload(error.to_string()))?;
241    if usize::try_from(stored.output_count).ok() != Some(outputs.len()) {
242        return Err(BatchPlanError::OutputCountMismatch {
243            stored: stored.output_count,
244            decoded: outputs.len(),
245        });
246    }
247    Ok(ExtractorBatch::new(stored.source.clone(), outputs))
248}
249
250/// Computes exact link neighborhoods affected by add, modify, and delete actions.
251///
252/// Modified sources contribute old and new keys, deleted sources contribute old keys, and added
253/// sources contribute new keys. Unchanged sources do not trigger relinking.
254///
255/// # Errors
256///
257/// Returns [`BatchPlanError`] for duplicate source batches or when a changed action lacks the
258/// required previous/current batch.
259pub fn affected_link_keys<T, K>(
260    plan: &ExtractorBatchPlan,
261    previous: &[ExtractorBatch<T>],
262    current: &[ExtractorBatch<T>],
263    link_key: impl Fn(&T) -> K,
264) -> Result<BTreeSet<K>, BatchPlanError>
265where
266    K: Ord,
267{
268    let previous = batch_map(previous, "previous")?;
269    let current = batch_map(current, "current")?;
270    let mut keys = BTreeSet::new();
271    for batch in plan.changed() {
272        match batch.action {
273            BatchAction::Add => {
274                extend_link_keys(
275                    &mut keys,
276                    required_batch(&current, batch, "current")?,
277                    &link_key,
278                );
279            }
280            BatchAction::Replace => {
281                extend_link_keys(
282                    &mut keys,
283                    required_batch(&previous, batch, "previous")?,
284                    &link_key,
285                );
286                extend_link_keys(
287                    &mut keys,
288                    required_batch(&current, batch, "current")?,
289                    &link_key,
290                );
291            }
292            BatchAction::Delete => {
293                extend_link_keys(
294                    &mut keys,
295                    required_batch(&previous, batch, "previous")?,
296                    &link_key,
297                );
298            }
299            BatchAction::Reuse => {}
300        }
301    }
302    Ok(keys)
303}
304
305fn batch_map<'a, T>(
306    batches: &'a [ExtractorBatch<T>],
307    side: &'static str,
308) -> Result<BTreeMap<ArtifactKey, &'a ExtractorBatch<T>>, BatchPlanError> {
309    let mut map = BTreeMap::new();
310    for batch in batches {
311        let key = batch.key();
312        if map.insert(key.clone(), batch).is_some() {
313            return Err(BatchPlanError::DuplicateBatch {
314                side,
315                extractor: key.extractor,
316                path: key.path.display,
317            });
318        }
319    }
320    Ok(map)
321}
322
323fn required_batch<'a, T>(
324    batches: &BTreeMap<ArtifactKey, &'a ExtractorBatch<T>>,
325    planned: &PlannedBatch,
326    side: &'static str,
327) -> Result<&'a ExtractorBatch<T>, BatchPlanError> {
328    batches
329        .get(&planned.key)
330        .copied()
331        .ok_or_else(|| BatchPlanError::MissingBatch {
332            side,
333            extractor: planned.key.extractor.clone(),
334            path: planned.key.path.display.clone(),
335        })
336}
337
338fn extend_link_keys<T, K>(
339    keys: &mut BTreeSet<K>,
340    batch: &ExtractorBatch<T>,
341    link_key: &impl Fn(&T) -> K,
342) where
343    K: Ord,
344{
345    keys.extend(batch.outputs.iter().map(link_key));
346}
347
348#[cfg(test)]
349mod tests {
350    use code_system_graph_model::{
351        ArtifactChange, ArtifactChangeKind, ArtifactFingerprint, CheckoutId, NativePath, NativePathEncoding, RepoId
352    };
353
354    use super::{
355        BatchAction, BatchPlanError, ExtractorBatch, affected_link_keys, load_extractor_batch, load_extractor_batch_with_budgets, plan_extractor_batches, store_extractor_batch
356    };
357    use crate::{ExtractionBudgets, ExtractionLimitExceeded, ExtractionTracker, IncrementalPlan};
358
359    fn tracker() -> ExtractionTracker {
360        ExtractionTracker::new("src/routes.rs", "test", &ExtractionBudgets::default())
361    }
362
363    fn path(value: &str) -> NativePath {
364        NativePath {
365            encoding: NativePathEncoding::Utf8,
366            bytes: value.as_bytes().to_vec(),
367            display: value.to_owned(),
368        }
369    }
370
371    fn change(source: &str, kind: ArtifactChangeKind) -> ArtifactChange {
372        ArtifactChange {
373            repo_id: RepoId::new("repo:api"),
374            checkout_id: CheckoutId::new("checkout:api"),
375            path: path(source),
376            extractor: "code-system-graph.http.openapi".to_owned(),
377            kind,
378        }
379    }
380
381    fn batch(source: &str, hash: &str, outputs: &[&str]) -> ExtractorBatch<String> {
382        ExtractorBatch::new(
383            ArtifactFingerprint {
384                repo_id: RepoId::new("repo:api"),
385                checkout_id: CheckoutId::new("checkout:api"),
386                path: path(source),
387                extractor: "code-system-graph.http.openapi".to_owned(),
388                content_hash: hash.to_owned(),
389                size_bytes: 1,
390            },
391            outputs.iter().map(|output| (*output).to_owned()).collect(),
392        )
393    }
394
395    #[test]
396    fn batch_plan_should_preserve_deterministic_source_actions() {
397        let plan = plan_extractor_batches(&IncrementalPlan {
398            changes: vec![
399                change("added.yaml", ArtifactChangeKind::Added),
400                change("deleted.yaml", ArtifactChangeKind::Deleted),
401                change("same.yaml", ArtifactChangeKind::Unchanged),
402            ],
403        });
404
405        assert_eq!(
406            plan.batches
407                .iter()
408                .map(|batch| batch.action)
409                .collect::<Vec<_>>(),
410            vec![BatchAction::Add, BatchAction::Delete, BatchAction::Reuse]
411        );
412    }
413
414    #[test]
415    fn affected_keys_should_include_old_and_new_modified_neighborhoods() {
416        let plan = plan_extractor_batches(&IncrementalPlan {
417            changes: vec![change("openapi.yaml", ArtifactChangeKind::Modified)],
418        });
419        let result = affected_link_keys(
420            &plan,
421            &[batch("openapi.yaml", "old", &["POST:/v1/orders"])],
422            &[batch("openapi.yaml", "new", &["POST:/v2/orders"])],
423            Clone::clone,
424        );
425
426        assert_eq!(
427            result,
428            Ok(["POST:/v1/orders".to_owned(), "POST:/v2/orders".to_owned()]
429                .into_iter()
430                .collect())
431        );
432    }
433
434    #[test]
435    fn affected_keys_should_require_deleted_previous_batch() {
436        let plan = plan_extractor_batches(&IncrementalPlan {
437            changes: vec![change("deleted.yaml", ArtifactChangeKind::Deleted)],
438        });
439        let previous: Vec<ExtractorBatch<String>> = Vec::new();
440        let current: Vec<ExtractorBatch<String>> = Vec::new();
441        let result = affected_link_keys(&plan, &previous, &current, Clone::clone);
442
443        assert!(matches!(
444            result,
445            Err(BatchPlanError::MissingBatch {
446                side: "previous",
447                ..
448            })
449        ));
450    }
451
452    #[test]
453    fn stored_batch_should_round_trip_without_source_text() {
454        let original = batch("src/routes.rs", "hash", &["GET:/orders", "POST:/orders"]);
455        let result = store_extractor_batch(&original, &mut tracker(), false)
456            .and_then(|stored| load_extractor_batch::<String>(&stored));
457
458        assert_eq!(result, Ok(original));
459    }
460
461    #[test]
462    fn stored_batch_should_reject_inconsistent_output_count() {
463        let original = batch("src/routes.rs", "hash", &["GET:/orders"]);
464        let result =
465            store_extractor_batch(&original, &mut tracker(), false).and_then(|mut stored| {
466                stored.output_count = 2;
467                load_extractor_batch::<String>(&stored)
468            });
469
470        assert!(matches!(
471            result,
472            Err(BatchPlanError::OutputCountMismatch {
473                stored: 2,
474                decoded: 1
475            })
476        ));
477    }
478
479    #[test]
480    fn stored_batch_should_check_payload_limit_before_decoding() {
481        let original = batch("src/routes.rs", "hash", &["GET:/orders"]);
482        let stored = store_extractor_batch(&original, &mut tracker(), false).expect("stored batch");
483        let exact = u64::try_from(stored.payload.len()).expect("payload length");
484        let exact_budgets = ExtractionBudgets {
485            max_serialized_output_bytes_per_artifact: exact,
486            ..ExtractionBudgets::default()
487        };
488        let below_budgets = ExtractionBudgets {
489            max_serialized_output_bytes_per_artifact: exact - 1,
490            ..ExtractionBudgets::default()
491        };
492
493        assert_eq!(
494            load_extractor_batch_with_budgets::<String>(&stored, &exact_budgets),
495            Ok(original)
496        );
497        assert!(matches!(
498            load_extractor_batch_with_budgets::<String>(&stored, &below_budgets),
499            Err(BatchPlanError::ExtractionLimit(ExtractionLimitExceeded {
500                resource: crate::ExtractionResource::SerializedOutputBytes,
501                observed,
502                maximum,
503                ..
504            })) if observed == exact && maximum == exact - 1
505        ));
506    }
507
508    #[test]
509    fn stored_batch_should_check_observation_limit_before_decoding() {
510        let original = batch("src/routes.rs", "hash", &["GET:/orders", "POST:/orders"]);
511        let stored = store_extractor_batch(&original, &mut tracker(), false).expect("stored batch");
512        let budgets = ExtractionBudgets {
513            max_observations_per_artifact: 1,
514            ..ExtractionBudgets::default()
515        };
516
517        assert!(matches!(
518            load_extractor_batch_with_budgets::<String>(&stored, &budgets),
519            Err(BatchPlanError::ExtractionLimit(ExtractionLimitExceeded {
520                resource: crate::ExtractionResource::Observations,
521                observed: 2,
522                maximum: 1,
523                ..
524            }))
525        ));
526    }
527}
code_system_graph_core/batch.rs

code_system_graph_core/
batch.rs
