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

use std::collections::{HashMap, hash_map::Entry};
use std::rc::Rc;
use std::fmt;

use crate::tokenizing::Location;
use crate::ast;
use crate::ast::{ItemId, TypeBody, Variant, TypeSingletonName, Declaration, Statement, StatementKind, ModuleSpecifier};
use crate::typechecking::TypeName;

mod resolver;
mod symbol_trie;
use symbol_trie::SymbolTrie;
mod test;

/// Fully-qualified symbol name
#[derive(Debug, Clone, Eq, PartialEq, Hash, PartialOrd, Ord)]
pub struct Fqsn {
  //TODO Fqsn's need to be cheaply cloneable
  scopes: Vec<Scope>, //TODO rename to ScopeSegment
}

impl Fqsn {
  fn from_scope_stack(scopes: &[Scope], new_name: Rc<String>) -> Self {
    let mut v = Vec::new();
    for s in scopes {
      v.push(s.clone());
    }
    v.push(Scope::Name(new_name));
    Fqsn { scopes: v }
  }

  #[cfg(test)]
  fn from_strs(strs: &[&str]) -> Fqsn {
    let mut scopes = vec![];
    for s in strs {
      scopes.push(Scope::Name(Rc::new(s.to_string())));
    }
    Fqsn {
      scopes
    }
  }
}



//TODO eventually this should use ItemId's to avoid String-cloning
/// One segment within a scope.
#[derive(Debug, Clone, Eq, PartialEq, Hash, PartialOrd, Ord)]
enum Scope {
  Name(Rc<String>)
}

#[allow(dead_code)]
#[derive(Debug, Clone)]
pub enum SymbolError {
  DuplicateName {
    prev_name: Fqsn,
    location: Location
  },
  DuplicateRecord {
    type_name: Fqsn,
    location: Location,
    member: String,
  }
}

#[allow(dead_code)]
#[derive(Debug)]
struct NameSpec<K> {
  location: Location,
  kind: K
}

#[derive(Debug)]
enum NameKind {
  Module,
  Function,
  Binding,
}

#[derive(Debug)]
struct TypeKind;

/// Keeps track of what names were used in a given namespace.
struct NameTable<K> {
  table: HashMap<Fqsn, NameSpec<K>>
}

impl<K> NameTable<K> {
  fn new() -> Self {
    Self { table: HashMap::new() }
  }

  fn register(&mut self, name: Fqsn, spec: NameSpec<K>) -> Result<(), SymbolError> {
    match self.table.entry(name.clone()) {
      Entry::Occupied(o) => {
        Err(SymbolError::DuplicateName { prev_name: name, location: o.get().location })
      },
      Entry::Vacant(v) => {
        v.insert(spec);
        Ok(())
      }
    }
  }
}

//cf. p. 150 or so of Language Implementation Patterns
pub struct SymbolTable {
  /// Used for import resolution.
  symbol_trie: SymbolTrie,

  /// These tables are responsible for preventing duplicate names.
  fq_names: NameTable<NameKind>, //Note that presence of two tables implies that a type and other binding with the same name can co-exist
  types: NameTable<TypeKind>,

  /// A map of the `ItemId`s of instances of use of names to their fully-canonicalized Fqsn form.
  /// Updated by the item id resolver.
  id_to_fqsn: HashMap<ItemId, Fqsn>,

  /// A map of the Fqsn of an AST definition to a Symbol data structure, which contains
  /// some basic information about what that symbol is and (ideally) references to other tables
  /// (e.g. typechecking tables) with more information about that symbol.
  fqsn_to_symbol: HashMap<Fqsn, Symbol>,
}

impl SymbolTable {
  pub fn new() -> SymbolTable {
    SymbolTable {
      symbol_trie: SymbolTrie::new(),
      fq_names: NameTable::new(),
      types: NameTable::new(),
      id_to_fqsn: HashMap::new(),
      fqsn_to_symbol: HashMap::new(),
    }
  }

  /// The main entry point into the symbol table. This will traverse the AST in several
  /// different ways and populate subtables with information that will be used further in the
  /// compilation process.
  pub fn process_ast(&mut self, ast: &ast::AST) -> Result<(), Vec<SymbolError>> {

    let errs = self.populate_name_tables(ast);
    if !errs.is_empty() {
      return Err(errs);
    }
    self.resolve_symbol_ids(ast);
    Ok(())
  }

  pub fn lookup_symbol(&self, id: &ItemId) -> Option<&Symbol> {
    let fqsn = self.id_to_fqsn.get(id);
    fqsn.and_then(|fqsn| self.fqsn_to_symbol.get(fqsn))
  }
}

#[allow(dead_code)]
#[derive(Debug)]
pub struct Symbol {
  pub local_name: Rc<String>,
  //fully_qualified_name: FullyQualifiedSymbolName,
  pub spec: SymbolSpec,
}

impl fmt::Display for Symbol {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    write!(f, "<Local name: {}, Spec: {}>", self.local_name, self.spec)
  }
}

#[derive(Debug)]
pub enum SymbolSpec {
  Func(Vec<TypeName>),
  DataConstructor {
    index: usize,
    arity: usize,
    type_name: TypeName, //TODO this eventually needs to be some kind of ID
  },
  RecordConstructor {
    index: usize,
    members: HashMap<Rc<String>, TypeName>,
    type_name: TypeName,
  },
  Binding,
}

impl fmt::Display for SymbolSpec {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    use self::SymbolSpec::*;
    match self {
      Func(type_names) => write!(f, "Func({:?})", type_names),
      DataConstructor { index, type_name, arity } => write!(f, "DataConstructor(idx: {}, arity: {}, type: {})", index, arity, type_name),
      RecordConstructor { type_name, index, ..} => write!(f, "RecordConstructor(idx: {})(<members> -> {})", index, type_name),
      Binding => write!(f, "Binding"),
    }
  }
}


impl SymbolTable {
  /* note: this adds names for *forward reference* but doesn't actually create any types. solve that problem
   * later */

  /// Register a new mapping of a fully-qualified symbol name (e.g. `Option::Some`)
  /// to a Symbol, a descriptor of what that name refers to.
  fn add_symbol(&mut self, fqsn: Fqsn, symbol: Symbol) {
    self.symbol_trie.insert(&fqsn);
    self.fqsn_to_symbol.insert(fqsn, symbol);
  }

  /// Walks the AST, matching the ID of an identifier used in some expression to
  /// the corresponding Symbol.
  fn resolve_symbol_ids(&mut self, ast: &ast::AST) {
    let mut resolver = resolver::Resolver::new(self);
    resolver.resolve(ast);
  }

  /// This function traverses the AST and adds symbol table entries for
  /// constants, functions, types, and modules defined within. This simultaneously
  /// checks for dupicate definitions (and returns errors if discovered), and sets
  /// up name tables that will be used by further parts of the compiler
  fn populate_name_tables(&mut self, ast: &ast::AST) -> Vec<SymbolError> {
    let mut scope_stack = vec![];
    self.add_from_scope(ast.statements.as_ref(), &mut scope_stack)
  }

  fn add_from_scope<'a>(&'a mut self, statements: &[Statement], scope_stack: &mut Vec<Scope>) -> Vec<SymbolError> {
    let mut errors = vec![];

    for statement in statements {
      let Statement { id: _, kind, location } = statement;  //TODO I'm not sure if I need to do anything with this ID
      let location = *location;
      if let Err(err) = self.add_single_statement(kind, location, scope_stack) {
        errors.push(err);
      } else { // If there's an error with a name, don't recurse into subscopes of that name
        let recursive_errs = match kind {
          StatementKind::Declaration(Declaration::FuncDecl(signature, body)) => {
            let new_scope = Scope::Name(signature.name.clone());
            scope_stack.push(new_scope);
            let output = self.add_from_scope(body.as_ref(), scope_stack);
            scope_stack.pop();
            output
          }
          StatementKind::Module(ModuleSpecifier { name, contents }) => {
            let new_scope = Scope::Name(name.clone());
            scope_stack.push(new_scope);
            let output = self.add_from_scope(contents.as_ref(), scope_stack);
            scope_stack.pop();
            output
          }
          StatementKind::Declaration(Declaration::TypeDecl { name, body, mutable }) => {
            self.add_type_members(name, body, mutable, location, scope_stack)
          }
          _ => vec![]
        };
        errors.extend(recursive_errs.into_iter());
      }
    }

    errors
  }

  fn add_single_statement(&mut self, kind: &StatementKind, location: Location, scope_stack: &[Scope]) -> Result<(), SymbolError> {
      match kind {
        StatementKind::Declaration(Declaration::FuncSig(signature)) => {
          let fq_function = Fqsn::from_scope_stack(scope_stack, signature.name.clone());
          self.fq_names.register(fq_function.clone(), NameSpec { location, kind: NameKind::Function })?;
          self.types.register(fq_function.clone(), NameSpec { location, kind: TypeKind } )?;

          self.add_symbol(fq_function, Symbol {
            local_name: signature.name.clone(),
            spec: SymbolSpec::Func(vec![]), //TODO does this inner vec need to exist at all?
          });
        }
        StatementKind::Declaration(Declaration::FuncDecl(signature, ..)) => {
          let fn_name = &signature.name;
          let fq_function = Fqsn::from_scope_stack(scope_stack, fn_name.clone());
          self.fq_names.register(fq_function.clone(), NameSpec { location, kind: NameKind::Function })?;
          self.types.register(fq_function.clone(), NameSpec { location, kind: TypeKind } )?;

          self.add_symbol(fq_function, Symbol {
            local_name: signature.name.clone(),
            spec: SymbolSpec::Func(vec![]), //TODO does this inner vec need to exist at all?
          });
        },
        StatementKind::Declaration(Declaration::TypeDecl { name, .. }) => {
          let fq_type = Fqsn::from_scope_stack(scope_stack, name.name.clone());
          self.types.register(fq_type, NameSpec { location, kind: TypeKind } )?;
        },
        StatementKind::Declaration(Declaration::Binding { name, .. }) => {
          let fq_binding = Fqsn::from_scope_stack(scope_stack, name.clone());
          self.fq_names.register(fq_binding.clone(), NameSpec { location, kind: NameKind::Binding })?;
          self.add_symbol(fq_binding, Symbol {
            local_name: name.clone(),
            spec: SymbolSpec::Binding,
          });
        }
        StatementKind::Module(ModuleSpecifier { name, .. }) => {
          let fq_module = Fqsn::from_scope_stack(scope_stack, name.clone());
          self.fq_names.register(fq_module, NameSpec { location, kind: NameKind::Module })?;
        },
        _ => (),
      }
      Ok(())
  }

  fn add_type_members(&mut self, type_name: &TypeSingletonName, type_body: &TypeBody, _mutable: &bool, location: Location, scope_stack: &mut Vec<Scope>) -> Vec<SymbolError> {
    let mut member_errors = vec![];
    let mut errors = vec![];

    let mut register = |fqsn: Fqsn, spec: SymbolSpec| {
      let name_spec = NameSpec { location, kind: TypeKind };
      if let Err(err) = self.types.register(fqsn.clone(), name_spec) {
        errors.push(err);
      } else {
        let local_name = match spec {
          SymbolSpec::DataConstructor { ref type_name, ..} | SymbolSpec::RecordConstructor { ref type_name, .. } => type_name.clone(),
          _ => panic!("This should never happen"),
        };
        let symbol = Symbol { local_name, spec };
        self.add_symbol(fqsn, symbol);
      };
    };

    let TypeBody(variants) = type_body;
    let new_scope = Scope::Name(type_name.name.clone());
    scope_stack.push(new_scope);

    for (index, variant) in variants.iter().enumerate() {
      match variant {
        Variant::UnitStruct(name) => {
          let fq_name = Fqsn::from_scope_stack(scope_stack.as_ref(), name.clone());
          let spec = SymbolSpec::DataConstructor {
            index,
            arity: 0,
            type_name: name.clone(),
          };
          register(fq_name, spec);
        },
        Variant::TupleStruct(name, items) => {
          let fq_name = Fqsn::from_scope_stack(scope_stack.as_ref(), name.clone());
          let spec = SymbolSpec::DataConstructor {
            index,
            arity: items.len(),
            type_name: name.clone(),
          };
          register(fq_name, spec);
        },
        Variant::Record { name, members } => {
          let fq_name = Fqsn::from_scope_stack(scope_stack.as_ref(), name.clone());

          let mut seen_members = HashMap::new();
          for (member_name, _) in members.iter() {
            match seen_members.entry(member_name.as_ref()) {
              Entry::Occupied(o) => {
                let location = *o.get();
                member_errors.push(SymbolError::DuplicateRecord {
                  type_name: fq_name.clone(),
                  location,
                  member: member_name.as_ref().to_string(),
                });
              },
              //TODO eventually this should track meaningful locations
              Entry::Vacant(v) => { v.insert(Location::default()); }
            }
          }

          let spec = SymbolSpec::RecordConstructor {
            index,
            type_name: name.clone(),
            members: members.iter()
              .map(|(_, _)| (Rc::new("DUMMY_FIELD".to_string()), Rc::new("DUMMY_TYPE_ID".to_string()))).collect()
          };
          register(fq_name, spec);
        }
      }
    }

    scope_stack.pop();
    errors.extend(member_errors.into_iter());
    errors
  }
}