collet_treemap/

lib.rs

use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};
use thiserror::Error;

/// Errors that can occur during map generation
#[derive(Error, Debug)]
pub enum MapError {
    #[error("IO error: {0}")]
    Io(#[from] std::io::Error),

    #[error("File not found: {0}")]
    FileNotFound(PathBuf),

    #[error("Unsupported language: {0}")]
    UnsupportedLanguage(String),

    #[error("Parse error: {0}")]
    ParseError(String),
}

/// Supported programming languages
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Hash)]
pub enum Language {
    Rust,
    Python,
    JavaScript,
}

impl Language {
    pub fn from_extension(ext: &str) -> Option<Self> {
        match ext {
            "rs" => Some(Language::Rust),
            "py" => Some(Language::Python),
            "js" | "jsx" | "ts" | "tsx" => Some(Language::JavaScript),
            _ => None,
        }
    }

    pub fn as_str(&self) -> &'static str {
        match self {
            Language::Rust => "rust",
            Language::Python => "python",
            Language::JavaScript => "javascript",
        }
    }
}

/// Configuration for map generation
#[derive(Debug, Clone, Copy)]
pub struct Config {
    pub format: OutputFormat,
    pub language: Option<Language>,
    pub max_depth: usize,
}

impl Default for Config {
    fn default() -> Self {
        Self {
            format: OutputFormat::Json,
            language: None,
            max_depth: 10,
        }
    }
}

/// Output format for generated maps
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum OutputFormat {
    Json,
    Text,
}

/// A code symbol in the repository
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Symbol {
    pub name: String,
    pub kind: SymbolKind,
    pub range: FileRange,
    pub file: PathBuf,
}

/// Kind of code symbol
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum SymbolKind {
    Module,
    Function,
    Class,
    Struct,
    Interface,
    Method,
    Property,
    Variable,
    Constant,
    Enum,
    Trait,
    Impl,
    Other(String),
}

impl SymbolKind {
    pub fn as_str(&self) -> &str {
        match self {
            SymbolKind::Module => "module",
            SymbolKind::Function => "function",
            SymbolKind::Class => "class",
            SymbolKind::Struct => "struct",
            SymbolKind::Interface => "interface",
            SymbolKind::Method => "method",
            SymbolKind::Property => "property",
            SymbolKind::Variable => "variable",
            SymbolKind::Constant => "constant",
            SymbolKind::Enum => "enum",
            SymbolKind::Trait => "trait",
            SymbolKind::Impl => "impl",
            SymbolKind::Other(s) => s,
        }
    }
}

/// File position information
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct FileRange {
    pub start_line: usize,
    pub start_col: usize,
    pub end_line: usize,
    pub end_col: usize,
}

/// Repository code map
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RepoMap {
    pub root: PathBuf,
    pub symbols: Vec<Symbol>,
    pub files: Vec<FileInfo>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FileInfo {
    pub path: PathBuf,
    pub language: Option<String>,
    pub lines: usize,
}

/// Generate a code map for the given directory
pub fn generate_map(path: &Path, config: Config) -> Result<RepoMap, MapError> {
    if !path.exists() {
        return Err(MapError::FileNotFound(path.to_path_buf()));
    }

    if !path.is_dir() {
        return Err(MapError::ParseError("Path must be a directory".to_string()));
    }

    let mut files = Vec::new();

    // Walk directory and collect symbols
    for entry in walkdir::WalkDir::new(path)
        .into_iter()
        .filter_map(Result::ok)
        .filter(|e| e.path().is_file())
    {
        let file_path = entry.path();

        // Skip common non-source directories
        if let Some(parent) = file_path.parent() {
            let parent_str = parent.to_string_lossy();
            if parent_str.contains("node_modules")
                || parent_str.contains(".git")
                || parent_str.contains("target")
                || parent_str.contains("__pycache__")
            {
                continue;
            }
        }

        let ext = file_path.extension().and_then(|e| e.to_str()).unwrap_or("");

        let language = Language::from_extension(ext);

        // Skip if language filter is set and doesn't match
        if let Some(filter_lang) = config.language {
            if language != Some(filter_lang) {
                continue;
            }
        }

        if let Some(lang) = language {
            let file_info = FileInfo {
                path: file_path.to_path_buf(),
                language: Some(lang.as_str().to_string()),
                lines: count_lines(file_path).unwrap_or(0),
            };
            files.push(file_info);
        }
    }

    Ok(RepoMap {
        root: path.to_path_buf(),
        symbols: Vec::new(),
        files,
    })
}

fn count_lines(path: &Path) -> Result<usize, MapError> {
    use std::io::BufRead;
    let file = std::fs::File::open(path)?;
    let reader = std::io::BufReader::new(file);
    Ok(reader.lines().count())
}

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

    #[test]
    fn test_language_from_extension() {
        assert_eq!(Language::from_extension("rs"), Some(Language::Rust));
        assert_eq!(Language::from_extension("py"), Some(Language::Python));
        assert_eq!(Language::from_extension("js"), Some(Language::JavaScript));
        assert_eq!(Language::from_extension("unknown"), None);
    }

    #[test]
    fn test_symbol_kind_as_str() {
        assert_eq!(SymbolKind::Function.as_str(), "function");
        assert_eq!(SymbolKind::Class.as_str(), "class");
    }

    #[test]
    fn test_config_default() {
        let cfg = Config::default();
        assert_eq!(cfg.format, OutputFormat::Json);
        assert_eq!(cfg.max_depth, 10);
    }
}