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

use std::rc::Rc;

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

type FqsnPrefix = Vec<Scope>;

#[derive(Debug)]
enum NameType {
    //TODO eventually this needs to support closures
    Param(u8), //TODO handle implications of functions being limited to 255 params
    LocalVariable,
    Import(Fqsn),
}

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

pub struct ScopeResolver<'a> {
    symbol_table: &'a mut super::SymbolTable,
    lexical_scopes: ScopeStack<'a, Rc<String>, NameType, ScopeType>,
}

impl<'a> ScopeResolver<'a> {
    pub fn new(symbol_table: &'a mut SymbolTable) -> Self {
        let lexical_scopes = ScopeStack::new(None);
        Self {
            symbol_table,
            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 lookup_name_in_scope(&mut self, name: &QualifiedName) {
        let QualifiedName { id, components } = name;
        //TODO handle a "partial" qualified name
        if components.len() == 1 {
            let local_name: Rc<String> = components[0].clone();
            let name_type = self.lexical_scopes.lookup(&local_name);
            match name_type {
                Some(NameType::Import(fqsn)) => {
                    let symbol = self.symbol_table.fqsn_to_symbol.get(&fqsn);
                    if let Some(symbol) = symbol {
                        self.symbol_table.id_to_symbol.insert(id.clone(), symbol.clone());
                    }
                },
                Some(NameType::Param(n)) => {
                    let spec = SymbolSpec::FunctionParam(*n);
                    let fqsn = Fqsn { scopes: vec![Scope::Name(local_name.clone())] };
                    self.symbol_table.add_symbol(id, fqsn, spec);
                }
                Some(NameType::LocalVariable) => {
                    let spec = SymbolSpec::LocalVariable;
                    let fqsn = Fqsn { scopes: vec![Scope::Name(local_name.clone())] };
                    self.symbol_table.add_symbol(id, fqsn, spec);
                },
                None => {
                    //TODO see if I can reduce this duplicate code
                    let fqsn = Fqsn { scopes: vec![Scope::Name(local_name.clone())] };
                    let symbol = self.symbol_table.fqsn_to_symbol.get(&fqsn);
                    if let Some(symbol) = symbol {
                        self.symbol_table.id_to_symbol.insert(id.clone(), symbol.clone());
                    }
                }
            }
        } else {
            let fqsn = Fqsn { scopes: components.iter().map(|name| Scope::Name(name.clone())).collect() };
            let symbol = self.symbol_table.fqsn_to_symbol.get(&fqsn);
            if let Some(symbol) = symbol {
                self.symbol_table.id_to_symbol.insert(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 fqsn in members.into_iter() {
                    self.lexical_scopes.insert( fqsn.local_name(), NameType::Import(fqsn));
                }
            }
            ImportedNames::LastOfPath => {
                let fqsn = Fqsn {
                    scopes: path_components
                    .iter()
                    .map(|c| Scope::Name(c.clone()))
                    .collect()
                };
                self.lexical_scopes.insert(fqsn.local_name(), NameType::Import(fqsn));
            }
            ImportedNames::List(ref names) => {
                let fqsn_prefix: Vec<Scope> = path_components
                    .iter()
                    .map(|c| Scope::Name(c.clone()))
                    .collect();
                for name in names.iter() {
                    let mut scopes = fqsn_prefix.clone();
                    scopes.push(Scope::Name(name.clone()));
                    let fqsn = Fqsn { scopes };
                    self.lexical_scopes.insert(fqsn.local_name(), NameType::Import(fqsn));
                }
            }
        };
        Recursion::Continue
    }

    fn declaration(&mut self, declaration: &Declaration) -> Recursion {
        let is_function_scope = matches!(self.lexical_scopes.get_name(), Some(ScopeType::Function { .. }));
        match declaration {
            Declaration::FuncDecl(signature, block) => {
                let param_names = signature.params.iter().map(|param| param.name.clone());
                //TODO I'm 90% sure this is right, until I get to closures
                //let mut new_scope = self.lexical_scopes.new_scope(Some(ScopeType::Function { name: signature.name.clone() }));
                let mut new_scope = ScopeStack::new(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,
                    lexical_scopes: new_scope,
                };
                walk_block(&mut new_resolver, block);
                Recursion::Stop
            }
            Declaration::Binding { name, .. } if is_function_scope => {
                self.lexical_scopes.insert(name.clone(), NameType::LocalVariable);
                Recursion::Continue
            }
            _ => Recursion::Continue
        }
    }

    fn expression(&mut self, expression: &Expression) -> Recursion {
        use ExpressionKind::*;
        match &expression.kind {
            Value(name) => {
                self.lookup_name_in_scope(name);
            },
            NamedStruct { name, fields: _ } => {
                self.lookup_name_in_scope(name);
            },
            Lambda { params, body, .. } => {
                let param_names = params.iter().map(|param| param.name.clone());
                //TODO need to properly handle closure scope, this is currently broken
                //let mut new_scope = self.lexical_scopes.new_scope(Some(ScopeType::Function { name: signature.name.clone() }));
                let mut new_scope = ScopeStack::new(Some(ScopeType::Lambda));

                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,
                    lexical_scopes: new_scope,
                };
                walk_block(&mut new_resolver, body);
                return Recursion::Stop;
            }
            _ => (),
        }
        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.lookup_name_in_scope(name);
            }
        };
        Recursion::Continue
    }
}