aster/aster-core/src/analysis.rs

use crate::ast::{Expr, Stmt};
use crate::lexer::{Token, TokenKind};

/// A symbol collected from walking the AST.
#[derive(Debug, Clone)]
pub struct Symbol {
    pub name: String,
    pub kind: SymbolKind,
}

#[derive(Debug, Clone)]
pub enum SymbolKind {
    Variable,
    Function { params: Vec<String> },
}

/// Collect all declared symbols from the given statements.
/// Walks into nested scopes but does NOT perform scope-aware filtering —
/// that's the caller's job (e.g. the LSP completion handler).
pub fn collect_symbols(stmts: &[Stmt]) -> Vec<Symbol> {
    let mut symbols = Vec::new();
    collect_stmts(stmts, &mut symbols);
    symbols
}

fn collect_stmts(stmts: &[Stmt], symbols: &mut Vec<Symbol>) {
    for stmt in stmts {
        collect_stmt(stmt, symbols);
    }
}

fn collect_stmt(stmt: &Stmt, symbols: &mut Vec<Symbol>) {
    match stmt {
        Stmt::Let { name, initializer, .. } => {
            symbols.push(Symbol {
                name: name.clone(),
                kind: SymbolKind::Variable,
            });
            collect_expr(initializer, symbols);
        }
        Stmt::Function { name, params, body } => {
            symbols.push(Symbol {
                name: name.clone(),
                kind: SymbolKind::Function {
                    params: params.clone(),
                },
            });
            for p in params {
                symbols.push(Symbol {
                    name: p.clone(),
                    kind: SymbolKind::Variable,
                });
            }
            collect_stmts(body, symbols);
        }
        Stmt::Block(body) => collect_stmts(body, symbols),
        Stmt::If {
            condition,
            then_branch,
            else_branch,
        } => {
            collect_expr(condition, symbols);
            collect_stmt(then_branch, symbols);
            if let Some(else_b) = else_branch {
                collect_stmt(else_b, symbols);
            }
        }
        Stmt::While { condition, body } => {
            collect_expr(condition, symbols);
            collect_stmt(body, symbols);
        }
        Stmt::For {
            initializer,
            condition,
            step,
            body,
        } => {
            if let Some(init) = initializer {
                collect_stmt(init, symbols);
            }
            if let Some(cond) = condition {
                collect_expr(cond, symbols);
            }
            if let Some(step_expr) = step {
                collect_expr(step_expr, symbols);
            }
            collect_stmt(body, symbols);
        }
        Stmt::ForIn {
            var_name,
            iterable,
            body,
        } => {
            symbols.push(Symbol {
                name: var_name.clone(),
                kind: SymbolKind::Variable,
            });
            collect_expr(iterable, symbols);
            collect_stmt(body, symbols);
        }
        Stmt::ExprStmt(expr) | Stmt::Return(Some(expr)) => {
            collect_expr(expr, symbols);
        }
        Stmt::Return(None) | Stmt::Break | Stmt::Continue => {}
    }
}

fn collect_expr(expr: &Expr, symbols: &mut Vec<Symbol>) {
    match expr {
        Expr::Literal(_) | Expr::Variable(_) => {}
        Expr::Assign { value, .. } => collect_expr(value, symbols),
        Expr::Get { object, .. } => collect_expr(object, symbols),
        Expr::Set { object, value, .. } => {
            collect_expr(object, symbols);
            collect_expr(value, symbols);
        }
        Expr::ObjectLiteral { properties } => {
            for (_, val) in properties {
                collect_expr(val, symbols);
            }
        }
        Expr::ArrayLiteral { elements } => {
            for elem in elements {
                collect_expr(elem, symbols);
            }
        }
        Expr::IndexGet { array, index } => {
            collect_expr(array, symbols);
            collect_expr(index, symbols);
        }
        Expr::IndexSet {
            array,
            index,
            value,
            ..
        } => {
            collect_expr(array, symbols);
            collect_expr(index, symbols);
            collect_expr(value, symbols);
        }
        Expr::Unary { right, .. } => collect_expr(right, symbols),
        Expr::Binary { left, right, .. } | Expr::Logical { left, right, .. } => {
            collect_expr(left, symbols);
            collect_expr(right, symbols);
        }
        Expr::Ternary {
            condition,
            then_branch,
            else_branch,
        } => {
            collect_expr(condition, symbols);
            collect_expr(then_branch, symbols);
            collect_expr(else_branch, symbols);
        }
        Expr::Call {
            callee,
            arguments,
        } => {
            collect_expr(callee, symbols);
            for arg in arguments {
                collect_expr(arg, symbols);
            }
        }
        Expr::Lambda { params, body } => {
            for p in params {
                symbols.push(Symbol {
                    name: p.clone(),
                    kind: SymbolKind::Variable,
                });
            }
            collect_stmts(body, symbols);
        }
    }
}

// ── Token-based position lookups ──────────────────────────────────

/// Scan tokens for the declaration site of `name`.
/// Recognises `let name`, `const name`, `fn name`, and `for (name in …)`.
/// Returns 1-based (line, column).
pub fn find_declaration(tokens: &[Token], name: &str) -> Option<(usize, usize)> {
    let len = tokens.len();
    let mut i = 0;
    while i < len {
        match &tokens[i].kind {
            TokenKind::Let | TokenKind::Const => {
                if let Some(tok) = tokens.get(i + 1) {
                    if let TokenKind::Identifier(n) = &tok.kind {
                        if n == name {
                            return Some((tok.line, tok.column));
                        }
                    }
                }
            }
            TokenKind::Fn => {
                // fn name( ... or fn name { ...
                if let Some(tok) = tokens.get(i + 1) {
                    if let TokenKind::Identifier(n) = &tok.kind {
                        if n == name {
                            return Some((tok.line, tok.column));
                        }
                    }
                }
            }
            TokenKind::For => {
                // for ( name in ... )
                // Skip past LeftParen
                let mut j = i + 1;
                if j < len && tokens[j].kind == TokenKind::LeftParen {
                    j += 1;
                }
                if j < len {
                    if let TokenKind::Identifier(n) = &tokens[j].kind {
                        if n == name {
                            return Some((tokens[j].line, tokens[j].column));
                        }
                    }
                }
            }
            _ => {}
        }
        i += 1;
    }
    None
}

/// Scan tokens for every occurrence of `Identifier(name)`.
/// Returns 1-based (line, column) for each match.
pub fn find_all_references(tokens: &[Token], name: &str) -> Vec<(usize, usize)> {
    tokens
        .iter()
        .filter_map(|tok| {
            if let TokenKind::Identifier(n) = &tok.kind {
                if n == name {
                    Some((tok.line, tok.column))
                } else {
                    None
                }
            } else {
                None
            }
        })
        .collect()
}