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>;
type NameScopeStack<'t> = ScopeStack<'t, Rc<String>, FQSNPrefix>;

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> {
    self.block(&mut ast.statements)?;
    Ok(())
  }


  //TODO this is incomplete
  fn lookup_name_in_scope(&self, 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)
  }

  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(ref spec) => self.import(spec),
      }?;
    }
    Ok(())
  }

  fn expr(&mut self, expr: &Expression) -> Result<(), String> {
    use ExpressionKind::*;
    match &expr.kind {
      ExpressionKind::Value(qualified_name) => {
        let fqsn = self.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 = self.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)?;
        }
      },
      IfExpression { discriminator, body } => {
        match discriminator {
          None => (),
          Some(ref expr) => self.expr(expr)?
        };
        self.if_expr_body(body)?;
      },
      Lambda { params, body, .. } => {
        self.block(&body)?;
        for param in params.iter() {
          if let Some(ref expr) = param.default {
            self.expr(expr)?;
          }
        }
      },
      _ => ()
    };
    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 if_expr_body(&mut self, body: &IfExpressionBody) -> Result<(), String> {
    use IfExpressionBody::*;
    match body {
      SimpleConditional { then_case, else_case } => {
        self.block(then_case)?;
        if let Some(block) = else_case {
          self.block(block)?;
        }
      },
      SimplePatternMatch { pattern, then_case, else_case } => {
        self.pattern(pattern)?;
        self.block(then_case)?;
        if let Some(block) = else_case {
          self.block(block)?;
        }
      },
      CondList(conditions) => {
        for arm in conditions.iter() {
          if let Condition::Pattern(ref pat) = arm.condition {
            self.pattern(pat)?;
          }
          self.block(&arm.body)?;
        }
      },
    }
    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 = self.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);
    }
  }
}

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

  }
}