Fix bug around Lazy caching and update Lazy model

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,139 @@
+package smithy4s
+
+import munit.FunSuite
+
+import smithy4s.example.{Tree, TreeNode, LeafNode, Foo}
+import cats.Hash
+import smithy4s.schema._
+import scala.annotation.tailrec
+import smithy4s.internals.maps.MMap
+import smithy4s.interopcats.SchemaVisitorHash
+import smithy4s.schema.Schema.recursive
+import smithy4s.schema.Schema._
+
+class RecursiveSpec extends FunSuite {
+
+  case class Recurse(n: Option[Recurse])
+
+  object Recurse {
+    implicit val schema: Schema[Recurse] = recursive {
+      struct(
+        schema.optional[Recurse]("n", _.n)
+      )(Recurse.apply)
+    }
+  }
+
+  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
+
+      assertEquals(
+        sizeAfterHashing,
+        sizeAfterInitializing,
+        "cache store size has grown after initialization"
+      )
+    }
+
+  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]
+    )
+  }
+
+  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")) {

modules/core/src/smithy4s/Lazy.scala

@@ -16,17 +16,24 @@
 
 package smithy4s
 
-final class Lazy[A](make: () => A) {
-  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 Lazy[A] {
+  def value: A
+  final def map[B](f: A => B): Lazy[B] = Lazy.Mapped(this, f)
 }
 
 object Lazy {
-  def apply[A](a: => A): Lazy[A] = new Lazy(() => a)
+  def apply[A](a: => A): Lazy[A] = new Root(() => a)
+
+  final class Root[A](make: () => A) extends Lazy[A] {
+    protected var thunk: () => A = make
+    lazy val value: A = {
+      val result = thunk()
+      thunk = null
+      result
+    }
+  }
+
+  final case class Mapped[A, B](left: Lazy[A], f: A => B) extends Lazy[B] {
+    lazy val value = f(left.value)
+  }
 }

modules/core/src/smithy4s/schema/Schema.scala

@@ -17,6 +17,7 @@
 package smithy4s
 package schema
 
+import smithy4s.internals.maps.MMap
 import Schema._
 import scala.reflect.ClassTag
 
@@ -214,6 +215,7 @@ object Schema {
   private final class TransitiveCompiler(
       underlying: Schema ~> Schema
   ) extends (Schema ~> Schema) {
+    val lazyCompileCache: MMap[Any, Any] = MMap.empty
 
     def apply[A](
         fa: Schema[A]
@@ -224,8 +226,14 @@ object Schema {
         underlying(u.copy(alternatives = u.alternatives.map(handleAlt(_))))
       case BijectionSchema(s, bijection) =>
         underlying(BijectionSchema(this(s), bijection))
-      case LazySchema(suspend) =>
-        underlying(LazySchema(suspend.map(this.apply)))
+      case l @ LazySchema(suspend) => {
+        lazyCompileCache
+          .getOrElseUpdate(
+            l,
+            LazySchema(suspend.map(this.apply))
+          )
+          .asInstanceOf[Schema[A]]
+      }
       case RefinementSchema(s, refinement) =>
         underlying(RefinementSchema(this(s), refinement))
       case c: CollectionSchema[c, a] =>

sampleSpecs/recursive.smithy

@@ -15,3 +15,20 @@ structure RecursiveListWrapper {
     @required
     items: RecursiveList
 }
+
+union Tree {
+  tree: TreeNode
+  leaf: LeafNode
+}
+
+structure TreeNode {
+  @required
+  left: Tree
+  @required
+  right: Tree
+}
+
+structure LeafNode {
+  @required
+  value: Integer
+}