reset/shift in scala

reset-shift.scala

class Shift[+A, -B, +C](val fun: (A => B) => C) {
  def map[A1](f: (A => A1)): Shift[A1, B, C] = {
    new Shift((k: A1 => B) => fun { (x: A) => k(f(x)) })
  }

  def flatMap[A1, B1, C1<:B](f: A => Shift[A1, B1, C1]): Shift[A1, B1, C] = {
    new Shift((k: A1 => B1) =>
      fun { (x: A) => f(x).fun(k) })
  }
}

def reset[A, C](c: Shift[A, A, C]) = c.fun((x: A) => x)


// create a shift, and maping.
val ctx = for {
  x <- new Shift((k: Int => Int) => k(k(k(7))))
} yield x + 1

reset(ctx) // => 10