WIP: Functional Cached Schema Visitors

modules/bootstrapped/src/generated/smithy4s/example/LeafNode.scala

@@ -0,0 +1,23 @@
+package smithy4s.example
+
+import smithy4s.Hints
+import smithy4s.Schema
+import smithy4s.ShapeId
+import smithy4s.ShapeTag
+import smithy4s.schema.Schema.int
+import smithy4s.schema.Schema.struct
+
+final case class LeafNode(value: Int)
+
+object LeafNode extends ShapeTag.Companion[LeafNode] {
+  val id: ShapeId = ShapeId("smithy4s.example", "LeafNode")
+
+  val hints: Hints = Hints.empty
+
+  // constructor using the original order from the spec
+  private def make(value: Int): LeafNode = LeafNode(value)
+
+  implicit val schema: Schema[LeafNode] = struct(
+    int.required[LeafNode]("value", _.value),
+  )(make).withId(id).addHints(hints)
+}

modules/bootstrapped/src/generated/smithy4s/example/Tree.scala

@@ -0,0 +1,67 @@
+package smithy4s.example
+
+import smithy4s.Hints
+import smithy4s.Schema
+import smithy4s.ShapeId
+import smithy4s.ShapeTag
+import smithy4s.schema.Schema.bijection
+import smithy4s.schema.Schema.recursive
+import smithy4s.schema.Schema.union
+
+sealed trait Tree extends scala.Product with scala.Serializable { self =>
+  @inline final def widen: Tree = this
+  def $ordinal: Int
+
+  object project {
+    def tree: Option[TreeNode] = Tree.TreeCase.alt.project.lift(self).map(_.tree)
+    def leaf: Option[LeafNode] = Tree.LeafCase.alt.project.lift(self).map(_.leaf)
+  }
+
+  def accept[A](visitor: Tree.Visitor[A]): A = this match {
+    case value: Tree.TreeCase => visitor.tree(value.tree)
+    case value: Tree.LeafCase => visitor.leaf(value.leaf)
+  }
+}
+object Tree extends ShapeTag.Companion[Tree] {
+
+  def tree(tree: TreeNode): Tree = TreeCase(tree)
+  def leaf(leaf: LeafNode): Tree = LeafCase(leaf)
+
+  val id: ShapeId = ShapeId("smithy4s.example", "Tree")
+
+  val hints: Hints = Hints.empty
+
+  final case class TreeCase(tree: TreeNode) extends Tree { final def $ordinal: Int = 0 }
+  final case class LeafCase(leaf: LeafNode) extends Tree { final def $ordinal: Int = 1 }
+
+  object TreeCase {
+    val hints: Hints = Hints.empty
+    val schema: Schema[Tree.TreeCase] = bijection(TreeNode.schema.addHints(hints), Tree.TreeCase(_), _.tree)
+    val alt = schema.oneOf[Tree]("tree")
+  }
+  object LeafCase {
+    val hints: Hints = Hints.empty
+    val schema: Schema[Tree.LeafCase] = bijection(LeafNode.schema.addHints(hints), Tree.LeafCase(_), _.leaf)
+    val alt = schema.oneOf[Tree]("leaf")
+  }
+
+  trait Visitor[A] {
+    def tree(value: TreeNode): A
+    def leaf(value: LeafNode): A
+  }
+
+  object Visitor {
+    trait Default[A] extends Visitor[A] {
+      def default: A
+      def tree(value: TreeNode): A = default
+      def leaf(value: LeafNode): A = default
+    }
+  }
+
+  implicit val schema: Schema[Tree] = recursive(union(
+    Tree.TreeCase.alt,
+    Tree.LeafCase.alt,
+  ){
+    _.$ordinal
+  }.withId(id).addHints(hints))
+}

modules/bootstrapped/src/generated/smithy4s/example/TreeNode.scala

@@ -0,0 +1,24 @@
+package smithy4s.example
+
+import smithy4s.Hints
+import smithy4s.Schema
+import smithy4s.ShapeId
+import smithy4s.ShapeTag
+import smithy4s.schema.Schema.recursive
+import smithy4s.schema.Schema.struct
+
+final case class TreeNode(left: Tree, right: Tree)
+
+object TreeNode extends ShapeTag.Companion[TreeNode] {
+  val id: ShapeId = ShapeId("smithy4s.example", "TreeNode")
+
+  val hints: Hints = Hints.empty
+
+  // constructor using the original order from the spec
+  private def make(left: Tree, right: Tree): TreeNode = TreeNode(left, right)
+
+  implicit val schema: Schema[TreeNode] = recursive(struct(
+    Tree.schema.required[TreeNode]("left", _.left),
+    Tree.schema.required[TreeNode]("right", _.right),
+  )(make).withId(id).addHints(hints))
+}

modules/bootstrapped/test/src/smithy4s/RecursiveSpec.scala

@@ -0,0 +1,215 @@
+package smithy4s
+
+import munit.FunSuite
+import smithy4s.example.{Foo, LeafNode, Tree, TreeNode}
+import cats.Hash
+import cats.effect.{IO, Sync}
+import smithy4s.capability.{CacheWrite, SyncLike}
+import smithy4s.internals.DocumentKeyEncoder.OptDocumentKeyEncoder
+import smithy4s.internals.{DocumentEncoder, DocumentKeyEncoder, EffectfulDocumentEncoderVisitor, TreeBasedTraversal}
+import smithy4s.schema._
+
+import scala.annotation.tailrec
+import smithy4s.internals.maps.MMap
+import smithy4s.interopcats.{SchemaVisitorHash, TreeBasedHashVisitor, monadThrowShim}
+import smithy4s.kinds.KleisliK
+import smithy4s.kinds.KleisliK.Unwrapped
+import smithy4s.schema.Schema.recursive
+import smithy4s.schema.Schema._
+
+class RecursiveSpec extends FunSuite {
+
+  implicit def syncShim[F[_]: Sync]: SyncLike[F] = new SyncLike[F] {
+    override def delay[A](thunk: => A): F[A] = Sync[F].delay(thunk)
+  }
+
+  case class Recurse(n: Option[Recurse])
+
+  object Recurse {
+    implicit val schema: Schema[Recurse] = recursive {
+      struct(
+        schema.optional[Recurse]("n", _.n)
+      )(Recurse.apply).withId("test","Recurse")
+    }
+  }
+
+  def buildTree(size: Int): Tree = {
+    val seed = Math.round(Math.random() * 1000).toInt
+    val nodes = (1 to size).map(i => Tree.leaf(LeafNode(i * seed))).toList
+
+    @tailrec()
+    def recursiveFold(els: List[Tree]): Tree =
+      els match {
+        case head :: Nil => head
+        case x => {
+          val joined: List[Tree] = x
+            .sliding(2, 2)
+            .flatMap {
+              case left :: right :: Nil =>
+                List(Tree.tree(TreeNode(left, right)))
+              case x => x
+            }
+            .toList
+          recursiveFold(joined)
+        }
+      }
+
+    recursiveFold(nodes)
+  }
+
+  def buildRecursive(size: Int): Recurse =
+    (1 to size).foldLeft(Recurse(None))((tail, i) => Recurse(Some(tail)))
+
+  def useLazyTestCache[F[_]](store: MMap[Any, Any]) = new CompilationCache[F] {
+    override def getOrElseUpdate[A](
+        schema: Schema[A],
+        fetch: Schema[A] => F[A]
+    ): F[A] = {
+      store.getOrElseUpdate(schema, fetch(schema)).asInstanceOf[F[A]]
+    }
+  }
+
+  def ignoreLazyTestCache[F[_]](store: MMap[Any, Any]) =
+    new CompilationCache[F] {
+      override def getOrElseUpdate[A](
+          schema: Schema[A],
+          fetch: Schema[A] => F[A]
+      ): F[A] = {
+        if (schema.isInstanceOf[Schema.LazySchema[_]]) { fetch(schema) }
+        else store.getOrElseUpdate(schema, fetch(schema)).asInstanceOf[F[A]]
+      }
+    }
+
+  def runTest[A](
+      caseString: String,
+      value: Int => A,
+      buildCache: MMap[Any, Any] => CompilationCache[Hash],
+      transformSchema: Schema[A] => Schema[A] = (x: Schema[A]) => x
+  )(implicit schema: Schema[A]) =
+    test(s"$caseString") {
+      val store: MMap[Any, Any] = MMap.empty
+
+      val updatedSchema = transformSchema(schema)
+
+      val hashVisitor: Hash[A] =
+        SchemaVisitorHash.fromSchema(updatedSchema, buildCache(store))
+
+      // Invoke hash with a size that will have some recursion so that the build cache an be materialized
+      hashVisitor.hash(value(2))
+      val sizeAfterInitializing = store.size
+
+      val sizes = List(10, 100, 256)
+      sizes.foreach(i => hashVisitor.hash(value(i)))
+      val sizeAfterHashing = store.size
+      println(s"$caseString: $sizeAfterInitializing $sizeAfterHashing")
+      assertEquals(
+        sizeAfterHashing,
+        sizeAfterInitializing,
+        "cache store size has grown after initialization"
+      )
+    }
+
+  def runTreeVisitorTest[A](
+      caseString: String,
+      value: Int => A,
+      transformSchema: Schema[A] => Schema[A] = (x: Schema[A]) => x
+  )(implicit schema: Schema[A]): Unit = test(s"treeVisitor - $caseString") {
+    import cats.effect.unsafe.implicits.global
+    val makeHashVisitor: IO[TreeBasedTraversal[IO, Hash]] =
+      TreeBasedTraversal.make[IO, Hash](TreeBasedHashVisitor, CacheWrite.make[IO,Hash[_]])
+    val updatedSchema: Schema[A] = transformSchema(schema)
+
+    val makeDocKeyVisitor = TreeBasedTraversal.make[IO, DocumentKeyEncoder.OptDocumentKeyEncoder](DocumentKeyEncoder.effectfulVisitor, CacheWrite.make[IO, DocumentKeyEncoder.OptDocumentKeyEncoder[_]])
+
+    def makeDocVisitor(deps: Schema ~> EffectfulDocumentEncoderVisitor.Dependencies) = TreeBasedTraversal.make[IO, EffectfulDocumentEncoderVisitor.Dependencies, DocumentEncoder](EffectfulDocumentEncoderVisitor(FieldFilter.Default), CacheWrite.make[IO,DocumentEncoder[_]], deps)
+
+    def makeDeps(keyVisitor: Schema ~> DocumentKeyEncoder.OptDocumentKeyEncoder): Schema ~> EffectfulDocumentEncoderVisitor.Dependencies = new (Schema ~> EffectfulDocumentEncoderVisitor.Dependencies) {
+      override def apply[A0](fa: Schema[A0]): EffectfulDocumentEncoderVisitor.Dependencies[A0] = EffectfulDocumentEncoderVisitor.Dependencies(keyVisitor(fa),fa)
+    }
+
+    val documentEffect: IO[Unit] = for {
+      keyVisitor <- makeDocKeyVisitor
+      unsafeKeyVisitor = new (Schema ~> DocumentKeyEncoder.OptDocumentKeyEncoder) {
+        override def apply[A0](fa: Schema[A0]): OptDocumentKeyEncoder[A0] = keyVisitor.visit(fa).unsafeRunSync()
+      }
+      docVisitor <- makeDocVisitor(makeDeps(unsafeKeyVisitor))
+      sizeBeforeInitializing <- docVisitor.cache.size
+      //_ <- keyVisitor.visit(updatedSchema)
+      //_ <- IO.println(s"Visited key encoder for $caseString")
+      updatedEncoder <- docVisitor.visit(updatedSchema)
+      _ <- IO.println(s"Visited doc encoder for $caseString")
+      sizeAfterInitializing <- docVisitor.cache.size
+      doc = updatedEncoder(value(2))
+      _ <- IO.println(doc.toString())
+      sizeAfterEncoding <- docVisitor.cache.size
+      _ = List(10, 100, 256).foreach(sz => updatedEncoder(value(sz)))
+      sizeAfterLargeEncoding <- docVisitor.cache.size
+      _ <- IO.println(s"$caseString: $sizeBeforeInitializing, $sizeAfterInitializing, $sizeAfterEncoding, $sizeAfterLargeEncoding")
+    } yield ()
+
+    val hashEffect: IO[Unit] = for {
+      hashVisitor <- makeHashVisitor
+      sizeBeforeInitializing <- hashVisitor.cache.size
+      updatedHash <- hashVisitor.visit(updatedSchema)
+      _ <- IO.println(s"Visited $caseString")
+      sizeAfterInitializing <- hashVisitor.cache.size
+      _ = updatedHash.hash(value(2))
+      sizeAfterHashing <- hashVisitor.cache.size
+      _ = List(10, 100, 256).foreach(sz => updatedHash.hash(value(sz)))
+      sizeAfterLargeHashing <- hashVisitor.cache.size
+      _ <- IO.println(s"$caseString: $sizeBeforeInitializing, $sizeAfterInitializing, $sizeAfterHashing, $sizeAfterLargeHashing")
+    } yield {
+      assertEquals(sizeBeforeInitializing, 0)
+      assertEquals(sizeAfterInitializing, sizeAfterHashing)
+      assertEquals(sizeAfterHashing, sizeAfterLargeHashing)
+    }
+    (hashEffect *> documentEffect).unsafeRunSync() // TODO: CatsEffectSuite
+  }
+
+  def addHints[A](schema: Schema[A]): Schema[A] = {
+    schema.transformHintsTransitively(
+      _.add(
+        smithy.api.Documentation("Adding some hints")
+      )
+    )
+  }
+
+  def runTestCases[A: Schema](caseString: String, value: Int => A) = {
+    runTest(
+      s"$caseString: current cache, hints unchanged",
+      value,
+      buildCache = ignoreLazyTestCache
+    )
+    runTest(
+      s"$caseString: current cache, hints transformed",
+      value,
+      buildCache = ignoreLazyTestCache,
+      transformSchema = addHints[A]
+    )
+    runTest(
+      s"$caseString: updated cache, hints unchanged",
+      value,
+      buildCache = useLazyTestCache
+    )
+    runTest(
+      s"$caseString: updated cache, hints transformed",
+      value,
+      buildCache = useLazyTestCache,
+      transformSchema = addHints[A]
+    )
+    runTreeVisitorTest(
+      s"tree based $caseString: hints unchanged",
+      value
+    )
+    runTreeVisitorTest(
+      s"tree based $caseString: hintsTransformed",
+      value,
+      transformSchema = addHints[A]
+    )
+  }
+
+  runTestCases("Foo", Foo.int)
+  runTestCases("Tree", buildTree)
+  runTestCases("Recurse", buildRecursive)
+
+}

modules/bootstrapped/test/src/smithy4s/schematests/HintsTransformationSpec.scala

@@ -127,9 +127,14 @@ class HintsTransformationSpec() extends FunSuite {
         struct(foos)(Foo.apply)
       }
     }
+
+    def buildFoo(size: Int): Foo =
+      (1 until size).foldLeft(Foo(None))((foo, _) => Foo(Some(foo)))
+
     checkSchema(Foo(None), 1)
     checkSchema(Foo(Some(Foo(None))), 2)
     checkSchema(Foo(Some(Foo(Some(Foo(None))))), 3)
+    checkSchema(buildFoo(256), 256)
   }
 
   test(header("nullable")) {