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(
85 "code: the code model has no reading for extension {0:?} (rs, py, js, mjs, cjs, jsx, c, h)"
86 )]
87 Language(String),
88 #[error(transparent)]
89 Backend(#[from] quarb_tree_sitter::TreeSitterError),
90 #[error("code: {0}")]
91 Encoding(String),
92}
93
94#[derive(Debug)]
101pub struct Decl {
102 pub parent: Option<usize>,
104 pub construct: &'static str,
106 pub name: Option<String>,
108 pub traits: &'static [&'static str],
110 pub kind: String,
112 pub span: (usize, usize),
114 pub lines: (usize, usize),
116 pub signature: Option<String>,
118 pub doc: Option<String>,
120 pub callee: Option<String>,
122 pub n_params: Option<i64>,
124}
125
126struct Node {
127 parent: Option<NodeId>,
128 children: Vec<NodeId>,
129 construct: &'static str,
130 name: Option<String>,
131 traits: &'static [&'static str],
132 kind: String,
133 span: (usize, usize),
134 lines: (usize, usize),
135 signature: Option<String>,
136 doc: Option<String>,
137 callee: Option<String>,
138 n_params: Option<i64>,
139}
140
141pub struct CodeModel {
143 source: String,
144 lang: Lang,
145 nodes: Vec<Node>,
146 defs_by_ident: HashMap<String, Vec<NodeId>>,
152 calls_by_ident: HashMap<String, Vec<NodeId>>,
155}
156
157const ALIASED: &[&str] = &[
162 "kind",
163 "construct",
164 "start-line",
165 "end-line",
166 "lang",
167 "n-params",
168];
169
170impl CodeModel {
171 pub fn build(source: String, lang: Lang, decls: Vec<Decl>) -> Self {
175 let mut nodes = Vec::with_capacity(decls.len() + 1);
176 nodes.push(Node {
178 parent: None,
179 children: Vec::new(),
180 construct: "",
181 name: None,
182 traits: &[],
183 kind: String::new(),
184 span: (0, source.len()),
185 lines: (1, source.lines().count().max(1)),
186 signature: None,
187 doc: None,
188 callee: None,
189 n_params: None,
190 });
191 for d in decls {
192 let id = NodeId(nodes.len() as u64);
193 let parent = NodeId(d.parent.map_or(0, |p| p as u64 + 1));
194 nodes.push(Node {
195 parent: Some(parent),
196 children: Vec::new(),
197 construct: d.construct,
198 name: d.name,
199 traits: d.traits,
200 kind: d.kind,
201 span: d.span,
202 lines: d.lines,
203 signature: d.signature,
204 doc: d.doc,
205 callee: d.callee,
206 n_params: d.n_params,
207 });
208 nodes[parent.0 as usize].children.push(id);
209 }
210 let mut model = CodeModel {
211 source,
212 lang,
213 nodes,
214 defs_by_ident: HashMap::new(),
215 calls_by_ident: HashMap::new(),
216 };
217 model.link_definitions();
218 model
219 }
220
221 fn link_definitions(&mut self) {
226 for (i, n) in self.nodes.iter().enumerate() {
227 if matches!(n.construct, "function" | "type")
228 && let Some(name) = &n.name
229 {
230 self.defs_by_ident
231 .entry(name.clone())
232 .or_default()
233 .push(NodeId(i as u64));
234 }
235 if let Some(callee) = &n.callee
236 && let Some(ident) = trailing_ident(callee)
237 {
238 self.calls_by_ident
239 .entry(ident.to_string())
240 .or_default()
241 .push(NodeId(i as u64));
242 }
243 }
244 self.calls_by_ident
246 .retain(|ident, _| self.defs_by_ident.contains_key(ident));
247 }
248
249 pub fn parse(text: &str, ext: &str) -> Result<Self, CodeError> {
254 let ext = ext.to_ascii_lowercase();
255 let lang = lang_for_ext(&ext).ok_or_else(|| CodeError::Language(ext.clone()))?;
256 let ts = quarb_tree_sitter::TreeSitterAdapter::parse(text, &ext)?;
257 let decls = lower::lower(&ts, lang);
258 Ok(Self::build(text.to_string(), lang, decls))
259 }
260
261 pub fn open(path: &std::path::Path) -> Result<Self, CodeError> {
263 Self::open_with_decoding(path, &quarb::source::SourceDecoding::default())
264 }
265
266 pub fn open_with_decoding(
270 path: &std::path::Path,
271 decoding: &quarb::source::SourceDecoding,
272 ) -> Result<Self, CodeError> {
273 let ext = path
274 .extension()
275 .and_then(|e| e.to_str())
276 .unwrap_or("")
277 .to_ascii_lowercase();
278 let bytes = std::fs::read(path)?;
279 let name = path.display().to_string();
280 let text = decoding
281 .decode(&name, &bytes)
282 .map_err(|e| CodeError::Encoding(format!("{name}: {e}")))?
283 .text;
284 Self::parse(&text, &ext)
285 }
286
287 pub fn locator(&self, node: NodeId) -> String {
291 let mut parts = Vec::new();
292 let mut cur = node;
293 while let Some(parent) = self.nodes[cur.0 as usize].parent {
294 parts.push(self.segment(parent, cur));
295 cur = parent;
296 }
297 parts.reverse();
298 format!("/{}", parts.join("/"))
299 }
300
301 fn label(&self, node: NodeId) -> &str {
302 let n = &self.nodes[node.0 as usize];
303 n.name.as_deref().unwrap_or(n.construct)
304 }
305
306 fn segment(&self, parent: NodeId, child: NodeId) -> String {
307 let label = self.label(child);
308 let same: Vec<NodeId> = self.nodes[parent.0 as usize]
309 .children
310 .iter()
311 .copied()
312 .filter(|&c| self.label(c) == label)
313 .collect();
314 if same.len() > 1 {
315 let pos = same.iter().position(|&c| c == child).unwrap() + 1;
316 format!("{label}[{pos}]")
317 } else {
318 label.to_string()
319 }
320 }
321
322 fn text_of(&self, n: &Node) -> &str {
323 &self.source[n.span.0.min(self.source.len())..n.span.1.min(self.source.len())]
324 }
325
326 pub fn source(&self) -> &str {
330 &self.source
331 }
332
333 pub fn lang(&self) -> Lang {
335 self.lang
336 }
337
338 pub fn ident(&self, node: NodeId) -> Option<&str> {
343 self.nodes[node.0 as usize].name.as_deref()
344 }
345
346 pub fn construct(&self, node: NodeId) -> &str {
348 self.nodes[node.0 as usize].construct
349 }
350
351 pub fn span(&self, node: NodeId) -> (usize, usize) {
353 self.nodes[node.0 as usize].span
354 }
355
356 pub fn line_span(&self, node: NodeId) -> (usize, usize) {
358 self.nodes[node.0 as usize].lines
359 }
360}
361
362fn qualifier(callee: &str) -> Option<&str> {
370 let end = callee.trim_end_matches(['!', '?']);
371 let last = trailing_ident(end)?;
372 let head = end[..end.len() - last.len()]
373 .trim_end_matches("::")
374 .trim_end_matches('.');
375 if head.len() == end.len() - last.len() {
376 return None;
377 }
378 let q = trailing_ident(head)?;
379 (!matches!(q, "self" | "this" | "super" | "crate" | "Self")).then_some(q)
380}
381
382fn trailing_ident(callee: &str) -> Option<&str> {
383 let end = callee.trim_end_matches(['!', '?']);
384 let start = end
385 .char_indices()
386 .rev()
387 .take_while(|(_, c)| c.is_alphanumeric() || *c == '_' || *c == '$')
388 .last()
389 .map(|(i, _)| i)?;
390 Some(&end[start..])
391}
392
393impl AstAdapter for CodeModel {
394 fn resolves_references(&self) -> bool {
396 false
397 }
398
399 fn root(&self) -> NodeId {
400 NodeId(0)
401 }
402
403 fn walks_for_references(&self) -> bool {
406 true
407 }
408
409 fn children(&self, node: NodeId) -> Vec<NodeId> {
410 self.nodes[node.0 as usize].children.clone()
411 }
412
413 fn name(&self, node: NodeId) -> Option<String> {
416 let n = &self.nodes[node.0 as usize];
417 n.parent?;
418 Some(n.name.clone().unwrap_or_else(|| n.construct.to_string()))
419 }
420
421 fn parent(&self, node: NodeId) -> Option<NodeId> {
422 self.nodes[node.0 as usize].parent
423 }
424
425 fn traits(&self, node: NodeId) -> Vec<String> {
426 self.nodes[node.0 as usize]
427 .traits
428 .iter()
429 .map(|t| t.to_string())
430 .collect()
431 }
432
433 fn property(&self, node: NodeId, name: &str) -> Option<Value> {
438 let n = &self.nodes[node.0 as usize];
439 match name {
440 "signature" => n.signature.clone().map(Value::Str),
441 "doc" => n.doc.clone().map(Value::Str),
442 "callee" => n.callee.clone().map(Value::Str),
443 _ => None,
444 }
445 }
446
447 fn default_value(&self, node: NodeId) -> Option<Value> {
449 Some(Value::Str(
450 self.text_of(&self.nodes[node.0 as usize]).to_string(),
451 ))
452 }
453
454 fn metadata_catalogue(&self) -> &'static [&'static str] {
459
460 &["kind", "construct", "start-line", "end-line", "lang", "n-params", "function"]
461
462 }
463
464
465 fn metadata(&self, node: NodeId, key: &str) -> Option<Value> {
466 let n = &self.nodes[node.0 as usize];
467 match key {
468 "kind" => (!n.kind.is_empty()).then(|| Value::Str(n.kind.clone())),
471 "construct" => (!n.construct.is_empty()).then(|| Value::Str(n.construct.to_string())),
472 "start-line" => Some(Value::Int(n.lines.0 as i64)),
473 "end-line" => Some(Value::Int(n.lines.1 as i64)),
474 "lang" => Some(Value::Str(self.lang.name().to_string())),
475 "n-params" => n.n_params.map(Value::Int),
476 _ => None,
477 }
478 }
479
480 fn aliased_metadata(&self, _node: NodeId) -> &'static [&'static str] {
481 ALIASED
482 }
483
484 fn links(&self, node: NodeId) -> Vec<(String, NodeId)> {
487 let n = &self.nodes[node.0 as usize];
488 let Some(callee) = n.callee.as_deref() else {
489 return Vec::new();
490 };
491 let Some(ident) = trailing_ident(callee) else {
492 return Vec::new();
493 };
494 let Some(targets) = self.defs_by_ident.get(ident) else {
495 return Vec::new();
496 };
497 let qualified: Vec<NodeId> = match qualifier(callee) {
503 Some(q) => targets
504 .iter()
505 .copied()
506 .filter(|&t| {
507 self.nodes[t.0 as usize]
508 .parent
509 .and_then(|p| self.nodes[p.0 as usize].name.as_deref())
510 == Some(q)
511 })
512 .collect(),
513 None => Vec::new(),
514 };
515 let chosen: &[NodeId] = if qualified.is_empty() {
516 targets
517 } else {
518 &qualified
519 };
520 chosen
521 .iter()
522 .map(|&t| ("definition".to_string(), t))
523 .collect()
524 }
525
526 fn backlinks(&self, node: NodeId) -> Vec<(String, NodeId)> {
529 let n = &self.nodes[node.0 as usize];
530 if !matches!(n.construct, "function" | "type") {
531 return Vec::new();
532 }
533 n.name
534 .as_deref()
535 .and_then(|name| self.calls_by_ident.get(name))
536 .map(|sources| {
537 sources
538 .iter()
539 .map(|&s| ("definition".to_string(), s))
540 .collect()
541 })
542 .unwrap_or_default()
543 }
544}