Data-Reify.scala

// see http://www.ittc.ku.edu/~andygill/papers/reifyGraph.pdf
// and https://hackage.haskell.org/package/data-reify

import cats.Applicative
import cats.syntax.all._
import cats.effect.IO

sealed abstract class Bit
object Bit {
  final class Xor(a: => Bit, b: => Bit) extends Bit {
    lazy val left: Bit = a
    lazy val right: Bit = b
  }
  final class Delay(a: => Bit) extends Bit {
    lazy val arg: Bit = a
  }
  final class Input(val values: List[Boolean]) extends Bit
  final class Var(val name: String) extends Bit
}

class Network {
  import Bit._
  
  lazy val input: Bit = new Var("input")
  lazy val parity: Bit = new Xor(new Delay(parity), input)
}

sealed abstract class BitF[+A]
object BitF {
  final case class Xor[A](a: A, b: A) extends BitF[A]
  final case class Delay[A](a: A) extends BitF[A]
  final case class Input(values: List[Boolean]) extends BitF[Nothing]
  final case class Var(name: String) extends BitF[Nothing]
}

val network: Network = new Network

println(Graph.reify(network.parity).unsafeRunSync)
// Graph(Map(1 -> Delay(0), 2 -> Var(input), 0 -> Xor(1,2)),0)

/////////////////////////////////////////////////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////////////////////////////////////////////////

type Unique = BigInt
final case class Graph[E[_]](nodes: Map[Unique, E[Unique]], id: Unique)

trait StableName[A] {
  type Id
  def stableName(a: A): IO[Id]
}
object StableName {
  final class AnyRefId[A <: AnyRef](val value: A) {
    override def equals(that : Any): Boolean = that match {
      case that : AnyRefId[b] => value eq that.value
      case _ => false
    }
    
    override def hashCode: Int = System.identityHashCode(value)
  }
  
  implicit def anyRefStableName[A <: AnyRef]: StableName[A] { type Id = AnyRefId[A] } =
    new StableName[A] {
      type Id = AnyRefId[A]
      def stableName(a: A): IO[Id] = IO { new AnyRefId(a) }
    } 
}

trait MuRef[A] {
  type DeRef[X]
  def mapDeRef[F[_], U](a: A)(f: A => F[U])(implicit F: Applicative[F]): F[DeRef[U]]
}

implicit val bit: MuRef[Bit] { type DeRef[X] = BitF[X] } =
  new MuRef[Bit] {
    type DeRef[X] = BitF[X]
    def mapDeRef[F[_], U](bit: Bit)(f: Bit => F[U])(implicit F: Applicative[F]): F[DeRef[U]] = bit match {
      case bit : Bit.Xor   => (f(bit.left), f(bit.right)).mapN(BitF.Xor.apply)
      case bit : Bit.Delay => f(bit.arg).map(BitF.Delay.apply)
      case bit : Bit.Input => F.pure(BitF.Input(bit.values))
      case bit : Bit.Var   => F.pure(BitF.Var(bit.name))  
    }
  }

object Graph {
  def reify[T](t: T)(implicit T: MuRef[T], S: StableName[T]): IO[Graph[T.DeRef]] = for {
    ctx <- Context.make[S.Id, T.DeRef]
    r   <- reifyWithContext[T, S.Id, T.DeRef](t, ctx)(T, S)
  } yield r
  
  private[this] final class Context[RefId, F[_]] private () {
    private[this] var _idMap: Map[RefId, BigInt] = Map.empty
    private[this] var _pfMap: Map[BigInt, F[BigInt]] = Map.empty
    private[this] var _uniqueId: BigInt = BigInt(0)
    
    def pfMap: IO[Map[BigInt, F[BigInt]]] = IO { _pfMap }
    
    def lookupRefId(id: RefId): IO[Option[BigInt]] = IO { _idMap.get(id) }
    
    def newUnique: IO[BigInt] = IO { 
      val result = _uniqueId
      _uniqueId += 1
      result
    }
    
    def addRefId(id: RefId, gid: BigInt): IO[Unit] = IO { _idMap += id -> gid }
    
    def addPF(from: BigInt, ref: F[BigInt]): IO[Unit] = IO { _pfMap += from -> ref }
  }
  private[this] object Context {
    def make[RefId, F[_]]: IO[Context[RefId, F]] = IO { new Context[RefId, F] }
  }
  
  private[this] def reifyWithContext[T, I, F[_]](t: T, ctx: Context[I, F])
  (implicit T: MuRef[T] { type DeRef[X] = F[X] }, S: StableName[T] { type Id = I }): IO[Graph[F]] = for {
    root  <- findNodes[T, I, F](t, ctx, Set())
    pfMap <- ctx.pfMap
  } yield Graph(pfMap, root)
  
  private[this] def findNodes[T, I, F[_]](j: T, ctx: Context[I, F], s: Set[BigInt])
  (implicit T: MuRef[T] { type DeRef[X] = F[X] }, S: StableName[T] { type Id = I }): IO[BigInt] = for {
    refId  <- S.stableName(j)
    optId  <- ctx.lookupRefId(refId)
    root   <- optId match {
      case Some(id) => 
        if (s.contains(id)) IO.pure(id)
        else for {
          res <- T.mapDeRef(j)(n => findNodes[T, I, F](n, ctx, s + id))
          _   <- ctx.addPF(id, res)
        } yield id
      case None => for {
        id  <- ctx.newUnique
        _   <- ctx.addRefId(refId, id)
        res <- T.mapDeRef(j)(n => findNodes[T, I, F](n, ctx, s + id))
        _   <- ctx.addPF(id, res)
      } yield id
    }
  } yield root 
}