Refactoring: Traits and Partial functions

In this post I was planning to add Citizens to my simulation. I did, and it turned out to be as exciting as watching termites mate.

So instead, in this post I will add Citizens to my simulation, clean up my codebase, factor out all duplicate code using Traits and refactor the act/loop/react by introducing partial functions.

Citizens, first attempt

The pattern for extending my program should be clear be now — I started out by writing a small spec:

"A Citizen" should {
  "be connected to a Harvester and react to Act messages" in {...}
  "still work when not connected to a Harvester" in {...}
}

After writing the full specs, I added the Act, Working and Idle messages and implemented my citizen:

class Citizen(val name: String) {
  var workplace: Actor = _

  def workAt(w: Actor) {
    workplace = w
  }

  def act() {
    loop {
      react {
        case Act if workplace == null => {
          println("[%s] Cannot act, no workplace".format(name))
          reply(Idle)
        }

        case Act => {
          val caller = sender
          println("[%s] Harvesting from %s".format(name, workplace))
          workplace ! Harvest
          react {
            case Harvested(_, _) => {
              caller ! Working
            }
          }
        }

        case Harvested(_, _) => { /* ignore */ }

        case Stop => exit()

        case msg => println("[%s] Ignoring %s".format(name, msg))
      }
    }
  }
}

The code compiled, my specs worked, and I felt thoroughly dissatisfied.

Stop, I've seen this before

When I wrote my case Stop => exit() handler for the third time, it was clear some refactoring was required. As it turns out, every actor in my simulation shares the following features:

• They have a name;
• They respond to Stop messages;
• They have a fallback for unknown messages;

I started out by extracting a common superclass for all my actors, called ScalasimActor:

abstract class ScalasimActor(val name: String) extends Actor {
}

This took care of factoring out the name, as seen for example in my improved Citizen:

class Citizen(name: String) extends ScalasimActor(name) {...}

Note the syntax: name is no longer a val in the Citizen constructor; in the extends clause, the name is passed on to the superclass, which in turn exposes the name using val. This was the easy part.

The difficult part was refactoring the react code. As it turns out, it is only difficult if you don't know how to do it. After reading an excellent post by Jim McBeath on refactoring Scala Actors, it was actually a matter of minutes to do the actual refactoring.

In his post, Jim McBeath explains how the code block after react is actually an anonymous partial function, and he explains how partial functions can be combined by using the orElse operator. To keep it practical, I added the following method to my ScalasimActor:

def handleDefault(): PartialFunction[Any, Unit] = {
  case Stop => exit
  case msg => println("[%s] Ignoring %s".format(name, msg))
}

In all my other actors, I replaced the actor-specific code by a separate function, and I could replace the standard code by calling handleDefault. For example, take a look at the new Producer class:

class Producer(name: String, val goods: String) extends ScalasimActor(name) {
  
  def act {
    loop {
      react (handleProduce orElse handleDefault)
    }
  }

  /** Handler for Produce messages. */
  def handleProduce(): PartialFunction[Any, Unit] = {
    case Produce => {
      val amount = 1
      println("[%s] Producing %d %s".format(name, amount, goods))
      reply(Produced(goods, amount))
    }
  }
}

Impressive. I got rid of a lot of duplicate code, and the readability of the code actually improved. But there is more to come.

Getting rid of the println

I found another sore point in my code: all the println calls were following the same format. However, putting a generic println-like function in the superclass felt... well, unScalaish. So I remembered what I read about Traits, and read even more about it in Jonar Bonér's post on traits and their relation to dependency injection. This was what I was looking for.

trait LoggingActor { self: ScalasimActor =>

  def debug(msg: String, args: Any*) {
    println("[DEBUG] [%s] %s".format(name, msg.format(args: _*)))
  }

  def info(msg: String, args: Any*) {
    println("[INFO]  [%s] %s".format(name, msg.format(args: _*)))
  }

  def warn(msg: String, args: Any*) {
    println("[WARN]  [%s] %s".format(name, msg.format(args: _*)))
  }

  def error(msg: String, args: Any*) {
    println("[ERROR] [%s] %s".format(name, msg.format(args: _*)))
  }

}

This must seem familiar to anyone that ever used a Java logging framework. Four methods that signify different log-levels. As a cute extra, the name of the calling actor is automatically included in the printed message. For now, the code is somewhat naive — there is no way to set the log level, and arguments are always evaluated. However, the code is contained into this single trait, so I will probably extend on this in the near future.

As an aside, please note the self: ScalasimActor => fragment at the beginning of this Trait. It is called a self type, and it ensures that the LoggingActor trait can only be applied to ScalasimActor classes, subclasses and instances. This in turn makes it possible to reference the name, as this field is present in every ScalasimActor.

After a short look at the new ScalasimActor:

abstract class ScalasimActor(val name: String)
  extends Actor
  with LoggingActor {

  /** Handler for Stop messages and unmatched messages. */
  def handleDefault(): PartialFunction[Any, Unit] = {
    case Stop => exit
    case msg => warn("[%s] Ignoring %s".format(name, msg))
  }
}

and one of its subclasses:

class Producer(name: String, val goods: String) extends ScalasimActor(name) {
  
  def act {
    loop {
      react (handleProduce orElse handleDefault)
    }
  }

  /** Handler for Produce messages. */
  def handleProduce(): PartialFunction[Any, Unit] = {
    case Produce => {
      val amount = 1
      info("Producing %d %s", amount, goods)
      reply(Produced(goods, amount))
    }
  }
}

it is clear that my refactorings had the desired effects: a lot of duplicate code has been eliminated, and separation of concerns has taken place. I was happy again.

Lessons learned

Refactoring my Scala-code did not take a lot of effort; actually, it was mostly about reading through different blogs that explained the various concepts. The posts I linked to above were the most helpful, but Googling around quickly turns up lots of interesting Scala blogs.

For posterity, here are the most important lessons:

• I implemented my first superclass;
• I implemented my first trait;
• I learned about visibility scoping in constructor arguments when subclassing;
• I learned about self types;
• And last but not least, I learned about Partial functions, composing them, and using them to keep my Actors lean and mean;

Get the source

As usual, you can find the code at Google Code.