_diffctx/edges/document/
mod.rs

1use std::path::Path;
2
3use once_cell::sync::Lazy;
4use regex::Regex;
5use rustc_hash::FxHashMap;
6
7use crate::config::weights::EDGE_WEIGHTS;
8use crate::types::{Fragment, FragmentId, FragmentKind};
9
10use super::EdgeDict;
11use super::base::{EdgeBuilder, add_edge};
12
13static HEADING_PREFIX_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"^#+\s*").unwrap());
14static MD_INTERNAL_LINK_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"\[.*?\]\(#([^)]+)\)").unwrap());
15static CITATION_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"\[@([^\]]+)\]").unwrap());
16
17fn slugify(text: &str) -> String {
18    let lower = text.to_lowercase();
19    let mut result = String::with_capacity(lower.len());
20    for ch in lower.chars() {
21        if ch.is_alphanumeric() || ch == '-' {
22            result.push(ch);
23        } else if ch.is_whitespace() || ch == '_' {
24            result.push('-');
25        }
26    }
27    result.trim_matches('-').to_string()
28}
29
30fn is_document_fragment(kind: FragmentKind) -> bool {
31    matches!(kind, FragmentKind::Section | FragmentKind::Chunk)
32}
33
34pub struct DocumentStructureEdgeBuilder;
35
36impl EdgeBuilder for DocumentStructureEdgeBuilder {
37    fn build(&self, fragments: &[Fragment], _repo_root: Option<&Path>) -> EdgeDict {
38        let weight = EDGE_WEIGHTS["doc_structure"].forward;
39        let reverse_factor = EDGE_WEIGHTS["doc_structure"].reverse_factor;
40
41        let mut by_path: FxHashMap<&str, Vec<&Fragment>> = FxHashMap::default();
42        for f in fragments {
43            if is_document_fragment(f.kind) {
44                by_path.entry(f.path()).or_default().push(f);
45            }
46        }
47
48        let mut edges: EdgeDict = FxHashMap::default();
49
50        for (_path, frags) in &mut by_path {
51            frags.sort_by_key(|f| f.start_line());
52            for pair in frags.windows(2) {
53                add_edge(&mut edges, &pair[0].id, &pair[1].id, weight, reverse_factor);
54            }
55        }
56
57        edges
58    }
59}
60
61pub struct AnchorLinkEdgeBuilder;
62
63impl AnchorLinkEdgeBuilder {
64    fn build_anchor_index<'a>(
65        &self,
66        fragments: &'a [Fragment],
67    ) -> FxHashMap<String, &'a FragmentId> {
68        let mut index: FxHashMap<String, &FragmentId> = FxHashMap::default();
69        for f in fragments {
70            if f.kind == FragmentKind::Section {
71                let first_line = f.content.lines().next().unwrap_or("");
72                let heading = HEADING_PREFIX_RE.replace(first_line, "");
73                let slug = slugify(heading.trim());
74                if !slug.is_empty() {
75                    index.entry(slug).or_insert(&f.id);
76                }
77            }
78        }
79        index
80    }
81}
82
83impl EdgeBuilder for AnchorLinkEdgeBuilder {
84    fn build(&self, fragments: &[Fragment], _repo_root: Option<&Path>) -> EdgeDict {
85        let weight = EDGE_WEIGHTS["anchor_link"].forward;
86        let reverse_factor = EDGE_WEIGHTS["anchor_link"].reverse_factor;
87
88        let anchor_index = self.build_anchor_index(fragments);
89        let mut edges: EdgeDict = FxHashMap::default();
90
91        for f in fragments {
92            for cap in MD_INTERNAL_LINK_RE.captures_iter(&f.content) {
93                let target_slug = slugify(&cap[1]);
94                if let Some(target_id) = anchor_index.get(&target_slug) {
95                    if **target_id != f.id {
96                        add_edge(&mut edges, &f.id, target_id, weight, reverse_factor);
97                    }
98                }
99            }
100        }
101
102        edges
103    }
104}
105
106pub struct CitationEdgeBuilder;
107
108impl EdgeBuilder for CitationEdgeBuilder {
109    fn build(&self, fragments: &[Fragment], _repo_root: Option<&Path>) -> EdgeDict {
110        let weight = EDGE_WEIGHTS["citation"].forward;
111
112        let mut citation_to_frags: FxHashMap<String, Vec<&FragmentId>> = FxHashMap::default();
113        for f in fragments {
114            for cap in CITATION_RE.captures_iter(&f.content) {
115                citation_to_frags
116                    .entry(cap[1].to_string())
117                    .or_default()
118                    .push(&f.id);
119            }
120        }
121
122        let mut edges: EdgeDict = FxHashMap::default();
123
124        for (_cit, frag_ids) in &citation_to_frags {
125            if frag_ids.len() < 2 {
126                continue;
127            }
128            let hub = frag_ids[0];
129            for other in &frag_ids[1..] {
130                let key_fwd = (hub.clone(), (*other).clone());
131                let existing_fwd = edges.get(&key_fwd).copied().unwrap_or(0.0);
132                if weight > existing_fwd {
133                    edges.insert(key_fwd, weight);
134                }
135                let key_rev = ((*other).clone(), hub.clone());
136                let existing_rev = edges.get(&key_rev).copied().unwrap_or(0.0);
137                if weight > existing_rev {
138                    edges.insert(key_rev, weight);
139                }
140            }
141        }
142
143        edges
144    }
145}
146
147pub fn get_document_builders() -> Vec<Box<dyn EdgeBuilder>> {
148    vec![
149        Box::new(DocumentStructureEdgeBuilder),
150        Box::new(AnchorLinkEdgeBuilder),
151        Box::new(CitationEdgeBuilder),
152    ]
153}
_diffctx/edges/document/mod.rs

_diffctx/edges/document/
mod.rs
