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

use std::rc::Rc;

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

#[derive(Debug)]
enum NameType {
    //TODO eventually this needs to support closures
    Param(u8),
    LocalVariable(ItemId),
    LocalFunction(ItemId),
    Import(Fqsn),
}

type LexScope<'a> = ScopeStack<'a, Rc<String>, NameType, ScopeType>;

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

pub struct ScopeResolver<'a> {
    symbol_table: &'a mut super::SymbolTable,
    //TODO maybe this shouldn't be a scope stack, b/c the recursion behavior comes from multiple
    //instances of ScopeResolver
    lexical_scopes: LexScope<'a>,
}

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);
    }

    /// This method correctly modifies the id_to_def table (ItemId) to have the appropriate
    /// mappings.
    fn lookup_name_in_scope(&mut self, name: &QualifiedName) {
        let QualifiedName { id, components } = name;

        let local_name = components.first().unwrap().clone();
        let name_type = self.lexical_scopes.lookup(&local_name);
        let fqsn = Fqsn { scopes: components.iter().map(|name| ScopeSegment::Name(name.clone())).collect() };
        let def_id = self.symbol_table.symbol_trie.lookup(&fqsn);

        //TODO handle a "partial" qualified name, and also handle it down in the pattern-matching
        //section
        if components.len() == 1 {
            match name_type {
                Some(NameType::Import(fqsn)) => {
                    let def_id = self.symbol_table.symbol_trie.lookup(fqsn);

                    if let Some(def_id) = def_id {
                        self.symbol_table.id_to_def.insert(*id, def_id);
                    }
                }
                Some(NameType::Param(n)) => {
                    let spec = SymbolSpec::FunctionParam(*n);
                    //TODO need to come up with a better solution for local variable FQSNs
                    let lscope = ScopeSegment::Name(Rc::new("<local-param>".to_string()));
                    let fqsn = Fqsn { scopes: vec![lscope, ScopeSegment::Name(local_name.clone())] };
                    self.symbol_table.add_symbol(id, fqsn, spec);
                }
                Some(NameType::LocalFunction(item_id)) => {
                    let def_id = self.symbol_table.id_to_def.get(item_id);
                    if let Some(def_id) = def_id {
                        let def_id = *def_id;
                        self.symbol_table.id_to_def.insert(*id, def_id);
                    }
                }
                Some(NameType::LocalVariable(item_id)) => {
                    let def_id = self.symbol_table.id_to_def.get(item_id);
                    if let Some(def_id) = def_id {
                        let def_id = *def_id;
                        self.symbol_table.id_to_def.insert(*id, def_id);
                    }
                }
                None =>
                    if let Some(def_id) = def_id {
                        self.symbol_table.id_to_def.insert(*id, def_id);
                    },
            }
        } else if let Some(def_id) = def_id {
            self.symbol_table.id_to_def.insert(*id, def_id);
        }
    }
}

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| ScopeSegment::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.last_elem(), NameType::Import(fqsn));
                }
            }
            ImportedNames::LastOfPath => {
                let fqsn =
                    Fqsn { scopes: path_components.iter().map(|c| ScopeSegment::Name(c.clone())).collect() };
                self.lexical_scopes.insert(fqsn.last_elem(), NameType::Import(fqsn));
            }
            ImportedNames::List(ref names) => {
                let fqsn_prefix: Vec<ScopeSegment> =
                    path_components.iter().map(|c| ScopeSegment::Name(c.clone())).collect();
                for name in names.iter() {
                    let mut scopes = fqsn_prefix.clone();
                    scopes.push(ScopeSegment::Name(name.clone()));
                    let fqsn = Fqsn { scopes };
                    self.lexical_scopes.insert(fqsn.last_elem(), NameType::Import(fqsn));
                }
            }
        };
        Recursion::Continue
    }

    fn declaration(&mut self, declaration: &Declaration, id: &ItemId) -> Recursion {
        let cur_function_name = match self.lexical_scopes.get_name() {
            //TODO this needs to be a fqsn
            Some(ScopeType::Function { name }) => Some(name.clone()),
            _ => None,
        };
        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() }));
                //TODO this will recurse unwantedly into scopes; need to pop an outer function
                //scope off first before going into a non-closure scope
                let mut new_scope =
                    ScopeStack::new(Some(ScopeType::Function { name: signature.name.clone() }));

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

                self.lexical_scopes.insert(signature.name.clone(), NameType::LocalFunction(*id));

                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 let Some(fn_name) = cur_function_name {
                    // We are within a function scope
                    let fqsn =
                        Fqsn { scopes: vec![ScopeSegment::Name(fn_name), ScopeSegment::Name(name.clone())] };
                    self.symbol_table.add_symbol(id, fqsn, SymbolSpec::LocalVariable);
                    self.lexical_scopes.insert(name.clone(), NameType::LocalVariable(*id));
                }
                Recursion::Continue
            }
            Declaration::Impl { block, .. } => {
                let mut new_scope = ScopeStack::new(Some(ScopeType::ImplBlock));
                let mut new_resolver =
                    ScopeResolver { symbol_table: self.symbol_table, lexical_scopes: new_scope };
                for stmt in block.iter() {
                    walk_declaration(&mut new_resolver, &stmt.kind, &stmt.id);
                }
                Recursion::Stop
            }
            _ => 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() {
                    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;
            }
            IfExpression { discriminator, body } => {
                if let Some(d) = discriminator.as_ref() {
                    walk_expression(self, d);
                }
                let mut resolver = ScopeResolver {
                    lexical_scopes: self.lexical_scopes.new_scope(Some(ScopeType::PatternMatch)),
                    symbol_table: self.symbol_table,
                };
                walk_if_expr_body(&mut resolver, body);
                return Recursion::Stop;
            }
            _ => (),
        }
        Recursion::Continue
    }

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

        match pat {
            Literal(..) | Ignored | TuplePattern(..) => (),
            TupleStruct(name, _) | Record(name, _) => {
                self.lookup_name_in_scope(name);
            }
            //TODO this isn't really the right syntax for a VarOrName
            VarOrName(QualifiedName { id, components }) => {
                if components.len() == 1 {
                    //TODO need a better way to construct a FQSN from a QualifiedName
                    let local_name: Rc<String> = components[0].clone();
                    let lscope = ScopeSegment::Name(Rc::new("<local-case-match>".to_string()));
                    let fqsn = Fqsn { scopes: vec![lscope, ScopeSegment::Name(local_name.clone())] };
                    self.symbol_table.add_symbol(id, fqsn, SymbolSpec::LocalVariable);
                    self.lexical_scopes.insert(local_name, NameType::LocalVariable(*id));
                } else {
                    let fqsn = Fqsn {
                        scopes: components.iter().map(|name| ScopeSegment::Name(name.clone())).collect(),
                    };
                    let def_id = self.symbol_table.symbol_trie.lookup(&fqsn);

                    if let Some(def_id) = def_id {
                        self.symbol_table.id_to_def.insert(*id, def_id);
                    }
                }
            }
        };
        Recursion::Continue
    }
}