prometeu-runtime/crates/compiler/prometeu-compiler/src/analysis/symbols/mod.rs

use crate::common::diagnostics::{Diagnostic, Severity};
use crate::common::spans::{Span, FileId};
use crate::frontends::pbs::ast::AstArena;
use prometeu_analysis::NodeId;
use prometeu_analysis::{NameId, SymbolId, ModuleId};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
pub enum SymbolKind {
    Type,
    Value,
    Service,
    Function,
    Struct,
    Contract,
    ErrorType,
    Local,
}

#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)]
pub enum Namespace {
    Type,
    Value,
    Service,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Symbol {
    pub name: NameId,
    pub kind: SymbolKind,
    pub exported: bool,
    pub module: ModuleId,
    pub decl_span: Span,
}

#[derive(Debug, Default, Clone, Serialize, Deserialize)]
pub struct SymbolArena {
    pub symbols: Vec<Symbol>,
}

#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
pub struct DefKey {
    pub module: ModuleId,
    pub name: NameId,
    pub namespace: Namespace,
}

#[derive(Debug, Default, Clone)]
pub struct DefIndex {
    symbols: HashMap<DefKey, SymbolId>,
}

#[derive(Debug, Default, Clone, Serialize, Deserialize)]
pub struct RefIndex {
    refs: Vec<Vec<Span>>,
}

#[derive(Debug, Default, Clone, Serialize, Deserialize)]
pub struct NodeToSymbol {
    map: Vec<Option<SymbolId>>,
}

impl SymbolArena {
    pub fn new() -> Self {
        Self { symbols: Vec::new() }
    }

    pub fn insert(&mut self, symbol: Symbol) -> SymbolId {
        let id = SymbolId(self.symbols.len() as u32);
        self.symbols.push(symbol);
        id
    }

    pub fn get(&self, id: SymbolId) -> &Symbol {
        &self.symbols[id.0 as usize]
    }
}

impl DefIndex {
    pub fn new() -> Self {
        Self {
            symbols: HashMap::new(),
        }
    }

    pub fn insert_symbol(&mut self, key: DefKey, symbol_id: SymbolId) -> Result<(), Diagnostic> {
        if self.symbols.contains_key(&key) {
            return Err(Diagnostic {
                severity: Severity::Error,
                code: "E_RESOLVE_DUPLICATE_SYMBOL".to_string(),
                message: "Duplicate symbol in the same module and namespace".to_string(),
                // Placeholder span; callers should overwrite with accurate span when known.
                span: Span::new(FileId(0), 0, 0),
                related: Vec::new(),
            });
        }

        self.symbols.insert(key, symbol_id);
        Ok(())
    }

    pub fn get(&self, key: DefKey) -> Option<SymbolId> {
        self.symbols.get(&key).copied()
    }

    /// Lookup by name/namespace ignoring module. Returns the first match with its module id.
    pub fn get_by_name_any_module(&self, name: NameId, namespace: Namespace) -> Option<(ModuleId, SymbolId)> {
        for (k, v) in &self.symbols {
            if k.name == name && k.namespace == namespace {
                return Some((k.module, *v));
            }
        }
        None
    }
}

impl RefIndex {
    pub fn new() -> Self {
        Self { refs: Vec::new() }
    }

    pub fn ensure_symbol(&mut self, symbol_id: SymbolId) {
        let index = symbol_id.0 as usize;
        if index >= self.refs.len() {
            self.refs.resize_with(index + 1, Vec::new);
        }
    }

    pub fn record_ref(&mut self, symbol_id: SymbolId, span: Span) {
        self.ensure_symbol(symbol_id);
        self.refs[symbol_id.0 as usize].push(span);
    }

    pub fn refs_of(&self, symbol_id: SymbolId) -> &[Span] {
        const EMPTY: [Span; 0] = [];
        self.refs
            .get(symbol_id.0 as usize)
            .map(|refs| refs.as_slice())
            .unwrap_or(&EMPTY)
    }
}

impl NodeToSymbol {
    pub fn new() -> Self {
        Self { map: Vec::new() }
    }

    pub fn bind_node(&mut self, node_id: NodeId, symbol_id: SymbolId) {
        self.ensure(node_id);
        self.map[node_id.0 as usize] = Some(symbol_id);
    }

    pub fn get(&self, node_id: NodeId) -> Option<SymbolId> {
        self.map.get(node_id.0 as usize).and_then(|opt| *opt)
    }

    pub fn ensure(&mut self, node_id: NodeId) {
        let index = node_id.0 as usize;
        if index >= self.map.len() {
            self.map.resize(index + 1, None);
        }
    }

    pub fn resize_to_fit(&mut self, arena: &AstArena) {
        self.map.resize(arena.nodes.len(), None);
    }
}

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

    fn sample_symbol(name: NameId, module: ModuleId) -> Symbol {
        Symbol {
            name,
            kind: SymbolKind::Function,
            exported: false,
            module,
            decl_span: Span::new(FileId(0), 1, 2),
        }
    }

    #[test]
    fn insert_returns_incremental_ids() {
        let mut arena = SymbolArena::new();
        let id0 = arena.insert(sample_symbol(NameId(0), ModuleId(0)));
        let id1 = arena.insert(sample_symbol(NameId(1), ModuleId(0)));

        assert_eq!(id0, SymbolId(0));
        assert_eq!(id1, SymbolId(1));
    }

    #[test]
    fn get_returns_correct_symbol() {
        let mut arena = SymbolArena::new();
        let symbol = sample_symbol(NameId(7), ModuleId(3));
        let id = arena.insert(symbol.clone());

        assert_eq!(arena.get(id).name, symbol.name);
        assert_eq!(arena.get(id).kind, symbol.kind);
        assert_eq!(arena.get(id).module, symbol.module);
        assert_eq!(arena.get(id).decl_span, symbol.decl_span);
    }

    #[test]
    fn def_index_duplicate_in_same_namespace_errors() {
        let mut index = DefIndex::new();
        let key = DefKey {
            module: ModuleId(1),
            name: NameId(10),
            namespace: Namespace::Type,
        };

        assert!(index.insert_symbol(key, SymbolId(0)).is_ok());
        let err = index.insert_symbol(key, SymbolId(1)).unwrap_err();

        assert_eq!(err.code, "E_RESOLVE_DUPLICATE_SYMBOL");
    }

    #[test]
    fn def_index_allows_same_name_in_different_namespace() {
        let mut index = DefIndex::new();
        let name = NameId(11);
        let type_key = DefKey {
            module: ModuleId(2),
            name,
            namespace: Namespace::Type,
        };
        let value_key = DefKey {
            module: ModuleId(2),
            name,
            namespace: Namespace::Value,
        };

        assert!(index.insert_symbol(type_key, SymbolId(0)).is_ok());
        assert!(index.insert_symbol(value_key, SymbolId(1)).is_ok());
        assert_eq!(index.get(type_key), Some(SymbolId(0)));
        assert_eq!(index.get(value_key), Some(SymbolId(1)));
    }

    #[test]
    fn ref_index_records_refs_per_symbol() {
        let mut index = RefIndex::new();
        let span_a1 = Span::new(FileId(0), 1, 2);
        let span_a2 = Span::new(FileId(0), 3, 4);
        let span_b1 = Span::new(FileId(1), 10, 12);

        index.record_ref(SymbolId(2), span_a1.clone());
        index.record_ref(SymbolId(2), span_a2.clone());
        index.record_ref(SymbolId(5), span_b1.clone());

        assert_eq!(index.refs_of(SymbolId(2)), &[span_a1.clone(), span_a2.clone()]);
        assert_eq!(index.refs_of(SymbolId(5)), &[span_b1.clone()]);
        assert!(index.refs_of(SymbolId(9)).is_empty());
    }

    #[test]
    fn node_to_symbol_bind_and_get() {
        let mut map = NodeToSymbol::new();
        let nid = NodeId(10);
        let sid = SymbolId(5);

        map.bind_node(nid, sid);
        assert_eq!(map.get(nid), Some(sid));
        assert_eq!(map.get(NodeId(0)), None);
    }

    #[test]
    fn node_to_symbol_expands_automatically() {
        let mut map = NodeToSymbol::new();
        let nid_high = NodeId(100);
        let sid = SymbolId(1);

        map.bind_node(nid_high, sid);
        assert_eq!(map.get(nid_high), Some(sid));
        assert!(map.map.len() > 100);
    }
}