schala/schala-lang/language/src/scope_resolution.rs

use std::rc::Rc;

use crate::symbol_table::{SymbolTable, ScopeSegment, ScopeSegmentKind, FullyQualifiedSymbolName};
use crate::ast::*;
use crate::util::ScopeStack;

type FQSNPrefix = Vec<ScopeSegment>;

pub struct ScopeResolver<'a> {
  symbol_table: &'a mut SymbolTable,
  name_scope_stack: ScopeStack<'a, Rc<String>, FQSNPrefix>,
}

impl<'a> ScopeResolver<'a> {
  pub fn new(symbol_table: &'a mut SymbolTable) -> ScopeResolver {
    let name_scope_stack = ScopeStack::new(None);
    ScopeResolver { symbol_table, name_scope_stack }
  }
  pub fn resolve(&mut self, ast: &mut AST) -> Result<(), String> {
    for statement in ast.statements.iter() {
      match statement.kind {
        StatementKind::Declaration(ref decl) => self.decl(decl),
        StatementKind::Expression(ref expr) => self.expr(expr),
        StatementKind::Import(ref spec) => self.import(spec),
      }?;
    }
    Ok(())
  }

  fn import(&mut self, import_spec: &ImportSpecifier) -> Result<(), String> {
    let ImportSpecifier { ref path_components, ref imported_names, .. } = &import_spec;
    match imported_names {
      ImportedNames::All => unimplemented!(),
      ImportedNames::LastOfPath => {
        let name = path_components.last().unwrap(); //TODO handle better
        let fqsn_prefix = path_components.iter().map(|c| ScopeSegment {
          name: c.clone(), kind: ScopeSegmentKind::Type
        }).collect();
        self.name_scope_stack.insert(name.clone(), fqsn_prefix);
        ()
      }
      ImportedNames::List(ref names) => {
        let fqsn_prefix: FQSNPrefix = path_components.iter().map(|c| ScopeSegment {
          name: c.clone(), kind: ScopeSegmentKind::Type
        }).collect();
        for name in names.iter() {
          self.name_scope_stack.insert(name.clone(), fqsn_prefix.clone());
        }
      }
    };
    Ok(())
  }

  fn decl(&mut self, decl: &Declaration) -> Result<(), String> {
    use Declaration::*;
    match decl {
      Binding { expr, .. } => self.expr(expr),
      FuncDecl(_, block) => self.block(block),
      _ => Ok(()),
    }
  }
  fn block(&mut self, block: &Block) -> Result<(), String> {
    for statement in block.iter() {
      match statement.kind {
        StatementKind::Declaration(ref decl) => self.decl(decl),
        StatementKind::Expression(ref expr) => self.expr(expr),
        StatementKind::Import(_) => Ok(())
      }?;
    }
    Ok(())
  }

  fn expr(&mut self, expr: &Expression) -> Result<(), String> {
    use ExpressionKind::*;
    match &expr.kind {
      ExpressionKind::Value(qualified_name) => {
        let fqsn = lookup_name_in_scope(&qualified_name);
        let ref id = qualified_name.id;
        self.symbol_table.map_id_to_fqsn(id, fqsn);
      },
      NamedStruct {  name, .. } => {
        let ref id = name.id;
        let fqsn = lookup_name_in_scope(&name);
        self.symbol_table.map_id_to_fqsn(id, fqsn);
      },
      BinExp(_, ref lhs, ref rhs) => {
        self.expr(lhs)?;
        self.expr(rhs)?;
      },
      PrefixExp(_, ref arg) => {
        self.expr(arg)?;
      },
      TupleLiteral(exprs) => {
        for expr in exprs.iter() {
          self.expr(expr)?;
        }
      },
      Call { f, arguments } => {
        self.expr(&f)?;
        for arg in arguments.iter() {
          self.invoc(arg)?;
        }
      },
      Lambda { params, body, .. } => { 
        self.block(&body)?;
        for param in params.iter() {
          if let Some(ref expr) = param.default {
            self.expr(expr)?;
          }
        }
      },
      IfExpression { ref body, ref discriminator } => {
        match &**discriminator {
          Discriminator::Simple(expr) | Discriminator::BinOp(expr, _) => self.expr(expr)?
        };

        match &**body {
          IfExpressionBody::SimplePatternMatch(ref pat, ref alt1, ref alt2) => {
            self.pattern(pat)?;
            self.block(alt1)?;
            if let Some(alt) = alt2 {
              self.block(alt)?;
            }
          },
          IfExpressionBody::GuardList(guardarms) => {
            for arm in guardarms.iter() {
              if let Guard::Pat(ref pat) = arm.guard {
                self.pattern(pat)?;
              }
              self.block(&arm.body)?;
            }
          }
          _ => ()
        }
      },
      _ => ()
    };
    Ok(())
  }

  fn invoc(&mut self, invoc: &InvocationArgument) -> Result<(), String> {
    use InvocationArgument::*;
    match invoc {
      Positional(expr) => self.expr(expr),
      Keyword { expr, .. } => self.expr(expr),
      _ => Ok(())
    }
  }

  fn pattern(&mut self, pat: &Pattern) -> Result<(), String> {
    use Pattern::*;
    match pat {
      Ignored => (),
      TuplePattern(patterns) => {
        for pat in patterns {
          self.pattern(pat)?;
        }
      },
      Literal(_) => (),
      TupleStruct(name, patterns) => {
        self.qualified_name_in_pattern(name);
        for pat in patterns {
          self.pattern(pat)?;
        }
      },
      Record(name, key_patterns) => {
        self.qualified_name_in_pattern(name);
        for (_, pat) in key_patterns {
          self.pattern(pat)?;
        }
      },
      VarOrName(name) => {
        self.qualified_name_in_pattern(name);
      },
    };
    Ok(())
  }

  /// this might be a variable or a pattern. if a variable, set to none
  fn qualified_name_in_pattern(&mut self, qualified_name: &QualifiedName) {
    let ref id = qualified_name.id;
    let fqsn = lookup_name_in_scope(qualified_name);
    if self.symbol_table.lookup_by_fqsn(&fqsn).is_some() {
      self.symbol_table.map_id_to_fqsn(&id, fqsn);
    }
  }
}

//TODO this is incomplete
fn lookup_name_in_scope(sym_name: &QualifiedName) -> FullyQualifiedSymbolName {
  let QualifiedName { components: vec, .. } = sym_name;
  let len = vec.len();
  let new_vec: Vec<ScopeSegment> = vec.iter().enumerate().map(|(i, name)| {
    let kind = if i == (len - 1) {
      ScopeSegmentKind::Terminal
    } else {
      ScopeSegmentKind::Type
    };
    ScopeSegment { name: name.clone(), kind }
  }).collect();
   FullyQualifiedSymbolName(new_vec)
}

#[cfg(test)]
mod tests {
  #[test]
  fn basic_scope() {

  }
}