1use quarb::{AstAdapter, NodeId, Value};
37
38mod lower;
39
40#[derive(Debug, Clone, Copy, PartialEq, Eq)]
42pub enum Lang {
43 Rust,
44 Python,
45 Javascript,
46 C,
47}
48
49impl Lang {
50 pub fn name(self) -> &'static str {
52 match self {
53 Lang::Rust => "rust",
54 Lang::Python => "python",
55 Lang::Javascript => "javascript",
56 Lang::C => "c",
57 }
58 }
59}
60
61pub fn lang_for_ext(ext: &str) -> Option<Lang> {
63 match ext {
64 "rs" => Some(Lang::Rust),
65 "py" => Some(Lang::Python),
66 "js" | "mjs" | "cjs" | "jsx" => Some(Lang::Javascript),
67 "c" | "h" => Some(Lang::C),
68 _ => None,
69 }
70}
71
72pub fn supported(ext: &str) -> bool {
75 lang_for_ext(ext).is_some()
76}
77
78#[derive(Debug, thiserror::Error)]
80pub enum CodeError {
81 #[error("code: {0}")]
82 Io(#[from] std::io::Error),
83 #[error("code: no code-level support for extension {0:?} (rs, py, js, mjs, cjs, jsx, c, h)")]
84 Language(String),
85 #[error(transparent)]
86 Backend(#[from] quarb_tree_sitter::TreeSitterError),
87}
88
89#[derive(Debug)]
96pub struct Decl {
97 pub parent: Option<usize>,
99 pub construct: &'static str,
101 pub name: Option<String>,
103 pub traits: &'static [&'static str],
105 pub kind: String,
107 pub span: (usize, usize),
109 pub lines: (usize, usize),
111 pub signature: Option<String>,
113 pub doc: Option<String>,
115 pub callee: Option<String>,
117 pub n_params: Option<i64>,
119}
120
121struct Node {
122 parent: Option<NodeId>,
123 children: Vec<NodeId>,
124 construct: &'static str,
125 name: Option<String>,
126 traits: &'static [&'static str],
127 kind: String,
128 span: (usize, usize),
129 lines: (usize, usize),
130 signature: Option<String>,
131 doc: Option<String>,
132 callee: Option<String>,
133 n_params: Option<i64>,
134 links: Vec<NodeId>,
136 backlinks: Vec<NodeId>,
138}
139
140pub struct CodeModel {
142 source: String,
143 lang: Lang,
144 nodes: Vec<Node>,
145}
146
147const ALIASED: &[&str] = &[
152 "kind",
153 "construct",
154 "start-line",
155 "end-line",
156 "lang",
157 "n-children",
158 "n-params",
159];
160
161impl CodeModel {
162 pub fn build(source: String, lang: Lang, decls: Vec<Decl>) -> Self {
166 let mut nodes = Vec::with_capacity(decls.len() + 1);
167 nodes.push(Node {
169 parent: None,
170 children: Vec::new(),
171 construct: "",
172 name: None,
173 traits: &[],
174 kind: String::new(),
175 span: (0, source.len()),
176 lines: (1, source.lines().count().max(1)),
177 signature: None,
178 doc: None,
179 callee: None,
180 n_params: None,
181 links: Vec::new(),
182 backlinks: Vec::new(),
183 });
184 for d in decls {
185 let id = NodeId(nodes.len() as u64);
186 let parent = NodeId(d.parent.map_or(0, |p| p as u64 + 1));
187 nodes.push(Node {
188 parent: Some(parent),
189 children: Vec::new(),
190 construct: d.construct,
191 name: d.name,
192 traits: d.traits,
193 kind: d.kind,
194 span: d.span,
195 lines: d.lines,
196 signature: d.signature,
197 doc: d.doc,
198 callee: d.callee,
199 n_params: d.n_params,
200 links: Vec::new(),
201 backlinks: Vec::new(),
202 });
203 nodes[parent.0 as usize].children.push(id);
204 }
205 let mut model = CodeModel {
206 source,
207 lang,
208 nodes,
209 };
210 model.link_definitions();
211 model
212 }
213
214 fn link_definitions(&mut self) {
219 let mut by_name: std::collections::HashMap<&str, Vec<NodeId>> =
220 std::collections::HashMap::new();
221 for (i, n) in self.nodes.iter().enumerate() {
222 if matches!(n.construct, "function" | "type")
223 && let Some(name) = &n.name
224 {
225 by_name
226 .entry(name.as_str())
227 .or_default()
228 .push(NodeId(i as u64));
229 }
230 }
231 let mut links: Vec<(NodeId, Vec<NodeId>)> = Vec::new();
232 for (i, n) in self.nodes.iter().enumerate() {
233 if let Some(callee) = &n.callee
234 && let Some(ident) = trailing_ident(callee)
235 && let Some(targets) = by_name.get(ident)
236 {
237 links.push((NodeId(i as u64), targets.clone()));
238 }
239 }
240 for (call, targets) in links {
241 for t in &targets {
242 self.nodes[t.0 as usize].backlinks.push(call);
243 }
244 self.nodes[call.0 as usize].links = targets;
245 }
246 }
247
248 pub fn parse(text: &str, ext: &str) -> Result<Self, CodeError> {
253 let ext = ext.to_ascii_lowercase();
254 let lang = lang_for_ext(&ext).ok_or_else(|| CodeError::Language(ext.clone()))?;
255 let ts = quarb_tree_sitter::TreeSitterAdapter::parse(text, &ext)?;
256 let decls = lower::lower(&ts, lang);
257 Ok(Self::build(text.to_string(), lang, decls))
258 }
259
260 pub fn open(path: &std::path::Path) -> Result<Self, CodeError> {
262 let ext = path
263 .extension()
264 .and_then(|e| e.to_str())
265 .unwrap_or("")
266 .to_ascii_lowercase();
267 let text = std::fs::read_to_string(path)?;
268 Self::parse(&text, &ext)
269 }
270
271 pub fn locator(&self, node: NodeId) -> String {
275 let mut parts = Vec::new();
276 let mut cur = node;
277 while let Some(parent) = self.nodes[cur.0 as usize].parent {
278 parts.push(self.segment(parent, cur));
279 cur = parent;
280 }
281 parts.reverse();
282 format!("/{}", parts.join("/"))
283 }
284
285 fn label(&self, node: NodeId) -> &str {
286 let n = &self.nodes[node.0 as usize];
287 n.name.as_deref().unwrap_or(n.construct)
288 }
289
290 fn segment(&self, parent: NodeId, child: NodeId) -> String {
291 let label = self.label(child);
292 let same: Vec<NodeId> = self.nodes[parent.0 as usize]
293 .children
294 .iter()
295 .copied()
296 .filter(|&c| self.label(c) == label)
297 .collect();
298 if same.len() > 1 {
299 let pos = same.iter().position(|&c| c == child).unwrap() + 1;
300 format!("{label}[{pos}]")
301 } else {
302 label.to_string()
303 }
304 }
305
306 fn text_of(&self, n: &Node) -> &str {
307 &self.source[n.span.0.min(self.source.len())..n.span.1.min(self.source.len())]
308 }
309
310 pub fn source(&self) -> &str {
314 &self.source
315 }
316
317 pub fn lang(&self) -> Lang {
319 self.lang
320 }
321
322 pub fn ident(&self, node: NodeId) -> Option<&str> {
327 self.nodes[node.0 as usize].name.as_deref()
328 }
329
330 pub fn construct(&self, node: NodeId) -> &str {
332 self.nodes[node.0 as usize].construct
333 }
334
335 pub fn span(&self, node: NodeId) -> (usize, usize) {
337 self.nodes[node.0 as usize].span
338 }
339
340 pub fn line_span(&self, node: NodeId) -> (usize, usize) {
342 self.nodes[node.0 as usize].lines
343 }
344}
345
346fn trailing_ident(callee: &str) -> Option<&str> {
350 let end = callee.trim_end_matches(['!', '?']);
351 let start = end
352 .char_indices()
353 .rev()
354 .take_while(|(_, c)| c.is_alphanumeric() || *c == '_' || *c == '$')
355 .last()
356 .map(|(i, _)| i)?;
357 Some(&end[start..])
358}
359
360impl AstAdapter for CodeModel {
361 fn root(&self) -> NodeId {
362 NodeId(0)
363 }
364
365 fn children(&self, node: NodeId) -> Vec<NodeId> {
366 self.nodes[node.0 as usize].children.clone()
367 }
368
369 fn name(&self, node: NodeId) -> Option<String> {
372 let n = &self.nodes[node.0 as usize];
373 n.parent?;
374 Some(n.name.clone().unwrap_or_else(|| n.construct.to_string()))
375 }
376
377 fn parent(&self, node: NodeId) -> Option<NodeId> {
378 self.nodes[node.0 as usize].parent
379 }
380
381 fn traits(&self, node: NodeId) -> Vec<String> {
382 self.nodes[node.0 as usize]
383 .traits
384 .iter()
385 .map(|t| t.to_string())
386 .collect()
387 }
388
389 fn property(&self, node: NodeId, name: &str) -> Option<Value> {
394 let n = &self.nodes[node.0 as usize];
395 match name {
396 "signature" => n.signature.clone().map(Value::Str),
397 "doc" => n.doc.clone().map(Value::Str),
398 "callee" => n.callee.clone().map(Value::Str),
399 _ => None,
400 }
401 }
402
403 fn default_value(&self, node: NodeId) -> Option<Value> {
405 Some(Value::Str(
406 self.text_of(&self.nodes[node.0 as usize]).to_string(),
407 ))
408 }
409
410 fn metadata(&self, node: NodeId, key: &str) -> Option<Value> {
411 let n = &self.nodes[node.0 as usize];
412 match key {
413 "kind" => (!n.kind.is_empty()).then(|| Value::Str(n.kind.clone())),
416 "construct" => (!n.construct.is_empty()).then(|| Value::Str(n.construct.to_string())),
417 "start-line" => Some(Value::Int(n.lines.0 as i64)),
418 "end-line" => Some(Value::Int(n.lines.1 as i64)),
419 "lang" => Some(Value::Str(self.lang.name().to_string())),
420 "n-children" => Some(Value::Int(n.children.len() as i64)),
421 "n-params" => n.n_params.map(Value::Int),
422 _ => None,
423 }
424 }
425
426 fn aliased_metadata(&self, _node: NodeId) -> &'static [&'static str] {
427 ALIASED
428 }
429
430 fn links(&self, node: NodeId) -> Vec<(String, NodeId)> {
433 self.nodes[node.0 as usize]
434 .links
435 .iter()
436 .map(|&t| ("definition".to_string(), t))
437 .collect()
438 }
439
440 fn backlinks(&self, node: NodeId) -> Vec<(String, NodeId)> {
443 self.nodes[node.0 as usize]
444 .backlinks
445 .iter()
446 .map(|&s| ("definition".to_string(), s))
447 .collect()
448 }
449}