sarah 0.1.7

Count tokens, lines and function sizes per file to help AI agents decide what to refactor
use std::collections::HashSet;
use std::path::Path;

use serde::Serialize;
use tree_sitter_tags::TagsContext;

use super::lang::Lang;
use super::nesting::NestingHelpers;
use crate::count::Tokenizer;
use crate::report::WarningKind;
use crate::units::{Depth, LineNo, Lines, Tokens};

#[derive(Debug, Clone, Serialize)]
pub struct FnStats {
    pub name: String,
    pub kind: String,
    pub start_line: LineNo,
    pub end_line: LineNo,
    pub lines: Lines,
    pub tokens: Tokens,
    /// Deepest control-flow nesting inside the function.
    pub depth: Depth,
    /// Warning kinds a `sarah-allow` comment above the function turns off.
    #[serde(skip_serializing_if = "Vec::is_empty")]
    pub allow: Vec<WarningKind>,
}

impl FnStats {
    /// Extract function/method definitions with their sizes. Returns `None` for unsupported languages.
    pub fn extract(path: &Path, source: &str, tokenizer: Tokenizer) -> Option<Vec<Self>> {
        let lang = Lang::from_path(path)?;
        let config = lang.config()?;
        let tree = NestingHelpers::parse(lang, source)?;
        let mut ctx = TagsContext::new();
        let (tags, _) = ctx.generate_tags(config, source.as_bytes(), None).ok()?;

        let mut seen = HashSet::new();
        let mut out = Vec::new();
        for tag in tags.filter_map(Result::ok) {
            let kind = config.syntax_type_name(tag.syntax_type_id);
            if !tag.is_definition || !matches!(kind, "function" | "method") {
                continue;
            }
            // Rust methods match both the method and function patterns; keep the first.
            if !seen.insert(tag.range.clone()) {
                continue;
            }
            // Tag ranges come from the parser and land on UTF-8 boundaries; skip any that don't.
            let (Some(body), Some(name), Some(before)) = (
                source.get(tag.range.clone()),
                source.get(tag.name_range),
                source.get(..tag.range.start),
            ) else {
                continue;
            };
            let start_line = LineNo(1).below(before.matches('\n').count());
            let end_line = start_line.below(body.matches('\n').count());
            out.push(Self {
                name: name.to_string(),
                kind: kind.to_string(),
                start_line,
                end_line,
                lines: start_line.span_to(end_line),
                tokens: tokenizer.count(body),
                depth: NestingHelpers::depth_of(lang, &tree, tag.range.clone()),
                allow: Vec::new(),
            });
        }
        Some(out)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    struct Sample {
        path: &'static str,
        src: &'static str,
    }

    fn names(sample: Sample) -> Vec<(String, String, usize, usize)> {
        let Sample { path, src } = sample;
        FnStats::extract(Path::new(path), src, Tokenizer::O200k)
            .unwrap()
            .into_iter()
            .map(|f| (f.name, f.kind, f.start_line.0, f.end_line.0))
            .collect()
    }

    #[test]
    fn rust_functions_and_methods() {
        let src = "fn a() {\n    1;\n}\n\nstruct S;\nimpl S {\n    fn b(&self) {}\n}\n";
        let got = names(Sample { path: "x.rs", src });
        assert!(got.contains(&("a".into(), "function".into(), 1, 3)));
        assert!(got.iter().any(|f| f.0 == "b" && f.2 == 7 && f.3 == 7));
        assert_eq!(got.len(), 2);
    }

    #[test]
    fn python_functions() {
        let src = "def f():\n    return 1\n\nclass C:\n    def m(self):\n        pass\n";
        let got = names(Sample { path: "x.py", src });
        assert!(got.iter().any(|f| f.0 == "f" && f.2 == 1 && f.3 == 2));
        assert!(got.iter().any(|f| f.0 == "m" && f.2 == 5 && f.3 == 6));
    }

    #[test]
    fn typescript_loads() {
        let src = "function f(a: number): number {\n  return a;\n}\nclass C {\n  m() {}\n}\n";
        let got = names(Sample { path: "x.ts", src });
        assert!(got.iter().any(|f| f.0 == "f"));
        assert!(got.iter().any(|f| f.0 == "m"));
        assert!(!names(Sample { path: "x.tsx", src }).is_empty());
    }

    #[test]
    fn new_languages() {
        let samples = [
            ("x.java", "class C {\n  void m() {\n  }\n}\n", "m"),
            ("x.c", "int f(void) {\n  return 1;\n}\n", "f"),
            (
                "x.cpp",
                "struct S { void m() {} };\nint f() {\n  return 1;\n}\n",
                "f",
            ),
            ("x.cs", "class C {\n  C() {}\n  void M() {\n  }\n}\n", "M"),
            ("x.rb", "class C\n  def m\n  end\nend\n", "m"),
            ("x.php", "<?php\nfunction f() {\n  return 1;\n}\n", "f"),
        ];
        for (path, src, name) in samples {
            let got = names(Sample { path, src });
            assert!(got.iter().any(|f| f.0 == name), "{path}: {got:?}");
        }
    }

    #[test]
    fn unsupported_language() {
        assert!(FnStats::extract(Path::new("x.txt"), "hello", Tokenizer::O200k).is_none());
    }
}