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

use std::rc::Rc;

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

type FqsnPrefix = Vec<Scope>;

pub struct Resolver<'a> {
    symbol_table: &'a mut super::SymbolTable,
    name_scope_stack: ScopeStack<'a, Rc<String>, FqsnPrefix>,
}

impl<'a> Resolver<'a> {
    pub fn new(symbol_table: &'a mut SymbolTable) -> Self {
        let name_scope_stack: ScopeStack<'a, Rc<String>, FqsnPrefix> = ScopeStack::new(None);
        Self {
            symbol_table,
            name_scope_stack,
        }
    }
    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 might be a variable or a pattern, depending on whether this symbol
    // already exists in the table.
    fn qualified_name_in_pattern(&mut self, qualified_name: &QualifiedName) {
        let fqsn = self.lookup_name_in_scope(qualified_name);
        if self.symbol_table.fqsn_to_symbol.get(&fqsn).is_some() {
            self.symbol_table
                .id_to_fqsn
                .insert(qualified_name.id.clone(), fqsn); //TODO maybe set this to an explicit value if none?
        }
    }
}

impl<'a> ASTVisitor for Resolver<'a> {
    //TODO need to un-insert these - maybe need to rethink visitor
    fn import(&mut self, import_spec: &ImportSpecifier) {
        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());
                }
            }
        };
    }

    fn qualified_name(&mut self, qualified_name: &QualifiedName) {
        let fqsn = self.lookup_name_in_scope(qualified_name);
        self.symbol_table
            .id_to_fqsn
            .insert(qualified_name.id.clone(), fqsn);
    }

    fn named_struct(
        &mut self,
        qualified_name: &QualifiedName,
        _fields: &[(Rc<String>, Expression)],
    ) {
        let fqsn = self.lookup_name_in_scope(qualified_name);
        self.symbol_table
            .id_to_fqsn
            .insert(qualified_name.id.clone(), fqsn);
    }

    fn pattern(&mut self, pat: &Pattern) {
        use Pattern::*;
        match pat {
            //TODO I think not handling TuplePattern is an oversight
            TuplePattern(_) => (),
            Literal(_) | Ignored => (),
            TupleStruct(name, _) | Record(name, _) | VarOrName(name) => {
                self.qualified_name_in_pattern(name)
            }
        };
    }
}