Skip to main content

pelicode_core/
language.rs

1//! Language adapters: the per-language rules for function extraction and normalization.
2
3use std::fmt;
4use std::path::Path;
5use std::str::FromStr;
6
7use serde::{Deserialize, Serialize};
8
9/// A programming language supported by pelicode.
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
11#[serde(rename_all = "lowercase")]
12pub enum Language {
13    /// Go (`.go`).
14    Go,
15    /// Java (`.java`).
16    Java,
17    /// Kotlin (`.kt`, `.kts`).
18    Kotlin,
19    /// Python (`.py`).
20    Python,
21    /// Ruby (`.rb`).
22    Ruby,
23    /// Rust (`.rs`).
24    Rust,
25}
26
27impl Language {
28    /// Every supported language.
29    pub const ALL: [Language; 6] = [
30        Language::Go,
31        Language::Java,
32        Language::Kotlin,
33        Language::Python,
34        Language::Ruby,
35        Language::Rust,
36    ];
37
38    /// Lower-case name of the language, as used in databases and reports.
39    ///
40    /// # Examples
41    ///
42    /// ```
43    /// use pelicode_core::Language;
44    ///
45    /// assert_eq!(Language::Kotlin.name(), "kotlin");
46    /// assert_eq!("kotlin".parse::<Language>(), Ok(Language::Kotlin));
47    /// ```
48    pub fn name(self) -> &'static str {
49        match self {
50            Language::Go => "go",
51            Language::Java => "java",
52            Language::Kotlin => "kotlin",
53            Language::Python => "python",
54            Language::Ruby => "ruby",
55            Language::Rust => "rust",
56        }
57    }
58
59    /// File extensions of the language, without the leading dot.
60    ///
61    /// # Examples
62    ///
63    /// ```
64    /// use pelicode_core::Language;
65    ///
66    /// assert_eq!(Language::Kotlin.extensions(), ["kt", "kts"]);
67    /// ```
68    pub fn extensions(self) -> &'static [&'static str] {
69        match self {
70            Language::Go => &["go"],
71            Language::Java => &["java"],
72            Language::Kotlin => &["kt", "kts"],
73            Language::Python => &["py"],
74            Language::Ruby => &["rb"],
75            Language::Rust => &["rs"],
76        }
77    }
78
79    /// The language of a file extension (without the dot); `None` when unsupported.
80    ///
81    /// # Examples
82    ///
83    /// ```
84    /// use pelicode_core::Language;
85    ///
86    /// assert_eq!(Language::from_extension("rs"), Some(Language::Rust));
87    /// assert_eq!(Language::from_extension("c"), None);
88    /// ```
89    pub fn from_extension(ext: &str) -> Option<Language> {
90        Language::ALL
91            .into_iter()
92            .find(|lang| lang.extensions().contains(&ext))
93    }
94
95    /// The language of a file, from its extension; `None` when unsupported.
96    ///
97    /// # Examples
98    ///
99    /// ```
100    /// use std::path::Path;
101    /// use pelicode_core::Language;
102    ///
103    /// assert_eq!(Language::from_path(Path::new("src/Main.java")), Some(Language::Java));
104    /// assert_eq!(Language::from_path(Path::new("README.md")), None);
105    /// ```
106    pub fn from_path(path: &Path) -> Option<Language> {
107        path.extension()
108            .and_then(|ext| ext.to_str())
109            .and_then(Language::from_extension)
110    }
111
112    /// The Tree-sitter grammar of the language.
113    ///
114    /// # Examples
115    ///
116    /// ```
117    /// use pelicode_core::Language;
118    ///
119    /// let mut parser = tree_sitter::Parser::new();
120    /// parser.set_language(&Language::Go.ts_language()).unwrap();
121    /// let tree = parser.parse("package main", None).unwrap();
122    /// assert_eq!(tree.root_node().kind(), "source_file");
123    /// ```
124    pub fn ts_language(self) -> tree_sitter::Language {
125        match self {
126            Language::Go => tree_sitter_go::LANGUAGE.into(),
127            Language::Java => tree_sitter_java::LANGUAGE.into(),
128            Language::Kotlin => tree_sitter_kotlin_ng::LANGUAGE.into(),
129            Language::Python => tree_sitter_python::LANGUAGE.into(),
130            Language::Ruby => tree_sitter_ruby::LANGUAGE.into(),
131            Language::Rust => tree_sitter_rust::LANGUAGE.into(),
132        }
133    }
134
135    pub(crate) fn spec(self) -> &'static LanguageSpec {
136        match self {
137            Language::Go => &GO,
138            Language::Java => &JAVA,
139            Language::Kotlin => &KOTLIN,
140            Language::Python => &PYTHON,
141            Language::Ruby => &RUBY,
142            Language::Rust => &RUST,
143        }
144    }
145}
146
147impl fmt::Display for Language {
148    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
149        f.write_str(self.name())
150    }
151}
152
153impl FromStr for Language {
154    type Err = String;
155
156    fn from_str(s: &str) -> Result<Self, Self::Err> {
157        Language::ALL
158            .into_iter()
159            .find(|lang| lang.name().eq_ignore_ascii_case(s))
160            .ok_or_else(|| format!("unsupported language: {s}"))
161    }
162}
163
164/// Node kinds that drive extraction and normalization for one language.
165pub(crate) struct LanguageSpec {
166    /// Nodes treated as a function unit.
167    pub functions: &'static [&'static str],
168    /// Comment nodes (dropped from every level).
169    pub comments: &'static [&'static str],
170    /// Variable-like identifiers, abstracted to `$ID` from N1.
171    pub identifiers: &'static [&'static str],
172    /// Type-like identifiers, kept at N1 and abstracted to `$TYPE` from N2.
173    pub type_identifiers: &'static [&'static str],
174    /// Ancestors that turn a plain identifier into a type identifier (e.g. Python's `type`).
175    pub type_contexts: &'static [&'static str],
176    /// String-like literals, emitted as one atomic token (`$STR` from N1).
177    pub strings: &'static [&'static str],
178    /// Numeric literals (`$NUM` from N1).
179    pub numbers: &'static [&'static str],
180    /// Whether the first string statement of a function body is a docstring to drop.
181    pub docstrings: bool,
182    /// Whether `;` is an optional statement terminator (layout, not structure) and is dropped.
183    pub optional_semicolons: bool,
184}
185
186static GO: LanguageSpec = LanguageSpec {
187    functions: &["function_declaration", "method_declaration", "func_literal"],
188    comments: &["comment"],
189    identifiers: &["identifier", "field_identifier", "label_name"],
190    type_identifiers: &["type_identifier", "package_identifier"],
191    type_contexts: &[],
192    strings: &[
193        "interpreted_string_literal",
194        "raw_string_literal",
195        "rune_literal",
196    ],
197    numbers: &["int_literal", "float_literal", "imaginary_literal"],
198    docstrings: false,
199    optional_semicolons: true,
200};
201
202static JAVA: LanguageSpec = LanguageSpec {
203    functions: &[
204        "method_declaration",
205        "constructor_declaration",
206        "compact_constructor_declaration",
207        "lambda_expression",
208    ],
209    comments: &["line_comment", "block_comment"],
210    identifiers: &["identifier"],
211    type_identifiers: &["type_identifier"],
212    type_contexts: &[],
213    strings: &["string_literal", "character_literal"],
214    numbers: &[
215        "decimal_integer_literal",
216        "hex_integer_literal",
217        "octal_integer_literal",
218        "binary_integer_literal",
219        "decimal_floating_point_literal",
220        "hex_floating_point_literal",
221    ],
222    docstrings: false,
223    optional_semicolons: false,
224};
225
226static KOTLIN: LanguageSpec = LanguageSpec {
227    functions: &[
228        "function_declaration",
229        "secondary_constructor",
230        "anonymous_function",
231        "lambda_literal",
232        "getter",
233        "setter",
234    ],
235    comments: &["line_comment", "block_comment"],
236    identifiers: &["identifier"],
237    type_identifiers: &[],
238    type_contexts: &["user_type"],
239    strings: &[
240        "string_literal",
241        "multiline_string_literal",
242        "character_literal",
243    ],
244    numbers: &["number_literal", "float_literal"],
245    docstrings: false,
246    optional_semicolons: true,
247};
248
249static PYTHON: LanguageSpec = LanguageSpec {
250    functions: &["function_definition", "lambda"],
251    comments: &["comment"],
252    identifiers: &["identifier"],
253    type_identifiers: &[],
254    type_contexts: &["type"],
255    strings: &["string", "concatenated_string"],
256    numbers: &["integer", "float"],
257    docstrings: true,
258    optional_semicolons: true,
259};
260
261static RUBY: LanguageSpec = LanguageSpec {
262    functions: &["method", "singleton_method", "lambda", "block", "do_block"],
263    comments: &["comment"],
264    identifiers: &[
265        "identifier",
266        "instance_variable",
267        "class_variable",
268        "global_variable",
269    ],
270    type_identifiers: &["constant"],
271    type_contexts: &[],
272    strings: &[
273        "string",
274        "chained_string",
275        "character",
276        "simple_symbol",
277        "delimited_symbol",
278        "regex",
279        "heredoc_body",
280        "string_array",
281        "symbol_array",
282    ],
283    numbers: &["integer", "float", "rational", "complex"],
284    docstrings: false,
285    optional_semicolons: true,
286};
287
288static RUST: LanguageSpec = LanguageSpec {
289    // Trait methods without a body (`function_signature_item`) are not functions.
290    functions: &["function_item", "closure_expression"],
291    comments: &["line_comment", "block_comment"],
292    identifiers: &[
293        "identifier",
294        "field_identifier",
295        "shorthand_field_identifier",
296    ],
297    type_identifiers: &["type_identifier"],
298    type_contexts: &[],
299    strings: &["string_literal", "raw_string_literal", "char_literal"],
300    numbers: &["integer_literal", "float_literal"],
301    docstrings: false,
302    // `;` turns an expression into a statement, so it is structure.
303    optional_semicolons: false,
304};
305
306#[cfg(test)]
307mod tests {
308    use super::*;
309
310    #[test]
311    fn detects_language_from_extension() {
312        assert_eq!(
313            Language::from_path(Path::new("a/b/Main.kt")),
314            Some(Language::Kotlin)
315        );
316        assert_eq!(
317            Language::from_path(Path::new("x.py")),
318            Some(Language::Python)
319        );
320        assert_eq!(Language::from_path(Path::new("x.rs")), Some(Language::Rust));
321        assert_eq!(Language::from_path(Path::new("x.c")), None);
322    }
323
324    #[test]
325    fn grammars_load() {
326        for lang in Language::ALL {
327            let mut parser = tree_sitter::Parser::new();
328            parser.set_language(&lang.ts_language()).unwrap();
329        }
330    }
331}