1#[cfg(test)]
3mod tests;
4use super::references::{ScopeReference, document_references};
5use super::{
6 BTreeMap, BTreeSet, DocumentAddress, DocumentEdge, DocumentEdgeKind, DocumentFrontier,
7 DocumentLoader, DocumentScope, DocumentSelector, LoadError, LoadPolicy, LoadSpec,
8 LoadedDocumentScope, MAX_SCOPE_CONTENT_BYTES, ResolvedContent, ResolvedDocumentScope,
9 ScopeLoadError, ScopedDocument, TraversalLimit, UnresolvedDocument, VecDeque, Write,
10 validate_document_scope,
11};
12
13impl DocumentLoader {
14 pub fn resolve_scope(
21 &self,
22 query: &DocumentScope,
23 ) -> Result<LoadedDocumentScope, ScopeLoadError> {
24 validate_document_scope(query)?;
25 let mut resolution = ScopeResolution::new(query);
26 resolution.resolve_roots(self);
27 if resolution.documents.is_empty() {
28 return Err(ScopeLoadError::NoResolvedDocuments {
29 reasons: resolution
30 .graph
31 .unresolved
32 .iter()
33 .map(|failure| failure.reason.clone())
34 .collect(),
35 });
36 }
37 if query.traversal.follow_links {
38 resolution.follow_links(self);
39 }
40 Ok(resolution.finish())
41 }
42
43 fn resolve_selector(
44 &self,
45 selector: &DocumentSelector,
46 policy: LoadPolicy,
47 ) -> Result<ResolvedContent, LoadError> {
48 self.load(
49 LoadSpec::Document {
50 selector: &selector.selector,
51 source: selector.source.as_deref(),
52 manual_section: selector.manual_section.as_deref(),
53 },
54 policy,
55 )
56 }
57}
58
59struct ScopeResolution {
60 graph: ResolvedDocumentScope,
61 documents: Vec<ResolvedContent>,
62 positions: BTreeMap<DocumentAddress, usize>,
63 queue: VecDeque<usize>,
64 content_bytes: u64,
65 failures: ResolutionFailures,
66 unresolved_keys: BTreeSet<UnresolvedKey>,
67}
68
69#[derive(PartialEq, Eq, PartialOrd, Ord)]
70struct ResolutionKey {
71 policy: u8,
72 selector: String,
73 source: Option<String>,
74 manual_section: Option<String>,
75}
76#[derive(PartialEq, Eq, PartialOrd, Ord)]
77struct UnresolvedKey {
78 from: Option<DocumentAddress>,
79 selector: String,
80 source: Option<String>,
81 manual_section: Option<String>,
82 reason: String,
83}
84
85#[derive(Default)]
88struct ResolutionFailures(BTreeMap<ResolutionKey, String>);
89
90impl ResolutionFailures {
91 fn resolve<T>(
92 &mut self,
93 selector: &DocumentSelector,
94 policy: LoadPolicy,
95 load: impl FnOnce() -> Result<T, String>,
96 ) -> Result<T, String> {
97 let key = ResolutionKey {
98 policy: match policy {
99 LoadPolicy::Combined => 0,
100 LoadPolicy::ManualOnly => 1,
101 LoadPolicy::TldrOnly => 2,
102 },
103 selector: selector.selector.clone(),
104 source: selector.source.clone(),
105 manual_section: selector.manual_section.clone(),
106 };
107 if let Some(reason) = self.0.get(&key) {
108 return Err(reason.clone());
109 }
110 let result = load();
111 if let Err(reason) = &result {
112 self.0.insert(key, reason.clone());
113 }
114 result
115 }
116}
117
118impl ScopeResolution {
119 fn new(query: &DocumentScope) -> Self {
120 Self {
121 graph: ResolvedDocumentScope {
122 query: query.clone(),
123 documents: Vec::new(),
124 edges: Vec::new(),
125 frontier: Vec::new(),
126 unresolved: Vec::new(),
127 reference_limits: Vec::new(),
128 },
129 documents: Vec::new(),
130 positions: BTreeMap::new(),
131 queue: VecDeque::new(),
132 content_bytes: 0,
133 failures: ResolutionFailures::default(),
134 unresolved_keys: BTreeSet::new(),
135 }
136 }
137
138 fn resolve_roots(&mut self, resolver: &DocumentLoader) {
139 for (root_index, selector) in self.graph.query.documents.clone().iter().enumerate() {
140 match self.failures.resolve(selector, LoadPolicy::Combined, || {
141 resolver
142 .resolve_selector(selector, LoadPolicy::Combined)
143 .map_err(|error| error.to_string())
144 }) {
145 Ok(bundle) => {
146 self.insert_root(bundle, selector, root_index);
147 }
148 Err(error) => self.record_unresolved(UnresolvedDocument {
149 from: None,
150 selector: selector.clone(),
151 reason: error,
152 }),
153 }
154 }
155 }
156
157 fn insert_root(
158 &mut self,
159 bundle: ResolvedContent,
160 selector: &DocumentSelector,
161 root_index: usize,
162 ) {
163 let Some(address) = bundle.address.clone() else {
164 self.record_unresolved(UnresolvedDocument {
165 from: None,
166 selector: selector.clone(),
167 reason: "selector did not resolve to a registered document".to_owned(),
168 });
169 return;
170 };
171 let root_index = u16::try_from(root_index).unwrap_or(u16::MAX);
172 if let Some(position) = self.positions.get(&address).copied() {
173 let roots = &mut self.graph.documents[position].root_indices;
174 if !roots.contains(&root_index) {
175 roots.push(root_index);
176 }
177 return;
178 }
179 if !self.commit_document(
180 bundle,
181 ScopedDocument {
182 address,
183 depth: 0,
184 root_indices: vec![root_index],
185 reached_from: Vec::new(),
186 },
187 None,
188 ) {
189 self.record_unresolved(UnresolvedDocument {
190 from: None,
191 selector: selector.clone(),
192 reason: format!(
193 "document exceeds the {} MiB aggregate scope content budget",
194 MAX_SCOPE_CONTENT_BYTES / (1024 * 1024)
195 ),
196 });
197 }
198 }
199
200 fn follow_links(&mut self, resolver: &DocumentLoader) {
201 while let Some(position) = self.queue.pop_front() {
202 let depth = self.graph.documents[position].depth;
203 if depth >= self.graph.query.traversal.effective_max_depth() {
204 self.record_depth_frontier(position);
205 continue;
206 }
207 let from = self.graph.documents[position].address.clone();
208 for reference in self.collect_outbound_references(position) {
209 self.follow_reference(resolver, &from, depth, &reference);
210 }
211 }
212 }
213
214 fn record_depth_frontier(&mut self, position: usize) {
215 let from = self.graph.documents[position].address.clone();
216 for reference in self.collect_outbound_references(position) {
217 if let Some(address) = reference.exact_address(&from) {
218 let edge = DocumentEdge {
219 from: from.clone(),
220 to: address,
221 kind: reference.kind,
222 };
223 if self.record_existing_edge(&edge) {
224 continue;
225 }
226 }
227 self.record_frontier(&from, &reference, TraversalLimit::MaxDepth);
228 }
229 }
230
231 fn collect_outbound_references(&mut self, position: usize) -> Vec<ScopeReference> {
232 let collected = document_references(&self.documents[position]);
233 if !collected.report.complete() {
234 self.graph
235 .reference_limits
236 .push(mant_protocol::ScopeReferenceLimit {
237 document: self.graph.documents[position].address.clone(),
238 coverage: mant_protocol::ReferenceCoverage::from_report(collected.report),
239 retention_limit: collected.retention_limit,
240 });
241 }
242 collected.references
243 }
244
245 fn follow_reference(
246 &mut self,
247 resolver: &DocumentLoader,
248 from: &DocumentAddress,
249 depth: u16,
250 reference: &ScopeReference,
251 ) {
252 if let Some(address) = reference.exact_address(from) {
253 let edge = DocumentEdge {
254 from: from.clone(),
255 to: address.clone(),
256 kind: reference.kind,
257 };
258 if self.record_existing_edge(&edge) {
259 return;
260 }
261 if self.at_document_limit() {
262 self.record_frontier(from, reference, TraversalLimit::MaxDocuments);
263 return;
264 }
265 } else if self.at_document_limit() {
266 self.record_frontier(from, reference, TraversalLimit::MaxDocuments);
267 return;
268 }
269
270 let Some(selector) = reference.selector(from) else {
271 self.record_unresolved(UnresolvedDocument {
272 from: Some(from.clone()),
273 selector: reference.fallback_selector(),
274 reason: "relative document link escapes its registered namespace".to_owned(),
275 });
276 return;
277 };
278 let policy = if reference.kind == DocumentEdgeKind::Manual {
279 LoadPolicy::ManualOnly
280 } else {
281 LoadPolicy::Combined
282 };
283 let bundle = match self.failures.resolve(&selector, policy, || {
284 resolver
285 .resolve_selector(&selector, policy)
286 .map_err(|error| error.to_string())
287 }) {
288 Ok(bundle) => bundle,
289 Err(error) => {
290 self.record_unresolved(UnresolvedDocument {
291 from: Some(from.clone()),
292 selector,
293 reason: error,
294 });
295 return;
296 }
297 };
298 let Some(address) = bundle.address.clone() else {
299 self.record_unresolved(UnresolvedDocument {
300 from: Some(from.clone()),
301 selector,
302 reason: "link did not resolve to a registered document".to_owned(),
303 });
304 return;
305 };
306 let edge = DocumentEdge {
307 from: from.clone(),
308 to: address.clone(),
309 kind: reference.kind,
310 };
311 if self.record_existing_edge(&edge) {
312 return;
313 }
314 if self.at_document_limit() {
315 self.record_frontier(from, reference, TraversalLimit::MaxDocuments);
316 return;
317 }
318 if !self.insert_linked(bundle, address, from, depth + 1, edge) {
319 self.record_frontier(from, reference, TraversalLimit::MaxContentBytes);
320 }
321 }
322
323 fn record_existing_edge(&mut self, edge: &DocumentEdge) -> bool {
324 let Some(position) = self.positions.get(&edge.to).copied() else {
325 return false;
326 };
327 if !self.graph.edges.contains(edge) {
328 self.graph.edges.push(edge.clone());
329 }
330 if edge.to != edge.from
331 && !self.graph.documents[position]
332 .reached_from
333 .contains(&edge.from)
334 {
335 self.graph.documents[position]
336 .reached_from
337 .push(edge.from.clone());
338 }
339 true
340 }
341
342 fn record_unresolved(&mut self, failure: UnresolvedDocument) {
343 let key = UnresolvedKey {
344 from: failure.from.clone(),
345 selector: failure.selector.selector.clone(),
346 source: failure.selector.source.clone(),
347 manual_section: failure.selector.manual_section.clone(),
348 reason: failure.reason.clone(),
349 };
350 if self.unresolved_keys.insert(key) {
351 self.graph.unresolved.push(failure);
352 }
353 }
354
355 fn insert_linked(
356 &mut self,
357 bundle: ResolvedContent,
358 address: DocumentAddress,
359 from: &DocumentAddress,
360 depth: u16,
361 edge: DocumentEdge,
362 ) -> bool {
363 self.commit_document(
364 bundle,
365 ScopedDocument {
366 address,
367 depth,
368 root_indices: Vec::new(),
369 reached_from: vec![from.clone()],
370 },
371 Some(edge),
372 )
373 }
374
375 fn commit_document(
381 &mut self,
382 bundle: ResolvedContent,
383 source: ScopedDocument,
384 edge: Option<DocumentEdge>,
385 ) -> bool {
386 debug_assert_eq!(bundle.address.as_ref(), Some(&source.address));
387 debug_assert!(!self.positions.contains_key(&source.address));
388 let bytes = normalized_content_bytes(&bundle);
389 let Some(total) = self.content_bytes.checked_add(bytes) else {
390 return false;
391 };
392 if total > MAX_SCOPE_CONTENT_BYTES {
393 return false;
394 }
395 let position = self.documents.len();
396 self.positions.insert(source.address.clone(), position);
397 self.documents.push(bundle);
398 self.graph.documents.push(source);
399 self.queue.push_back(position);
400 if let Some(edge) = edge
401 && !self.graph.edges.contains(&edge)
402 {
403 self.graph.edges.push(edge);
404 }
405 self.content_bytes = total;
406 true
407 }
408
409 fn at_document_limit(&self) -> bool {
410 u32::try_from(self.documents.len()).unwrap_or(u32::MAX)
411 >= self.graph.query.traversal.effective_max_documents()
412 }
413
414 fn record_frontier(
415 &mut self,
416 from: &DocumentAddress,
417 reference: &ScopeReference,
418 limit: TraversalLimit,
419 ) {
420 let frontier = DocumentFrontier {
421 from: from.clone(),
422 target: reference
423 .selector(from)
424 .unwrap_or_else(|| reference.fallback_selector()),
425 kind: reference.kind,
426 limit,
427 };
428 if !self.graph.frontier.contains(&frontier) {
429 self.graph.frontier.push(frontier);
430 }
431 }
432
433 fn finish(self) -> LoadedDocumentScope {
434 LoadedDocumentScope {
435 scope: self.graph,
436 documents: self.documents,
437 }
438 }
439}
440
441fn normalized_content_bytes(content: &ResolvedContent) -> u64 {
446 let mut counter = ByteCounter::default();
447 if let Some(document) = &content.document {
448 serde_json::to_writer(&mut counter, document)
449 .expect("writing normalized document bytes to a counter cannot fail");
450 }
451 if let Some(tldr) = &content.tldr {
452 serde_json::to_writer(&mut counter, tldr)
453 .expect("writing normalized tldr bytes to a counter cannot fail");
454 }
455 counter.0
456}
457
458#[derive(Default)]
459struct ByteCounter(u64);
460
461impl Write for ByteCounter {
462 fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
463 self.0 = self
464 .0
465 .saturating_add(u64::try_from(bytes.len()).unwrap_or(u64::MAX));
466 Ok(bytes.len())
467 }
468
469 fn flush(&mut self) -> std::io::Result<()> {
470 Ok(())
471 }
472}