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

use std::rc::Rc;

use crate::symbol_table::{SymbolTable, FQSN, Scope};
use crate::ast::*;
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: &Vec<(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),
    };
  }
}