build.sbt

scalaVersion := "2.11.8"

libraryDependencies ++= Seq(
  "com.chuusai" %% "shapeless" % "2.3.2",
  "org.typelevel" %% "cats" % "0.8.1",
  "org.scalacheck" %% "scalacheck" % "1.13.4"
)

DAGGen.scala

package net.tixxit.daggen

import cats.{Apply, Monad}
import cats.data.StateT
import cats.syntax.functor._

import shapeless.{::, HList, HNil}
import shapeless.ops.hlist._

import org.scalacheck.Arbitrary.arbitrary
import org.scalacheck.{Gen, rng}

object gen {
  implicit val genMonad: Monad[Gen] =
    new Monad[Gen] {
      def pure[A](x: A): Gen[A] = Gen.const(x)
      def flatMap[A, B](fa: Gen[A])(f: A => Gen[B]): Gen[B] = fa.flatMap(f)
      override def map[A, B](fa: Gen[A])(f: A => B): Gen[B] = fa.map(f)
      def tailRecM[A, B](a: A)(f: A => Gen[Either[A,B]]): Gen[B] =
        f(a).flatMap {
          case Left(a2) => tailRecM(a)(f)
          case Right(b) => Gen.const(b)
        }
    }
}

import gen._

trait GetNodeState[In <: HList] {
  type Out <: HList

  def empty: Out
}

object GetNodeState {
  type Aux[I <: HList, O <: HList] = GetNodeState[I] { type Out = O }

  implicit def hnil: Aux[HNil, HNil] = new GetNodeState[HNil] {
    type Out = HNil
    def empty: Out = HNil
  }

  implicit def cons[H, T <: HList, OT <: HList](implicit tail: Aux[T, OT]): Aux[H :: T, Set[H] :: OT] =
    new GetNodeState[H :: T] {
      type Out = Set[H] :: tail.Out
      def empty: Out = Set.empty[H] :: tail.empty
    }
}

trait DAG[L <: HList] {
  type NodeState <: HList
  type Node[A] = StateT[Gen, NodeState, A]

  def initialNodeState: NodeState

  def node0[A](genA: Gen[A])(implicit
    selector: Selector[NodeState, Set[A]],
    replacer: Replacer.Aux[NodeState, Set[A], Set[A], (Set[A], NodeState)]
  ): Node[A] = {
    StateT[Gen, NodeState, A] { origState =>
      val available = selector(origState)
      val gen: Gen[A] =
        if (available.isEmpty) genA
        else Gen.frequency(available.size -> Gen.oneOf(available.toSeq), 1 -> genA)
      gen.map { a =>
        val (_, nextState) = replacer(origState, available + a)
        (nextState, a)
      }
    }
  }

  def node1[A, B](na: => Node[A])(f: A => Gen[B])(implicit
    selector: Selector[NodeState, Set[B]],
    replacer: Replacer.Aux[NodeState, Set[B], Set[B], (Set[B], NodeState)]
  ): Node[B] = StateT[Gen, NodeState, B] { s =>
    na.run(s).flatMap { case (s0, a) =>
      node0(f(a)).run(s0)
    }
  }

  def node2[A, B, C](na: => Node[A], nb: => Node[B])(f: (A, B) => Gen[C])(implicit
    selector: Selector[NodeState, Set[C]],
    replacer: Replacer.Aux[NodeState, Set[C], Set[C], (Set[C], NodeState)]
  ): Node[C] = StateT[Gen, NodeState, C] { s =>
    na.run(s).flatMap { case (s0, a) =>
      nb.run(s0).flatMap { case (s1, b) =>
        node0(f(a, b)).run(s1)
      }
    }
  }

  def frequency[A](nodes: (Int, Node[A])*): Node[A] =
    StateT[Gen, NodeState, A] { s =>
      val gens = nodes.map { case (w, n) => (w, n.run(s)) }
      Gen.frequency(gens: _*)
    }

  def toGen[A](node: Node[A]): Gen[A] = node.runA(initialNodeState)
}

object DAG {
  type Aux[L <: HList, S <: HList] = DAG[L] { type NodeState = S }

  def apply[L <: HList](implicit getNodeState: GetNodeState[L]): Aux[L, getNodeState.Out] =
    new DAG[L] {
      type NodeState = getNodeState.Out
      def initialNodeState: NodeState = getNodeState.empty
    }
}

object Example extends App {
  // Define our bipartite DAG's ADT.
  case class Foo(x: Bar, y: Bar)
  sealed trait Bar
  case class Baz(foo: Foo) extends Bar
  case class Qux(n: Int) extends Bar

  // Define our DAG's structure.
  type Nodes = Foo :: Baz :: Qux :: HNil
  val dag = DAG[Nodes]
  lazy val qux: dag.Node[Qux] = dag.node0(arbitrary[Int].map(Qux(_)))
  lazy val baz: dag.Node[Baz] = dag.node1(foo)(Baz(_))
  lazy val bar: dag.Node[Bar] = dag.frequency[Bar](1 -> qux.widen, 1 -> baz.widen)
  lazy val foo: dag.Node[Foo] = dag.node2(bar, bar)(Foo(_, _))

  // Create the Scalacheck generator.
  val genFoo: Gen[Foo] = dag.toGen(foo)
  (1 to 10).foreach { _ =>
    println(genFoo(Gen.Parameters.default, rng.Seed.random()))
  }
}