Idiomatic Scala translation of Kiselyov's zippers?

Oleg Kiselyov showed how to make a zipper from any traversable by using delimited continuations. His Haskell code is quite short:

module ZipperTraversable where 

import qualified Data.Traversable as T
import qualified Data.Map as M


-- In the variant Z a k, a is the current, focused value
-- evaluate (k Nothing) to move forward
-- evaluate (k v)       to replace the current value with v and move forward.

data Zipper t a = ZDone (t a) 
                | Z a (Maybe a -> Zipper t a)

make_zipper :: T.Traversable t => t a -> Zipper t a
make_zipper t = reset $ T.mapM f t >>= return . ZDone
 where
 f a = shift (\k -> return $ Z a (k . maybe a id))

-- The Cont monad for delimited continuations, implemented here to avoid
-- importing conflicting monad transformer libraries

newtype Cont r a = Cont{runCont :: (a -> r) -> r}

instance Monad (Cont r) where
    return x = Cont $ \k -> k x
    m >>= f  = Cont $ \k -> runCont m (\v -> runCont (f v) k)

-- Two delimited control operators,
-- without answer-type polymorphism or modification
-- These features are not needed for the application at hand

reset :: Cont r r -> r
reset m = runCont m id

shift :: ((a -> r) -> Cont r r) -> Cont r a
shift e = Cont (\k -> reset (e k))

I've run into a few problems trying to implement it in Scala. I started off trying to use Scala's delimited continuations package, but even using Rompf's richIterable idea generalized to @cps[X] instead of @suspendable, it's impossible to have the function provided to shift return a different type than the function provided to reset.

I tried implementing the continuation monad following Kiselyov's definition, but Scala makes it hard to curry type parameters and I couldn't figure out how to turn Cont[R] into a monad in a way that scalaz's traverse method was happy with.

I'm a beginner at both Haskell and Scala, and would really appreciate help with this.

You can use the continuations plugin. After the plugin does its translation works, it has similarities to the Cont monad and the shift and reset from Oleg. The tricky part is to figure out the types. So here is my translation:

import util.continuations._
import collection.mutable.ListBuffer

sealed trait Zipper[A] { def zipUp: Seq[A] }
case class ZDone[A](val zipUp: Seq[A]) extends Zipper[A]
case class Z[A](current: A, forward: Option[A] => Zipper[A]) extends Zipper[A] {
  def zipUp = forward(None).zipUp
}

object Zipper {
  def apply[A](seq: Seq[A]): Zipper[A] = reset[Zipper[A], Zipper[A]] {
    val coll = ListBuffer[A]()
    val iter = seq.iterator
    while (iter.hasNext) {
      val a = iter.next()
      coll += shift { (k: A=>Zipper[A]) =>
        Z(a, (a1: Option[A]) => k(a1.getOrElse(a)))
      }
    }
    ZDone(coll.toList)
  }
}

The continuations plugin has support for while loop but not for map or flatMap, so I have made the choice of using while and a mutable ListBuffer to capture the possibly updated elements. The make_zipper function is translated into the companion Zipper.apply - a typical Scala location for creating new objects or collection. The data type is translated into a sealed trait with two case classes extended it. And I have put the zip_up function as a method of Zipper with different implementations for each case class. Also pretty typical.

Here it is in action:

object ZipperTest extends App {
  val sample = for (v <- 1 to 5) yield (v, (1 to v).reduceLeft(_ * _))
  println(sample) // Vector((1,1), (2,2), (3,6), (4,24), (5,120))

  def extract134[A](seq: Seq[A]) = {
    val Z(a1, k1) = Zipper(seq)
    val Z(a2, k2) = k1(None)
    val Z(a3, k3) = k2(None)
    val Z(a4, k4) = k3(None)
    List(a1, a3, a4)
  }
  println(extract134(sample)) // List((1,1), (3,6), (4,24))

  val updated34 = {
    val Z(a1, k1) = Zipper(sample)
    val Z(a2, k2) = k1(None)
    val Z(a3, k3) = k2(None)
    val Z(a4, k4) = k3(Some(42 -> 42))
    val z = k4(Some(88 -> 22))
    z.zipUp
  }
  println(updated34) // List((1,1), (2,2), (42,42), (88,22), (5,120))
}

How did I figure the types for shift, k and reset or how to translate T.mapM?

I looked at mapM and I know it will allow me to get a Cont, but I am not sure what is inside the Cont as it depends on the shift. So I start inside the shift. Ignoring the haskell return to construct a Cont, shift returns a Zipper. I also guess that I need to add an element of type A to my collection to build. So the shift will be in the "hole" where an element of type A is expected, therefore k will be a A=>? function. Let's assume this. I'm putting question marks after the types I wasn't so sure about. So I started with:

shift { (k: A?=>?) =>
  Z(a, ?)
}

Next I looked (hard) at (k . maybe a id). The function maybe a id will return an A, so that is consistent with what k takes as an argument. It is the equivalent of a1.getOrElse(a). Also because I need to fill in Z(a, ?), I need to figure out how to get a function from option to Zipper. The simplest way is to assume that k returns a Zipper. Also, looking at how the zipper is used k1(None) or k1(Some(a)), I know I have to give a way for the user to optionally replace the elements, which is what the forward function does. It continues with the original a or an updated a. It starts to make sense. So now I have:

shift { (k: A=>Zipper[A]) =>
  Z(a, (a1: Option[A]) => k(a1.getOrElse(a)))
}

Next I look at mapM again. I see that it is composed with return . ZDone. Ignoring return again (because it is just for the Cont monad), I see that ZDone will take the resulting collection. So that is perfect, I just need to put coll in it and by the time the program arrives there, it will have the updated elements. Also, the type of the expression inside the reset is now consistent with the return type of k which is Zipper[A].

Finally I fill in the type for reset which the compiler can infer, but when I guess right, it gives me a (false) sense of confidence I know what is going on.

My translation is not as general or as pretty as the Haskell one because it does not preserve the type from the collection and uses mutation but hopefully it is easier to understand.

Edit: Here is the version that preserves the type and uses an immutable list so that z.zipUp == z.zipUp:

import util.continuations._
import collection.generic.CanBuild
import collection.SeqLike

sealed trait Zipper[A, Repr] { def zipUp: Repr }
case class ZDone[A, Repr](val zipUp: Repr) extends Zipper[A, Repr]
case class Z[A, Repr](current: A,
    forward: Option[A] => Zipper[A, Repr]) extends Zipper[A, Repr] {
  def zipUp = forward(None).zipUp
}

object Zipper {
  def apply[A, Repr](seq: SeqLike[A, Repr])
                    (implicit cb: CanBuild[A, Repr]): Zipper[A, Repr] = {
    type ZAR = Zipper[A, Repr]

    def traverse[B](s: Seq[A])(f: A => B@cps[ZAR]): List[B]@cps[ZAR] =
      if (s.isEmpty) List()
      else f(s.head) :: traverse(s.tail)(f)

    reset {
      val list = traverse(seq.toSeq)(a => shift { (k: A=>ZAR) =>
        Z(a, (a1: Option[A]) => k(a1.getOrElse(a)))
      })
      val builder = cb() ++= list
      ZDone(builder.result): ZAR
    }
  }
}

By the way, here are additional resources on the continuation monad in scala:

http://blog.tmorris.net/continuation-monad-in-scala/
what I put myself through when I first tried it out: https://gist.github.com/huynhjl/4077185; it includes translating to Scala various Haskell continuation examples as well as some example from Tiark's paper (but not using the continuation plugin which points more clearly the similarity between the approach).
if you download scalaz, you may be able to make Tony Morris's cont a scalaz Monad instance and use scalaz traverse - then the translation to scala would then be a more literal one.

Recommended topics

Hot tags