Skip to main content

ailint_extractor/
lib.rs

1//! Lexical extraction of source-code comments as virtual guidance documents.
2//!
3//! `ailint-extractor` uses [`logos`](https://crates.io/crates/logos) to tokenize
4//! source files just enough to pull out line and block comments without building
5//! a full AST. The returned [`Comment`]s can be fed to `ailint-core` as
6//! "virtual documents" so semantic rules (vague instructions, negative overload,
7//! bloat) can catch AI slop that leaks into inline comments and docstrings.
8//!
9//! Supported languages: [`Language::Rust`], [`Language::TypeScript`],
10//! [`Language::JavaScript`], [`Language::Python`].
11//!
12//! ```
13//! use ailint_extractor::{extract, Language, CommentKind};
14//!
15//! let src = "fn main() {\n    // TODO: refactor this later\n}\n";
16//! let comments = extract(src, Language::Rust);
17//! assert_eq!(comments.len(), 1);
18//! assert_eq!(comments[0].kind, CommentKind::Line);
19//! assert_eq!(comments[0].line, 2);
20//! assert_eq!(comments[0].body().trim(), "TODO: refactor this later");
21//! ```
22
23#![warn(missing_docs)]
24
25use std::ops::Range;
26use std::path::Path;
27
28mod lexers;
29
30use lexers::{cs_comments, go_comments, java_comments, js_comments, py_comments, rust_comments};
31
32/// Programming languages supported by the extractor.
33#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
34pub enum Language {
35    /// Rust (`.rs`). Handles `//`, `///`, `//!`, `/* */`, `/** */`, `/*! */`,
36    /// plus regular, raw, and byte string literals.
37    Rust,
38    /// TypeScript (`.ts`, `.tsx`). Handles `//`, `/* */`, `/** */`, plus single,
39    /// double, and template string literals.
40    TypeScript,
41    /// JavaScript (`.js`, `.jsx`, `.mjs`, `.cjs`). Same handling as TypeScript.
42    JavaScript,
43    /// Python (`.py`). Handles `#` line comments and triple-quoted strings
44    /// (surfaced as [`CommentKind::Docstring`]).
45    Python,
46    /// Go (`.go`). Handles `//`, `/* */`, plus `"…"`, raw `` `…` ``, and
47    /// `'…'` rune literals. Go has no dedicated doc-comment syntax.
48    Go,
49    /// Java (`.java`). Handles `//`, `/* */`, `/** */` javadoc, plus `"…"`,
50    /// `"""…"""` text blocks (Java 15+), and `'…'` char literals.
51    Java,
52    /// C# (`.cs`). Handles `//`, `///` XML doc, `/* */`, plus `"…"`, `@"…"`
53    /// verbatim (with `""` escapes), `$"…"` interpolated (surface only),
54    /// and `'…'` char literals.
55    CSharp,
56}
57
58impl Language {
59    /// Detect a supported language from a file path's extension.
60    ///
61    /// Returns `None` for unsupported extensions. Extension matching is
62    /// case-insensitive.
63    pub fn from_path(path: &Path) -> Option<Self> {
64        let ext = path.extension()?.to_str()?.to_ascii_lowercase();
65        match ext.as_str() {
66            "rs" => Some(Self::Rust),
67            "ts" | "tsx" | "mts" | "cts" => Some(Self::TypeScript),
68            "js" | "jsx" | "mjs" | "cjs" => Some(Self::JavaScript),
69            "py" | "pyi" => Some(Self::Python),
70            "go" => Some(Self::Go),
71            "java" => Some(Self::Java),
72            "cs" => Some(Self::CSharp),
73            _ => None,
74        }
75    }
76
77    /// Short stable identifier for the language (`"rust"`, `"typescript"`, …).
78    pub fn as_str(&self) -> &'static str {
79        match self {
80            Self::Rust => "rust",
81            Self::TypeScript => "typescript",
82            Self::JavaScript => "javascript",
83            Self::Python => "python",
84            Self::Go => "go",
85            Self::Java => "java",
86            Self::CSharp => "csharp",
87        }
88    }
89}
90
91/// Classification of an extracted comment.
92#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
93pub enum CommentKind {
94    /// A single-line comment: Rust/TS/JS `//`, Python `#`.
95    Line,
96    /// A block comment: `/* … */`. Not nested.
97    Block,
98    /// A documentation comment:
99    /// - Rust: `///` (outer) or `//!` (inner), `/** … */`, `/*! … */`.
100    /// - TS/JS: `/** … */` (JSDoc/TSDoc).
101    Doc,
102    /// A Python triple-quoted string (`""" … """` or `''' … '''`).
103    /// Emitted regardless of syntactic position — the extractor does not
104    /// distinguish docstrings from ordinary triple-quoted literals.
105    Docstring,
106}
107
108/// A single extracted comment.
109#[derive(Debug, Clone, PartialEq, Eq)]
110pub struct Comment {
111    /// The raw comment slice, including delimiters (`// foo`, `/* foo */`).
112    pub raw: String,
113    /// Classification of the comment.
114    pub kind: CommentKind,
115    /// Byte offsets within the source file.
116    pub byte_range: Range<usize>,
117    /// 1-based line number of the comment's start.
118    pub line: usize,
119}
120
121impl Comment {
122    /// The comment content with delimiters stripped, but internal whitespace
123    /// preserved.
124    ///
125    /// - `// foo` -> `" foo"`
126    /// - `/// foo` -> `" foo"`
127    /// - `/* foo */` -> `" foo "`
128    /// - `""" foo """` -> `" foo "`
129    pub fn body(&self) -> &str {
130        strip_delimiters(&self.raw, self.kind)
131    }
132}
133
134fn strip_delimiters(raw: &str, kind: CommentKind) -> &str {
135    match kind {
136        CommentKind::Line => {
137            if let Some(rest) = raw.strip_prefix("///") {
138                rest
139            } else if let Some(rest) = raw.strip_prefix("//!") {
140                rest
141            } else if let Some(rest) = raw.strip_prefix("//") {
142                rest
143            } else if let Some(rest) = raw.strip_prefix('#') {
144                rest
145            } else {
146                raw
147            }
148        }
149        CommentKind::Doc => {
150            if let Some(rest) = raw.strip_prefix("///") {
151                rest
152            } else if let Some(rest) = raw.strip_prefix("//!") {
153                rest
154            } else if let Some(rest) = raw.strip_prefix("/**").and_then(|r| r.strip_suffix("*/")) {
155                rest
156            } else if let Some(rest) = raw.strip_prefix("/*!").and_then(|r| r.strip_suffix("*/")) {
157                rest
158            } else {
159                raw
160            }
161        }
162        CommentKind::Block => raw
163            .strip_prefix("/*")
164            .and_then(|r| r.strip_suffix("*/"))
165            .unwrap_or(raw),
166        CommentKind::Docstring => raw
167            .strip_prefix("\"\"\"")
168            .and_then(|r| r.strip_suffix("\"\"\""))
169            .or_else(|| raw.strip_prefix("'''").and_then(|r| r.strip_suffix("'''")))
170            .unwrap_or(raw),
171    }
172}
173
174/// Extract every comment from `source` for the given `language`.
175///
176/// Comments are returned in source order. Byte ranges refer to the input
177/// `source`. Line numbers are 1-based.
178pub fn extract(source: &str, language: Language) -> Vec<Comment> {
179    let raw_comments = match language {
180        Language::Rust => rust_comments(source),
181        Language::TypeScript | Language::JavaScript => js_comments(source),
182        Language::Python => py_comments(source),
183        Language::Go => go_comments(source),
184        Language::Java => java_comments(source),
185        Language::CSharp => cs_comments(source),
186    };
187    attach_lines(source, raw_comments)
188}
189
190/// A raw comment span produced by a language lexer, before line-number
191/// annotation. Kept crate-internal.
192pub(crate) struct RawComment {
193    pub raw: String,
194    pub kind: CommentKind,
195    pub byte_range: Range<usize>,
196}
197
198fn attach_lines(source: &str, raws: Vec<RawComment>) -> Vec<Comment> {
199    if raws.is_empty() {
200        return Vec::new();
201    }
202    let bytes = source.as_bytes();
203    let mut newline_offsets: Vec<usize> = bytes
204        .iter()
205        .enumerate()
206        .filter_map(|(i, b)| if *b == b'\n' { Some(i) } else { None })
207        .collect();
208    newline_offsets.push(bytes.len());
209
210    raws.into_iter()
211        .map(|r| {
212            let line = 1 + newline_offsets
213                .binary_search(&r.byte_range.start)
214                .unwrap_or_else(|idx| idx);
215            Comment {
216                raw: r.raw,
217                kind: r.kind,
218                byte_range: r.byte_range,
219                line,
220            }
221        })
222        .collect()
223}
224
225#[cfg(test)]
226mod tests {
227    use super::*;
228    use std::path::PathBuf;
229
230    #[test]
231    fn language_from_path_covers_common_extensions() {
232        for (path, lang) in [
233            ("src/lib.rs", Language::Rust),
234            ("src/App.tsx", Language::TypeScript),
235            ("index.ts", Language::TypeScript),
236            ("index.js", Language::JavaScript),
237            ("bundle.mjs", Language::JavaScript),
238            ("main.py", Language::Python),
239            ("types.pyi", Language::Python),
240            ("cmd/main.go", Language::Go),
241            ("App.java", Language::Java),
242            ("Program.cs", Language::CSharp),
243        ] {
244            assert_eq!(
245                Language::from_path(&PathBuf::from(path)),
246                Some(lang),
247                "path={path}"
248            );
249        }
250    }
251
252    #[test]
253    fn language_from_path_rejects_unknown() {
254        assert!(Language::from_path(&PathBuf::from("README.md")).is_none());
255        assert!(Language::from_path(&PathBuf::from("Cargo.toml")).is_none());
256        assert!(Language::from_path(&PathBuf::from("noext")).is_none());
257    }
258
259    #[test]
260    fn language_as_str_stable() {
261        assert_eq!(Language::Rust.as_str(), "rust");
262        assert_eq!(Language::TypeScript.as_str(), "typescript");
263        assert_eq!(Language::JavaScript.as_str(), "javascript");
264        assert_eq!(Language::Python.as_str(), "python");
265        assert_eq!(Language::Go.as_str(), "go");
266        assert_eq!(Language::Java.as_str(), "java");
267        assert_eq!(Language::CSharp.as_str(), "csharp");
268    }
269
270    #[test]
271    fn body_strips_line_markers() {
272        let c = Comment {
273            raw: "/// docstring".into(),
274            kind: CommentKind::Doc,
275            byte_range: 0..13,
276            line: 1,
277        };
278        assert_eq!(c.body(), " docstring");
279    }
280
281    #[test]
282    fn body_strips_block_markers() {
283        let c = Comment {
284            raw: "/* hi */".into(),
285            kind: CommentKind::Block,
286            byte_range: 0..8,
287            line: 1,
288        };
289        assert_eq!(c.body(), " hi ");
290    }
291
292    #[test]
293    fn body_strips_python_docstring() {
294        let c = Comment {
295            raw: "\"\"\"module\"\"\"".into(),
296            kind: CommentKind::Docstring,
297            byte_range: 0..12,
298            line: 1,
299        };
300        assert_eq!(c.body(), "module");
301    }
302
303    #[test]
304    fn empty_source_returns_empty_vec() {
305        assert!(extract("", Language::Rust).is_empty());
306        assert!(extract("", Language::Python).is_empty());
307        assert!(extract("", Language::TypeScript).is_empty());
308    }
309}