ast.scala

package juicy.source.ast

trait Expression extends Visitable
trait Statement extends Visitable

trait BinOp extends Expression {
  val lhs: Expression
  val rhs: Expression

  protected def rewriter[T <: BinOp]
      (ctor: (Expression, Expression) => T)
      (implicit rule: Rewriter) = {
    rule(
      ctor(
        lhs.rewrite.asInstanceOf[Expression],
        rhs.rewrite.asInstanceOf[Expression]
      ))
  }
}

case class Add(lhs: Expression, rhs: Expression) extends BinOp {
  def rewrite(implicit rule: Rewriter) = rewriter(Add.apply _)
}

case class Sub(lhs: Expression, rhs: Expression) extends BinOp {
  def rewrite(implicit rule: Rewriter) = rewriter(Sub.apply _)
}

case class WhileStmnt(
  cond: Expression,
  body: Statement
) extends Statement {
  def rewrite(implicit rule: Rewriter) =
    rule(
      WhileStmnt(
        cond.rewrite.asInstanceOf[Expression],
        body.rewrite.asInstanceOf[Statement]
      ))
}

case class ExprStmnt(
  expr: Expression
) extends Statement {
  def rewrite(implicit rule: Rewriter) =
    rule(
      ExprStmnt(
        expr.rewrite.asInstanceOf[Expression]
      ))
}

callsite.scala

      ast.rewrite(Rewriter {
        case stmnt@ExprStmnt(expr) =>
          WhileStmnt(
            expr,
            stmnt)

        case otherwise => otherwise
      })

visitor.scala

case class Rewriter(rule: Visitable => Visitable) {
  def apply(node: Visitable): Visitable = rule(node)
}


trait Visitable {
  def rewrite(implicit rule: Rewriter): Visitable
}