monadicDC.scala

package monadicDC

import scala.annotation.tailrec

enum Equal[A, B] {
  case EQ[A]() extends Equal[A, A]
  case NE()
}

opaque type Prompt[A] = Int

object Prompt {
  enum P[Ans] {
    case Done(value: Ans)
    case More[Ans, A](k: Prompt[A] => P[Ans]) extends P[Ans]

    def run: Ans = {
      @tailrec def go(n: Int, p: P[Ans]): Ans =
        p match {
          case Done(a) => a
          case More(k) => go(n + 1, k(n))
        }
      go(0, this)
    }
  }
}

extension [A](self: Prompt[A]) {
  def abort[B](e: CC[A]): CC[B] = withSubCont(_ => e)

  def withSubCont[B](f: SubCont[B, A] => CC[A]): CC[B] = CC.WithSubCont(self, f)

  def control[B](f: SubCont[B, A] => CC[A]): CC[B] =
    withSubCont(k => pushPrompt(self)(f(k)))

  def shift[B](f: SubCont[B, A] => CC[A]): CC[B] =
    withSubCont(k => pushPrompt(self)(f(k ++ SubCont.PushPrompt(self))))

  def ===[B](that: Prompt[B]): Equal[A, B] =
    if (self == that) Equal.EQ().asInstanceOf[Equal[A, B]] else Equal.NE()
}

enum Frame[A, B] {
  case Leaf(value: A => CC[B])
  case Node[A, B, C](left: Frame[A, B], right: Frame[B, C]) extends Frame[A, C]

  def :+[C](f: B => CC[C]): Frame[A, C] = Node(this, Leaf(f))

  def ++[C](that: Frame[B, C]): Frame[A, C] = Node(this, that)

  def apply(a: A): CC[B] = {
    @tailrec def go[A](frame: Frame[A, B], a: A): CC[B] =
      frame match {
        case Leaf(f) => f(a)
        case Node(Leaf(f), r) => CC.FlatMap(f(a), r)
        case Node(Node(ll, lr), r) => go(Node(ll, Node(lr, r)), a)
      }
    go(this, a)
  }
}

enum Cont[Ans, A] {
  case Empty[Ans]() extends Cont[Ans, Ans]
  case PushPrompt(prompt: Prompt[A], cont: Cont[Ans, A]) extends Cont[Ans, A]
  case PushFrame[Ans, A, B](frame: Frame[A, B], cont: Cont[Ans, B]) extends Cont[Ans, A]

  def split[B](prompt: Prompt[B]): (SubCont[A, B], Cont[Ans, B]) =
    this match {
      case Empty() => throw new NoSuchElementException("Prompt was not found on the stack")
      case PushPrompt(p, k) => p === prompt match {
        case Equal.EQ() => (SubCont.Empty(), k)
        case Equal.NE() => k.split(prompt) match {
          case (sk, k) => (SubCont.PushPrompt(p) ++ sk, k)
        }
      }
      case PushFrame(frame, k) => k.split(prompt) match {
        case (sk, k) => (SubCont.PushFrame(frame) ++ sk, k)
      }
    }

  def pushFrame[B](frame: Frame[B, A]): Cont[Ans, B] =
    this match {
      case PushFrame(f, k) => PushFrame(frame ++ f, k)
      case _ => PushFrame(frame, this)
    }

  def pushSubCont[B](sk: SubCont[B, A]): Cont[Ans, B] = {
    @tailrec def go[A](sk: SubCont[B, A], k: Cont[Ans, A]): Cont[Ans, B] =
      sk match {
        case SubCont.Empty() => k
        case SubCont.PushPrompt(p) => Cont.PushPrompt(p, k)
        case SubCont.PushFrame(f) => Cont.PushFrame(f, k)
        case SubCont.Append(l, SubCont.Empty()) => go(l, k)
        case SubCont.Append(l, SubCont.PushPrompt(p)) => go(l, Cont.PushPrompt(p, k))
        case SubCont.Append(l, SubCont.PushFrame(f)) => go(l, Cont.PushFrame(f, k))
        case SubCont.Append(l, SubCont.Append(rl, rr)) => go(SubCont.Append(SubCont.Append(l, rl), rr), k)
      }
    go(sk, this)
  }

  @tailrec
  final def apply(a: A): Prompt.P[Ans] =
    this match {
      case Empty() => Prompt.P.Done(a)
      case PushPrompt(_, k) => k(a)
      case PushFrame(f, k) => f(a)(k)
    }
}

enum SubCont[A, B] {
  case Empty[A]() extends SubCont[A, A]
  case PushPrompt[A](prompt: Prompt[A]) extends SubCont[A, A]
  case PushFrame(frame: Frame[A, B])
  case Append[A, B, C](f: SubCont[A, B], g: SubCont[B, C]) extends SubCont[A, C]

  def ++[C](that: SubCont[B, C]): SubCont[A, C] = Append(this, that)

  def apply(e: CC[A]): CC[B] = CC.PushSubCont(this, e)

  def apply(a: A): CC[B] = CC.PushSubCont(this, CC.Pure(a))
}

enum CC[A] {
  case Pure(value: A)
  case NewPrompt[A]() extends CC[Prompt[A]]
  case PushPrompt(prompt: Prompt[A], e: CC[A])
  case PushSubCont[A, B](sk: SubCont[A, B], e: CC[A]) extends CC[B]
  case WithSubCont[A, B](p: Prompt[A], f: SubCont[B, A] => CC[A]) extends CC[B]
  case FlatMap[A, B](e: CC[A], frame: Frame[A, B]) extends CC[B]

  def apply[Ans](k: Cont[Ans, A]): Prompt.P[Ans] = {
    @tailrec def go[A](e: CC[A], k: Cont[Ans, A]): Prompt.P[Ans] = e match {
      case Pure(a) => k(a)
      case _: NewPrompt[a] => Prompt.P.More[Ans, a](p => k(p))
      case PushPrompt(p, e) => go(e, Cont.PushPrompt(p, k))
      case PushSubCont(sk, e) => go(e, k.pushSubCont(sk))
      case WithSubCont(p, f) => k.split(p) match {
        case (sk, k) => go(f(sk), k)
      }
      case FlatMap(e, f) => go(e, k.pushFrame(f))
    }
    go(this, k)
  }

  def map[B](f: A => B): CC[B] = flatMap(a => Pure(f(a)))

  def flatMap[B](f: A => CC[B]): CC[B] =
    this match {
      case Pure(a) => f(a)
      case FlatMap(e, fs) => FlatMap(e, fs :+ f)
      case _ => FlatMap(this, Frame.Leaf(f))
    }

  def ap[B](f: CC[A => B]): CC[B] = flatMap(a => f.map(_(a)))

  def zipWith[B, C](that: CC[B])(f: (A, B) => C): CC[C] = ap(that.map(b => a => f(a, b)))

  def run: A = this match {
    case Pure(a) => a
    case _ => this(Cont.Empty()).run
  }
}

object CC {
  def apply[A](a: A): CC[A] = Pure(a)
}

def newPrompt[A]: CC[Prompt[A]] = CC.NewPrompt()

def pushPrompt[A](p: Prompt[A])(e: CC[A]): CC[A] = CC.PushPrompt(p, e)

def reset[A](f: Prompt[A] => CC[A]): CC[A] = newPrompt[A].flatMap(p => pushPrompt(p)(f(p)))