prometeu-runtime/crates/prometeu-compiler/src/frontends/pbs/collector.rs

use crate::common::diagnostics::{Diagnostic, DiagnosticBundle, Severity};
use crate::frontends::pbs::ast::*;
use crate::frontends::pbs::symbols::*;
use crate::semantics::export_surface::ExportSurfaceKind;
use prometeu_analysis::NameInterner;

pub struct SymbolCollector<'a> {
    interner: &'a NameInterner,
    type_symbols: SymbolTable,
    value_symbols: SymbolTable,
    diagnostics: Vec<Diagnostic>,
}

impl<'a> SymbolCollector<'a> {
    pub fn new(interner: &'a NameInterner) -> Self {
        Self {
            interner,
            type_symbols: SymbolTable::new(),
            value_symbols: SymbolTable::new(),
            diagnostics: Vec::new(),
        }
    }

    pub fn collect(
        &mut self,
        arena: &AstArena,
        root: NodeId,
    ) -> Result<(SymbolTable, SymbolTable), DiagnosticBundle> {
        let file = match arena.kind(root) {
            NodeKind::File(file) => file,
            _ => {
                return Err(DiagnosticBundle::error(
                    "E_COLLECT_INVALID_ROOT",
                    "Expected File node as root".to_string(),
                    arena.span(root),
                ))
            }
        };

        for decl in &file.decls {
            match arena.kind(*decl) {
                NodeKind::FnDecl(fn_decl) => self.collect_fn(arena, *decl, fn_decl),
                NodeKind::ServiceDecl(service_decl) => self.collect_service(arena, *decl, service_decl),
                NodeKind::TypeDecl(type_decl) => self.collect_type(arena, *decl, type_decl),
                _ => {}
            }
        }

        if !self.diagnostics.is_empty() {
            return Err(DiagnosticBundle {
                diagnostics: self.diagnostics.clone(),
            });
        }

        Ok((
            std::mem::replace(&mut self.type_symbols, SymbolTable::new()),
            std::mem::replace(&mut self.value_symbols, SymbolTable::new()),
        ))
    }

    fn collect_fn(&mut self, arena: &AstArena, id: NodeId, decl: &FnDeclNodeArena) {
        let vis = match decl.vis.as_deref() {
            Some("pub") => Visibility::Pub,
            Some("mod") => Visibility::Mod,
            _ => Visibility::FilePrivate,
        };

        let span = arena.span(id);
        self.check_export_eligibility(SymbolKind::Function, vis, span);

        let symbol = Symbol {
            name: decl.name,
            kind: SymbolKind::Function,
            namespace: Namespace::Value,
            visibility: vis,
            ty: None, // Will be resolved later
            is_host: false,
            span,
            origin: None,
        };
        self.insert_value_symbol(symbol);
    }

    fn collect_service(&mut self, arena: &AstArena, id: NodeId, decl: &ServiceDeclNodeArena) {
        let vis = match decl.vis.as_deref() {
            Some("pub") => Visibility::Pub,
            _ => Visibility::Mod, // Defaults to Mod
        };

        let span = arena.span(id);

        self.check_export_eligibility(SymbolKind::Service, vis, span);

        let symbol = Symbol {
            name: decl.name,
            kind: SymbolKind::Service,
            namespace: Namespace::Type, // Service is a type
            visibility: vis,
            ty: None,
            is_host: false,
            span,
            origin: None,
        };
        self.insert_type_symbol(symbol);
    }

    fn collect_type(&mut self, arena: &AstArena, id: NodeId, decl: &TypeDeclNodeArena) {
        let vis = match decl.vis.as_deref() {
            Some("pub") => Visibility::Pub,
            Some("mod") => Visibility::Mod,
            _ => Visibility::FilePrivate,
        };
        let kind = match decl.type_kind.as_str() {
            "struct" => SymbolKind::Struct,
            "contract" => SymbolKind::Contract,
            "error" => SymbolKind::ErrorType,
            _ => SymbolKind::Struct, // Default
        };

        let span = arena.span(id);

        self.check_export_eligibility(kind.clone(), vis, span);

        let symbol = Symbol {
            name: decl.name,
            kind,
            namespace: Namespace::Type,
            visibility: vis,
            ty: None,
            is_host: decl.is_host,
            span,
            origin: None,
        };
        self.insert_type_symbol(symbol);
    }

    fn insert_type_symbol(&mut self, symbol: Symbol) {
        // Check for collision in value namespace first
        if let Some(existing) = self.value_symbols.get(symbol.name) {
            let existing = existing.clone();
            self.error_collision(&symbol, &existing);
            return;
        }

        if let Err(()) = self.type_symbols.insert(symbol.clone()) {
            if let Some(existing) = self.type_symbols.get(symbol.name).cloned() {
                self.error_duplicate(&symbol, &existing);
            }
        }
    }

    fn insert_value_symbol(&mut self, symbol: Symbol) {
        // Check for collision in type namespace first
        if let Some(existing) = self.type_symbols.get(symbol.name) {
            let existing = existing.clone();
            self.error_collision(&symbol, &existing);
            return;
        }

        if let Err(()) = self.value_symbols.insert(symbol.clone()) {
            if let Some(existing) = self.value_symbols.get(symbol.name).cloned() {
                self.error_duplicate(&symbol, &existing);
            }
        }
    }

    fn error_duplicate(&mut self, symbol: &Symbol, existing: &Symbol) {
        self.diagnostics.push(Diagnostic {
            severity: Severity::Error,
            code: "E_RESOLVE_DUPLICATE_SYMBOL".to_string(),
            message: format!(
                "Duplicate symbol '{}' already defined at {:?}",
                self.interner.resolve(symbol.name),
                existing.span
            ),
            span: symbol.span,
            related: Vec::new(),
        });
    }

    fn error_collision(&mut self, symbol: &Symbol, existing: &Symbol) {
        self.diagnostics.push(Diagnostic {
            severity: Severity::Error,
            code: "E_RESOLVE_NAMESPACE_COLLISION".to_string(),
            message: format!(
                "DebugSymbol '{}' collides with another symbol in the {:?} namespace defined at {:?}",
                self.interner.resolve(symbol.name),
                existing.namespace,
                existing.span
            ),
            span: symbol.span,
            related: Vec::new(),
        });
    }

    fn check_export_eligibility(&mut self, kind: SymbolKind, vis: Visibility, span: crate::common::spans::Span) {
        if let Err(msg) = ExportSurfaceKind::validate_visibility(kind, vis) {
            self.diagnostics.push(Diagnostic {
                severity: Severity::Error,
                code: "E_SEMANTIC_EXPORT_RESTRICTION".to_string(),
                message: msg,
                span,
                related: Vec::new(),
            });
        }
    }
}