@examples trait for completions (#339)

* WIP: @examples trait for completions

* Less hacky brace stuff

* Move out of visitor

* refact

* Get rid of sortTextOverride

* Add some tests for example completions

* remove todo and outdated comment

modules/ast/src/main/scala/playground/smithyql/AST.scala

@@ -127,6 +127,12 @@ sealed trait InputNode[F[_]] extends AST[F] {
     fk: F ~> G
   ): InputNode[G]
 
+  def asStruct: Option[Struct[F]] =
+    this match {
+      case Struct(fields) => Some(Struct(fields))
+      case _              => None
+    }
+
 }
 
 final case class OperationName[F[_]](

modules/examples/src/main/smithy/demo.smithy

@@ -49,6 +49,28 @@ operation GetVersion {
 @documentation("""
 Create a hero.
 """)
+@examples([{
+    title: "Valid input"
+    documentation: "This is a valid input"
+    input: {
+        hero: {
+            good: {
+                howGood: 10
+            }
+        }
+    }
+}, {
+    title: "Valid input v2"
+    documentation: "This is also a valid input, but for a bad hero"
+    input: {
+        hero: {
+            bad: {
+                evilName: "Evil"
+                powerLevel: 10
+            }
+        }
+    }
+}])
 operation CreateHero {
     input: CreateHeroInput
     output: CreateHeroOutput

modules/formatter/src/main/scala/playground/smithyql/format/Formatter.scala

@@ -28,6 +28,13 @@ object Formatter {
   implicit val useClauseFormatter: Formatter[UseClause] = writeDoc
   implicit val preludeFormatter: Formatter[Prelude] = writeDoc
   implicit val qonFormatter: Formatter[QueryOperationName] = writeDoc
+
+  implicit val fieldsFormatter: Formatter[Struct.Fields] =
+    (
+      f,
+      w,
+    ) => FormattingVisitor.writeStructFields(WithSource.liftId(f)).renderTrim(w)
+
   implicit val inputNodeFormatter: Formatter[InputNode] = writeDoc
   implicit val structFormatter: Formatter[Struct] = writeDoc
   implicit val listedFormatter: Formatter[Listed] = writeDoc
@@ -108,9 +115,14 @@ private[format] object FormattingVisitor extends ASTVisitor[WithSource, Doc] { v
     query: WithSource[Query[WithSource]]
   ): Doc = printGeneric(query)
 
+  // no braces
+  def writeStructFields(
+    fields: WithSource[Struct.Fields[WithSource]]
+  ): Doc = writeCommaSeparated(fields.map(_.value))(writeField)
+
   override def struct(
     fields: WithSource[Struct.Fields[WithSource]]
-  ): Doc = writeBracketed(fields.map(_.value))(Doc.char('{'), Doc.char('}'))(writeField)
+  ): Doc = Doc.char('{') + writeStructFields(fields) + Doc.char('}')
 
   private def forceLineAfterTrailingComments[A](
     printer: WithSource[A] => Doc
@@ -177,19 +189,24 @@ private[format] object FormattingVisitor extends ASTVisitor[WithSource, Doc] { v
     // Force newlines between fields
     fields.map(renderField).intercalate(Doc.hardLine)
 
+  private def writeCommaSeparated[T](
+    items: WithSource[List[T]]
+  )(
+    renderItem: T => Doc
+  ): Doc =
+    Doc.hardLine +
+      printWithComments(items)(writeFields(_)(renderItem(_) + Doc.comma))
+        .indent(2) +
+      Doc.hardLine
+
   private def writeBracketed[T](
     items: WithSource[List[T]]
   )(
     before: Doc,
     after: Doc,
   )(
     renderItem: T => Doc
-  ): Doc =
-    before + Doc.hardLine +
-      printWithComments(items)(writeFields(_)(renderItem(_) + Doc.comma))
-        .indent(2) +
-      Doc.hardLine +
-      after
+  ): Doc = before + writeCommaSeparated(items)(renderItem) + after
 
   def writeIdent(
     ident: QualifiedIdentifier

modules/language-support/src/main/scala/playground/language/CompletionProvider.scala

@@ -17,6 +17,8 @@ import playground.smithyql.SourceFile
 import playground.smithyql.WithSource
 import playground.smithyql.parser.SourceParser
 import playground.smithyql.syntax.*
+import smithy.api.Examples
+import smithy4s.Hints
 import smithy4s.dynamic.DynamicSchemaIndex
 
 trait CompletionProvider {
@@ -105,7 +107,12 @@ object CompletionProvider {
           .service
           .endpoints
           .map { endpoint =>
-            OperationName[Id](endpoint.name) -> endpoint.input.compile(CompletionVisitor)
+            OperationName[Id](endpoint.name) -> endpoint
+              .input
+              .addHints(
+                endpoint.hints.get(Examples).map(Hints(_)).getOrElse(Hints.empty)
+              )
+              .compile(CompletionVisitor)
           }
           .toMap
       }

modules/language-support/src/main/scala/playground/language/CompletionVisitor.scala

@@ -3,6 +3,7 @@ package playground.language
 import cats.Id
 import cats.kernel.Eq
 import cats.syntax.all.*
+import playground.NodeEncoder
 import playground.ServiceNameExtractor
 import playground.TextUtils
 import playground.language.CompletionItem.InsertUseClause.NotRequired
@@ -23,6 +24,7 @@ import playground.smithyql.WithSource
 import playground.smithyql.format.Formatter
 import smithy.api
 import smithy4s.Bijection
+import smithy4s.Document
 import smithy4s.Endpoint
 import smithy4s.Hints
 import smithy4s.Lazy
@@ -68,6 +70,8 @@ final case class CompletionItem(
     insertText = InsertText.JustString(s"$label = ")
   )
 
+  def withSortText(text: String): CompletionItem = copy(sortText = text.some)
+
 }
 
 sealed trait TextEdit extends Product with Serializable
@@ -383,6 +387,73 @@ object CompletionItem {
     )
   }
 
+  // Examples for operation inputs.
+  // TODO: currently only works inside the struct (and assumes that by rendering only fields, no braces).
+  // If/when we ever have graceful parsing in completions, we should handle other contexts, such as being outside of the struct.
+  def forInputExamples[S](
+    schema: Schema[S]
+  ): List[CompletionItem] = {
+    val documentDecoder = Document.Decoder.fromSchema(schema)
+    val nodeEncoder = NodeEncoder.derive(schema)
+
+    case class Sample(
+      name: String,
+      documentation: Option[String],
+      inputObject: Struct[Id],
+    )
+
+    def decodeSample(
+      example: api.Example
+    ): Option[Sample] =
+      for {
+        input <- example.input
+        decoded <- documentDecoder.decode(input).toOption
+        // note: we could've transcoded from Document to Node directly, without the intermediate decoding
+        // but the examples we suggest should be valid, and this is the only way to ensure that.
+        encoded = nodeEncoder.toNode(decoded)
+
+        // we're only covering inputs, and operation inputs must be structures.
+        asObject <- encoded.asStruct
+      } yield Sample(
+        name = example.title,
+        documentation = example.documentation,
+        inputObject = asObject,
+      )
+
+    def completionForSample(
+      sample: Sample,
+      index: Int,
+    ): CompletionItem = {
+      val text = Formatter[Struct.Fields]
+        .format(
+          sample
+            .inputObject
+            .fields
+            .mapK(WithSource.liftId),
+          Int.MaxValue,
+        )
+
+      CompletionItem
+        .fromHints(
+          kind = CompletionItemKind.Constant,
+          label = s"Example: ${sample.name}",
+          insertText = InsertText.JustString(text),
+          schema = schema.addHints(
+            sample.documentation.map(api.Documentation(_)).map(Hints(_)).getOrElse(Hints.empty)
+          ),
+        )
+        .withSortText(s"0_$index")
+    }
+
+    schema
+      .hints
+      .get(api.Examples)
+      .foldMap(_.value)
+      .flatMap(decodeSample)
+      .zipWithIndex
+      .map(completionForSample.tupled)
+  }
+
   def deprecationString(
     info: api.Deprecated
   ): String = {
@@ -570,11 +641,16 @@ object CompletionVisitor extends SchemaVisitor[CompletionResolver] {
     fields: Vector[Field[S, ?]],
     make: IndexedSeq[Any] => S,
   ): CompletionResolver[S] = {
+    // Artificial schema resembling this one. Should be pretty much equivalent.
+    val schema = Schema.struct(fields)(make).addHints(hints).withId(shapeId)
+
     val compiledFields = fields.map(field => (field, field.schema.compile(this)))
 
+    val examples = CompletionItem.forInputExamples(schema)
+
     structLike(
       inBody =
-        fields
+        examples ++ fields
           // todo: filter out present fields
           .sortBy(field => (field.isRequired && !field.hasDefaultValue, field.label))
           .map(CompletionItem.fromField)

modules/language-support/src/test/scala/playground/language/CompletionItemTests.scala

@@ -1,5 +1,8 @@
 package playground.language
 
+import demo.smithy.CreateHeroInput
+import demo.smithy.DemoServiceOperation.CreateHero
+import demo.smithy.Good
 import demo.smithy.Hero
 import demo.smithy.Subscription
 import playground.Assertions.*
@@ -8,6 +11,9 @@ import playground.smithyql.QualifiedIdentifier
 import playground.smithyql.syntax.*
 import playground.std.ClockGen
 import playground.std.ClockOperation
+import smithy.api.Documentation
+import smithy.api.Examples
+import smithy4s.Hints
 import smithy4s.schema.Schema
 import weaver.*
 
@@ -178,6 +184,19 @@ object CompletionItemTests extends FunSuite {
     )
   }
 
+  test("CompletionItem.forInputExamples: documentation already present is ignored") {
+    val schema = CreateHeroInput
+      .schema
+      .addHints(CreateHero.hints.get(Examples).map(a => a: Hints.Binding).toList*)
+      .addHints(Documentation("hello"))
+
+    val results = CompletionItem.forInputExamples(schema)
+
+    val containsDocs = results.exists(_.docs.exists(_.contains("hello")))
+
+    assert(!containsDocs)
+  }
+
   test("describeSchema: recursive struct") {
     val result = CompletionItem.describeSchema(Subscription.schema)()
 

modules/language-support/src/test/scala/playground/language/CompletionProviderTests.scala

@@ -1,5 +1,6 @@
 package playground.language
 
+import demo.smithy.DemoService
 import demo.smithy.DemoServiceGen
 import demo.smithy.DeprecatedServiceGen
 import playground.Assertions.*
@@ -210,6 +211,69 @@ object CompletionProviderTests extends SimpleIOSuite {
     assertNoDiff(result, expected)
   }
 
+  pureTest("completing operation examples - input examples get used") {
+    val provider = CompletionProvider.forServices(List(wrapService(DemoService)))
+
+    val input =
+      s"""use service ${DemoService.id}
+         |
+         |CreateHero {}""".stripMargin
+
+    val results = provider
+      .provide(
+        input,
+        input.positionOf("}"),
+      )
+      .filter(_.label.startsWith("Example:"))
+
+    val expected = List(
+      CompletionItem(
+        kind = CompletionItemKind.Constant,
+        label = "Example: Valid input",
+        insertText = InsertText.JustString(
+          """
+            |  hero: {
+            |    good: {
+            |      howGood: 10,
+            |    },
+            |  },
+            |""".stripMargin
+        ),
+        detail = ": structure CreateHeroInput",
+        description = Some("demo.smithy"),
+        deprecated = false,
+        docs = Some("This is a valid input"),
+        extraTextEdits = Nil,
+        sortText = Some("0_0"),
+      ),
+      CompletionItem(
+        kind = CompletionItemKind.Constant,
+        label = "Example: Valid input v2",
+        insertText = InsertText.JustString(
+          """
+            |  hero: {
+            |    bad: {
+            |      evilName: "Evil",
+            |      powerLevel: 10,
+            |    },
+            |  },
+            |""".stripMargin
+        ),
+        detail = ": structure CreateHeroInput",
+        description = Some("demo.smithy"),
+        deprecated = false,
+        docs = Some("This is also a valid input, but for a bad hero"),
+        extraTextEdits = Nil,
+        sortText = Some("0_1"),
+      ),
+    )
+
+    assertNoDiff(
+      results,
+      expected,
+    )
+  }
+
   // needs: completions inside prelude (entire use clauses)
   // needs: completions inside use clause (only service id)
   // https://github.com/kubukoz/smithy-playground/issues/163