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

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

pub struct SymbolCollector {
    type_symbols: SymbolTable,
    value_symbols: SymbolTable,
    diagnostics: Vec<Diagnostic>,
}

impl SymbolCollector {
    pub fn new() -> Self {
        Self {
            type_symbols: SymbolTable::new(),
            value_symbols: SymbolTable::new(),
            diagnostics: Vec::new(),
        }
    }

    pub fn collect(&mut self, file: &FileNode) -> Result<(SymbolTable, SymbolTable), DiagnosticBundle> {
        for decl in &file.decls {
            match decl {
                Node::FnDecl(fn_decl) => self.collect_fn(fn_decl),
                Node::ServiceDecl(service_decl) => self.collect_service(service_decl),
                Node::TypeDecl(type_decl) => self.collect_type(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, decl: &FnDeclNode) {
        let vis = match decl.vis.as_deref() {
            Some("pub") => Visibility::Pub,
            Some("mod") => Visibility::Mod,
            _ => Visibility::FilePrivate,
        };
        
        self.check_export_eligibility(SymbolKind::Function, vis, decl.span);

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

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

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

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

    fn collect_type(&mut self, decl: &TypeDeclNode) {
        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
        };

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

        let symbol = Symbol {
            name: decl.name.clone(),
            kind,
            namespace: Namespace::Type,
            visibility: vis,
            ty: None,
            is_host: decl.is_host,
            span: decl.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(existing) = self.type_symbols.insert(symbol.clone()) {
            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(existing) = self.value_symbols.insert(symbol.clone()) {
            self.error_duplicate(&symbol, &existing);
        }
    }

    fn error_duplicate(&mut self, symbol: &Symbol, existing: &Symbol) {
        self.diagnostics.push(Diagnostic {
            level: DiagnosticLevel::Error,
            code: Some("E_RESOLVE_DUPLICATE_SYMBOL".to_string()),
            message: format!("Duplicate symbol '{}' already defined at {:?}", symbol.name, existing.span),
            span: Some(symbol.span),
        });
    }

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

    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 {
                level: DiagnosticLevel::Error,
                code: Some("E_SEMANTIC_EXPORT_RESTRICTION".to_string()),
                message: msg,
                span: Some(span),
            });
        }
    }
}