Question

Here is an Expr class from the stairway book.

abstract class Expr
case class Var(name: String) extends Expr
case class Number(num: Double) extends Expr
case class UnOp(operator: String, arg: Expr) extends Expr
case class BinOp(operator: String, left: Expr, right: Expr) extends Expr

Now, I want a function to rename a variable in a expression. Here is my first attempt.

def renameVar(expr: Expr, varName: String, newName: String): Expr = expr match {
    case Var(name) if name == varName => Var(newName)
    case Number(_) => expr
    case UnOp(operator, arg) => UnOp(operator, renameVar(arg, varName, newName))
    case BinOp(operator, left, right) => BinOp(operator, renameVar(left, varName, newName), renameVar(right, varName, newName))
}

val anExpr = BinOp("+", Number(1), Var("x"))
val anExpr2 = renameVar(anExpr, "x", "y")

This works but is tedious (the actual class I am working with has several case subclasses). Also, I might need several similar transformations. Is there a better alternative (possibly using higher order functions)?

Answer 1

So your version of renameVar has to know about two separate things: it has to know how to recurse the tree and it has to know how to rename a variable.

One solution might be to separate these two concerns. You could use the visitor design pattern to give each class control over how it does recursion; the visit method is only concerned with how to traverse the tree. As it traverses, it can pass through a function that handles the actual work (renaming a variable in your case).

Here is a simple implementation that passes a transformation function (that operates on Expr and returns an Expr). The fact that it uses a PartialFunction allows you to pattern-match the expression in the tree to operate on. Any expressions not covered by the cases just fall back to the normal recursion (as specified by doVisit).

Depending on the variety of different tasks, you might need to have a more complex visit method. But this should give you an idea of the direction:

// Class Hierarchy
abstract class Expr {
  def visit(f: PartialFunction[Expr, Expr]): Expr = if (f.isDefinedAt(this)) f(this) else doVisit(f)
  protected def doVisit(f: PartialFunction[Expr, Expr]): Expr
}
case class Var(name: String) extends Expr {
  protected def doVisit(f: PartialFunction[Expr, Expr]) = this
}
case class Number(num: Double) extends Expr {
  protected def doVisit(f: PartialFunction[Expr, Expr]) = this
}
case class UnOp(operator: String, arg: Expr) extends Expr {
  protected def doVisit(f: PartialFunction[Expr, Expr]) = UnOp(operator, arg.visit(f))
}
case class BinOp(operator: String, left: Expr, right: Expr) extends Expr {
  protected def doVisit(f: PartialFunction[Expr, Expr]) = BinOp(operator, left.visit(f), right.visit(f))
}

// Transformation Functions
def renameVar(expr: Expr, varName: String, newName: String): Expr = {
  expr.visit { case Var(`varName`) => Var(newName) }
}

Now you can introduce a new class, like TernaryOp(String, Expr, Expr, Expr), define its doVisit method in a similar way, and it will work without you having to modify renameVar (or any other transformation functions like renameVar).