Ambassador to the Computers

Logic programming in Scala, part 3: unification and state

2011-06-08T20:41:00.000-07:00

In this post I want to build on the backtracking logic monad we covered last time by adding unification, yielding an embedded DSL for Prolog-style logic programming.

Prolog

Here is a small Prolog example, the rough equivalent of List.contains in Scala:

  member(X, [X|T]). 
  member(X, [_|T]) :- member(X, T).

Member doesn’t return a boolean; instead it succeeds or fails (in the same way as the logic monad). The goal member(1, [1,2,3]) succeeds; the goal member(4, [1,2,3]) fails. (What happens for member(1, [1,1,3])?)

A Prolog predicate is defined by one or more clauses (each ending in a period), made up of a head (the predicate and arguments before the :-) and zero or more subgoals (goals after the :-, separated by commas; if there are no subgoals the :- is omitted). To solve a goal, we unify it (match it) with each clause head, then solve each subgoal in the clause. If a subgoal fails we backtrack and try the next matching head; if there is no matching head the goal fails. A goal may succeed more than once.

For member we have two clauses: the first says that member succeeds if X is the head of the list ([X|T] is the same as x::t in Scala); the second says that member succeeds if X is a member of the tail of the list, regardless of the head. There is no clause where the list is empty (written []); a goal with an empty list fails since there is no matching clause head.

Prolog unification is more expressive than pattern matching as found in Scala, OCaml, etc. Both sides of a unification may contain variables; unification attempts to instantiate them so that the two sides are equal. Variables are instantiated by terms, which themselves may contain variables; unification finds the most general instantiation which makes the sides equal.

As a small example of this expressivity, we can run member “backwards”: the goal member(X, [1,2,3]) succeeds once for each element of the list, with X bound to the element.

There is much more on Prolog and logic programming in Frank Pfenning’s course notes, which I recommend highly.

Unification

For each type we want to use in unification we’ll define a corresponding type of terms, which have the same structure as the underlying type but can also contain variables. These aren’t Scala variables (which of course can’t be stored in a data structure) but “existential variables”, or evars. Evars are just tags; computations will carry an environment mapping evars to terms, which may be updated after a successful unification.

import scala.collection.immutable.{Map,HashMap} 
 
class Evar[A](val name: String) 
object Evar { def apply[A](name: String) = new Evar[A](name) } 
 
trait Term[A] { 
  // invariant: on call to unify, this and t have e substituted 
  def unify(e: Env, t: Term[A]): Option[Env] 
 
  def occurs[B](v: Evar[B]): Boolean 
  def subst(e: Env): Term[A] 
  def ground: A 
}

The important property of an evar is that it is distinct from every other evar; the name attached to it is just a label. An evar is indexed by a phantom type indicating the underlying type of terms which may be bound to it.

A term is indexed by its underlying type. So Int becomes Term[Int], String becomes Term[String], and so on; an evar of type Evar[A] may only be bound to a term of type Term[A]. (Prolog is dynamically typed, but this statically-typed treatment of evars and terms fits better with Scala.)

The unify method unifies a term with another term of the same type, taking an environment and returning an updated environment (or None if the unification fails). Occurs checks if an evar occurs in a term (as we will see this is used to prevent circular bindings). Subst substitutes the variables in a term with their bindings in an environment, and ground returns the underlying Scala value represented by the term (provided the term contains no evars).

class Env(m: Map[Evar[Any],Term[Any]]) { 
  def apply[A](v: Evar[A]) = 
    m(v.asInstanceOf[Evar[Any]]).asInstanceOf[Term[A]] 
  def get[A](v: Evar[A]): Option[Term[A]] = 
    m.get(v.asInstanceOf[Evar[Any]]).asInstanceOf[Option[Term[A]]] 
  def updated[A](v: Evar[A], t: Term[A]): Env = { 
    val v2 = v.asInstanceOf[Evar[Any]] 
    val t2 = t.asInstanceOf[Term[Any]] 
    val e2 = Env(Map(v2 -> t2)) 
    val m2 = m.mapValues(_.subst(e2)) 
    Env(m2.updated(v2, t2)) 
  } 
} 
object Env { 
  def apply(m: Map[Evar[Any],Term[Any]]) = new Env(m) 
  def empty = new Env(HashMap()) 
}

An environment is just a map from evars to terms. Because we need to store evars and terms of different types in the same environment, we cast them to and from Any; this is safe because of the phantom type on Evar. For simplicity we maintain the invariant that the term bound to each evar is already substituted by the rest of the environment.

case class VarTerm[A](v: Evar[A]) extends Term[A] { 
  def unify(e: Env, t: Term[A]) = 
    t match { 
      case VarTerm(v2) if (v2 == v) => Some(e) 
      case _ => 
        if (t.occurs(v)) None 
        else Some(e.updated(v, t)) 
    } 
 
  def occurs[B](v2: Evar[B]) = v2 == v 
 
  def subst(e: Env) = 
    e.get(v) match { 
      case Some(t) => t 
      case None => this 
    } 
 
  def ground = 
    throw new IllegalArgumentException("not ground") 
 
  override def toString = { v.name  } 
}

The VarTerm class represents terms consisting of an evar. To unify a VarTerm with another VarTerm containing the same evar, we just return the environment unchanged (since there is no new information). Otherwise we check that the evar doesn’t appear in the term (since a unification x =:= List(x) would create a circular term) then return the updated environment.

To substitute a VarTerm we return the term bound to the evar in the environment if one exists, otherwise the unsubstituted VarTerm. A VarTerm is never ground (we assume ground is called only on terms which are already substituted by the environment).

case class LitTerm[A](a: A) extends Term[A] { 
  def unify(e: Env, t: Term[A]) = 
    t match { 
      case LitTerm(a2) => if (a == a2) Some(e) else None 
      case _: VarTerm[_] => t.unify(e, this) 
      case _ => None 
    } 
 
  def occurs[B](v: Evar[B]) = false 
  def subst(e: Env) = this 
  def ground = a 
 
  override def toString = { a.toString } 
}

LitTerm represents terms of literal Scala values. A LitTerm unifies with another LitTerm containing an equal value, but that adds nothing to the environment. Then we have two cases which we need for every term type—to unify with a VarTerm call unify back on it; otherwise fail.

case class NilTerm[A]() extends Term[List[A]] { 
  def unify(e: Env, t: Term[List[A]]) = 
    t match { 
      case NilTerm() => Some(e) 
      case _: VarTerm[_] => t.unify(e, this) 
      case _ => None 
    } 
 
  def occurs[B](v: Evar[B]) = false 
  def subst(e: Env) = this 
  def ground = Nil 
 
  override def toString = { Nil.toString } 
} 
 
case class ConsTerm[A](hd: Term[A], tl: Term[List[A]]) 
  extends Term[List[A]] 
{ 
  def unify(e: Env, t: Term[List[A]]) = 
    t match { 
      case ConsTerm(hd2, tl2) => 
        for { 
          e1 <- hd.unify(e, hd2) 
          e2 <- tl.subst(e1).unify(e1, tl2.subst(e1)) 
        } yield e2 
      case _: VarTerm[_] => t.unify(e, this) 
      case _ => None 
    } 
 
  def occurs[C](v: Evar[C]) = hd.occurs(v) || tl.occurs(v) 
  def subst(e: Env) = ConsTerm(hd.subst(e), tl.subst(e)) 
  def ground = hd.ground :: tl.ground 
 
  override def toString = { hd.toString + " :: " + tl.toString } 
}

NilTerm and ConsTerm represent the Nil and :: constructors for lists. Nil is sort of like a literal, so the methods for NilTerm are similar to those for LitTerm. For ConsTerm we unify by unifying the heads and tails, calling subst on the tails since unifying the heads may have added bindings to the environment. (Here it’s convenient to use a for-comprehension on the Option[Env] type since either unification may fail.) Similarly we implement occurs, subst, and ground by calling them on the head and tail.

object Term { 
  implicit def var2Term[A](v: Evar[A]): Term[A] = VarTerm(v) 
  //implicit def lit2term[A](a: A): Term[A] = LitTerm(a) 
  implicit def int2Term(a: Int): Term[Int] = LitTerm(a) 
  implicit def list2Term[A](l: List[Term[A]]): Term[List[A]] = 
    l match { 
      case Nil => NilTerm[A] 
      case hd :: tl => ConsTerm(hd, list2Term(tl)) 
    } 
}

Finally we have some implicit conversions to make it a little easier to build Term values. The lit2term conversion turned out to be a bad idea; in particular you don’t want a LitTerm[List[A]] since it doesn’t unify with a ConsTerm[A] or NilTerm[A].

State

In order to combine unification with backtracking, we need to keep track of the environment along each branch of the tree of choices. We don’t want the environments from different branches to interfere, so it’s convenient to use a purely functional environment representation; we pass the current environment down the tree as computation proceeds. However, we can hide this state passing in the monad interface:

trait LogicState { L => 
  type T[S,A] 
  // as before 
  def split[S,A](s: S, t: T[S,A]): Option[(S,A,T[S,A])] 
 
  def get[S]: T[S,S] 
  def set[S](s: S): T[S, Unit] 
 
  case class Syntax[S,A](t: T[S,A]) { 
    // as before 
    def &[B](t2: => T[S,B]): T[S,B] = L.bind(t, { _: A => t2 }) 
  } 
}

LogicState is mostly the same as Logic, except that the type of choices has an extra parameter for the type of the state. The get and set functions get and set the current state. To split we need an initial state to get things started, and each result includes an updated state. Finally we add the syntax & to sequence two computations, ignoring the value of the first. We’ll use this to sequence goals, since we care only about the updated environment.

The simplest implementation of LogicState builds on Logic:

trait LogicStateT extends LogicState { 
  val Logic: Logic 
 
  type T[S,A] = S => Logic.T[(S, A)]

We embed state-passing in a Logic.T as a function from an initial state to a choice of alternatives, where each alternative includes an updated state along with its value.

  def fail[S,A] = { s: S => Logic.fail } 
  def unit[S,A](a: A) = { s: S => Logic.unit((s, a)) } 
 
  def or[S,A](t1: T[S,A], t2: => T[S,A]) = 
    { s: S => Logic.or(t1(s), t2(s)) } 
 
  def bind[S,A,B](t: T[S,A], f: A => T[S,B]) = { 
    val f2: ((S,A)) => Logic.T[(S,B)] = { case (s, a) => f(a)(s) } 
    { s: S => Logic.bind(t(s), f2) } 
  } 
 
  def apply[S,A,B](t: T[S,A], f: A => B) = { 
    val f2: ((S,A)) => ((S,B)) = { case (s, a) => (s, f(a)) } 
    { s: S => Logic.apply(t(s), f2) } 
  } 
 
  def filter[S,A](t: T[S,A], p: A => Boolean) = { 
    val p2: ((S,A)) => Boolean = { case (_, a) => p(a) } 
    { s: S => Logic.filter(t(s), p2) } 
  }

All of these operations pass the state through unchanged. Note that or passes the same state to both alternatives—different branches of the tree cannot interfere with one another’s state.

  def split[S,A](s: S, t: T[S,A]) = { 
    Logic.split(t(s)) match { 
      case None => None 
      case Some(((s, a), t)) => Some((s, a, { _ => t })) 
    } 
  } 
 
  def get[S] = { s: S => Logic.unit((s,s)) } 
  def set[S](s: S) = { _: S => Logic.unit((s,())) } 
}

In split we pass the given state to the underlying Logic.T, and for each alternative we unpack the pair of state and value. The choice of remaining alternatives t encapsulates the current state, so when we return it we ignore the input state. In get and set we return and replace the current state.

Another approach is to pass state explicitly through LogicSFK:

object LogicStateSFK extends LogicState { 
  type FK[R] = () => R 
  type SK[S,A,R] = (S, A, FK[R]) => R 
 
  trait T[S,A] { def apply[R](s: S, sk: SK[S,A,R], fk: FK[R]): R }

This is not really any different from LogicStateT applied to LogicSFK—we have just uncurried the state argument. We can take the same path as last time and defunctionalize this into a tail-recursive implementation (see the full code) although LogicStateT applied to LogicSFKDefuncTailrec inherits tail-recursiveness from the underlying Logic monad.

Scrolog

Finally we can put the pieces together into a Prolog-like embedded DSL:

trait Scrolog { 
  val LogicState: LogicState 
  import LogicState._ 
 
  type G = T[Env,Unit]

From our point of view, a goal is a stateful choice among alternatives, where we don’t care about the value returned, only the environment.

  class TermSyntax[A](t: Term[A]) { 
    def =:=(t2: Term[A]): G = 
      for { 
        env <- get 
        env2 <- { 
          t.subst(env).unify(env, t2.subst(env)) match { 
            case None => fail[Env,Unit] 
            case Some(e) => set(e) 
          } 
        } 
      } yield env2 
  } 
 
  implicit def termSyntax[A](t: Term[A]) = new TermSyntax(t) 
  implicit def syntax[A](t: G) = LogicState.syntax(t)

We connect term unification to the stateful logic monad with a wrapper class defining a =:= operator. To unify terms in the monad, we get the current environment, substitute it into the two terms (to satisfy the invariant above), then call unify; if it fails we fail the computation, else we set the new state.

  def run[A](t: G, n: Int, tm: Term[A]): List[Term[A]] = 
    LogicState.run(Env.empty, t, n) 
      .map({ case (e, _) => tm.subst(e) }) 
}

The run function solves a goal, taking as arguments the goal, the maximum number of solutions to find, and a term to be evaluated in the environment of each solution.

Examples

First we need to set up Scrolog:

val Scrolog = 
  new Scrolog { val LogicState = 
    new LogicStateT { val Logic = LogicSFKDefuncTailrec } 
  } 
import Scrolog._

Here is a translation of the member predicate:

  def member[A](x: Term[A], l: Term[List[A]]): G = { 
    val hd = Evar[A]("hd"); val tl = Evar[List[A]]("tl") 
    ConsTerm(x, tl) =:= l | 
    (ConsTerm(hd, tl) =:= l & member(x, tl)) 
  }

We implement predicates by functions, and goals by function calls. To implement matching the clause head, we explicitly unify the input arguments against each clause head, and combine the clauses with |. Subgoals are sequenced with &. Finally, we must create local evars explicitly, since they are fresh for each call (just as local variables are in Scala).

Finally we can run the goal above:

scala> val x = Evar[Int]("x") 
scala> run(member(x, List[Term[Int]](1, 2, 3)), 3, x) 
res6: List[Term[Int]] = List(1, 2, 3)

As another example, we can implement addition over unary natural numbers. In Prolog this would be

  sum(z, N, N). 
  sum(s(M), N, s(P)) :- sum(M, N, P).

In Prolog we can just invent symbols like s and z; in Scala we need first to define a type of natural numbers, then terms over that type:

  sealed trait Nat 
  case object Z extends Nat 
  case class S(n: Nat) extends Nat 
 
  case object ZTerm extends Term[Nat] { 
    // like NilTerm 
 
  case class STerm(n: Term[Nat]) extends Term[Nat] { 
    // like ConsTerm

Then we can define sum, again separating the clauses by | and explicitly unifying the clause heads:

  def sum(m: Term[Nat], n: Term[Nat], p: Term[Nat]): G = { 
    val m2 = Evar[Nat]("m"); val p2 = Evar[Nat]("p") 
    (m =:= Z & n =:= p) | 
    (m =:= STerm(m2) & p =:= STerm(p2) & sum(m2, n, p2)) 
  }

We can use sum to do addition:

scala> val x = Evar[Nat]("x"); val y = Evar[Nat]("y") 
scala> run(sum(S(Z), S(S(Z)), x), 1, x) 
res8: List[Term[Nat]] = List(S(S(S(Z))))

or subtraction:

scala> run(sum(x, S(S(Z)), S(S(S(Z)))), 1, x) 
res10: List[Term[Nat]] = List(S(Z)) 
 
scala> run(sum(S(Z), x, S(S(S(Z)))), 1, x) 
res11: List[Term[Nat]] = List(S(S(Z)))

or even to find all the pairs of naturals which sum to 3:

scala> run(sum(x, y, S(S(S(Z)))), 10, List[Term[Nat]](x, y)) 
res14: List[Term[List[Nat]]] = 
  List(Z :: S(S(S(Z))) :: List(), 
       S(Z) :: S(S(Z)) :: List(), 
       S(S(Z)) :: S(Z) :: List(), 
       S(S(S(Z))) :: Z :: List())

although the printing of Term[List] could be better.

This is only a small taste of the expressivity of Prolog-style logic programming. Again let me recommend Frank Pfenning’s course notes, which explore the semantics of Prolog in a “definitional interpreters” style, by gradually refining an interpreter to expose more of the machinery of the language.

See the full code.

Logic programming in Scala, part 2: backtracking

2011-04-29T22:07:00.000-07:00

In the previous post we saw how to write computations in a logic monad, where a “value” is a choice among alternatives, and operating on a value means operating on all the alternatives.

Our first implementation of the logic monad represents a choice among alternatives as a list, and operating on a value means running the operation for each alternative immediately (to produce a new list of alternatives). If we imagine alternatives as leaves of a tree (with | indicating branching), the first implementation explores the tree breadth-first.

This is OK for some problems, but we run into trouble when there are a large or infinite number of alternatives. For example, a choice among the natural numbers:

scala> import LogicList._ 
import LogicList._ 
 
scala> val nat: T[Int] = unit(0) | nat.map(_ + 1) 
java.lang.NullPointerException 
        ...

This goes wrong because even though the right-hand argument to | is by-name, we immediately try to use it, and fail because nat is not yet defined.

scala> def nat: T[Int] = unit(0) | nat.map(_ + 1) 
nat: LogicList.T[Int] 
scala> run(nat, 10) 
java.lang.StackOverflowError 
        ...

With def we can successfully define nat, because the right-hand side isn’t evaluated until nat is used in the call to run, but we overflow the stack trying to compute all the natural numbers.

Let’s repair this with a fancier implementation of the logic monad, translated from Kiselyov et al.’s Backtracking, Interleaving, and Terminating Monad Transformers. This implementation will explore the tree depth-first.

Success and failure continuations

The idea is to represent a choice of alternatives by a function, which takes as arguments two functions: a success continuation and a failure continuation. The success continuation is just a function indicating what to do next with each alternative; the failure continuation is what to do next when there are no more alternatives.

For success, what we do next is either return the alternative (when we have reached a leaf of the tree), or perform some operation on it (possibly forming new branches rooted at the alternative). For failure, what we do next is back up to the last branch point and succeed with the next alternative. If there are no more alternatives at the previous branch point we back up again, and so on until we can succeed or finally run out of alternatives. In other words, we do depth-first search on the tree, except that the tree isn’t a materialized data structure—it’s created on the fly.

(In the jargon of logic programming, a branch point is called a “choice point”, and going back to an earlier choice point is called “backtracking”.)

object LogicSFK extends Logic { 
  type FK[R] = () => R 
  type SK[A,R] = (A, FK[R]) => R 
 
  trait T[A] { def apply[R](sk: SK[A,R], fk: FK[R]): R }

The continuations can return a result of some arbitrary type R. This means that the function representing a choice has a “rank-2” polymorphic type—it takes functions which are themselves polymorphic—which is not directly representable in Scala. But we can encode it by making the representation function a method on a trait.

The success continuation takes a value of the underlying type (i.e. an alternative), and also a failure continuation, to call in case this branch of the tree eventually fails (by calling fail, or filter when no alternative satisfies the predicate). The failure continuation is also called to succeed with the next alternative after returning a leaf (see split).

  def fail[A] = 
    new T[A] { 
      def apply[R](sk: SK[A,R], fk: FK[R]) = fk() 
    } 
 
  def unit[A](a: A) = 
    new T[A] { 
      def apply[R](sk: SK[A,R], fk: FK[R]) = sk(a, fk) 
    }

To fail, just call the failure continuation. To succeed with one alternative, call the success continuation with the single alternative and the passed-in failure continuation—there are no more alternatives to try, so if this branch fails the unit fails.

  def or[A](t1: T[A], t2: => T[A]) = 
    new T[A] { 
      def apply[R](sk: SK[A,R], fk: FK[R]) = 
        t1(sk, { () => t2(sk, fk) }) 
    }

Or creates a choice point. We want to explore the alternatives in both t1 and t2, so we pass the success continuation to t1 (which calls it on each alternative); when t1 is exhausted we pass the success continuation to t2; finally we fail with the caller’s failure continuation—that is, we backtrack.

  def bind[A,B](t: T[A], f: A => T[B]) = 
    new T[B] { 
      def apply[R](sk: SK[B,R], fk: FK[R]) = 
        t(({ (a, fk) => f(a)(sk, fk) }: SK[A,R]), fk) 
    } 
 
  def apply[A,B](t: T[A], f: A => B) = 
    new T[B] { 
      def apply[R](sk: SK[B,R], fk: FK[R]) = 
        t(({ (a, fk) => sk(f(a), fk) }: SK[A,R]), fk) 
    }

For bind we extend each branch by calling f on the current leaf. To succeed we call f on the alternative a. Now f(a) returns a choice of alternatives, so we pass it the original success continuation (which says what to do next with alternatives resulting from the bind), and the failure continuation in force at the point a was generated (which succeeds with the next available alternative from f(a)).

For apply things are simpler, since f(a) returns a single value rather than a choice of alternatives: we succeed immediately with the returned value.

  def filter[A](t: T[A], p: A => Boolean) = 
    new T[A] { 
      def apply[R](sk: SK[A,R], fk: FK[R]) = { 
        val sk2: SK[A,R] = 
          { (a, fk) => if (p(a)) sk(a, fk) else fk() } 
        t(sk2, fk) 
      } 
    }

To filter a choice of alternatives, each time we succeed with a value we see if it satisfies the predicate p; if it does, we succeed with that value (extending the branch), otherwise we fail (pruning the branch).

  def split[A](t: T[A]) = { 
    def unsplit(fk: FK[Option[(A,T[A])]]): T[A] = 
      fk() match { 
        case None => fail 
        case Some((a, t)) => or(unit(a), t) 
      } 
    def sk : SK[A,Option[(A,T[A])]] = 
      { (a, fk) => Some((a, bind(unit(fk), unsplit))) } 
    t(sk, { () => None }) 
  } 
}

The point of split is to pull a single alternative from a choice, returning along with it a choice of the remaining alternatives. In the list implementation we just returned the head and tail of the list. In this implementation, the alternatives are computed on demand; we want to be careful to do only as much computation as needed to pull the first alternative

The failure continuation we pass to t just returns None when there are no more alternatives. The success continuation sk returns the first alternative and a choice of the remaining alternatives (wrapped in Some).

The tricky part is the choice of remaining alternatives. We’re given the failure continuation fk; calling it calls sk on the next alternative, which ultimately returns Some(a, t) where a is the next alternative, or None if there are no more alternatives. We repackage this Option as a choice of alternatives with unsplit. So that we don’t call fk too soon, we call unsplit via bind, which defers the call until the resulting choice of alternatives is actually used.

Now we can write infinite choices:

scala> import LogicSFK._ 
import LogicSFK._ 
 
scala> val nat: T[Int] = unit(0) | nat.map(_ + 1) 
nat: LogicSFK.T[Int] = LogicSFK$$anon$3@27aea0c1 
 
scala> run(nat, 10) 
res1: List[Int] = List(0, 1, 2, 3, 4, 5, 6, 7, 8, 9)

Well, this is a pretty complicated way to generate the natural numbers up to 10…

While nat looks like a lazy stream (as you might write in Haskell), no results are memoized (as they are in Haskell). To compute each successive number all the previous ones must be recomputed, and the running time of run(nat, N) is O(N²).

Defunctionalization

The code above is a fairly direct translation of the Haskell code from the paper. But its use of continuation-passing style doesn’t map well to Scala, because Scala doesn’t implement tail-call elimination (because the JVM doesn’t). Every call to a success or failure continuation adds a frame to the stack, even though all we ever do with the result is return it (i.e. the call is in tail-position), so the stack frame could be eliminated.

Surprisingly, we run out of memory before we run out of stack:

scala> run(nat, 2000) 
java.lang.OutOfMemoryError: Java heap space 
 ...

A little heap profiling shows that we’re using quadratic space as well as quadratic time. It turns out that the implementation of Logic.run (from the previous post) has a space leak. The call to run is not tail-recursive, so the stack frame hangs around, and although t is dead after split(t), there’s still a reference to it on the stack.

We can rewrite run with an accumulator to be tail-recursive:

  def run[A](t: T[A], n: Int): List[A] = { 
    def runAcc(t: T[A], n: Int, acc: List[A]): List[A] = 
      if (n <= 0) acc.reverse else 
        split(t) match { 
          case None => acc.reverse 
          case Some((a, t)) => runAcc(t, n - 1, a :: acc) 
        } 
    runAcc(t, n, Nil) 
  }

Now scalac compiles runAcc as a loop, so there are no stack frames holding on to dead values of t, and we get the expected:

scala> run(nat, 9000) 
java.lang.StackOverflowError 
 ...

To address the stack overflow we turn to defunctionalization. The idea (from John Reynold’s classic paper Definitional Interpreters for Higher-Order Programming Languages) is to replace functions and their applications with data constructors (we’ll use case classes) and an apply function, which matches the data constructor and does whatever the corresponding function body does. If a function captures variables, the data constructor must capture the same variables.

After defunctionalization we’re left with three mutually recursive apply functions (one for each of T, FK, and SK) where each recursive call is in tail position. In theory the compiler could transform these into code that takes only constant stack space (since they are local functions private to split). But in fact it will do so only for single recursive functions, so we will need to do this transformation by hand.

There is one hitch: the original code is not completely tail-recursive, because of unsplit, which calls a failure continuation then matches on the result. To fix this we need to add yet another continuation, which represents what to do after returning a result from a success or failure continuation.

object LogicSFKDefunc extends Logic { 
  type O[A] = Option[(A,T[A])] 
 
  sealed trait T[A] 
  case class Fail[A]() extends T[A] 
  case class Unit[A](a: A) extends T[A] 
  case class Or[A](t1: T[A], t2: () => T[A]) extends T[A] 
  case class Bind[A,B](t: T[A], f: A => T[B]) extends T[B] 
  case class Apply[A,B](t: T[A], f: A => B) extends T[B] 
  case class Filter[A](t: T[A], p: A => Boolean) extends T[A] 
  case class Unsplit[A](fk: FK[O[A]]) extends T[A] 
 
  def fail[A] = Fail() 
  def unit[A](a: A) = Unit(a) 
  def or[A](t1: T[A], t2: => T[A]) = Or(t1, { () => t2 }) 
  def bind[A,B](t: T[A], f: A => T[B]) = Bind(t, f) 
  def apply[A,B](t: T[A], f: A => B) = Apply(t, f) 
  def filter[A](t: T[A], p: A => Boolean) = Filter(t, p)

A choice of alternatives T[A] is now represented symbolically by case classes, and the functions which operate on choices just return the corresponding case. The cases capture the same variables that were captured in the original functions.

We have an additional case Unsplit which represents the bind(unit(fk), unsplit) combination from split. And we use O[A] as a convenient abbreviation.

  sealed trait FK[R] 
  case class FKOr[A,R](t: () => T[A], sk: SK[A,R], fk: FK[R]) 
    extends FK[R] 
  case class FKSplit[R](r: R) extends FK[R] 
 
  sealed trait SK[A,R] 
  case class SKBind[A,B,R](f: A => T[B], sk: SK[B,R]) 
    extends SK[A,R] 
  case class SKApply[A,B,R](f: A => B, sk: SK[B,R]) 
    extends SK[A,R] 
  case class SKFilter[A,R](p: A => Boolean, sk: SK[A,R]) 
    extends SK[A,R] 
  case class SKSplit[A,R](r: (A, FK[R]) => R) extends SK[A,R] 
 
  sealed trait K[R,R2] 
  case class KReturn[R]() extends K[R,R] 
  case class KUnsplit[A,R,R2](sk: SK[A,R], fk:FK[R], k: K[R,R2]) 
    extends K[O[A],R2]

Each case for FK (respectively SK) corresponds to a success (respectively failure) continuation function in the original code—it’s easy to match them up.

The K cases are for the new return continuation. They are defunctionalized from functions R => R2; we can either return a value directly, or match on whether it is Some or None as in unsplit. (If K is hard to understand you might try “refunctionalizing” it by replacing the cases with functions.)

We see that case classes are more powerful than variants in OCaml, without GADTs at least. Cases can have “input” type variables (appearing in arguments) which do not appear in the “output” (the type the case extends). When we match on the case these are treated as existentials. And the output type of a case can be more restrictive than type it extends; when we match on the case we can make more restrictive assumptions about types in that branch of the match. More on this in Emir, Odersky, and Williams’ Matching Objects with Patterns.

  def split[A](t: T[A]) = { 
 
    def applyT[A,R,R2] 
      (t: T[A], sk: SK[A,R], fk: FK[R], k: K[R,R2]): R2 = 
      t match { 
        case Fail() => applyFK(fk, k) 
        case Unit(a) => applySK(sk, a, fk, k) 
        case Or(t1, t2) => applyT(t1, sk, FKOr(t2, sk, fk), k) 
        case Bind(t, f) => applyT(t, SKBind(f, sk), fk, k) 
        case Apply(t, f) => applyT(t, SKApply(f, sk), fk, k) 
        case Filter(t, p) => applyT(t, SKFilter(p, sk), fk, k) 
        case Unsplit(fk2) => applyFK(fk2, KUnsplit(sk, fk, k)) 
      } 
 
    def applyFK[R,R2](fk: FK[R], k: K[R,R2]): R2 = 
      fk match { 
        case FKOr(t, sk, fk) => applyT(t(), sk, fk, k) 
        case FKSplit(r) => applyK(k, r) 
      } 
 
    def applySK[A,R,R2] 
      (sk: SK[A,R], a: A, fk: FK[R], k: K[R,R2]): R2 = 
      sk match { 
        case SKBind(f, sk) => applyT(f(a), sk, fk, k) 
        case SKApply(f, sk) => applySK(sk, f(a), fk, k) 
        case SKFilter(p, sk) => 
          if (p(a)) applySK(sk, a, fk, k) else applyFK(fk, k) 
        case SKSplit(rf) => applyK(k, rf(a, fk)) 
      }

Again, each of these cases corresponds directly to a function in the original code, and again it is easy to match them up (modulo the extra return continuation argument) to see that all we have done is separated the data part of the function (i.e. the captured variables) from the code part.

The exception is Unsplit, which again corresponds to bind(unit(fk), unsplit). To apply it, we apply fk (which collapses unit(fk), bind, and the application of fk in unsplit) with KUnsplit as continuation, capturing sk, fk, and k (corresponding to their capture in the success continuation of bind).

    def applyK[R,R2](k: K[R,R2], r: R): R2 = 
      k match { 
        case KReturn() => r.asInstanceOf[R2] 
        case KUnsplit(sk, fk, k) => { 
          r match { 
            case None => applyFK(fk, k) 
            case Some((a, t)) => applyT(or(unit(a), t), sk, fk, k) 
          } 
        } 
      }

For KReturn we just return the result. Although KReturn extends K[R,R], Scala doesn’t deduce from this that R = R2, so we must coerce the result. For KUnsplit we do the same match as unsplit, then apply the resulting T (for the None case we call the failure continuation directly instead of applying fail). Here Scala deduces from the return type of KUnsplit that is safe to treat r as an Option.

    applyT[A,O[A],O[A]]( 
      t, 
      SKSplit((a, fk) => Some((a, Unsplit(fk)))), 
      FKSplit(None), 
      KReturn()) 
  } 
}

Finally we apply the input T in correspondence to the original split.

Tail call elimination

(This section has been revised; you can see the original here.)

To eliminate the stack frames from tail calls, we next rewrite the four mutually-recursive functions into a single recursive function (which Scala compiles as a loop). To do this we have to abandon some type safety (but only in the implementation of the Logic monad; we’ll still present the same safe interface).

object LogicSFKDefuncTailrec extends Logic { 
  type O[A] = Option[(A,T[A])] 
 
  type T[A] = I 
 
  sealed trait I 
  case class Fail() extends I 
  case class Unit(a: Any) extends I 
  case class Or(t1: I, t2: () => I) extends I 
  case class Bind(t: I, f: Any => I) extends I 
  case class Apply(t: I, f: Any => Any) extends I 
  case class Filter(t: I, p: Any => Boolean) extends I 
  case class Unsplit(fk: I) extends I 
 
  case class FKOr(t: () => I, sk: I, fk: I) extends I 
  case class FKSplit(r: O[Any]) extends I 
 
  case class SKBind(f: Any => I, sk: I) extends I 
  case class SKApply(f: Any => Any, sk: I) extends I 
  case class SKFilter(p: Any => Boolean, sk: I) extends I 
  case class SKSplit(r: (Any, I) => O[Any]) extends I 
 
  case object KReturn extends I 
  case class KUnsplit(sk: I, fk: I, k: I) extends I

This is all pretty much as before except that we erase all the type parameters. Having done so we can combine the four defunctionalized types into a single type I (for “instruction” perhaps), which will allow us to write a single recursive apply function. The type parameter in T[A] is then a phantom type since it does not appear on the right-hand side of the definition; it is used only to enforce constraints outside the module.

  def fail[A]: T[A] = Fail() 
  def unit[A](a: A): T[A] = Unit(a) 
  def or[A](t1: T[A], t2: => T[A]): T[A] = Or(t1, { () => t2 }) 
  def bind[A,B](t: T[A], f: A => T[B]): T[B] = 
    Bind(t, f.asInstanceOf[Any => I]) 
  def apply[A,B](t: T[A], f: A => B): T[B] = 
    Apply(t, f.asInstanceOf[Any => I]) 
  def filter[A](t: T[A], p: A => Boolean): T[A] = 
    Filter(t, p.asInstanceOf[Any => Boolean])

The functions for building T[A] values are mostly the same. We have to cast passed-in functions since Any is not a subtype of arbitrary A. The return type annotations don’t seem necessary but I saw some strange type errors without them (possibly related to the phantom type?) when using the Logic.Syntax wrapper.

def split[A](t: T[A]): O[A] = { 
  def apply(i: I, a: Any, r: O[Any], sk: I, fk: I, k: I): O[Any] = 
    i match { 
      case Fail() => apply(fk, null, null, null, null, k) 
      case Unit(a) => apply(sk, a, null, null, fk, k) 
      case Or(t1, t2) => 
        apply(t1, null, null, sk, FKOr(t2, sk, fk), k) 
      case Bind(t, f) => 
        apply(t, null, null, SKBind(f, sk), fk, k) 
      case Apply(t, f) => 
        apply(t, null, null, SKApply(f, sk), fk, k) 
      case Filter(t, p) => 
        apply(t, null, null, SKFilter(p, sk), fk, k) 
      case Unsplit(fk2) => 
        apply(fk2, null, null, null, null, KUnsplit(sk, fk, k)) 
 
      case FKOr(t, sk, fk) => apply(t(), null, null, sk, fk, k) 
      case FKSplit(r) => apply(k, null, r, null, null, null) 
 
      case SKBind(f, sk) => apply(f(a), null, null, sk, fk, k) 
      case SKApply(f, sk) => apply(sk, f(a), null, null, fk, k) 
      case SKFilter(p, sk) => 
        if (p(a)) 
          apply(sk, a, null, null, fk, k) 
        else 
          apply(fk, null, null, null, null, k) 
      case SKSplit(rf) => 
        apply(k, null, rf(a, fk), null, null, null) 
 
      case KReturn => r 
      case KUnsplit(sk, fk, k) => { 
        r match { 
          case None => apply(fk, null, null, null, null, k) 
          case Some((a, t)) => 
            apply(or(unit(a), t), null, null, sk, fk, k) 
        } 
      } 
    } 
 
  apply(t, 
        null, 
        null, 
        SKSplit((a, fk) => Some((a, Unsplit(fk)))), 
        FKSplit(None), 
        KReturn).asInstanceOf[O[A]] 
}

The original functions took varying arguments; the single function takes all the arguments which the original ones did. We pass null for unused arguments in each call, but otherwise the cases are the same as before.

Now we can evaluate nat to large N without running out of stack (but since the running time is quadratic it takes longer than I care to wait to complete):

scala> run(nat, 100000) 
^C

See the complete code here.

Next time we’ll thread state through this backtracking logic monad, and use it to implement unification.

Logic programming in Scala, part 1

2011-04-06T22:03:00.000-07:00

I got a new job where I am hacking some Scala. I thought I would learn something by translating some functional code into Scala, and a friend had recently pointed me to Kiselyov et al.’s Backtracking, Interleaving, and Terminating Monad Transformers, which provides a foundation for Prolog-style logic programming. Of course, a good translation should use the local idiom. So in this post (and the next) I want to explore an embedded domain-specific language for logic programming in Scala.

A search problem

Here is a problem I sometimes give in interviews:

Four people need to cross a rickety bridge, which can hold only two people at a time. It’s a moonless night, so they need a light to cross; they have one flashlight with a battery which lasts 60 minutes. Each person crosses the bridge at a different speed: Alice takes 5 minutes, Bob takes 10, Candace takes 20 minutes, and Dave 25. How do they get across?

I’m not interested in the answer—I’m interviewing programmers, not law school applicants—but rather in how to write a program to find the answer.

The basic shape of the solution is to represent the state of the world (where are the people, where is the flashlight, how much battery is left), write a function to compute from any particular state the set of possible next states, then search for an answer (a path from the start state to the final state) in the tree formed by applying the next state function transitively to the start state. (Here is a paper describing solutions in Prolog and Haskell.)

Here is a first solution in Scala:

object Bridge0 { 
  object Person extends Enumeration { 
    type Person = Value 
    val Alice, Bob, Candace, Dave = Value 
    val all = List(Alice, Bob, Candace, Dave) // values is broken 
  } 
  import Person._ 
 
  val times = Map(Alice -> 5, Bob -> 10, Candace -> 20, Dave -> 25) 
 
  case class State(left: List[Person], 
                   lightOnLeft: Boolean, 
                   timeRemaining: Int)

We define an enumeration of people (the Enumeration class is a bit broken in Scala 2.8.1), a map of the time each takes to cross, and a case class to store the state of the world: the list of people on the left side of the bridge (the right side is just the complement); whether the flashlight is on the left; and how much time remains in the flashlight.

  def chooseTwo(list: List[Person]): List[(Person,Person)] = { 
    val init: List[(Person, Person)] = Nil 
    list.foldLeft(init) { (pairs, p1) => 
      list.foldLeft(pairs) { (pairs, p2) => 
        if (p1 < p2) (p1, p2) :: pairs else pairs 
      } 
    } 
  }

This function returns the list of pairs of people from the input list. We use foldLeft to do a double loop over the input list, accumulating pairs (p1, p2) where p1 < p2; this avoids returning (Alice, Bob) and also (Bob, Alice). The use of foldLeft is rather OCamlish, and if you know Scala you will complain that foldLeft is not idiomatic—we will repair this shortly.

In Scala, Nil doesn’t have type 'a list like in OCaml and Haskell, but rather List[Nothing]. The way local type inference works, the type variable in the type of foldLeft is instantiated with the type of the init argument, so you have to ascribe a type to init (or explicitly instantiate the type variable with foldLeft[List[(Person, Person)]]) or else you get a type clash between List[Nothing] and List[(Person, Person)].

  def next(state: State): List[State] = { 
    if (state.lightOnLeft) { 
      val init: List[State] = Nil 
      chooseTwo(state.left).foldLeft(init) { 
        case (states, (p1, p2)) => 
          val timeRemaining = 
            state.timeRemaining - math.max(times(p1), times(p2)) 
          if (timeRemaining >= 0) { 
            val left = 
              state.left.filterNot { p => p == p1 || p == p2 } 
            State(left, false, timeRemaining) :: states 
          } 
          else 
            states 
      } 
    } else { 
      val right = Person.all.filterNot(state.left.contains) 
      val init: List[State] = Nil 
      right.foldLeft(init) { (states, p) => 
        val timeRemaining = state.timeRemaining - times(p) 
        if (timeRemaining >= 0) 
          State(p :: state.left, true, timeRemaining) :: states 
        else 
          states 
      } 
    } 
  }

Here we compute the set of successor states for a state. We make a heuristic simplification: when the flashlight is on the left (the side where everyone begins) we move two people from the left to the right; when it is on the right we move only one. I don’t have a proof that an answer must take this form, but I believe it, and it makes the code shorter.

So when the light is on the left we fold over all the pairs of people still on the left, compute the time remaining if they were to cross, and if it is not negative build a new state where they and the flashlight are moved to the right and the time remaining updated.

If the light is on the right we do the same in reverse, but choose only one person to move.

  def tree(path: List[State]): List[List[State]] = 
    next(path.head). 
      map(s => tree(s :: path)). 
        foldLeft(List(path)) { _ ++ _ } 
 
  def search: List[List[State]] = { 
    val start = List(State(Person.all, true, 60)) 
    tree(start).filter { _.head.left == Nil } 
  } 
}

A list of successive states is a path (with the starting state at the end and the most recent state at the beginning); the state tree is a set of paths. The tree rooted at a path is the set of paths with the input path as a suffix. To compute this tree, we find the successor states of the head of the path, augment the path with each state in turn, recursively find the tree rooted at each augmented path, then append them all (including the input path).

Then to find an answer, we generate the state tree rooted at the path consisting only of the start state (everybody and the flashlight on the left, 60 minutes remaining on the light), then filter out the paths which end in a final state (everybody on the right).

For-comprehensions

To make the code above more idiomatic Scala (and more readable), we would of course use for-comprehensions, for example:

  def chooseTwo(list: List[Person]): List[(Person,Person)] = 
    for { p1 <- list; p2 <- list; if p1 < p2 } yield (p1, p2)

Just as before, we do a double loop over the input list, returning pairs where p1 < p2. (However, under the hood the result list is constructed by appending to a ListBuffer rather than with ::, so the pairs are returned in the reverse order.)

The for-comprehension syntax isn’t specific to lists. It’s syntactic sugar which translates to method calls, so we can use it on any objects which implement the right methods. The methods we need are

  def filter(p: A => Boolean): T[A] 
  def map[B](f: A => B): T[B] 
  def flatMap[B](f: A => T[B]): T[B] 
  def withFilter(p: A => Boolean): T[A]

where T is some type constructor, like List. For List, filter and map have their ordinary meaning, and flatMap is a map (where the result type must be a list) which concatenates the resulting lists (that is, it flattens the list of lists).

WithFilter is like filter but should be implemented as a “virtual” filter for efficiency—for List it doesn’t build a new filtered list, but instead just keeps track of the filter function; this way multiple adjacent filters can be combined and the result produced with a single pass over the list.

The details of the translation are in the Scala reference manual, section 6.19. Roughly speaking, <- becomes flatMap, if becomes filter, and yield becomes map. So another way to write chooseTwo is:

  def chooseTwo(list: List[Person]): List[(Person,Person)] = 
    list.flatMap(p1 => 
      list.filter(p2 => p1 < p2).map(p2 => (p1, p2)))

The logic monad

So far we have taken a concrete view of the choices that arise in searching the state tree, by representing a choice among alternatives as a list. For example, in the chooseTwo function we returned a list of alternative pairs. I want now to take a more abstract view, and define an abstract type T[A] to represent a choice among alternatives of type A, along with operations on the type, packaged into a trait:

trait Logic { L => 
  type T[A] 
 
  def fail[A]: T[A] 
  def unit[A](a: A): T[A] 
  def or[A](t1: T[A], t2: => T[A]): T[A] 
  def apply[A,B](t: T[A], f: A => B): T[B] 
  def bind[A,B](t: T[A], f: A => T[B]): T[B] 
  def filter[A](t: T[A], p: A => Boolean): T[A] 
  def split[A](t: T[A]): Option[(A,T[A])]

A fail value is a choice among no alternatives. A unit(a) is a choice of a single alternative. The value or(t1, t2) is a choice among the alternatives represented by t1 together with the alternatives represented by t2.

The meaning of applying a function to a choice of alternatives is a choice among the results of applying the function to each alternative; that is, if t represents a choice among 1, 2, and 3, then apply(t, f) represents a choice among f(1), f(2), and f(3).

Bind is the same except the function returns a choice of alternatives, so we must combine all the alternatives in the result; that is, if t is a choice among 1, 3, and 5, and f is { x => or(unit(x), unit(x + 1)) }, then bind(t, f) is a choice among 1, 2, 3, 4, 5, and 6.

A filter of a choice of alternatives by a predicate is a choice among only the alternatives which pass the the predicate.

Finally, split is a function which returns the first alternative in a choice of alternatives (if there is at least one) along with a choice among the remaining alternatives.

  def or[A](as: List[A]): T[A] = 
    as.foldRight(fail[A])((a, t) => or(unit(a), t)) 
 
  def run[A](t: T[A], n: Int): List[A] = 
    if (n <= 0) Nil else 
      split(t) match { 
        case None => Nil 
        case Some((a, t)) => a :: run(t, n - 1) 
      }

As a convenience, or(as: List[A]) means a choice among the elements of as. And run returns a list of the first n alternatives in a choice, picking them off one by one with split; this is how we get answers out of a T[A].

  case class Syntax[A](t: T[A]) { 
    def map[B](f: A => B): T[B] = L.apply(t, f) 
    def filter(p: A => Boolean): T[A] = L.filter(t, p) 
    def flatMap[B](f: A => T[B]): T[B] = L.bind(t, f) 
    def withFilter(p: A => Boolean): T[A] = L.filter(t, p) 
 
    def |(t2: => T[A]): T[A] = L.or(t, t2) 
  } 
 
  implicit def syntax[A](t: T[A]) = Syntax(t) 
}

Here we hook into the for-comprehension notation, by wrapping values of type T[A] in an object with the methods we need (and | as an additional bit of syntactic sugar), which methods just delegate to the functions defined above. We arrange with an implicit conversion for these wrappers to spring into existence when we need them.

The bridge puzzle with the logic monad

Now we can rewrite the solution in terms of the Logic trait:

class Bridge(Logic: Logic) { 
  import Logic._

We pass an implementation of the logic monad in, then open it so the implicit conversion is available (we can also use T[A] and the Logic functions without qualification).

The Person, times, and State definitions are as before.

  private def chooseTwo(list: List[Person]): T[(Person,Person)] = 
    for { p1 <- or(list); p2 <- or(list); if p1 < p2 } 
    yield (p1, p2)

As we saw, we can write chooseTwo more straightforwardly using a for-comprehension. In the previous version we punned on list as a concrete list and as a choice among alternatives; here we convert one to the other explicitly.

  private def next(state: State): T[State] = { 
    if (state.lightOnLeft) { 
      for { 
        (p1, p2) <- chooseTwo(state.left) 
        timeRemaining = 
          state.timeRemaining - math.max(times(p1), times(p2)) 
        if timeRemaining >= 0 
      } yield { 
        val left = 
          state.left.filterNot { p => p == p1 || p == p2 } 
        State(left, false, timeRemaining) 
      } 
    } else { // ...

This is pretty much as before, except with for-comprehensions instead of foldLeft and explicit consing. (You can easily figure out the branch for the flashlight on the right.)

  private def tree(path: List[State]): T[List[State]] = 
    unit(path) | 
      (for { 
         state <- next(path.head) 
         path <- tree(state :: path) 
       } yield path) 
 
  def search(n: Int): List[List[State]] = { 
    val start = List(State(Person.all, true, 60)) 
    val t = 
      for { path <- tree(start); if path.head.left == Nil } 
      yield path 
    run(t, n) 
  } 
}

In tree we use | to adjoin the input path (previously we gave it in the initial value of foldLeft). In search we need to actually run the Logic.T[A] value rather than returning it, because it’s an abstract type and can’t escape the module (see the Postscript for an alternative); this is why the other methods must be private.

Implementing the logic monad with lists

We can recover the original solution by implementing Logic with lists:

object LogicList extends Logic { 
  type T[A] = List[A] 
 
  def fail[A] = Nil 
  def unit[A](a: A) = a :: Nil 
  def or[A](t1: List[A], t2: => List[A]) = t1 ::: t2 
  def apply[A,B](t: List[A], f: A => B) = t.map(f) 
  def bind[A,B](t: List[A], f: A => List[B]) = t.flatMap(f) 
  def filter[A](t: List[A], p: A => Boolean) = t.filter(p) 
  def split[A](t: List[A]) = 
    t match { 
      case Nil => None 
      case h :: t => Some(h, t) 
    } 
}

A choice among alternatives is just a List of the alternatives, so the semantics we sketched above are realized in a very direct way.

The downside to the List implementation is that we compute all the alternatives, even if we only care about one of them. (In the bridge problem any path to the final state is a satisfactory answer, but our program computes all such paths, even if we pass an argument to search requesting only one answer.) We might even want to solve problems with an infinite number of solutions.

Next time we’ll repair this downside by implementing the backtracking monad from the paper by Kiselyov et al.

See the complete code here.

Postscript: modules in Scala

I got the idea of implementing the for-comprehension methods as an implict wrapper from Edward Kmett’s functorial library. It’s nice that T[A] remains completely abstract, and the for-comprehension notation is just sugar. I also tried an implementation where T[A] is bounded by a trait containing the methods:

trait Monadic[T[_], A] { 
  def map[B](f: A => B): T[B] 
  def filter(p: A => Boolean): T[A] 
  def flatMap[B](f: A => T[B]): T[B] 
  def withFilter(p: A => Boolean): T[A] 
 
  def |(t: => T[A]): T[A] 
  def split: Option[(A,T[A])] 
} 
 
trait Logic { 
  type T[A] <: Monadic[T, A] 
  // no Syntax class needed

This works too but the type system hackery is a bit ugly, and it constrains implementations of Logic more than is necessary.

Another design choice is whether T[A] is an abstract type (as I have it) or a type parameter of Logic:

trait Logic[T[_]] { L => 
  // no abstract type T[A] but otherwise as before 
}

Neither alternative provides the expressivity of OCaml modules (but see addendum below): with abstract types, consumers of Logic cannot return values of T[A] (as we saw above); with a type parameter, they can, but the type is no longer abstract.

In OCaml we would write

module type Logic = 
sig 
  type 'a t 
 
  val unit : 'a -> 'a t 
  (* and so on *) 
end 
 
module Bridge(L : Logic) = 
struct 
  type state = ... 
  val search : state list L.t 
end

and get both the abstract type and the ability to return values of the type.

Addendum

Jorge Ortiz points out in the comments that it is possible to keep T[A] abstract and also return its values from Bridge, by making the Logic argument a (public) val. We can then remove the privates, and write search as just:

  def search: T[List[State]] = { 
    val start = List(State(Person.all, true, 60)) 
    for { path <- tree(start); if path.head.left == Nil } 
    yield path 
  }

instead of baking run into it. Now, if we write val b = new Bridge(LogicList) then b.search has type b.Logic.T[List[b.State]], and we can call b.Logic.run to evaluate it.

This is only a modest improvement; what’s still missing, compared to the OCaml version, is the fact that LogicList and b.Logic are the same module. So we can’t call LogicList.run(b.search) directly. Worse, we can’t compose modules which use the same Logic implementation, because they each have their own incompatibly-typed Logic member.

I thought there might be a way out of this using singleton types—the idea is that a match of a value v against a typed pattern where the type is w.type succeeds when v eq w (section 8.2 in the reference manual). So we can define

def run[A]( 
  Logic: Logic, 
  b: Bridge, 
  t: b.Logic.T[A], 
  n: Int): List[A] = 
{ 
  Logic match { 
    case l : b.Logic.type => l.run(t, n) 
  } 
}

which is accepted, but sadly

scala> run[List[b.State]](LogicList, b, b.search, 2) 
<console>:8: error: type mismatch; 
 found   : b.Logic.T[List[b.State]] 
 required: b.Logic.T[List[b.State]] 
       run[List[b.State]](LogicList, b, b.search, 2) 
                                          ^

Addendum addendum

Some further advice from Jorge Ortiz: the specific type of Logic (not just Logic.type) can be exposed outside Bridge either through polymorphism:

class Bridge[L <: Logic](val Logic: L) { 
  ... 
} 
 
val b = new Bridge(LogicList)

or by defining an abstract value (this works the same if Bridge is a trait):

abstract class Bridge { 
  val Logic: Logic 
  ... 
}

So we can compose uses of T but it remains abstract.

Three uses for a binary heap

2010-11-24T20:58:00.000-08:00

Lately I have been interviewing for jobs, so doing a lot of whiteboard programming, and binary heaps keep arising in the solutions to these interview problems. There is nothing new or remarkable about these applications (binary heaps and their uses are covered in any undergraduate algorithms class), but I thought I would write them down because they are cute, and in the hope that they might be useful to someone else who (like me) gets by most days as a working programmer with no algorithm fancier than quicksort or binary search.

Binary heaps

Here’s a signature for a binary heap module Heap:

module type OrderedType = 
sig 
  type t 
  val compare : t -> t -> int 
end 
 
module type S = sig 
  type elt 
  type t 
  val make : unit -> t 
  val add : t -> elt -> unit 
  val peek_min : t -> elt option 
  val take_min : t -> elt 
  val size : t -> int 
end 
 
module Make (O : OrderedType) : S with type elt = O.t

We start with a signature for ordered types (following the Set and Map modules in the standard library), so we can provide a type-specific comparison function.

From an ordered type we can make a heap which supports adding elements, peeking the smallest element (None if there are no elements) without removing it, removing and returning the smallest element (raising Not_found if the heap is empty), and returning the number of elements.

We’ll work out the asymptotic running times of the algorithms below, so it will be useful to know that the worst-case running time of the add and take_min functions is O(log n) where n is the number of elements in the heap.

Finding the k smallest elements in a list

Here’s a simple one. To find the smallest element in a list, we could sort the list then take the first element in the sorted list, at a cost of O(log n). Or we could just take a pass over the list keeping a running minimum, at a cost of O(n).

What if we want the k smallest elements? Again, we could sort the list, but if k < n we can do better by generalizing the single-pass solution. The idea is to keep the k smallest elements we’ve seen so far in a binary heap. For each element in the list we add it to the heap, then (if there were already k elements in the heap) remove the largest element in the heap, leaving the k smallest.

The running time is O(n log k) since we do an add and a take_min in a heap of size k for each of n elements in the list. Here’s the code:

let kmin (type s) k l = 
  let module OT = struct 
    type t = s 
    let compare e1 e2 = compare e2 e1 
  end in 
  let module H = Heap.Make(OT) in 
 
  let h = H.make () in 
  List.iter 
    (fun e -> 
       H.add h e; 
       if H.size h > k 
       then ignore (H.take_min h)) 
    l; 
  let rec loop mins = 
    match H.peek_min h with 
      | None -> mins 
      | _ -> loop (H.take_min h :: mins) in 
  loop []

Here we make good use of OCaml 3.12’s new feature for explicitly naming type variables in a polymorphic function to make a structure matching OrderedType. The heap has the same element type as the list, but we reverse the comparison since we want to remove the largest rather than smallest element from the heap in the loop. At the end of kmin we drain the heap to build a list of the k smallest elements.

Merging k lists

Suppose we want to merge k lists. We could merge them pairwise until there is only one list, but that would take k - 1 passes, for a worst-case running time of O(n * (k - 1)). Instead we can merge them all in one pass, using a binary heap so we can find the next smallest element of k lists in O(log k) time, for a running time of O(n log k). Here’s the code:

let merge (type s) ls = 
  let module OT = struct 
    type t = s list 
    let compare e1 e2 = 
      compare (List.hd e1) (List.hd e2) 
  end in 
  let module H = Heap.Make(OT) in 
 
  let h = H.make () in 
  let add = function 
    | [] -> () 
    | l -> H.add h l in 
  List.iter add ls; 
  let rec loop () = 
    match H.peek_min h with 
      | None -> [] 
      | _ -> 
          match H.take_min h with 
            | [] -> assert false 
            | m :: t -> 
                add t; 
                m :: loop () in 
  loop ()

We store the lists in the heap, and compare them by comparing their head element (we’re careful not to put an empty list in the heap). When we take the smallest list from the heap, its head becomes the next element in the output list, and we return its tail (if it is not empty) to the heap.

Computing a skyline

The next problem was told to me in terms of computing the skyline of a set of buildings. A building has a height and a starting and ending x-coordinate; buildings may overlap. The skyline of a set of buildings is a list of (x, y) pairs (in ascending x order), describing a sequence of horizontal line segments (each starting at (x, y) and ending at the subsequent x), such that at any x there is no space between the line segment and the tallest building. (Here’s another description with diagrams.)

I googled a bit to see what this is useful for, and didn’t find much. One application is to extract a monophonic line from polyphonic music, where x is time and height is some metric on notes, like pitch or volume. It might be useful for searching data which is only intermittently applicable—say, to compute a schedule over time of the nearest open restaurant.

The algorithm scans the building start and end points in ascending x order, keeping the “active” buildings (those which overlap the current x) in a binary heap. The height of the skyline can only change at a building start or end point. We can determine the tallest building at a point by calling peek_min on the heap.

When we hit a start point we add the building to the heap; for an end point we do nothing (the heap has no operation to remove an element). So we may have inactive buildings in the heap. We remove them lazily—before checking the height of the highest building, we call take_min to remove any higher inactive buildings.

The worst-case running time is O(n log n), since we do some heap operations for each building, and we might end up with all the buildings in the heap.

Here’s the code:

type building = int * int * int (* x0, x1, h *) 
 
let skyline bs = 
  let module OT = struct 
    type t = int * building 
    let compare (x1, _) (x2, _) = compare x1 x2 
  end in 
  let module Events = Heap.Make(OT) in 
  let events = Events.make () in 
  List.iter 
    (fun ((x0,x1,_) as b) -> 
       Events.add events (x0, b); 
       Events.add events (x1, b)) 
    bs; 
 
  let module OT = struct 
    type t = building 
    let compare (_,_,h1) (_,_,h2) = compare h2 h1 
  end in 
  let module Heights = Heap.Make(OT) in 
  let heights = Heights.make () in 
 
  let rec loop last = 
    match Events.peek_min events with 
      | None -> [] 
      | _ -> 
          let (x, (x0,_,h as b)) = Events.take_min events in 
          if x = x0 then Heights.add heights b; 
          while (match Heights.peek_min heights with 
                   | Some (_,x1,_) -> x1 <= x 
                   | _ -> false) do 
            ignore (Heights.take_min heights) 
          done; 
          let h = 
            match Heights.peek_min heights with 
              | Some (_,_,h) -> h 
              | None -> 0 in 
          match last with 
            | Some h' when h = h' -> loop last 
            | _ -> (x, h) :: loop (Some h) in 
  loop None

We use a second heap events to store the “events” (the start and end points of all the buildings), in order to process them in ascending x order. (This use is not dynamic—we do not add new elements to the heap while processing them—so we could just as well use another means of sorting the points.) In this heap we store the x coordinate and the building (we can tell whether we have a start or end point by comparing the x coordinate to the building’s start point), and compare elements by comparing just the x coordinates.

The main heap heights stores buildings, and we compare them by comparing heights (reversed, so peek_min peeks the tallest building). While there are still events, we add the building to heights if the event is a start point, clear out inactive buildings, then return the pair (x, y) where x is the point we’re processing and y is the height of the tallest active building. Additionally we filter out adjacent pairs with the same height; these can arise when a shorter building starts or ends while a taller building is active.

Implementing binary heaps

The following implementation is derived from the one in Daniel Bünzli’s React library (edited a little bit for readability). The Wikipedia article on binary heaps explains the standard technique well, so I won’t repeat it.

The only piece of trickiness is the use of Obj.magic 0 for unused elements of the array, so we can grow it by doubling the size rather than adding a single element each time, and thereby amortize the cost of blitting the old array.

module Make (O : OrderedType) : S with type elt = O.t = 
struct 
  type elt = O.t 
  type t = { mutable arr : elt array; mutable len : int } 
 
  let make () = { arr = [||]; len = 0; } 
 
  let compare h i1 i2 = O.compare h.arr.(i1) h.arr.(i2) 
 
  let swap h i1 i2 = 
    let t = h.arr.(i1) in 
    h.arr.(i1) <- h.arr.(i2); 
    h.arr.(i2) <- t 
 
  let rec up h i = 
    if i = 0 then () 
    else 
      let p = (i - 1) / 2 in 
      if compare h i p < 0 then begin 
        swap h i p; 
        up h p 
      end 
 
  let rec down h i = 
    let l = 2 * i + 1 in 
    let r = 2 * i + 2 in 
    if l >= h.len then () 
    else 
      let child = 
        if r >= h.len then l 
        else if compare h l r < 0 then l else r in 
      if compare h i child > 0 then begin 
        swap h i child; 
        down h child 
      end 
 
  let add h e = 
    if h.len = Array.length h.arr 
    then begin 
      let len = 2 * h.len + 1 in 
      let arr' = Array.make len (Obj.magic 0) in 
      Array.blit h.arr 0 arr' 0 h.len; 
      h.arr <- arr' 
    end; 
    h.arr.(h.len) <- e; 
    up h h.len; 
    h.len <- h.len + 1 
 
  let peek_min h = 
    match h.len with 
      | 0 -> None 
      | _ -> Some h.arr.(0) 
 
  let take_min h = 
    match h.len with 
      | 0 -> raise Not_found 
      | 1 -> 
          let m = h.arr.(0) in 
          h.arr.(0) <- (Obj.magic 0); 
          h.len <- 0; 
          m 
      | k -> 
          let m = h.arr.(0) in 
          let k = k - 1 in 
          h.arr.(0) <- h.arr.(k); 
          h.arr.(k) <- (Obj.magic 0); 
          h.len <- k; 
          down h 0; 
          m 
 
  let size h = h.len 
end

(Complete code is here.)

Reading Camlp4, part 11: syntax extensions

2010-09-10T17:16:00.000-07:00

In this final (?) post in my series on Camlp4, I want at last to cover syntax extensions. A nontrivial syntax extension involves almost all the topics we have previously covered, so it seems fitting that we treat them last.

Extending grammars

In the post on parsing we covered Camlp4 grammars but stopped short of explaining how to extend them. Well, this is not completely true: we used the EXTEND form to extend an empty grammar, and we can also use it to extend non-empty grammars. We saw a small example of this when implementing quotations, where we extended the JSON grammar with a new json_eoi entry (which refered to an entry in the original grammar). Rules and levels may also be added to existing entries, and rules may be deleted.

Let’s look at a complete syntax extension, which demonstrates modifying Camlp4’s OCaml grammar. The purpose of the extension is to change the precedence of the method call operator # to make “method chaining” read better. For example, if the foo method returns an object, you can write

  obj#foo "bar" #baz

to call the baz method, rather than needing

  (obj#foo "bar")#baz

(I originally wrote this for use with the jQuery binding for ocamljs; method chaining is common with jQuery.)

Here is the extension:

  open Camlp4 
  
  module Id : Sig.Id = 
  struct 
    let name = "pa_jquery" 
    let version = "0.1" 
  end 
  
  module Make (Syntax : Sig.Camlp4Syntax) = 
  struct 
    open Sig 
    include Syntax 
  
    DELETE_RULE Gram expr: SELF; "#"; label END; 
  
    EXTEND Gram 
      expr: BEFORE "apply" 
        [ "#" LEFTA 
          [ e = SELF; "#"; lab = label -> 
              <:expr< $e$ # $lab$ >> ] 
        ]; 
    END 
  end 
  
  module M = Register.OCamlSyntaxExtension(Id)(Make)

To make sense of a syntax extension it’s helpful to refer to Camlp4OCamlRevisedParser.ml (which defines the revised syntax grammar) and Camlp4OCamlParser.ml (which defines the original syntax as an extension of the revised syntax). There we see that the # operator is parsed in the expr entry, in a level called ”.” (which includes other dereferencing operators), and that this level appears below the apply level, which parses function application. Recall from the parsing post that operators in lower levels bind more tightly. So to get the effect we want, we need to move the # rule above the apply level in the grammar.

First we delete the rule from its original location: DELETE_RULE takes the grammar, the entry, and the symbols on the left-hand side of the rule, followed by END; you don’t have to say in what level it appears. Then we add the rule at a new location: we create a new level # containing the rule from the original grammar, and add it before the level named apply.

There are several ways to specify where a level is inserted: BEFORE level and AFTER level put it before or after some other level; LEVEL level adds rules to an existing level (you will be warned but not stopped from changing the label or associativity of the level); FIRST and LAST put the level before or after all other levels. If you don’t specify, rules are added to the topmost level in the entry. The resulting grammar works just as if you had given it all at once, making the insertions in the specified places. (However, it is not very clear from the code how ordering works when inserting rules into an existing level; it is perhaps best not to rely on the order of rules in a level anyway.)

Finally we register the extension. The Make argument to OCamlSyntaxExtension returns a Sig.Camlp4Syntax for some reason (in Register.ml it is just ignored) so we include Syntax to provide it.

(The complete code for this example is here.)

Transforming the AST

Let’s do a slightly more complicated example involving some transformation of the parsed AST. It often comes up that we want to let-bind the value of an expression to a name, trapping exceptions, then evaluate the body of the let outside the scope of the exception handler. This is a bit painful to write in stock OCaml; we can only straightforwardly express trapping exceptions in the whole let expression:

  try let x = e1 in e2 
  with e -> h

A nice alternative is to use thunks to delay the evaluation of the body, doing it outside the scope of the try/with:

  (try let x = e1 in fun () -> e2 
   with e -> fun () -> h)()

(We must thunkify the exception handler to make the types work out.) This is simple enough to do by hand, but let’s give it some syntactic sugar:

  let try x = e1 in e2 
  with e -> h

which should expand to the thunkified version above. (The idea and syntax are taken from Martin Jambon’s micmatch extension.)

Let’s look at the existing rules in Camlp4OCamlRevisedParser.ml for let and try to get an idea of how to parse the let/try form:

  [ "let"; r = opt_rec; bi = binding; "in"; x = SELF -> 
      <:expr< let $rec:r$ $bi$ in $x$ >> 
  ... 
  | "try"; e = sequence; "with"; a = match_case -> 
      <:expr< try $mksequence' _loc e$ with [ $a$ ] >>

For let, the opt_rec entry parses an optional rec keyword (we see there is a special antiquotation for interpolating rec). Binding parses a group of bindings separated by and. SELF is just expr. For try, sequence is a sequence of expressions separated by ;, and match_case is a group of match cases separated by |. (These entries are both a little different in the original syntax, to account for the different semicolon rules and the [] delimiters around the match cases.) Recall that Camlp4OCamlRevisedParser.ml uses the revised syntax quotations, so we have [] around the match cases. The call to mksequence' just wraps a do {} around a sequence if necessary; more on this below.

The parsing rule we want is a combination of these. Here is the extension:

  EXTEND Gram 
    expr: LEVEL "top" [ 
      [ "let"; "try"; r = opt_rec; bi = binding; "in"; 
        e = sequence; "with"; a = match_case -> 
          let a = 
            List.map 
              (function 
                 | <:match_case< $p$ when $w$ -> $e$ >> -> 
                     <:match_case< 
                       $p$ when $w$ -> fun () -> $e$ 
                     >> 
                 | mc -> mc) 
              (Ast.list_of_match_case a []) in 
          <:expr< 
            (try let $rec:r$ $bi$ in fun () -> do { $e$ } 
             with [ $list:a$ ])() 
          >> 
      ] 
    ]; 
  END

We put rec after try (following micmatch), which is a little weird ~~, but if we put it before we would need to look ahead to disambiguate `let` from `let try`; once we parse `opt_rec` we are committed to one rule or the other~~ ; instead we could start the rule "let"; r = opt_rec; "try", which has no ambiguity with the ordinary let rule because the "let"; opt_rec prefix is factored out; the parser doesn’t choose between the rules until it tries to parse try. After in we parse sequence rather than SELF; this seems like a good choice because there is a with to end the sequence.

Now, to transform the AST, we map over the match cases. The match_case entry returns a list of cases separated by Ast.McOr; we call list_of_match_case to get an ordinary list. For each case, we match the pattern, when clause, and expression on the right-hand side (these are packaged in an Ast.McArr, where the when clause field is Ast.ExNil if there is no when clause), and return it with the expression thunkified. Then we return the whole let inside try, with the body sequence thunkified.

We have to add a do {} around the body, creating an Ast.ExSeq node, because that’s what is expected by Camlp4Ast2OCamlAst.ml—recall from the filters post that the Camlp4 AST is translated to an OCaml AST and marshalled to the compiler. If we forget this (and “we” often forget these idiosyncrasies) then we get the error ”expr; expr: not allowed here, use do {...} or [|...|] to surround them”, which is pretty helpful as these errors go.

(The complete code for this example is here.)

Extending pattern matching

As a final example, let’s extend OCaml’s pattern syntax. In the quotations post we noted that JSON quotations in a pattern are not very useful, because we would usually like a pattern to match even if the fields of an object come in a different order or there are extra fields. To keep the code short let’s abstract the problem a little and consider matching association lists: if we write a match case

  | alist [ "foo", x; "bar", y ] -> e

we would like it to match association lists with "foo" and "bar" keys, in any order, with any extra pairs in the list. Our translation looks like this:

  | __pa_alist_patt_1 when 
      (match ((try Some (List.assoc "foo" __pa_alist_patt_1) 
               with | Not_found -> None), 
              (try Some (List.assoc "bar" __pa_alist_patt_1) 
               with | Not_found -> None)) 
       with 
       | (Some x, Some y) -> true 
       | _ -> false) 
      -> 
      (match ((try Some (List.assoc "foo" __pa_alist_patt_1) 
               with | Not_found -> None), 
              (try Some (List.assoc "bar" __pa_alist_patt_1) 
               with | Not_found -> None)) 
       with 
       | (Some x, Some y) -> e 
       | _ -> assert false)

This might seem overcomplicated, and it is true that we could simplify it for this case. But the built-in pattern syntax is complicated, and it is tricky handling all the cases to make things work smoothly; the strategy that produces the code above will handle some (but not all) of the complications. (We’ll consider some improvements below.)

The basic idea is that when we come to an alist we replace it with a new fresh name, then do further matching in a when clause, so if it fails we can continue to the next case by returning false. In the when clause we look up the keys, putting them in options, then match on the options; we handle nested patterns (to the right of a key) by embedding them in the when clause match. The when clause match also binds variables appearing in the original pattern, so they are available to the when clause of the original case (if there is one). Finally, we do the whole thing over again in the match case body to provide bindings to the original body.

In order to implement this we’ll use both a syntax extension and a filter. The reason is that we’d like to extend the patt entry, but to do the AST transformation we sketched above we need to transform match_cases. We could replace the match_case part of the parser as well but that seems needlessly hairy, and generally when writing a syntax extension we’d like to touch as little of the parser as possible so it interoperates well with other extensions.

First, here is the syntax extension:

  EXTEND Gram 
    patt: LEVEL "simple" 
    [[ 
       "alist"; "["; 
         l = 
           LIST0 
             [ e = expr LEVEL "simple"; ","; 
               p = patt LEVEL "simple" -> 
                 Ast.PaOlbi (_loc, "", p, e) ] 
             SEP ";"; 
       "]" -> 
         <:patt< $uid:"alist"$ $Ast.paSem_of_list l$ >> 
    ]]; 
  END

We extend the simple level of the patt entry, which parses primitive patterns. Inside alist [] we parse a list of expr / patt pairs; we parse expr at the simple level or else it would parse the whole pair as an expr, and the same for patt just in case. Then we return the pair of expression and pattern in an Ast.PaOlbi (ordinarily used for optional argument defaults in function definitions). Why? Well, we need to return something of type patt, but we need somehow to get the expr to our filter, and this is the only patt constructor that holds an expr. (As an alternative we could parse a patt instead of an expr, but then we’d need to translate it to an expr at the point we use it.) Finally we return the list wrapped in a data constructor so we can recognize it easily in the filter; because it is lower-case, we can be sure that “alist” is not the identifier of a real data constructor.

Now let’s look at the filter. First, some helper functions:

  let fresh = 
    let id = ref 0 in 
    fun () -> 
      incr id; 
      "__pa_alist_patt_"  ^ string_of_int !id 
 
  let expr_tup_of_list _loc = function 
    | [] -> <:expr< () >> 
    | [ v ] -> v 
    | vs -> <:expr< $tup:Ast.exCom_of_list vs$ >> 
 
  let patt_tup_of_list _loc = function 
    | [] -> <:patt< () >> 
    | [ p ] -> p 
    | ps -> <:patt< $tup:Ast.paCom_of_list ps$ >>

We have a function to generate fresh names, a function to turn a list of expressions into a tuple, and a similar function for patterns. The reason we need these latter two is that a tuple with 0 or 1 elements is not allowed by Camlp4Ast2OCamlAst.ml (the empty “tuple” is actually a special identifier in the Camlp4 AST). Next, the main rewrite function:

  let rewrite _loc p w e = 
    let k = ref (fun s f -> s) in

The function takes the parts of an Ast.McArr (that is, a match case). We’re going to map over the pattern p, building up a function k as we encounter nested alist forms. We want to generate the same matching code in the when clause and the body, so k is parameterized with an expression in case of success (the original when clause or the body) and in case of failure (false or assert false). We will build k from the inside out, starting with a function that just returns the success expression.

    let map = 
      object 
        inherit Ast.map as super 
 
        method patt p = 
          match super#patt p with 
            | <:patt< $uid:"alist"$ $l$ >> -> 
                let id = fresh () in 
                let l = 
                  List.map 
                    (function 
                       | Ast.PaOlbi (_, _, p, e) -> p, e 
                       | _ -> assert false) 
                    (Ast.list_of_patt l []) in 
                let vs = 
                  List.map 
                    (fun (_, e) -> 
                       <:expr< 
                         try Some (List.assoc $e$ $lid:id$) 
                         with Not_found -> None 
                       >>) 
                    l in 
                let ps = 
                  List.map 
                    (fun (p, _) -> <:patt< Some $p$ >>) 
                    l in 
                let k' = !k in 
                k := 
                  (fun s f -> 
                     <:expr< 
                       match $expr_tup_of_list _loc vs$ with 
                         | $patt_tup_of_list _loc ps$ -> $k' s f$ 
                         | _ -> $f$ 
                     >>); 
                <:patt< $lid:id$ >> 
            | p -> p 
      end in

The Ast.map object provides methods to map each syntactic class of the AST, along with default implementations which return the node unchanged. We extend it to walk over the pattern, leaving it unchanged except when we come to our special alist constructor. In that case we generate a fresh name, which we return as the value of the function. Then we extract the expr / patt pairs and map them to try Some (List.assoc ... expressions and Some patterns. Finally we extend k by matching all the expressions against all the patterns; if the match succeeds we call the previous k, otherwise the failure expression. Since we build k from the inside out, we transform subpatterns first (by matching over super#patt p).

    let p = map#patt p in 
    let w = match w with <:expr< >> -> <:expr< true >> | _ -> w in 
    let w = !k w <:expr< false >> in 
    let e = !k e <:expr< assert false >> in 
    <:match_case< $p$ when $w$ -> $e$ >>

We call map#patt on p to replace special alist constructor nodes with fresh names and build up k, then call the resulting k on the when clause (if there is no when clause we replace it with true) and body, and finally return the result as a match_case, completing the rewrite function.

  let filter = 
    let map = 
      object 
        inherit Ast.map as super 
 
        method match_case mc = 
          match super#match_case mc with 
            | <:match_case@_loc< $p$ when $w$ -> $e$ >> -> 
                rewrite _loc p w e 
            | e -> e 
      end in 
    map#str_item 
 
  let _ = AstFilters.register_str_item_filter filter

We extend Ast.map again to call the rewrite function on each match_case, then register the resulting filter.

The code above handles when clauses and nested alist patterns, and interacts properly with ordinary OCaml patterns. However, it completely falls down on nested pattern alternatives. If we write

match x with 
  | alist [ "foo", f ] 
  | alist [ "fooz", f ] -> e

we get this mess:

  | __pa_alist_patt_1 | __pa_alist_patt_2 when 
      (match try Some (List.assoc "fooz" __pa_alist_patt_2) 
             with | Not_found -> None 
       with 
       | Some f -> 
           (match try Some (List.assoc "foo" __pa_alist_patt_1) 
                  with | Not_found -> None 
            with 
            | Some f -> true 
            | _ -> false) 
       | _ -> false) 
      -> 
      ... (* the same mess for the body *)

The first problem is that both arms of an alternative must bind the same names, but we have replaced them with two different fresh names. The second problem is that we have blindly treated alternative alist patterns as being nested one inside the other. A solution to both these problems is to split nested alternatives into separate cases, at the cost of duplicating the when clause and body in each.

Jeremy Yallop’s patterns framework (see here for an update that works with OCaml 3.12.0) allows multiple pattern extensions to coexist, and provides some common facilities to make them easier to write. In particular it splits nested alternatives into separate cases. Another deficiency in the code above is that it duplicates the match expression (built in k) in the when clause and body. This can be avoided by computing the body within the when clause, setting a reference, and dereferencing it in the body. However, the reference must be bound outside the match_case to be visible both in the when clause and the body, so this approach must transform each AST node that contains match_cases in order to bind the refs in the right place. The patterns framework handles this as well.

(The complete code for this example is here. A version using the patterns framework is here.)

ocamljs 0.3

2010-08-26T14:45:00.001-07:00

I am happy to announce version 0.3 of ocamljs. Ocamljs is a system for compiling OCaml to Javascript. It includes a Javascript back-end for the OCaml compiler, as well as several support libraries, such as bindings to the browser DOM. Ocamljs also works with orpc for RPC over HTTP, and froc for functional reactive browser programming.

Changes since version 0.2 include:

support for OCaml 3.11.x and 3.12.0
jQuery binding (contributed by Dave Benjamin)
full support for OCaml objects (interoperable with Javascript objects)
Lwt 2.x support
ocamllex and ocamlyacc support
better interoperability with Javascript
many small fixes and improvements

Development of ocamljs has moved from Google Code to Github; see

project page: http://github.com/jaked/ocamljs
documentation: http://jaked.github.com/ocamljs
downloads: http://github.com/jaked/ocamljs/downloads

Comparison to js_of_ocaml

Since I last did an ocamljs release, a new OCaml-to-Javascript system has arrived, js_of_ocaml. I want to say a little about how the two systems compare:

Ocamljs is a back-end to the existing OCaml compiler; it translates the “lambda” intermediate language to Javascript. (This is also where the bytecode and native code back-ends connect to the common front-end.) Js_of_ocaml post-processes ordinary OCaml bytecode (compiled and linked with the ordinary OCaml bytecode compiler) into Javascript. With ocamljs you need a special installation of the compiler (and special support for ocamlbuild and ocamlfind), you need to recompile libraries, and you need the OCaml source to build it. With js_of_ocaml you don’t need any of this.

Since ocamljs recompiles libraries, it’s possible to special-case code for the Javascript build to take advantage of Javascript facilities. For example, ocamljs implements the Buffer module on top of Javascript arrays instead of strings, for better performance. Similarly, it implements CamlinternalOO to use Javascript method dispatch directly instead of layering OCaml method dispatch on top. Js_of_ocaml can’t do this (or at least it would be necessary to recognize the compiled bytecode and replace it with the special case).

Because js_of_ocaml works from bytecode, it can’t always know the type of values (at the bytecode level, ints, bools, and chars all have the same representation, for example). This makes interoperating with native Javascript more difficult: you usually need conversion functions between the OCaml and Javascript representation of values when you call a Javascript function from OCaml. Ocamljs has more information to work with, and can represent OCaml bools as Javascript bools, for example, so you can usually call a Javascript function from OCaml without conversions.

Ocamljs has a mixed representation of strings: literal strings and the result of ^, Buffer.contents, and Printf.sprintf are all immutable Javascript strings; strings created with String.create are mutable strings implemented by Javascript arrays (with a toString method which returns the represented string). This is good for interoperability—you can usually pass a string directly to Javascript—but it doesn’t match regular OCaml’s semantics, and it can cause runtime failures (e.g. if you try to mutate an immutable string). Js_of_ocaml implements only mutable strings, so you need conversions when calling Javascript, but the semantics match regular OCaml.

With ocamljs, Javascript objects can be called from OCaml using the ordinary OCaml method-call syntax, and objects written in OCaml can be called using the ordinary Javascript syntax. With js_of_ocaml, a special syntax is needed to call Javascript objects, and OCaml objects can’t easily be called from Javascript. However, there is an advantage to having a special call syntax: with ocamljs it is not possible to partially apply calls to native Javascript methods, but this is not caught by the compiler, so there can be a runtime failure.

Ocamljs supports inline Javascript, while js_of_ocaml does not. I think it might be possible for js_of_ocaml to do so using the same approach that ocamljs takes: use Camlp4 quotations to embed a syntax tree, then convert the syntax tree from its OCaml representation (as lambda code or bytecode) into Javascript. However, you would still need conversion functions between OCaml and Javascript values.

I haven’t compared the performance of the two systems. It seems like there must be a speed penalty to translating from bytecode compared to translating from lambda code. On the other hand, while ocamljs is very naive in its translation, js_of_ocaml makes several optimization passes. With many programs it doesn’t matter, since most of the time is spent in browser code. (For example, the planet example seems to run at the same speed in ocamljs and js_of_ocaml.) It would be interesting to compare them on something computationally intensive like Andrej Bauer’s random-art.org.

Js_of_ocaml is more complete and careful in its implementation of OCaml (e.g. it supports int64s), and it generates much more compact code than ocamljs. I hope to close the gap in these areas, possibly by borrowing some code and good ideas from js_of_ocaml.

Mixing monadic and direct-style code with delimited continuations

2010-08-20T17:50:00.000-07:00

The Lwt library is a really nice way to write concurrent programs. A big downside, however, is that you can’t use direct-style libraries with it. Suppose we’re writing an XMPP server, and we want to parse XML as it arrives over a network connection, using Daniel Bünzli’s nice xmlm library. Xmlm can read from a string, or from a Pervasives.in_channel, or you can give it a function of type (unit -> int) to return the next character of input. But there is no way to have it read from an Lwt thread; that is, we can’t give it a function of type (unit -> int Lwt.t), since it doesn’t know what do with an Lwt.t. To keep track of the parser state at the point the input blocks, the whole library would need to be rewritten in Lwt style (i.e. monadic style).

Now, Lwt does provide the Lwt_preemptive module, which gives you a way to spin off a preemptive thread (implemented as an ordinary OCaml thread) and wait for its result in the usual Lwt way with bind. This is useful, but has two drawbacks: preemptive threads are preemptive, so you’re back to traditional locking if you want to operate on shared data; and preemptive threads are threads, so they are much heavier than Lwt threads, and (continuing the XMPP hypothetical) it may not be feasible to use one per open connection.

Fibers

What we would really like is to be able spin off a cooperative, direct-style thread. The thread needs a way to block on Lwt threads, but when it blocks we need to be able to schedule another Lwt thread. As a cooperative thread it of course has exclusive access to the process state while it is running. A cooperative, direct-style thread is sometimes called a coroutine (although to me that word connotes a particular style of inter-thread communication as well, where values are yielded between coroutines), or a fiber.

Here’s an API for mixing Lwt threads with fibers:

  val start : (unit -> 'a) -> 'a Lwt.t 
  val await : 'a Lwt.t -> 'a

The start function spins off a fiber, returning an Lwt thread which is woken with the result of the fiber once it completes. The await function (which may be called only from within a fiber) blocks on the result of an Lwt thread, allowing another Lwt thread to be scheduled while it is waiting.

With this API we could implement our XMPP server by calling xmlm from within a fiber, and passing it a function that awaits the next character available on the network connection. But how do we implement it?

Delimited continuations

Oleg Kiselyov’s recent announcement of a native-code version of his Delimcc library for delimited continuations in OCaml reminded me of two things:

I should find out what delimited continuations are.
They sound useful for implementing fibers.

The paper describing the library, Delimited Control in OCaml, Abstractly and Concretely, has a pretty good overview of delimited continuations, and section 2 of A Monadic Framework for Delimited Continuations is helpful too.

The core API is small:

  type 'a prompt 
  type ('a,'b) subcont 
  
  val new_prompt   : unit -> 'a prompt 
  
  val push_prompt  : 'a prompt -> (unit -> 'a) -> 'a 
  val take_subcont : 
    'b prompt -> (('a,'b) subcont -> unit -> 'b) -> 'a 
  val push_subcont : ('a,'b) subcont -> (unit -> 'a) -> 'b

I find it easiest to think about these functions as operations on the stack. A prompt is an identifier used to mark a point on the stack (the stack can be marked more than once with the same prompt). The function new_prompt makes a new prompt which is not equal to any other prompt.

The call push_prompt p f marks the stack with p then runs f, so the stack, growing to the right, looks like

 
  ABCDpEFGH

where ABCD are stack frames in the continuation of the call to push_prompt, and EFGH are frames created while running f. If f returns normally (that is, without calling take_subcont) then its return value is returned by push_prompt, and we are back to the original stack ABCD.

If take_subcont p g is called while running f, the stack fragment EFGH is packaged up as an ('a,'b) subcont and passed to g. You can think of an ('a,'b) subcont as a function of type 'a -> 'b, where 'a is the return type of the call to take_subcont and 'b is the return type of the call to push_prompt. Take_subcont removes the fragment pEFGH from the stack, and there are some new frames IJKL from running g, so we have

 
  ABCDIJKL

Now g can make use of the passed-in subcont using push_subcont. (Thinking of a subcont as a function, push_subcont is just a weird function application operator, which takes the argument as a thunk). Then the stack becomes

 
  ABCDIJKLEFGH

Of course g can call the subcont as many times as you like.

A common pattern is to re-mark the stack with push_prompt before calling push_subcont (so take_subcont may be called again). There is an optimized version of this combination called push_delim_subcont, which produces the stack

 
  ABCDIJKLpEFGH

The idea that a subcont is a kind of function is realized by shift0, which is like take_subcont except that instead of passing a subcont to g it passes an ordinary function. The passed function just wraps a call to push_delim_subcont. (It is push_delim_subcont rather than push_subcont for historical reasons I think—see the Monadic Framework paper for a comparison of various delimited continuation primitives.)

Implementing fibers

To implement fibers, we want start f to mark the stack, then run f; and await t to unwind the stack back to the mark, wait for t to complete, then restore the stack. Here is start:

  let active_prompt = ref None 
  
  let start f = 
    let t, u = Lwt.wait () in 
    let p = Delimcc.new_prompt () in 
    active_prompt := Some p; 
  
    Delimcc.push_prompt p begin fun () -> 
      let r = 
        try Lwt.Return (f ()) 
        with e -> Lwt.Fail e in 
      active_prompt := None; 
      match r with 
        | Lwt.Return v -> Lwt.wakeup u v 
        | Lwt.Fail e -> Lwt.wakeup_exn u e 
        | Lwt.Sleep -> assert false 
    end; 
    t

We make a sleeping Lwt thread, and store a new prompt in a global (this is OK because we won’t yield control to another Lwt thread before using it; of course this is not safe with OCaml threads). Then we mark the stack with push_prompt and run the fiber. (The let r = ... match r with ... is to avoid calling Lwt.wakeup{,_exn} in the scope of the try; we use Lwt.state as a handy type to store either a result or an exception.) If the fiber completes without calling await then all we do is wake up the Lwt thread with the returned value or exception.

Here is await:

  let await t = 
    let p = 
      match !active_prompt with 
        | None -> failwith "await called outside start" 
        | Some p -> p in 
    active_prompt := None; 
 
    match Lwt.poll t with 
      | Some v -> v 
      | None -> 
          Delimcc.shift0 p begin fun k -> 
            let ready _ = 
              active_prompt := Some p; 
              k (); 
              Lwt.return () in 
            ignore (Lwt.try_bind (fun () -> t) ready ready) 
          end; 
          match Lwt.poll t with 
            | Some v -> v 
            | None -> assert false

We first check to be sure that we are in the scope of start, and that t isn’t already completed (in which case we just return its result). If we actually need to wait for t, we call shift0, which capture the stack fragment back to the push_prompt call in start (this continuation includes the subsequent match Lwt.poll t and everything after the call to await), then try_bind so we can restore the stack fragment when t completes (whether by success or failure). When t completes, the ready function restores the global active_prompt, in case the fiber calls await again, then restores the stack by calling k (recall that this also re-marks the stack with p, which is needed if the fiber calls await again).

It’s pretty difficult to follow what’s going on here, so let’s try it with stacks. After calling start we have

 
  ABCDpEFGH

where ABCD is the continuation of push_prompt in start (just the return of t) and EFGH are frames created by the thunk passed to start. Now, a call to await (on an uncompleted thread) calls shift0, which packs up EFGH as k and unwinds the stack to p. The function passed to shift0 stores k in ready but doesn’t call it, and control returns to start (since the stack has been unwound).

The program continues normally until t completes. Now control is in Lwt.run_waiters running threads that were waiting on t; one of them is our ready function. When it is called, the stack is re-marked and EFGH is restored, so we have

 
  QRSTpEFGH

where QRST is wherever we happen to be in the main program, ending in Lwt.run_waiters. Now, EFGH ends with the second call to match Lwt.poll in await, which returns the value of t and continues the thunk passed to start. The stack is now marked with p inside Lwt.run_waiters, so when await is called again control returns there.

Events vs. threads

We have seen that we can use fibers to write Lwt threads in direct style. Should we abandon Lwt’s monadic style entirely, and use Lwt only for its event handling?

First, how does each style perform? Every time a fiber blocks and resumes, we have to copy, unwind, and restore its entire stack. With Lwt threads, the “stack” is a bunch of linked closures in the heap, so we don’t need to do anything to block or resume. On the other hand, building and garbage-collecting the closures is more expensive than pushing and popping the stack. We can imagine that which style performs better depends on the thread: if it blocks infrequently enough, the amortized cost of copying and restoring the stack might be lower than the cost of building and garbage-collecting the closures. (We can also imagine that a different implementation of delimited continuations might change this tradeoff.)

Second, how does the code look? The paper Cooperative Task Management without Manual Stack Management considers this question in the context of the “events vs. threads” debate. Many of its points lose their force when translated to OCaml and Lwt—closures, the >>= operator, and Lwt’s syntax extension go a long way toward making Lwt code look like direct style—but some are still germane. In favor of fibers is that existing direct-style code need not be rewritten to work with Lwt (what motivated us in the first place). In favor of monadic style is that the type of a function reflects the possibility that it might block, yield control to another thread, and disturb state invariants.

Direct-style FRP

We could apply this idea, of replacing monadic style with direct style using delimited continuations, to other monads—in particular to the froc library for functional reactive programming. (The Scala.React FRP library also uses delimited continuations to implement direct style; see Deprecating the Observer Pattern for details.)

Here’s the API:

  val direct : (unit -> 'a) -> 'a Froc.behavior 
  val read : 'a Froc.behavior -> 'a

Not surprisingly, it’s just the same as for Lwt, but with a different monad and different names (I don’t know if direct is quite right but it is better than start). There is already a function Froc.sample with the same type as read, but it has a different meaning: sample takes a snapshot of a behavior but creates no dependency on it.

The implementation is very similar as well:

  let active_prompt = ref None 
  
  let direct f = 
    let t, u = Froc_ddg.make_changeable () in 
    let p = Delimcc.new_prompt () in 
    active_prompt := Some p; 
  
    Delimcc.push_prompt p begin fun () -> 
      let r = 
        try Froc_ddg.Value (f ()) 
        with e -> Froc_ddg.Fail e in 
      active_prompt := None; 
      Froc_ddg.write_result u r 
    end; 
    (Obj.magic t : _ Froc.behavior)

This is essentially the same code as start, modulo the change of monad. However, some of the functions we need aren’t exported from Froc, so we need to use the underlying Froc_ddg module and magic the result at the end. Froc_ddg.make_changeable is the equivalent of Lwt.wait: it returns an “uninitialized” monadic value along with a writer for that value. We use Froc_ddg.result instead of Lwt.state to store a value or exception, and Froc_ddg.write_result instead of the pattern match and Lwt.wakeup{,_exn}.

  
  let read t = 
    let p = 
      match !active_prompt with 
        | None -> failwith "read called outside direct" 
        | Some p -> p in 
    active_prompt := None; 
  
    Delimcc.shift0 p begin fun k -> 
      Froc.notify_result_b t begin fun _ -> 
        active_prompt := Some p; 
        k () 
      end 
    end; 
    Froc.sample t

And this is essentially the same code as await. A Froc.behavior always has a value, so we don’t poll it as we did with Lwt.t, but go straight to shift0. We have Froc.try_bind but it’s a little more compact to use use notify_result_b, which passes a result.

Monadic reflection

The similarity between these implementations suggests that we could use the same code to get a direct style version of any monad; we only need a way to create an uninitialized monadic value, then set it. The call to Lwt.poll in await is an optimization which we would have to forgo. (In both these examples we have a monad with failure, and try_bind, but we could do without it.)

A little googling turns up Andrzej Filinski’s paper Representing Monads, which reaches the same conclusion, with a lot more rigor. In that work start/direct are called reify, and await/read are called reflect. Reflect is close to the implementations above, but in reify the paper marks the stack inside a function passed to bind rather than creating an uninitialized monadic value and later setting it.

This makes sense—inside bind an uninitialized monadic value is created, then set from the result of the function passed to bind. So we are partially duplicating bind in the code above. If we mark the stack in the right place we should be able to use bind directly. It is hard to see how to make the details work out, however, since Lwt.bind and Froc.bind each have some cases where uninitialized values are not created.

(You can find the complete code for Lwt fibers here and direct-style froc here.)

(revised 10/22)