Codensity of Trees

codensity-tree.scala

sealed trait TreeLike[M[_]] {
  def node[A](f: M[A], g: M[A]): M[A]
}

sealed trait Tree[+A] { self =>
  def flatMap[B](f: A => Tree[B]): Tree[B] = self match {
    case Leaf(a) => f(a)
    case Node(l, r) => Node(l.flatMap(f), r.flatMap(f))
  }
}

case class Leaf[+A](val a: A) extends Tree[A]
case class Node[+A](val l: Tree[A], val r: Tree[A]) extends Tree[A]

sealed trait CTree[A] { self =>
  def apply[B](f: A => Tree[B]): Tree[B]
  def flatMap[B](f: A => CTree[B]): CTree[B] = {
    new CTree[B] {
      def apply[C](h: B => Tree[C]): Tree[C] = {
        self.apply { a: A => f(a).apply(h) }
      }
    }
  }
}

object CTree {
  def treeabs[A](t: CTree[A]): Tree[A] = t(a => Leaf(a))
  def treerep[A](t: Tree[A]): CTree[A] = t match {
    case Leaf(a) => new CTree[A] {
      def apply[B](k: A => Tree[B]): Tree[B] = k(a)
    }
    case Node(l, r) => new CTree[A] {
      def apply[B](k: A => Tree[B]): Tree[B] = {
        Node(l.flatMap(k), r.flatMap(k))
      }
    }
  }
}

object Main {

  implicit val treeIsTreeLike = new TreeLike[Tree] {
    def node[A](f: Tree[A], g: Tree[A]): Tree[A] = Node(f, g)
  }

  implicit val cTreeIsTreeLike = new TreeLike[CTree] {
    def node[A](p: CTree[A], q: CTree[A]): CTree[A] = {
      new CTree[A] {
        def apply[B](f: A => Tree[B]): Tree[B] = Node(p(f), q(f))
      }
    }
  }

  def sampleTree: Tree[Int] = {
    val f = { n: Int => Node(Leaf(n + 1), Leaf(n + 1)) }
    Leaf(0).flatMap(f).flatMap(f)
  }
}