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.entry(name.as_str()).or_default().push(NodeId(i as u64));
226 }
227 }
228 let mut links: Vec<(NodeId, Vec<NodeId>)> = Vec::new();
229 for (i, n) in self.nodes.iter().enumerate() {
230 if let Some(callee) = &n.callee
231 && let Some(ident) = trailing_ident(callee)
232 && let Some(targets) = by_name.get(ident)
233 {
234 links.push((NodeId(i as u64), targets.clone()));
235 }
236 }
237 for (call, targets) in links {
238 for t in &targets {
239 self.nodes[t.0 as usize].backlinks.push(call);
240 }
241 self.nodes[call.0 as usize].links = targets;
242 }
243 }
244
245 pub fn parse(text: &str, ext: &str) -> Result<Self, CodeError> {
250 let ext = ext.to_ascii_lowercase();
251 let lang = lang_for_ext(&ext).ok_or_else(|| CodeError::Language(ext.clone()))?;
252 let ts = quarb_tree_sitter::TreeSitterAdapter::parse(text, &ext)?;
253 let decls = lower::lower(&ts, lang);
254 Ok(Self::build(text.to_string(), lang, decls))
255 }
256
257 pub fn open(path: &std::path::Path) -> Result<Self, CodeError> {
259 let ext = path
260 .extension()
261 .and_then(|e| e.to_str())
262 .unwrap_or("")
263 .to_ascii_lowercase();
264 let text = std::fs::read_to_string(path)?;
265 Self::parse(&text, &ext)
266 }
267
268 pub fn locator(&self, node: NodeId) -> String {
272 let mut parts = Vec::new();
273 let mut cur = node;
274 while let Some(parent) = self.nodes[cur.0 as usize].parent {
275 parts.push(self.segment(parent, cur));
276 cur = parent;
277 }
278 parts.reverse();
279 format!("/{}", parts.join("/"))
280 }
281
282 fn label(&self, node: NodeId) -> &str {
283 let n = &self.nodes[node.0 as usize];
284 n.name.as_deref().unwrap_or(n.construct)
285 }
286
287 fn segment(&self, parent: NodeId, child: NodeId) -> String {
288 let label = self.label(child);
289 let same: Vec<NodeId> = self.nodes[parent.0 as usize]
290 .children
291 .iter()
292 .copied()
293 .filter(|&c| self.label(c) == label)
294 .collect();
295 if same.len() > 1 {
296 let pos = same.iter().position(|&c| c == child).unwrap() + 1;
297 format!("{label}[{pos}]")
298 } else {
299 label.to_string()
300 }
301 }
302
303 fn text_of(&self, n: &Node) -> &str {
304 &self.source[n.span.0.min(self.source.len())..n.span.1.min(self.source.len())]
305 }
306
307 pub fn source(&self) -> &str {
311 &self.source
312 }
313
314 pub fn lang(&self) -> Lang {
316 self.lang
317 }
318
319 pub fn ident(&self, node: NodeId) -> Option<&str> {
324 self.nodes[node.0 as usize].name.as_deref()
325 }
326
327 pub fn construct(&self, node: NodeId) -> &str {
329 self.nodes[node.0 as usize].construct
330 }
331
332 pub fn span(&self, node: NodeId) -> (usize, usize) {
334 self.nodes[node.0 as usize].span
335 }
336
337 pub fn line_span(&self, node: NodeId) -> (usize, usize) {
339 self.nodes[node.0 as usize].lines
340 }
341}
342
343fn trailing_ident(callee: &str) -> Option<&str> {
347 let end = callee.trim_end_matches(['!', '?']);
348 let start = end
349 .char_indices()
350 .rev()
351 .take_while(|(_, c)| c.is_alphanumeric() || *c == '_' || *c == '$')
352 .last()
353 .map(|(i, _)| i)?;
354 Some(&end[start..])
355}
356
357impl AstAdapter for CodeModel {
358 fn root(&self) -> NodeId {
359 NodeId(0)
360 }
361
362 fn children(&self, node: NodeId) -> Vec<NodeId> {
363 self.nodes[node.0 as usize].children.clone()
364 }
365
366 fn name(&self, node: NodeId) -> Option<String> {
369 let n = &self.nodes[node.0 as usize];
370 n.parent?;
371 Some(n.name.clone().unwrap_or_else(|| n.construct.to_string()))
372 }
373
374 fn parent(&self, node: NodeId) -> Option<NodeId> {
375 self.nodes[node.0 as usize].parent
376 }
377
378 fn traits(&self, node: NodeId) -> Vec<String> {
379 self.nodes[node.0 as usize]
380 .traits
381 .iter()
382 .map(|t| t.to_string())
383 .collect()
384 }
385
386 fn property(&self, node: NodeId, name: &str) -> Option<Value> {
391 let n = &self.nodes[node.0 as usize];
392 match name {
393 "signature" => n.signature.clone().map(Value::Str),
394 "doc" => n.doc.clone().map(Value::Str),
395 "callee" => n.callee.clone().map(Value::Str),
396 _ => None,
397 }
398 }
399
400 fn default_value(&self, node: NodeId) -> Option<Value> {
402 Some(Value::Str(
403 self.text_of(&self.nodes[node.0 as usize]).to_string(),
404 ))
405 }
406
407 fn metadata(&self, node: NodeId, key: &str) -> Option<Value> {
408 let n = &self.nodes[node.0 as usize];
409 match key {
410 "kind" => (!n.kind.is_empty()).then(|| Value::Str(n.kind.clone())),
413 "construct" => (!n.construct.is_empty()).then(|| Value::Str(n.construct.to_string())),
414 "start-line" => Some(Value::Int(n.lines.0 as i64)),
415 "end-line" => Some(Value::Int(n.lines.1 as i64)),
416 "lang" => Some(Value::Str(self.lang.name().to_string())),
417 "n-children" => Some(Value::Int(n.children.len() as i64)),
418 "n-params" => n.n_params.map(Value::Int),
419 _ => None,
420 }
421 }
422
423 fn aliased_metadata(&self, _node: NodeId) -> &'static [&'static str] {
424 ALIASED
425 }
426
427 fn links(&self, node: NodeId) -> Vec<(String, NodeId)> {
430 self.nodes[node.0 as usize]
431 .links
432 .iter()
433 .map(|&t| ("definition".to_string(), t))
434 .collect()
435 }
436
437 fn backlinks(&self, node: NodeId) -> Vec<(String, NodeId)> {
440 self.nodes[node.0 as usize]
441 .backlinks
442 .iter()
443 .map(|&s| ("definition".to_string(), s))
444 .collect()
445 }
446}