microKanren In Scala

uKanren.scala

/**
 * An implementation of microKanren (and probably most of miniKanren), with
 * a few extras. Currently, it supports:
 *
 *  - The essential core of microKanren: Unify, Fresh, Disjunction, Conjunction
 *    - Standard terms: Variables, Atoms, TermCons, EmptyTerm.
 *    - An implicit conversion from type T to Atom[T]. This makes writing programs
 *      much easier.
 *    - A decent reifier, which converts terms to strings.
 *
 *  - Goal.and and Goal.or operators (Conjunction and Disjunction, respectively)
 *    for easily composing Goals.
 *
 *  - Goal negation, although this comes with a caveat. Read the documentation for
 *    the Negate case class for an explanation.
 *
 *  - Relations, which are like standard facts in Prolog. They make the process
 *    of defining certain types of goals simpler.
 *
 * There are lots of parts, so I've tried my best to separate the parts of the
 * implementation and comment them when necessary.
 *
 * This source code file is covered under the Unlicense.
 */
import scala.annotation.tailrec
import scala.collection.mutable.HashSet

/**************** STREAMS ***************/

/**
 * We need to write our own Stream on top of Scala's stream, because Scala's
 * streams do not support lazy heads. The guys behind microLogic[1] explain
 * well the reasoning and structure - go read what they have to say about the
 * subject.
 * 
 * [1] http://mullr.github.io/micrologic/literate.html#sec-2-4
 */
sealed trait LazyStream[+T] {
  /**
   * Converts a possibly ImmatureStream into a stream that is either empty, or
   * has a head.
   */
  def force: RealizedStream[T]

  /**
   * Evaluates a function for each element of a stream.
   */
  def flatMap[U](fn: T => LazyStream[U]): LazyStream[U]

  /**
   * Joins together the elements of a stream. Note that this respects the fact
   * that the other stream may be infinite, by not allowing this stream to block
   * the others.
   */
  def interleave[U >: T](other: LazyStream[U]): LazyStream[U]
}

/**
 * An ImmatureStream is a non-RelaizedStream - it may or may not have a head,
 * depending upon how its thunk evaluates.
 *
 * Extra care must be taken when working with this, since this is often used
 * to represent computations that can take place indefinitely. You don't want
 * to see the StackOverflow exeptions that can result if this is handled poorly.
 */
class ImmatureStream[T](protected val fn: () => LazyStream[T]) extends LazyStream[T]{
  def force = {
    // We need to keep running the functions that we get back from forcing each
    // stream, since they may return a LazyStream rather than a RealizedStream.
    var value = fn()
    var done = false
    var result: RealizedStream[T] = null

    while (!done) {
      value match {
        case other: ImmatureStream[_] =>
          value = other.fn()
        case realized: RealizedStream[_] =>
          done = true
          result = realized
      }
    }
    result
  }

  def flatMap[U](mapFn: T => LazyStream[U]) = 
    new ImmatureStream(() => this.force.flatMap(mapFn))

  def interleave[U >: T](other: LazyStream[U]) =
    /*
     * Note the order reversal - this is necessary because we want 
     * an ImmatureStream to let the other stream do a unit of work before
     * this stream does anything after the initial force.
     */
    new ImmatureStream({ () =>
      var thisForced = force
      other.interleave(thisForced)
    })

  override def toString = s"%$fn"
}

/**
 * A RealizedStream has evaluated its head element - it is either the
 * case that the RealizedStream is empty (that is, it lacks a head) or
 * mature (that is, it has a head).
 */
sealed trait RealizedStream[+T] extends LazyStream[T] {
  def force = this

  /**
   * Checks to see if this is empty.
   */
  def isEmpty: Boolean

  /**
   * Converts this into a Stream, thus allowing the use of Stream's
   * collection methods.
   */
  def toStream: Stream[T]
}

/**
 * The EmptyLazyStream is a realized stream which has no content.
 */
case object EmptyLazyStream extends RealizedStream[Nothing] {
  def isEmpty = true
  def flatMap[U](fn: Nothing => LazyStream[U]) = this
  def interleave[U >: Nothing](other: LazyStream[U]) = other
  def toStream = Stream.empty

  override def toString = "%()"
}

/**
 * A MatureStream is a stream which has a realized head, and a lazily evaluated tail.
 */
class MatureStream[T](head: T, tail: () => LazyStream[T]) extends RealizedStream[T]{
  def isEmpty = false

  def flatMap[U](fn: T => LazyStream[U]) = 
    fn(head).interleave(
      new ImmatureStream(() => tail().flatMap(fn)))

  def interleave[U >: T](other: LazyStream[U]) =
    new MatureStream(head, () => new ImmatureStream(tail).interleave(other))

  def toStream = head #:: (tail().force.toStream)
  override def toString = s"%($head . $tail)"
}

/**
 * This interleaves a group of LazyStreams; it is useful, because interleaving
 * is a binary operation on LazyStream, while this makes it an n-ary operation.
 */
def interleave[T](streams: Seq[LazyStream[T]]): LazyStream[T] =
  streams.foldLeft(EmptyLazyStream: LazyStream[Nothing]) {
    case (interleaved: LazyStream[T], nextStream: LazyStream[T]) =>
      interleaved.interleave(nextStream)
  }

/**************** TERMS ****************/

/**
 * Bindings map variables to terms, which themselves may be variables.
 */
type Bindings = Map[Var, Term]

/**
 * A Term is any value that can appear when trying to evaluate a Goal.
 */
sealed trait Term {
  /**
   * Pulls out the binding corresponding to this term. Note that this only
   * does anything useful for Vars.
   */
  def read(bindings: Bindings): Term = this

  /**
   * Produces a readable String representation - it takes a fresh variable
   * counter, and returns the String representation and the next fresh
   * variable counter.
   */
  def reify(varCounter: Int): (String, Int)
}

/**
 * A Variable is a way to uniquely define a name that can be bound to.
 *
 * Note: This is a class because putting 'new' in front of it makes the fact
 * that this is generating a new logic variable more obvious.
 */
class Var extends Term {
  @tailrec
  override final def read(bindings: Bindings): Term =
    bindings.get(this) match {
      case Some(v: Var) => v.read(bindings)
      case Some(term)   => term
      case None         => this
    }

  def reify(varCounter: Int) =
    (s"Var#${varCounter}", varCounter + 1)
}

/**
 * An Atom is a value which doesn't contain any further Terms. Use this to
 * wrap Scala datatypes as Terms.
 * 
 * Note that an implicit conversion is defined from Any => Atom[T], so you
 * shouldn't have to use this constructor manually.
 */
case class Atom[T](value: T) extends Term {
  def reify(varCounter: Int) =
    (s"!${value}", varCounter)
}

/**
 * A TermCons is the equivalent of a List, but for terms.
 */
case class TermCons(head: Term, tail: Term) extends Term {
  def reify(varCounter: Int) = {
    val (headText, newCounter) = head.reify(varCounter)
    var (tailText, finalCounter) = tail.reify(newCounter)
    (s"($headText . $tailText)", finalCounter)
  }
}

/**
 * A Term without any value. Useful for defining lists along with TermCons.
 */
case object EmptyTerm extends Term {
  def reify(varCounter: Int) =
    ("?", varCounter)
}

/**
 * This makes code which unifies with non-Scala types easier to write, since 
 * it would otherwise require writing Atom(x) everywhere.
 */
import scala.language.implicitConversions
implicit def toAtom[T](value: T): Atom[T] = Atom(value)

/**
 * This converts a Scala List into a string of TermCons values.
 */
def listToCons(list: List[Term]): Term =
  list match {
    case head :: tail => TermCons(head, listToCons(tail))
    case Nil          => EmptyTerm
  }

/**
 * Converts a group of bindings to a string representation, by recursively
 * reifying its keys and values. Example output:
 *
 * {{{
 *  - ;; Var#0 -> !foo ;; Var#1 -> !bar
 *  - ;; Var#0 -> !baz
 * }}}
 */
def reifyBinds(seq: Seq[Bindings]): String =
  seq.foldLeft("") {
    case (buff: String, binds: Bindings) =>
      val text = binds.foldLeft(("", 0)) {
        case ((buff: String, counter: Int), (key: Var, value: Term)) =>
          val (keyText, keyCounter) = key.reify(counter)
          val (valueText, valueCounter) = value.reify(keyCounter)
          (buff + ";; " + s"$keyText -> $valueText", valueCounter)
      }._1

      buff + "\n - " + text
  }

/**************** GOALS ****************/

/**
 * This recursively unifies two expressions.
 */
def unify(a:Term, b: Term, bindings: Bindings): Option[Bindings] = {
  val finalA = a.read(bindings)
  val finalB = b.read(bindings)

  (finalA, finalB) match {
    // If either is a variable, then extend the bindings
    case (x: Var, y) =>
      Some(bindings + (x -> y))
    case (x, y: Var) =>
      Some(bindings + (y -> x))

    // If both side are conses, then recursively unify them
    case (TermCons(xHead, xTail), TermCons(yHead, yTail)) =>
      for {
        headBindings <- unify(xHead, yHead, bindings)
        tailBindings <- unify(xTail, yTail, headBindings)
      } yield tailBindings

    // Anything else requires equality to unify
    case _ =>
      if (finalA == finalB) Some(bindings)
      else                  None
  }
}

/**
 * A Goal is an operation than can either succeed or fail, and it generates a
 * group of new Bindings if it succeeds.
 */
sealed trait Goal {
  def attempt(bindings: Bindings): LazyStream[Bindings]
  def suspend = SuspendedGoal(this)

  def and(other: Goal) = Conjunction(this, other)
  def or(other: Goal) = Disjunction(this, other)
  def negate = Negate(this)
}

/**
 * AlwaysG always succeeds.
 */
case object SuccessG extends Goal {
  def attempt(bindings: Bindings): LazyStream[Bindings] =
    new MatureStream(bindings, () => EmptyLazyStream)
}

/**
 * FailG always fails.
 */
case object FailG extends Goal {
  def attempt(bindings: Bindings): LazyStream[Bindings] =
    EmptyLazyStream
}

/**
 * A SuspendedGoal is a Goal which is wraps the Goal in an ImmatureStream, to
 * prevent it from being evaluated immediately. This is necessary for goals
 * which are recursive.
 *
 * See also: Goal.suspend
 */
case class SuspendedGoal(goal: Goal) extends Goal {
  def attempt(bindings: Bindings): LazyStream[Bindings] =
    new ImmatureStream(() => goal.attempt(bindings))
}

/**
 * Unify succeeds if the two terms unify, and fails if they do not.
 */
case class Unify(a: Term, b: Term) extends Goal {
  def attempt(bindings: Bindings): LazyStream[Bindings] =
    unify(a, b, bindings) match {
      case Some(newBindings) => new MatureStream(newBindings, () => EmptyLazyStream)
      case None              => EmptyLazyStream
    }
}

/**
 * Fresh generates a new logic variable, and applies it to a function which
 * returns a Goal. It then executes this Goal.
 */
case class Fresh(fn: Var => Goal) extends Goal {
  def attempt(bindings: Bindings): LazyStream[Bindings] = {
    val fresh = new Var

    fn(fresh).attempt(bindings)
  }
}

/**
 * Disjunction tries each Goal, and succeeds as long as one of them does.
 *
 * See also: Goal.or
 */
case class Disjunction(goals: Goal*) extends Goal {
  def attempt(bindings: Bindings): LazyStream[Bindings] = {
    val bindingStreams = goals.map(_.attempt(bindings))
    interleave(bindingStreams)
  }
}

/**
 * Conjunction only succeeds if all of its Goals succeed.
 *
 * See also: Goal.and
 */
case class Conjunction(goals: Goal*) extends Goal {
  /*
   * It is easier to deal with Conjunction if we can reduce the Conjunction
   * down into simpler cases, since it allows a more natural use of recursion.
   *
   * In fact, microKanren defines them in exactly this way. I implemented 
   * Goal.and later, which is why I decided to make both Disjunction and
   * Conjunction n-ary rather than binary goals.
   */
  trait Junction
  case class Bijunction(a: Goal, b: Junction) extends Junction
  case object EmptyJunction extends Junction

  /**
   * Converts a group of Goals into a binary tree of Goals and Junctions.
   */
  private def toJunction(goals: List[Goal]): Junction =
    goals match {
      case head :: rest => Bijunction(head, toJunction(rest))
      case Nil          => EmptyJunction
    }

  private val goalJunctions = toJunction(goals.toList)
    
  @tailrec
  private def attemptJunction(
      bindStream: LazyStream[Bindings], junction: Junction): LazyStream[Bindings] =
    junction match {
      case Bijunction(head, tail) =>
        val headStream = bindStream.flatMap(head.attempt(_))
        attemptJunction(headStream, tail)
      case EmptyJunction => bindStream
    }

  def attempt(bindings: Bindings): LazyStream[Bindings] =
    attemptJunction(new MatureStream(bindings, () => EmptyLazyStream), 
                    goalJunctions)
}

/**
 * Negation is used to ensure that a Goal fails, rather than succeeds.
 *
 * Note, however, that this is more "imperative" than the rest of this language.
 * You cannot, for example, discern all immortal beings this way:
 * 
 * {{{
 *   val isMortal = (thing: Term) => ...
 *   val isImmortal = (thing: Term) => isMortal(thing).negate
 * }}}
 *
 * The reason for this is simple - in order to do this, you'd have to have some
 * representation of "all possible things" so that you could subtract the chunk
 * that are satisfied by isMortal. The problem, then, is that there is no such
 * "group of all possible things" - there are an infinite number of them!
 *
 * Negation can only be used to test a given thing - you can't use it for 
 * unification, since it doesn't produce useful results outside of a "this thing
 * is false" or "this thing is true."
 *
 * See also: Goal.negate
 */
case class Negate(goal: Goal) extends Goal {
  def attempt(bindings: Bindings): LazyStream[Bindings] = {
    val forcedResult = goal.attempt(bindings).force

    if (forcedResult.isEmpty) new MatureStream(bindings, () => EmptyLazyStream)
    else                      EmptyLazyStream
  }
} 

/**
 * This starts executing a Goal, returning a Stream of answers.
 */
def run(goal: Goal, bindings: Bindings = Map()) = {
  val result = goal.attempt(bindings)
  result.force.toStream
}

/**
 * Like run(), but it constrains the output and reifies it.
 */
def reifyRun(goal: Goal, bindings: Bindings = Map(), max: Int = 0) = {
  val result = run(goal, bindings)
  if (max == 0) reifyBinds(result)
  else          reifyBinds(result.take(max))
}

/**************** RELATIONS ****************/

/*
 * Relations are a group of fixed facts, which may need to be consulted in
 * order to compute a relation.
 *
 * Note that the simpler microKanren method would be to do this:
 *
 * {{{
 *   val isMan = Fresh { thing => Term
 *    Unify(thing, "Socrates") or
 *    Unify(thing, "Martin Odersky")
 *   }
 * }}}
 *
 * However, this isn't extensible - I want to be able to declare facts, and
 * add to them rather than encode them in this way. Using relations, I can do
 * this:
 *
 * {{{
 *  val man = new UnaryRelation
 *  man += "Socrates"
 *  man += "Martin Odersky"
 *  val isMan = (thing: Term) => man(thing)
 * }}}
 */

/**
 * This is capable of servicing all n-ary relations, by unifying each member of
 * the relation against its input.
 *
 * See also: UnaryRelation, BinaryRelation, TrinaryRelation
 */
case class RelationUnifier(terms: List[Term], facts: List[List[Term]]) extends Goal {
  val termsCons = listToCons(terms)

  def attempt(bindings: Bindings): LazyStream[Bindings] =
    unifyRecursively(bindings, facts)

  private def unifyRecursively(bindings: Bindings, facts: List[List[Term]]): LazyStream[Bindings] = 
    facts match {
      case head :: tail =>
        val factCons = listToCons(head)
        val unifyTail = () => unifyRecursively(bindings, tail)

        unify(termsCons, factCons, bindings) match {
          case Some(newBindings) => new MatureStream(newBindings, unifyTail)
          case None => new ImmatureStream(unifyTail)
        }
      case Nil =>
        EmptyLazyStream
    }
}

/**
 * A relation over a single value.
 */
class UnaryRelation extends (Term => Goal) {
  private val values = HashSet[Term]()
  override def toString = s"Relation[1]:\n    - ${values.mkString("\n    - ")}"

  /**
   * Adds a fact to this relation - a fact is simply an assertion that the
   * given element participates in this relation. Note that we operate under
   * the closed-world assumption - if the relation doesn't include X, then
   * the relation doesn't apply to X.
   */
  def +=(a: Term) = values += a

  /**
   * This unifies the given term against each member of the relations, by
   * using RelationUnifier to do the heavy lifting.
   */
  def apply(a: Term): Goal = 
    RelationUnifier(List(a), values.toList.map(List(_)))
}

/**
 * A relation over two values.
 */
class BinaryRelation extends ((Term, Term) => Goal) {
  private val values = HashSet[(Term, Term)]()
  override def toString = s"Relation[2]:\n    - ${values.mkString("\n    - ")}"

  def +=(a: Term, b: Term) = values += ((a, b))
  def apply(a: Term, b: Term): Goal =
    RelationUnifier(List(a, b), values.toList.map { case (x, y) => List(x, y) })
}

/**
 * A relation over three values.
 */
class TernaryRelation extends ((Term, Term, Term) => Goal) {
  private val values = HashSet[(Term, Term, Term)]()
  override def toString = s"Relation[3]:\n    - ${values.mkString("\n    - ")}"

  def +=(a: Term, b: Term, c: Term) = values += ((a, b, c))
  def apply(a: Term, b: Term, c: Term): Goal =
    RelationUnifier(List(a, b, c), values.toList.map { case (x, y, z) => List(x, y, z) })
}

/**************** TEST ****************/

/**
 * A small twist on the standard "Socrates is mortal" syllogism.
 * The one thing this doesn't test is recursive relations like:
 * 
 * {{{
 *  val fives = (thing: Term) => 
 *   Unify(thing, 5) or Fresh(fives).suspend
 *  
 *  run(Fresh(fives))
 * }}}
 * 
 * This should produce an infinite stream of fives. I've tested such
 * constructions (they were an integral part in getting some parts
 * correct), and they should work.
 */
val dragon = new UnaryRelation
dragon += "Numinex"
dragon += "Alduin"
dragon += "Dovahkiin" // Meh. It ensures that conjugation works.

val man = new UnaryRelation
man += "Socrates"
man += "Martin Odersky"
man += "Dovahkiin"

val explicitlyMortal = new UnaryRelation
explicitlyMortal += "Spongebob Squarepants" // Probably true.

// A thing is mortal iif:
// - It is a man, and not a dragon.
// - It is given as explicitly mortal.
val isMortal = (thing: Term) => 
  (man(thing) and dragon(thing).negate) or explicitlyMortal(thing)

println(s"All Mortals: ${reifyRun(Fresh(isMortal))}")