GraphWalk.scala

import org.scalacheck.Arbitrary
import org.scalacheck.Prop.forAll


object GraphTest extends Properties("GraphTest") {
  def find[K,V](walkfn: Set[K] => Map[K,Set[K]], endsOnly: Boolean)(
    s: Set[K], acc: Set[K] = Set[K](), visited: Set[K] = Set[K]()): Set[K] = {
    if(s.isEmpty) {
      acc
    }
    else {
      val node = walkfn(s)
      val nextAcc = if(endsOnly) {
        // Only add to the acc if K has no children
        s.foldLeft(acc) { (oldAcc, k) =>
          if(node.get(k).map { _.size }.getOrElse(0) == 0) (oldAcc + k) else oldAcc
        }
      }
      else {
        // Add the whole keyset:
        acc ++ s
      }
      // Whom have we visited?
      val nextVisited = visited ++ s
      // Whom to visit next?
      val nextVis: Set[K] = (node.flatMap { _._2 }.toSet -- nextVisited)
      //println(nextVis, nextAcc, nextVisited)
      find(walkfn, endsOnly)(nextVis, nextAcc, nextVisited)
    }
  }

  def multiGet[K](from: Map[K,Set[K]]) = { (s: Set[K]) =>
    s.flatMap { k => from.get(k).flatMap { v => Some((k,v)) } }.toMap
  }

  property("union commutes") = forAll { (g: Map[Int,Set[Int]], keys: Set[Int]) =>
    find(multiGet(g),  false)(keys) == (keys.flatMap { k => find(multiGet(g), false)(Set(k)) })
  }
  property("on-keyset gets all") = forAll { (g: Map[Int,Set[Int]]) =>
    find(multiGet(g),  false)(g.keySet) == (g.keySet ++ (g.flatMap { _._2 }.toSet))
  }
  property("dead-ends are real") = forAll { (g: Map[Int,Set[Int]], keys: Set[Int]) =>
    val deadEnds = find(multiGet(g), true)(keys)
    multiGet(g)(deadEnds).map { _._2.size }.sum == 0
  }
}