Attempt to do AST Transformation

ast_transformation.scala

#!/bin/sh
  exec scala "$0" "$@"
!#

/*
  Subset of our Simple Java AST 
*/

sealed abstract class SJStatement
case class SJAssert (assertion : SJExpression) extends SJStatement
case class SJWhile (test : SJExpression, body : List[SJStatement]) extends SJStatement
case class SJConditional (test : SJExpression, consequent : List[SJStatement], alternative : List[SJStatement]) extends SJStatement
case class SJAssignment (left : SJVariableAccess, right : SJExpression) extends SJStatement 
case class SJFieldWrite (variable : SJVariableAccess, field : String, value : SJExpression) extends SJStatement 
case class SJFieldRead (value : SJVariableAccess, variable : SJVariableAccess, field : String) extends SJStatement
case class SJReturn (ret : SJExpression) extends SJStatement 
case class SJCall (value : Option[SJVariableAccess], receiver : SJExpression, fun : String, arguments : List[SJExpression]) extends SJStatement 
case class SJNewExpression (value : SJVariableAccess, jtype : String, arguments : List[SJExpression]) extends SJStatement 

trait SJExpression extends SJStatement
case class SJBinaryExpression (operation : String, left : SJExpression, right : SJExpression) extends SJExpression 
case class SJUnaryExpression (operation : String, expr : SJExpression) extends SJExpression 
case class SJLiteral (value : String) extends SJExpression 
case class SJVariableAccess (variable : String) extends SJExpression

/*
  The transform methods
*/

// http://apocalisp.wordpress.com/2010/07/02/higher-rank-polymorphism-in-scala/

type Id[A] = A

trait ~>[F[_],G[_], -UpperBound] {
  def apply[A <: UpperBound](a: F[A]): G[A]
}

def transformExpr(expressions: List[SJExpression], f: ~>[Id,Id,SJExpression]): List[SJExpression] = {
  expressions map { _ match {
    case SJBinaryExpression(op,l,r)    => f(SJBinaryExpression(op, f(l),f(r)))
    case SJUnaryExpression (op,expr)   => f(SJUnaryExpression(op,f(expr))) 
    case a@SJLiteral(_)                => f(a)
    case a@SJVariableAccess(_)         => f(a)
    case x => f(x)
  }}
}

def transform(statements: List[SJStatement], f: ~>[Id,Id,SJStatement]): List[SJStatement] = {
  statements map { _ match {
    case SJAssert(a)                   => f(SJAssert(f(a)))
    case SJWhile(test, body)           => f(SJWhile(f(test),transform(body,f)))
    case SJConditional(test, c, a)     => f(SJConditional(f(test), transform(c,f), transform(a,f)))
    case SJAssignment(l,r)             => f(SJAssignment(f(l),f(r)))
    case SJFieldWrite(v, field, v2)    => f(SJFieldWrite(f(v),field,f(v2)))
    case SJFieldRead(v, v2, field)     => f(SJFieldRead(f(v),f(v2),field))
    case SJReturn(ret)                 => f(SJReturn(f(ret)))
    case SJCall(x,rec,fun,args)        => f(SJCall(x.map(f(_)),f(rec),fun,transformExpr(args,f)))
    case SJNewExpression(v, typ, args) => f(SJNewExpression(f(v),typ,transformExpr(args,f)))    
    case x: SJStatement => f(x)
  }}
}


val program = List(
   SJConditional(SJBinaryExpression(">=", SJVariableAccess("n"), SJLiteral("0")),
      List(SJCall(Some(SJVariableAccess("tmp_1")), SJVariableAccess("this"), "fac",
                            List(SJBinaryExpression("-", SJVariableAccess("n"), SJLiteral("1")))),
           SJAssignment(SJVariableAccess("x"), SJBinaryExpression("*", SJVariableAccess("n"), SJVariableAccess("tmp_1")))),
      List(SJAssignment(SJVariableAccess("x"), SJLiteral("1")))),
   SJReturn(SJVariableAccess("x")))

val rewritten = transform(program, new (~>[Id,Id,SJStatement]) {
  def apply[A <: SJStatement](a: A): A = a match {
    case SJVariableAccess("tmp_1") => SJVariableAccess("x") 
    case x: SJStatement => x
  }
})

println(rewritten)