use quarb::{AstAdapter, NodeId, Value};
mod lower;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Lang {
Rust,
Python,
Javascript,
C,
}
impl Lang {
pub fn name(self) -> &'static str {
match self {
Lang::Rust => "rust",
Lang::Python => "python",
Lang::Javascript => "javascript",
Lang::C => "c",
}
}
}
pub fn lang_for_ext(ext: &str) -> Option<Lang> {
match ext {
"rs" => Some(Lang::Rust),
"py" => Some(Lang::Python),
"js" | "mjs" | "cjs" | "jsx" => Some(Lang::Javascript),
"c" | "h" => Some(Lang::C),
_ => None,
}
}
pub fn supported(ext: &str) -> bool {
lang_for_ext(ext).is_some()
}
#[derive(Debug, thiserror::Error)]
pub enum CodeError {
#[error("code: {0}")]
Io(#[from] std::io::Error),
#[error("code: no code-level support for extension {0:?} (rs, py, js, mjs, cjs, jsx, c, h)")]
Language(String),
#[error(transparent)]
Backend(#[from] quarb_tree_sitter::TreeSitterError),
}
#[derive(Debug)]
pub struct Decl {
pub parent: Option<usize>,
pub construct: &'static str,
pub name: Option<String>,
pub traits: &'static [&'static str],
pub kind: String,
pub span: (usize, usize),
pub lines: (usize, usize),
pub signature: Option<String>,
pub doc: Option<String>,
pub callee: Option<String>,
pub n_params: Option<i64>,
}
struct Node {
parent: Option<NodeId>,
children: Vec<NodeId>,
construct: &'static str,
name: Option<String>,
traits: &'static [&'static str],
kind: String,
span: (usize, usize),
lines: (usize, usize),
signature: Option<String>,
doc: Option<String>,
callee: Option<String>,
n_params: Option<i64>,
links: Vec<NodeId>,
backlinks: Vec<NodeId>,
}
pub struct CodeModel {
source: String,
lang: Lang,
nodes: Vec<Node>,
}
const ALIASED: &[&str] = &[
"kind",
"construct",
"start-line",
"end-line",
"lang",
"n-children",
"n-params",
];
impl CodeModel {
pub fn build(source: String, lang: Lang, decls: Vec<Decl>) -> Self {
let mut nodes = Vec::with_capacity(decls.len() + 1);
nodes.push(Node {
parent: None,
children: Vec::new(),
construct: "",
name: None,
traits: &[],
kind: String::new(),
span: (0, source.len()),
lines: (1, source.lines().count().max(1)),
signature: None,
doc: None,
callee: None,
n_params: None,
links: Vec::new(),
backlinks: Vec::new(),
});
for d in decls {
let id = NodeId(nodes.len() as u64);
let parent = NodeId(d.parent.map_or(0, |p| p as u64 + 1));
nodes.push(Node {
parent: Some(parent),
children: Vec::new(),
construct: d.construct,
name: d.name,
traits: d.traits,
kind: d.kind,
span: d.span,
lines: d.lines,
signature: d.signature,
doc: d.doc,
callee: d.callee,
n_params: d.n_params,
links: Vec::new(),
backlinks: Vec::new(),
});
nodes[parent.0 as usize].children.push(id);
}
let mut model = CodeModel {
source,
lang,
nodes,
};
model.link_definitions();
model
}
fn link_definitions(&mut self) {
let mut by_name: std::collections::HashMap<&str, Vec<NodeId>> =
std::collections::HashMap::new();
for (i, n) in self.nodes.iter().enumerate() {
if matches!(n.construct, "function" | "type")
&& let Some(name) = &n.name
{
by_name.entry(name.as_str()).or_default().push(NodeId(i as u64));
}
}
let mut links: Vec<(NodeId, Vec<NodeId>)> = Vec::new();
for (i, n) in self.nodes.iter().enumerate() {
if let Some(callee) = &n.callee
&& let Some(ident) = trailing_ident(callee)
&& let Some(targets) = by_name.get(ident)
{
links.push((NodeId(i as u64), targets.clone()));
}
}
for (call, targets) in links {
for t in &targets {
self.nodes[t.0 as usize].backlinks.push(call);
}
self.nodes[call.0 as usize].links = targets;
}
}
pub fn parse(text: &str, ext: &str) -> Result<Self, CodeError> {
let ext = ext.to_ascii_lowercase();
let lang = lang_for_ext(&ext).ok_or_else(|| CodeError::Language(ext.clone()))?;
let ts = quarb_tree_sitter::TreeSitterAdapter::parse(text, &ext)?;
let decls = lower::lower(&ts, lang);
Ok(Self::build(text.to_string(), lang, decls))
}
pub fn open(path: &std::path::Path) -> Result<Self, CodeError> {
let ext = path
.extension()
.and_then(|e| e.to_str())
.unwrap_or("")
.to_ascii_lowercase();
let text = std::fs::read_to_string(path)?;
Self::parse(&text, &ext)
}
pub fn locator(&self, node: NodeId) -> String {
let mut parts = Vec::new();
let mut cur = node;
while let Some(parent) = self.nodes[cur.0 as usize].parent {
parts.push(self.segment(parent, cur));
cur = parent;
}
parts.reverse();
format!("/{}", parts.join("/"))
}
fn label(&self, node: NodeId) -> &str {
let n = &self.nodes[node.0 as usize];
n.name.as_deref().unwrap_or(n.construct)
}
fn segment(&self, parent: NodeId, child: NodeId) -> String {
let label = self.label(child);
let same: Vec<NodeId> = self.nodes[parent.0 as usize]
.children
.iter()
.copied()
.filter(|&c| self.label(c) == label)
.collect();
if same.len() > 1 {
let pos = same.iter().position(|&c| c == child).unwrap() + 1;
format!("{label}[{pos}]")
} else {
label.to_string()
}
}
fn text_of(&self, n: &Node) -> &str {
&self.source[n.span.0.min(self.source.len())..n.span.1.min(self.source.len())]
}
pub fn source(&self) -> &str {
&self.source
}
pub fn lang(&self) -> Lang {
self.lang
}
pub fn ident(&self, node: NodeId) -> Option<&str> {
self.nodes[node.0 as usize].name.as_deref()
}
pub fn construct(&self, node: NodeId) -> &str {
self.nodes[node.0 as usize].construct
}
pub fn span(&self, node: NodeId) -> (usize, usize) {
self.nodes[node.0 as usize].span
}
pub fn line_span(&self, node: NodeId) -> (usize, usize) {
self.nodes[node.0 as usize].lines
}
}
fn trailing_ident(callee: &str) -> Option<&str> {
let end = callee.trim_end_matches(['!', '?']);
let start = end
.char_indices()
.rev()
.take_while(|(_, c)| c.is_alphanumeric() || *c == '_' || *c == '$')
.last()
.map(|(i, _)| i)?;
Some(&end[start..])
}
impl AstAdapter for CodeModel {
fn root(&self) -> NodeId {
NodeId(0)
}
fn children(&self, node: NodeId) -> Vec<NodeId> {
self.nodes[node.0 as usize].children.clone()
}
fn name(&self, node: NodeId) -> Option<String> {
let n = &self.nodes[node.0 as usize];
n.parent?;
Some(n.name.clone().unwrap_or_else(|| n.construct.to_string()))
}
fn parent(&self, node: NodeId) -> Option<NodeId> {
self.nodes[node.0 as usize].parent
}
fn traits(&self, node: NodeId) -> Vec<String> {
self.nodes[node.0 as usize]
.traits
.iter()
.map(|t| t.to_string())
.collect()
}
fn property(&self, node: NodeId, name: &str) -> Option<Value> {
let n = &self.nodes[node.0 as usize];
match name {
"signature" => n.signature.clone().map(Value::Str),
"doc" => n.doc.clone().map(Value::Str),
"callee" => n.callee.clone().map(Value::Str),
_ => None,
}
}
fn default_value(&self, node: NodeId) -> Option<Value> {
Some(Value::Str(
self.text_of(&self.nodes[node.0 as usize]).to_string(),
))
}
fn metadata(&self, node: NodeId, key: &str) -> Option<Value> {
let n = &self.nodes[node.0 as usize];
match key {
"kind" => (!n.kind.is_empty()).then(|| Value::Str(n.kind.clone())),
"construct" => (!n.construct.is_empty()).then(|| Value::Str(n.construct.to_string())),
"start-line" => Some(Value::Int(n.lines.0 as i64)),
"end-line" => Some(Value::Int(n.lines.1 as i64)),
"lang" => Some(Value::Str(self.lang.name().to_string())),
"n-children" => Some(Value::Int(n.children.len() as i64)),
"n-params" => n.n_params.map(Value::Int),
_ => None,
}
}
fn aliased_metadata(&self, _node: NodeId) -> &'static [&'static str] {
ALIASED
}
fn links(&self, node: NodeId) -> Vec<(String, NodeId)> {
self.nodes[node.0 as usize]
.links
.iter()
.map(|&t| ("definition".to_string(), t))
.collect()
}
fn backlinks(&self, node: NodeId) -> Vec<(String, NodeId)> {
self.nodes[node.0 as usize]
.backlinks
.iter()
.map(|&s| ("definition".to_string(), s))
.collect()
}
}