1use quarb::{AstAdapter, NodeId, Value};
37use std::collections::HashMap;
38
39mod lower;
40
41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
43pub enum Lang {
44 Rust,
45 Python,
46 Javascript,
47 C,
48}
49
50impl Lang {
51 pub fn name(self) -> &'static str {
53 match self {
54 Lang::Rust => "rust",
55 Lang::Python => "python",
56 Lang::Javascript => "javascript",
57 Lang::C => "c",
58 }
59 }
60}
61
62pub fn lang_for_ext(ext: &str) -> Option<Lang> {
64 match ext {
65 "rs" => Some(Lang::Rust),
66 "py" => Some(Lang::Python),
67 "js" | "mjs" | "cjs" | "jsx" => Some(Lang::Javascript),
68 "c" | "h" => Some(Lang::C),
69 _ => None,
70 }
71}
72
73pub fn supported(ext: &str) -> bool {
76 lang_for_ext(ext).is_some()
77}
78
79#[derive(Debug, thiserror::Error)]
81pub enum CodeError {
82 #[error("code: {0}")]
83 Io(#[from] std::io::Error),
84 #[error("code: no code-level support for extension {0:?} (rs, py, js, mjs, cjs, jsx, c, h)")]
85 Language(String),
86 #[error(transparent)]
87 Backend(#[from] quarb_tree_sitter::TreeSitterError),
88 #[error("code: {0}")]
89 Encoding(String),
90}
91
92#[derive(Debug)]
99pub struct Decl {
100 pub parent: Option<usize>,
102 pub construct: &'static str,
104 pub name: Option<String>,
106 pub traits: &'static [&'static str],
108 pub kind: String,
110 pub span: (usize, usize),
112 pub lines: (usize, usize),
114 pub signature: Option<String>,
116 pub doc: Option<String>,
118 pub callee: Option<String>,
120 pub n_params: Option<i64>,
122}
123
124struct Node {
125 parent: Option<NodeId>,
126 children: Vec<NodeId>,
127 construct: &'static str,
128 name: Option<String>,
129 traits: &'static [&'static str],
130 kind: String,
131 span: (usize, usize),
132 lines: (usize, usize),
133 signature: Option<String>,
134 doc: Option<String>,
135 callee: Option<String>,
136 n_params: Option<i64>,
137}
138
139pub struct CodeModel {
141 source: String,
142 lang: Lang,
143 nodes: Vec<Node>,
144 defs_by_ident: HashMap<String, Vec<NodeId>>,
150 calls_by_ident: HashMap<String, Vec<NodeId>>,
153}
154
155const ALIASED: &[&str] = &[
160 "kind",
161 "construct",
162 "start-line",
163 "end-line",
164 "lang",
165 "n-children",
166 "n-params",
167];
168
169impl CodeModel {
170 pub fn build(source: String, lang: Lang, decls: Vec<Decl>) -> Self {
174 let mut nodes = Vec::with_capacity(decls.len() + 1);
175 nodes.push(Node {
177 parent: None,
178 children: Vec::new(),
179 construct: "",
180 name: None,
181 traits: &[],
182 kind: String::new(),
183 span: (0, source.len()),
184 lines: (1, source.lines().count().max(1)),
185 signature: None,
186 doc: None,
187 callee: None,
188 n_params: None,
189 });
190 for d in decls {
191 let id = NodeId(nodes.len() as u64);
192 let parent = NodeId(d.parent.map_or(0, |p| p as u64 + 1));
193 nodes.push(Node {
194 parent: Some(parent),
195 children: Vec::new(),
196 construct: d.construct,
197 name: d.name,
198 traits: d.traits,
199 kind: d.kind,
200 span: d.span,
201 lines: d.lines,
202 signature: d.signature,
203 doc: d.doc,
204 callee: d.callee,
205 n_params: d.n_params,
206 });
207 nodes[parent.0 as usize].children.push(id);
208 }
209 let mut model = CodeModel {
210 source,
211 lang,
212 nodes,
213 defs_by_ident: HashMap::new(),
214 calls_by_ident: HashMap::new(),
215 };
216 model.link_definitions();
217 model
218 }
219
220 fn link_definitions(&mut self) {
225 for (i, n) in self.nodes.iter().enumerate() {
226 if matches!(n.construct, "function" | "type")
227 && let Some(name) = &n.name
228 {
229 self.defs_by_ident
230 .entry(name.clone())
231 .or_default()
232 .push(NodeId(i as u64));
233 }
234 if let Some(callee) = &n.callee
235 && let Some(ident) = trailing_ident(callee)
236 {
237 self.calls_by_ident
238 .entry(ident.to_string())
239 .or_default()
240 .push(NodeId(i as u64));
241 }
242 }
243 self.calls_by_ident
245 .retain(|ident, _| self.defs_by_ident.contains_key(ident));
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 Self::open_with_decoding(path, &quarb::source::SourceDecoding::default())
263 }
264
265 pub fn open_with_decoding(
269 path: &std::path::Path,
270 decoding: &quarb::source::SourceDecoding,
271 ) -> Result<Self, CodeError> {
272 let ext = path
273 .extension()
274 .and_then(|e| e.to_str())
275 .unwrap_or("")
276 .to_ascii_lowercase();
277 let bytes = std::fs::read(path)?;
278 let name = path.display().to_string();
279 let text = decoding
280 .decode(&name, &bytes)
281 .map_err(|e| CodeError::Encoding(format!("{name}: {e}")))?
282 .text;
283 Self::parse(&text, &ext)
284 }
285
286 pub fn locator(&self, node: NodeId) -> String {
290 let mut parts = Vec::new();
291 let mut cur = node;
292 while let Some(parent) = self.nodes[cur.0 as usize].parent {
293 parts.push(self.segment(parent, cur));
294 cur = parent;
295 }
296 parts.reverse();
297 format!("/{}", parts.join("/"))
298 }
299
300 fn label(&self, node: NodeId) -> &str {
301 let n = &self.nodes[node.0 as usize];
302 n.name.as_deref().unwrap_or(n.construct)
303 }
304
305 fn segment(&self, parent: NodeId, child: NodeId) -> String {
306 let label = self.label(child);
307 let same: Vec<NodeId> = self.nodes[parent.0 as usize]
308 .children
309 .iter()
310 .copied()
311 .filter(|&c| self.label(c) == label)
312 .collect();
313 if same.len() > 1 {
314 let pos = same.iter().position(|&c| c == child).unwrap() + 1;
315 format!("{label}[{pos}]")
316 } else {
317 label.to_string()
318 }
319 }
320
321 fn text_of(&self, n: &Node) -> &str {
322 &self.source[n.span.0.min(self.source.len())..n.span.1.min(self.source.len())]
323 }
324
325 pub fn source(&self) -> &str {
329 &self.source
330 }
331
332 pub fn lang(&self) -> Lang {
334 self.lang
335 }
336
337 pub fn ident(&self, node: NodeId) -> Option<&str> {
342 self.nodes[node.0 as usize].name.as_deref()
343 }
344
345 pub fn construct(&self, node: NodeId) -> &str {
347 self.nodes[node.0 as usize].construct
348 }
349
350 pub fn span(&self, node: NodeId) -> (usize, usize) {
352 self.nodes[node.0 as usize].span
353 }
354
355 pub fn line_span(&self, node: NodeId) -> (usize, usize) {
357 self.nodes[node.0 as usize].lines
358 }
359}
360
361fn qualifier(callee: &str) -> Option<&str> {
369 let end = callee.trim_end_matches(['!', '?']);
370 let last = trailing_ident(end)?;
371 let head = end[..end.len() - last.len()]
372 .trim_end_matches("::")
373 .trim_end_matches('.');
374 if head.len() == end.len() - last.len() {
375 return None;
376 }
377 let q = trailing_ident(head)?;
378 (!matches!(q, "self" | "this" | "super" | "crate" | "Self")).then_some(q)
379}
380
381fn trailing_ident(callee: &str) -> Option<&str> {
382 let end = callee.trim_end_matches(['!', '?']);
383 let start = end
384 .char_indices()
385 .rev()
386 .take_while(|(_, c)| c.is_alphanumeric() || *c == '_' || *c == '$')
387 .last()
388 .map(|(i, _)| i)?;
389 Some(&end[start..])
390}
391
392impl AstAdapter for CodeModel {
393 fn resolves_references(&self) -> bool {
395 false
396 }
397
398 fn root(&self) -> NodeId {
399 NodeId(0)
400 }
401
402 fn walks_for_references(&self) -> bool {
405 true
406 }
407
408 fn children(&self, node: NodeId) -> Vec<NodeId> {
409 self.nodes[node.0 as usize].children.clone()
410 }
411
412 fn name(&self, node: NodeId) -> Option<String> {
415 let n = &self.nodes[node.0 as usize];
416 n.parent?;
417 Some(n.name.clone().unwrap_or_else(|| n.construct.to_string()))
418 }
419
420 fn parent(&self, node: NodeId) -> Option<NodeId> {
421 self.nodes[node.0 as usize].parent
422 }
423
424 fn traits(&self, node: NodeId) -> Vec<String> {
425 self.nodes[node.0 as usize]
426 .traits
427 .iter()
428 .map(|t| t.to_string())
429 .collect()
430 }
431
432 fn property(&self, node: NodeId, name: &str) -> Option<Value> {
437 let n = &self.nodes[node.0 as usize];
438 match name {
439 "signature" => n.signature.clone().map(Value::Str),
440 "doc" => n.doc.clone().map(Value::Str),
441 "callee" => n.callee.clone().map(Value::Str),
442 _ => None,
443 }
444 }
445
446 fn default_value(&self, node: NodeId) -> Option<Value> {
448 Some(Value::Str(
449 self.text_of(&self.nodes[node.0 as usize]).to_string(),
450 ))
451 }
452
453 fn metadata(&self, node: NodeId, key: &str) -> Option<Value> {
454 let n = &self.nodes[node.0 as usize];
455 match key {
456 "kind" => (!n.kind.is_empty()).then(|| Value::Str(n.kind.clone())),
459 "construct" => (!n.construct.is_empty()).then(|| Value::Str(n.construct.to_string())),
460 "start-line" => Some(Value::Int(n.lines.0 as i64)),
461 "end-line" => Some(Value::Int(n.lines.1 as i64)),
462 "lang" => Some(Value::Str(self.lang.name().to_string())),
463 "n-children" => Some(Value::Int(n.children.len() as i64)),
464 "n-params" => n.n_params.map(Value::Int),
465 _ => None,
466 }
467 }
468
469 fn aliased_metadata(&self, _node: NodeId) -> &'static [&'static str] {
470 ALIASED
471 }
472
473 fn links(&self, node: NodeId) -> Vec<(String, NodeId)> {
476 let n = &self.nodes[node.0 as usize];
477 let Some(callee) = n.callee.as_deref() else {
478 return Vec::new();
479 };
480 let Some(ident) = trailing_ident(callee) else {
481 return Vec::new();
482 };
483 let Some(targets) = self.defs_by_ident.get(ident) else {
484 return Vec::new();
485 };
486 let qualified: Vec<NodeId> = match qualifier(callee) {
492 Some(q) => targets
493 .iter()
494 .copied()
495 .filter(|&t| {
496 self.nodes[t.0 as usize]
497 .parent
498 .and_then(|p| self.nodes[p.0 as usize].name.as_deref())
499 == Some(q)
500 })
501 .collect(),
502 None => Vec::new(),
503 };
504 let chosen: &[NodeId] = if qualified.is_empty() {
505 targets
506 } else {
507 &qualified
508 };
509 chosen
510 .iter()
511 .map(|&t| ("definition".to_string(), t))
512 .collect()
513 }
514
515 fn backlinks(&self, node: NodeId) -> Vec<(String, NodeId)> {
518 let n = &self.nodes[node.0 as usize];
519 if !matches!(n.construct, "function" | "type") {
520 return Vec::new();
521 }
522 n.name
523 .as_deref()
524 .and_then(|name| self.calls_by_ident.get(name))
525 .map(|sources| {
526 sources
527 .iter()
528 .map(|&s| ("definition".to_string(), s))
529 .collect()
530 })
531 .unwrap_or_default()
532 }
533}