schala/schala-lang/language/src/symbol_table/resolver.rs

use std::rc::Rc;

use crate::ast::*;
use crate::symbol_table::{Fqsn, Scope, SymbolTable, DefId};
use crate::util::ScopeStack;

type FqsnPrefix = Vec<Scope>;

#[derive(Debug)]
enum NameType {
    //TODO eventually this needs to support closures
    Param(u8), //TODO functions limited to 255 params
    LocalVariable,
    ImportedDefinition(DefId),
}

#[derive(Debug)]
enum ScopeType {
    Function {
        name: Rc<String>
    },
    //TODO add some notion of a let-like scope?
}

pub struct ScopeResolver<'a> {
    symbol_table: &'a mut super::SymbolTable,
    /// Used for import resolution. TODO maybe this can also use the lexical scope table.
    name_scope_stack: ScopeStack<'a, Rc<String>, FqsnPrefix>,
    lexical_scopes: ScopeStack<'a, Rc<String>, NameType, ScopeType>
}

impl<'a> ScopeResolver<'a> {
    pub fn new(symbol_table: &'a mut SymbolTable) -> Self {
        let name_scope_stack = ScopeStack::new(None);
        let lexical_scopes = ScopeStack::new(None);
        Self {
            symbol_table,
            name_scope_stack,
            lexical_scopes,
        }
    }

    pub fn resolve(&mut self, ast: &AST) {
        walk_ast(self, ast);
    }

    fn lookup_name_in_scope(&self, sym_name: &QualifiedName) -> Fqsn {
        let QualifiedName { components, .. } = sym_name;
        let first_component = &components[0];
        match self.name_scope_stack.lookup(first_component) {
            None => Fqsn {
                scopes: components
                    .iter()
                    .map(|name| Scope::Name(name.clone()))
                    .collect(),
            },
            Some(fqsn_prefix) => {
                let mut full_name = fqsn_prefix.clone();
                let rest_of_name: FqsnPrefix = components[1..]
                    .iter()
                    .map(|name| Scope::Name(name.clone()))
                    .collect();
                full_name.extend_from_slice(&rest_of_name);

                Fqsn { scopes: full_name }
            }
        }
    }

    /// This method correctly modifies the id_to_symbol table (ItemId) to have the appropriate
    /// mappings.
    fn handle_qualified_name(&mut self, name: &QualifiedName) {
        println!("Handling qualified_name in resolver.rs: {:?}", name);
        let fqsn = self.lookup_name_in_scope(name);
        let symbol = self.symbol_table.fqsn_to_symbol.get(&fqsn);
        if let Some(symbol) = symbol {
            self.symbol_table.id_to_symbol.insert(name.id.clone(), symbol.clone());
        }
    }
}

impl<'a> ASTVisitor for ScopeResolver<'a> {
    // Import statements bring in a bunch of local names that all map to a specific FQSN.
    // FQSNs map to a Symbol (or this is an error), Symbols have a DefId. So for every
    // name we import, we map a local name (a string) to a NameType::ImportedDefinition(DefId).
    fn import(&mut self, import_spec: &ImportSpecifier) -> Recursion {
        let ImportSpecifier {
            ref path_components,
            ref imported_names,
            ..
        } = &import_spec;
        match imported_names {
            ImportedNames::All => {
                let prefix = Fqsn {
                    scopes: path_components
                        .iter()
                        .map(|c| Scope::Name(c.clone()))
                        .collect(),
                };
                let members = self.symbol_table.symbol_trie.get_children(&prefix);
                for member in members.into_iter() {
                    let Scope::Name(n) = member.scopes.last().unwrap();
                    let local_name = n.clone();
                    self.name_scope_stack.insert(local_name, member.scopes);
                }
            }
            ImportedNames::LastOfPath => {
                let name = path_components.last().unwrap(); //TODO handle better
                let fqsn_prefix = path_components
                    .iter()
                    .map(|c| Scope::Name(c.clone()))
                    .collect();
                self.name_scope_stack.insert(name.clone(), fqsn_prefix);
            }
            ImportedNames::List(ref names) => {
                let fqsn_prefix: FqsnPrefix = path_components
                    .iter()
                    .map(|c| Scope::Name(c.clone()))
                    .collect();
                for name in names.iter() {
                    self.name_scope_stack
                        .insert(name.clone(), fqsn_prefix.clone());
                }
            }
        };
        Recursion::Continue
    }

    fn declaration(&mut self, declaration: &Declaration) -> Recursion {
        if let Declaration::FuncDecl(signature, block) = declaration {
            let param_names = signature.params.iter().map(|param| param.name.clone());
            let mut new_scope = self.lexical_scopes.new_scope(Some(ScopeType::Function { name: signature.name.clone() }));

            for (n, param) in param_names.enumerate().into_iter() {
                new_scope.insert(param, NameType::Param(n as u8));
            }

            let mut new_resolver = ScopeResolver {
                symbol_table: self.symbol_table,
                name_scope_stack: self.name_scope_stack.new_scope(None),
                lexical_scopes: new_scope,
            };
            walk_block(&mut new_resolver, block);
            Recursion::Stop
        } else {
            Recursion::Continue
        }
    }

    fn expression(&mut self, expression: &Expression) -> Recursion {
        use ExpressionKind::*;
        match &expression.kind {
            Value(name) => {
                self.handle_qualified_name(name);
            },
            NamedStruct { name, fields: _ } => {
                self.handle_qualified_name(name);
            },
            _ => (),
        }
        Recursion::Continue
    }

    fn pattern(&mut self, pat: &Pattern) -> Recursion {
        use Pattern::*;

        match pat {
            //TODO I think not handling TuplePattern is an oversight
            TuplePattern(_) => (),
            Literal(_) | Ignored => (),
            TupleStruct(name, _) | Record(name, _) | VarOrName(name) => {
                self.handle_qualified_name(name);
            }
        };
        Recursion::Continue
    }
}