schala/schala-lang/language/src/ast/visitor.rs

use std::rc::Rc;
use crate::ast::*;
use crate::util::deref_optional_box;

//TODO maybe these functions should take closures that return a KeepRecursing | StopHere type,
//or a tuple of (T, <that type>)

pub trait ASTVisitor: Sized {
  fn ast(&mut self, _ast: &AST) {}
  fn block(&mut self, _statements: &Vec<Statement>) {}
  fn statement(&mut self, _statement: &Statement) {}
  fn declaration(&mut self, _declaration: &Declaration) {}
  fn signature(&mut self, _signature: &Signature) {}
  fn type_declaration(&mut self, _name: &TypeSingletonName, _body: &TypeBody, _mutable: bool) {}
  fn type_alias(&mut self, _alias: &Rc<String>, _original: &Rc<String>) {}
  fn binding(&mut self, _name: &Rc<String>, _constant: bool, _type_anno: Option<&TypeIdentifier>, _expr: &Expression) {}
  fn implemention(&mut self, _type_name: &TypeIdentifier, _interface_name: Option<&TypeSingletonName>, _block: &Vec<Declaration>) {}
  fn interface(&mut self, _name: &Rc<String>, _signatures: &Vec<Signature>) {}
  fn expression(&mut self, _expression: &Expression) {}
  fn expression_kind(&mut self, _kind: &ExpressionKind) {}
  fn type_annotation(&mut self, _type_anno: Option<&TypeIdentifier>) {}
  fn named_struct(&mut self, _name: &QualifiedName, _fields: &Vec<(Rc<String>, Expression)>) {}
  fn call(&mut self, _f: &Expression, _arguments: &Vec<InvocationArgument>) {}
  fn index(&mut self, _indexee: &Expression, _indexers: &Vec<Expression>) {}
  fn if_expression(&mut self, _discrim: Option<&Expression>, _body: &IfExpressionBody) {}
  fn condition_arm(&mut self, _arm: &ConditionArm) {}
  fn while_expression(&mut self, _condition: Option<&Expression>, _body: &Block) {}
  fn for_expression(&mut self, _enumerators: &Vec<Enumerator>, _body: &ForBody) {}
  fn lambda(&mut self, _params: &Vec<FormalParam>, _type_anno: Option<&TypeIdentifier>, _body: &Block) {}
  fn invocation_argument(&mut self, _arg: &InvocationArgument) {}
  fn formal_param(&mut self, _param: &FormalParam) {}
  fn import(&mut self, _import: &ImportSpecifier) {}
  fn module(&mut self, _module: &ModuleSpecifier) {}
  fn qualified_name(&mut self, _name: &QualifiedName) {}
  fn nat_literal(&mut self, _n: u64) {}
  fn float_literal(&mut self, _f: f64) {}
  fn string_literal(&mut self, _s: &Rc<String>) {}
  fn bool_literal(&mut self, _b: bool) {}
  fn binexp(&mut self, _op: &BinOp, _lhs: &Expression, _rhs: &Expression) {}
  fn prefix_exp(&mut self, _op: &PrefixOp, _arg: &Expression) {}
  fn pattern(&mut self, _pat: &Pattern) {}
}

pub trait ASTVisitorNew: Sized {
  fn expression(&mut self, _expression: &Expression) -> VisitResult { VisitResult::default() }
  fn expression_kind(&mut self, _kind: &ExpressionKind) -> VisitResult { VisitResult::default() }

  fn type_annotation(&mut self, _id: ItemId, _type_anno: Option<&TypeIdentifier>) -> VisitResult { VisitResult::default() }

  fn named_struct(&mut self, _name: &QualifiedName, _fields: &Vec<(Rc<String>, Expression)>) -> VisitResult { VisitResult::default() }

  fn nat_literal(&mut self, _id: ItemId, _n: u64)-> VisitResult { VisitResult::default() }
  fn float_literal(&mut self, _id: ItemId,  _f: f64) -> VisitResult { VisitResult::default() }
  fn string_literal(&mut self, _id: ItemId,  _s: &Rc<String>) -> VisitResult { VisitResult::default() }
  fn bool_literal(&mut self, _id: ItemId, _b: bool) -> VisitResult { VisitResult::default() }

  fn qualified_name(&mut self, _id: ItemId, _name: &QualifiedName) {}
}

enum VisitResult {
  KeepRecursing,
  Stop
}

impl std::default::Default for VisitResult {
  fn default() -> VisitResult { VisitResult::KeepRecursing }
}

pub fn traverse<V: ASTVisitorNew>(visitor: &mut V, expression: &Expression) {
  expresion(visitor, expresion)
}

fn expression<V: ASTVisitorNew>(visitor: &mut V, expression: &Expression) {
  use ExpressionKind::*;

  let Expression { id, kind, type_anno } = expression;

  let output = match kind {
    NatLiteral(n) => v.nat_literal(id, *n),
    FloatLiteral(f) => v.float_literal(id, *f),
    StringLiteral(s) => v.string_literal(id, s),
    BoolLiteral(b) => v.bool_literal(d, *b),
    BinExp(op, lhs, rhs) => {
      expression(v, lhs);
      expression(v, rhs);
    },
    PrefixExp(op, arg) => {
      expression(v, arg);
    }
    TupleLiteral(exprs) => {
      for expr in exprs {
        expression(v, expr);
      }
    },
    Value(name) => v.qualified_name(id, name),
    NamedStruct { name, fields } => {
      for (_, expr) in fields.iter() {
        v.expression(expr);
      }
    }
    Call { f, arguments } => {
      expression(v, f);
      for arg in args.iter() {
        match arg {
          InvocationArgument::Positional(expr) => {
            expression(v, expr);
          },
          InvocationArgument::Keyword { expr, .. } => {
            expression(v, expr);
          },
          Ignored => (),
        }
      }
    },
    Index { indexee, indexers } => {
      expression(v, indexee);
      for i in indexers.iter() {
        expression(v, i);
      }
    },
    IfExpression { discriminator, body } => {
      discriminator.as_ref().map(|d| expression(v, d));
      match body {
        IfExpressionBody::SimpleConditional { then_case, else_case } => {
          walk_block(v, then_case);
          else_case.as_ref().map(|block| walk_block(v, block));
        },
        IfExpressionBody::SimplePatternMatch { pattern, then_case, else_case } => {
          v.pattern(pattern);
          walk_pattern(v, pattern);
          walk_block(v, then_case);
          else_case.as_ref().map(|block| walk_block(v, block));
        },
        IfExpressionBody::CondList(arms) => {
          for arm in arms {
            v.condition_arm(arm);
            condition_arm(v, arm);
          }
        }
      }
    },
    WhileExpression { condition, body } => (),
    ForExpression { enumerators, body } => (),
    Lambda { params , type_anno, body } => {
      for param in params {
        formal_param(v, param);
      }
      v.type_annotation(type_anno);
    },
    ListLiteral(exprs) => {
      for expr in exprs {
        expression(v, expr);
      }
    },
  };
  output
}

fn formal_param<V: ASTVisitorNew>(v: &mut V, param: &FormalParam) {
  param.default.as_ref().map(|p| {
    expression(v, p);
  });
  v.type_annotation(param.anno.as_ref());
}