PureCaching.scala

type ShapeId = String                                                            //defined type alias ShapeId
type Hints = List[String]                                                        //defined type alias Hints
type Field[A] = (String, Schema[A])                                              //defined type alias Field
type Alt[U, A] = (String, Schema[A])                                             //defined type alias Alt

trait PolyF[F[_], G[_]] {                                                        //defined trait PolyF
  def apply[A](fa: F[A]): G[A]
}
type ~>[F[_], G[_]] = PolyF[F, G]                                                //defined type alias ~>

type Visitor[F[_]] = Schema ~> F                                                 //defined type alias Visitor

trait MonoidK[F[_]] {                                                            //defined trait MonoidK
  def appendK[A](a1: F[A], a2: F[A]): F[A]
}

final class Lazy[A](make: () => A) {                                             //defined class Lazy
  private[this] var thunk: () => A = make
  lazy val value: A = {
    val result = thunk()
    thunk = null
    result
  }
  def map[B](f: A => B): Lazy[B] = new Lazy(() => f(make()))
}

sealed trait Schema[A] {                                                         //defined trait Schema
  def shapeId: ShapeId
  def hints: Hints
}
object Schema {                                                                  //defined object Schema
  final case class Primitive(shapeId: ShapeId, hints: Hints) extends Schema[Int]
  final case class Struct[S](
      shapeId: ShapeId,
      hints: Hints,
      field: List[Field[_]]
  ) extends Schema[S]
  final case class Union[U](
      shapeId: ShapeId,
      hints: Hints,
      alts: List[Alt[U, _]]
  ) extends Schema[U]
  final case class LazySchema[S](suspend: Lazy[Schema[S]]) extends Schema[S] {
    override def shapeId: ShapeId = suspend.value.shapeId
    override def hints: Hints = suspend.value.hints
  }

  def lazily[A](s: => Schema[A]): Schema[A] = LazySchema(new Lazy(() => s))
}

case class SZip[A](focus: Set[Schema[_]], cache: Set[Schema[_]], state: A) {     //defined class SZip
  def map[B](f: A => B): SZip[B] = SZip(focus, cache, f(state))
  def extract: A = state
  def extend[B](f: SZip[A] => B): SZip[B] = SZip(focus, cache, f(this))

  // Return any unvisited children if empty
  def next: SZip[A] = {
    val nextSchemas: Set[Schema[_]] = focus.toList.flatMap {
      case Schema.Primitive(_, _) => List.empty
      case Schema.Struct(_, _, field) =>
        field.map(_._2).filterNot(ss => cache.contains(ss))
      case Schema.Union(_, _, alts) =>
        alts.map(_._2).filterNot(ss => cache.contains(ss))
      case Schema.LazySchema(suspend) =>
        List(suspend.value)
    }.toSet
    SZip(nextSchemas, cache ++ nextSchemas, state)
  }
}
object SZip {                                                                    //defined object SZip
  def apply[A](schema: Schema[_], a: A): SZip[A] = {
    val singleton = Set[Schema[_]](schema)
    SZip(singleton, singleton, a)
  }

  def visit[A](schema: Schema[_])(f: Schema[_] => A): List[A] = {
    val init = SZip(schema, f(schema))
    def go(zip: SZip[A]): List[A] = {
      if (zip.focus.isEmpty) {
        List.empty
      } else {
        val cacheStr = zip.cache.map(ce =>
          s"${ce.shapeId}: ${ce.getClass.getSimpleName.split('$').last}"
        )
        println(s"cache: $cacheStr")
        zip.focus.toList.map(f) ++ go(zip.next)
      }
    }
    go(init)
  }

  def visitF[F[_], A](
      schema: Schema[A]
  )(f: Schema ~> F)(implicit M: MonoidK[F]): F[A] = {
    val init = SZip(schema, f(schema))
    def go(zip: SZip[F[A]]): List[F[A]] = {
      if (zip.focus.isEmpty) {
        List.empty
      } else {
        println(
          zip.cache.map(ce =>
            s"${ce.shapeId}: ${ce.getClass.getSimpleName.split('$').last}"
          )
        )
        zip.focus.toList.map(s => f(s.asInstanceOf[Schema[A]])) ++ go(zip.next)
      }
    }
    go(init).reduce(M.appendK[A])
  }
}

//Visitors
type Hash[A] = Int                                                               //defined type alias Hash
lazy val hashVisitor: Visitor[Hash] = new Visitor[Hash] {                        //hashVisitor: Visitor[Hash] = $anon$1@36c832b4
  override def apply[A](fa: Schema[A]): Hash[A] = fa match {
    case Schema.Primitive(shapeId, hints) => shapeId.hashCode + hints.hashCode()
    case Schema.Struct(shapeId, hints, field) =>
      shapeId.hashCode + hints.hashCode()
    case Schema.Union(shapeId, hints, alts) =>
      shapeId.hashCode + hints.hashCode()
    case Schema.LazySchema(suspend) => hashVisitor(suspend.value)
  }
}
implicit val hashMonoidK: MonoidK[Hash] = new MonoidK[Hash] {                    //hashMonoidK: MonoidK[Hash] = $anon$2@b5b40a7
  override def appendK[A](a1: Hash[A], a2: Hash[A]) = a1 * a2
}

//Impl
sealed trait Tree                                                                //defined trait Tree
object Tree {                                                                    //defined object Tree
  final case class Node(left: Tree, right: Tree) extends Tree
  final case class Leaf(n: Int) extends Tree
}

lazy val nodeSchema: Schema[Tree.Node] = Schema.Struct(                          //nodeSchema: Schema[Tree.Node] = Struct(Node,List(),List((left,LazySchema(A$A11$A$A11$Lazy@7578513)), (right,LazySchema(A$A11$A$A11$Lazy@7578513))))
  "Node",
  List.empty,
  List(
    "left" -> treeSchema,
    "right" -> treeSchema
  )
)
lazy val leafSchema: Schema[Tree.Leaf] = Schema.Struct(                          //leafSchema: Schema[Tree.Leaf] = Struct(Leaf,List(),List((n,Primitive(n,List()))))
  "Leaf",
  List.empty,
  List("n" -> Schema.Primitive("n", List.empty))
)
lazy val treeSchema: Schema[Tree] = Schema.lazily {                              //treeSchema: Schema[Tree] = LazySchema(A$A11$A$A11$Lazy@7578513)
  Schema.Union[Tree](
    "Tree",
    List.empty,
    List[Alt[Tree, _]](
      "Node" -> nodeSchema,
      "Leaf" -> leafSchema
    )
  )
}

sealed trait Json                                                                //defined trait Json
object Json {                                                                    //defined object Json
  case class JObject(fields: Map[String, Json]) extends Json
  case class JArray(xs: List[Json]) extends Json
  case class JString(str: String) extends Json
  case class JNum(n: Int) extends Json
}

lazy val jObjectSchema: Schema[Json.JObject] = Schema.Struct(                    //jObjectSchema: Schema[Json.JObject] = Struct(object,List(),List((fields,LazySchema(A$A11$A$A11$Lazy@6544570f))))
  "object",
  List.empty,
  List(
    "fields" -> jSchema
  )
)
lazy val jArraySchema: Schema[Json.JArray] =                                     //jArraySchema: Schema[Json.JArray] = Struct(array,List(),List((xs,LazySchema(A$A11$A$A11$Lazy@6544570f))))
  Schema.Struct("array", List.empty, List("xs" -> jSchema))
lazy val jStringSchema: Schema[Json.JString] = Schema.Struct(                    //jStringSchema: Schema[Json.JString] = Struct(jString,List(),List((str,Primitive(String,List()))))
  "jString",
  List.empty,
  List("str" -> Schema.Primitive("String", List.empty))
)
lazy val jNumSchema: Schema[Json.JNum] = Schema.Struct(                          //jNumSchema: Schema[Json.JNum] = Struct(jNum,List(),List((num,Primitive(Number,List()))))
  "jNum",
  List.empty,
  List("num" -> Schema.Primitive("Number", List.empty))
)
lazy val jSchema: Schema[Json] = Schema.lazily {                                 //jSchema: Schema[Json] = LazySchema(A$A11$A$A11$Lazy@6544570f)
  Schema.Union[Json](
    "Json",
    List.empty,
    List[Alt[Json, _]](
      "object" -> jObjectSchema,
      "array" -> jArraySchema,
      "string" -> jStringSchema,
      "number" -> jNumSchema
    )
  )
}

SZip.visit(treeSchema) {                                                         //cache: Set(Tree: LazySchema)
  case Schema.Primitive(shapeId, _) => shapeId                                   //cache: Set(Tree: LazySchema, Tree: Union)
  case Schema.Struct(shapeId, _, _) => shapeId                                   //cache: Set(Tree: LazySchema, Tree: Union, Node: Struct, Leaf: Struct)
  case Schema.Union(shapeId, _, _)  => shapeId                                   //cache: Set(Leaf: Struct, Node: Struct, Tree: LazySchema, Tree: Union, n: Primitive)
  case Schema.LazySchema(_)         => "Lazy"                                    //res0: List[ShapeId] = List(Lazy, Tree, Node, Leaf, n)
}

SZip.visitF(treeSchema)(hashVisitor)                                             //Set(Tree: LazySchema)
                                                                                 //Set(Tree: LazySchema, Tree: Union)
                                                                                 //Set(Tree: LazySchema, Tree: Union, Node: Struct, Leaf: Struct)
                                                                                 //Set(Leaf: Struct, Node: Struct, Tree: LazySchema, Tree: Union, n: Primitive)
                                                                                 //res1: Hash[Tree] = 1356268128
SZip.visitF(jSchema)(hashVisitor)                                                //Set(Json: LazySchema)
                                                                                 //Set(Json: LazySchema, Json: Union)
                                                                                 //Set(array: Struct, jString: Struct, Json: LazySchema, object: Struct, jNum: Struct, Json: Union)
                                                                                 //Set(array: Struct, jString: Struct, Json: LazySchema, object: Struct, jNum: Struct, String: Primitive, Json: Union, Number: Primitive)
                                                                                 //res2: Hash[Json] = 2008906496
SZip.visitF(jObjectSchema)(hashVisitor)                                          //Set(object: Struct)
                                                                                 //Set(object: Struct, Json: LazySchema)
                                                                                 //Set(object: Struct, Json: LazySchema, Json: Union)
                                                                                 //Set(array: Struct, jString: Struct, Json: LazySchema, object: Struct, jNum: Struct, Json: Union)
                                                                                 //Set(array: Struct, jString: Struct, Json: LazySchema, object: Struct, jNum: Struct, String: Primitive, Json: Union, Number: Primitive)
                                                                                 //res3: Hash[Json.JObject] = 2008906496