OXYGENHUB
◆ v0.7.0 oxygenhub / crates / hub / src / repos / functions.rs
4.7 KiB · Rust · last changed Function history on the file page 1 month ago source blame history raw
1 
//! Function boundaries on file pages (hub spec §6.6): tree-sitter finds
2 
//! the functions and methods of a Rust, TypeScript, Python or shell file,
3 
//! with the lines they span, so a function's history (Oxygen spec §27) is
4 
//! one click away.
5
6 
/// A function or method and the lines it spans (1-based, inclusive).
7 
#[derive(Clone, Debug, PartialEq, Eq)]
8 
pub struct Function {
9 
    /// Qualified by its enclosing type or class where there is one: `Repo::log`.
10 
    pub name: String,
11 
    pub start: usize,
12 
    pub end: usize,
13 
}
14
15 
/// The grammar for a path, by its extension, and the node kinds that are
16 
/// functions in it.
17 
fn grammar(path: &str) -> Option<(tree_sitter::Language, &'static [&'static str])> {
18 
    let ext = path.rsplit_once('.').map(|(_, e)| e.to_ascii_lowercase())?;
19 
    Some(match ext.as_str() {
20 
        "rs" => (tree_sitter_rust::LANGUAGE.into(), &["function_item"][..]),
21 
        "ts" | "mts" | "cts" => (
22 
            tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(),
23 
            &["function_declaration", "method_definition", "generator_function_declaration"][..],
24 
        ),
25 
        "tsx" => (
26 
            tree_sitter_typescript::LANGUAGE_TSX.into(),
27 
            &["function_declaration", "method_definition", "generator_function_declaration"][..],
28 
        ),
29 
        "py" => (tree_sitter_python::LANGUAGE.into(), &["function_definition"][..]),
30 
        "sh" | "bash" => (tree_sitter_bash::LANGUAGE.into(), &["function_definition"][..]),
31 
        _ => return None,
32 
    })
33 
}
34
35 
/// Node kinds whose name qualifies the functions inside them.
36 
const CONTAINERS: &[&str] = &["impl_item", "trait_item", "class_declaration", "class_definition", "mod_item"];
37
38 
/// The functions of `text`, in file order; empty for a language without a
39 
/// grammar here or text that does not parse.
40 
pub fn functions(path: &str, text: &str) -> Vec<Function> {
41 
    let Some((language, kinds)) = grammar(path) else { return vec![] };
42 
    let mut parser = tree_sitter::Parser::new();
43 
    if parser.set_language(&language).is_err() {
44 
        return vec![];
45 
    }
46 
    let Some(tree) = parser.parse(text, None) else { return vec![] };
47 
    let mut out = Vec::new();
48 
    walk(tree.root_node(), text.as_bytes(), kinds, &mut Vec::new(), &mut out);
49 
    out
50 
}
51
52 
fn name_of(node: tree_sitter::Node, src: &[u8]) -> Option<String> {
53 
    let n = node.child_by_field_name("name").or_else(|| node.child_by_field_name("type"))?;
54 
    n.utf8_text(src).ok().map(|s| s.split_whitespace().collect::<Vec<_>>().join(" "))
55 
}
56
57 
fn walk(node: tree_sitter::Node, src: &[u8], kinds: &[&str], scope: &mut Vec<String>, out: &mut Vec<Function>) {
58 
    let kind = node.kind();
59 
    let mut pushed = false;
60 
    if kinds.contains(&kind) {
61 
        if let Some(name) = name_of(node, src) {
62 
            let mut full = scope.clone();
63 
            full.push(name);
64 
            out.push(Function { name: full.join("::"), start: node.start_position().row + 1, end: node.end_position().row + 1 });
65 
        }
66 
    } else if CONTAINERS.contains(&kind) {
67 
        if let Some(name) = name_of(node, src) {
68 
            scope.push(name);
69 
            pushed = true;
70 
        }
71 
    }
72 
    let mut cursor = node.walk();
73 
    for child in node.children(&mut cursor) {
74 
        walk(child, src, kinds, scope, out);
75 
    }
76 
    if pushed {
77 
        scope.pop();
78 
    }
79 
}
80
81 
#[cfg(test)]
82 
mod tests {
83 
    use super::*;
84
85 
    fn names(path: &str, text: &str) -> Vec<(String, usize, usize)> {
86 
        functions(path, text).into_iter().map(|f| (f.name, f.start, f.end)).collect()
87 
    }
88
89 
    #[test]
90 
    fn rust_functions_and_methods() {
91 
        let src = "use x;\n\nfn free() {\n    a();\n}\n\nimpl Repo {\n    pub fn log(&self) -> u8 {\n        1\n    }\n}\n\ntrait T {\n    fn t(&self);\n}\n";
92 
        assert_eq!(names("src/lib.rs", src), [("free".to_string(), 3, 5), ("Repo::log".to_string(), 8, 10)]);
93 
        let generic = "impl<T> Wrap<T> {\n    fn get(&self) {}\n}\n";
94 
        assert_eq!(names("a.rs", generic), [("Wrap<T>::get".to_string(), 2, 2)]);
95 
    }
96
97 
    #[test]
98 
    fn typescript_python_and_shell() {
99 
        let ts = "export function activate(ctx: X): void {\n  go();\n}\nclass Repo {\n  refresh(): void {\n    x();\n  }\n}\n";
100 
        assert_eq!(names("src/extension.ts", ts), [("activate".to_string(), 1, 3), ("Repo::refresh".to_string(), 5, 7)]);
101 
        let py = "def top():\n    pass\n\nclass Recipe:\n    def main(self):\n        return 1\n";
102 
        assert_eq!(names("images/rust/recipe.py", py), [("top".to_string(), 1, 2), ("Recipe::main".to_string(), 5, 6)]);
103 
        let sh = "#!/bin/sh\nbuild() {\n  cargo build\n}\n";
104 
        assert_eq!(names("build.sh", sh), [("build".to_string(), 2, 4)]);
105 
    }
106
107 
    #[test]
108 
    fn other_files_have_none() {
109 
        assert!(functions("README.md", "# fn main() {}\n").is_empty());
110 
        assert!(functions("Makefile", "all:\n").is_empty());
111 
    }
112 
}