fix: unbreak enum-based waiters and render null-aware acceptors (#842)

Author: lucix-aws

.changes/fe9d51c3-98cd-4c45-aed3-a797d278a36c.json

@@ -0,0 +1,8 @@
+{
+    "id": "fe9d51c3-98cd-4c45-aed3-a797d278a36c",
+    "type": "bugfix",
+    "description": "Fix broken enum-based waiters.",
+    "issues": [
+        "awslabs/aws-sdk-kotlin#892"
+    ]
+}

codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/waiters/AcceptorGenerator.kt

@@ -11,12 +11,14 @@ import software.amazon.smithy.kotlin.codegen.core.KotlinWriter
 import software.amazon.smithy.kotlin.codegen.core.RuntimeTypes
 import software.amazon.smithy.kotlin.codegen.core.withBlock
 import software.amazon.smithy.kotlin.codegen.utils.dq
+import software.amazon.smithy.model.shapes.ShapeId
+import software.amazon.smithy.model.shapes.StructureShape
 import software.amazon.smithy.waiters.*
 
 /**
  * Renders an individual acceptor for a waiter.
  */
-private fun KotlinWriter.renderAcceptor(acceptor: Acceptor) {
+private fun KotlinWriter.renderAcceptor(wi: WaiterInfo, acceptor: Acceptor) {
     addImport(RuntimeTypes.Core.Retries.Policy.RetryDirective)
 
     val directive = when (acceptor.state!!) {
@@ -39,8 +41,8 @@ private fun KotlinWriter.renderAcceptor(acceptor: Acceptor) {
             write("ErrorTypeAcceptor(RetryDirective.#L, #L),", directive, matcher.value.dq())
         }
 
-        is Matcher.InputOutputMember -> renderPathAcceptor(directive, true, matcher.value)
-        is Matcher.OutputMember -> renderPathAcceptor(directive, false, matcher.value)
+        is Matcher.InputOutputMember -> renderPathAcceptor(wi, directive, true, matcher.value)
+        is Matcher.OutputMember -> renderPathAcceptor(wi, directive, false, matcher.value)
         else -> throw CodegenException("""Unknown matcher type "${matcher::class}"""")
     }
 }
@@ -58,38 +60,49 @@ internal fun KotlinWriter.renderAcceptorList(wi: WaiterInfo, asValName: String)
         wi.inputSymbol,
         wi.outputSymbol,
     ) {
-        wi.waiter.acceptors.forEach(::renderAcceptor)
+        wi.waiter.acceptors.forEach { renderAcceptor(wi, it) }
     }
 }
 
 /**
  * Render a path-based acceptor (i.e., one that uses an output or inputOutput matcher).
  */
-private fun KotlinWriter.renderPathAcceptor(directive: String, includeInput: Boolean, matcher: PathMatcher) {
+private fun KotlinWriter.renderPathAcceptor(wi: WaiterInfo, directive: String, includeInput: Boolean, matcher: PathMatcher) {
     val acceptorType = if (includeInput) {
-        addImport(RuntimeTypes.Core.Retries.Policy.InputOutputAcceptor)
-        "InputOutputAcceptor"
+        RuntimeTypes.Core.Retries.Policy.InputOutputAcceptor
     } else {
-        addImport(RuntimeTypes.Core.Retries.Policy.OutputAcceptor)
-        "OutputAcceptor"
+        RuntimeTypes.Core.Retries.Policy.OutputAcceptor
     }
 
-    withBlock("#L(RetryDirective.#L) {", "},", acceptorType, directive) {
-        val visitor = KotlinJmespathExpressionVisitor(this)
+    val visitedShape = if (includeInput) {
+        buildSyntheticInputOutputShape(wi.op.id.name, wi.input.id, wi.output.id)
+    } else {
+        wi.output
+    }
+
+    withBlock("#T(RetryDirective.#L) {", "},", acceptorType, directive) {
+        val visitor = KotlinJmespathExpressionVisitor(wi.ctx, this, visitedShape)
         val expression = JmespathExpression.parse(matcher.path)
         val actual = expression.accept(visitor)
 
         val expected = matcher.expected
         val comparison = when (matcher.comparator!!) {
-            PathComparator.STRING_EQUALS -> "$actual?.toString() == ${expected.dq()}"
-            PathComparator.BOOLEAN_EQUALS -> "$actual == ${expected.toBoolean()}"
-            PathComparator.ANY_STRING_EQUALS -> "$actual?.any { it?.toString() == ${expected.dq()} } ?: false"
+            PathComparator.STRING_EQUALS -> "${actual.identifier} == ${expected.dq()}"
+            PathComparator.BOOLEAN_EQUALS -> "${actual.identifier} == ${expected.toBoolean()}"
+            PathComparator.ANY_STRING_EQUALS -> "${actual.identifier}?.any { it == ${expected.dq()} } ?: false"
 
-            // NOTE: the size > 0 check is necessary because the waiter spec says that `allStringEquals` requires
+            // NOTE: the isNotEmpty check is necessary because the waiter spec says that `allStringEquals` requires
             // at least one value unlike Kotlin's `all` which returns true if the collection is empty
             PathComparator.ALL_STRING_EQUALS ->
-                "$actual != null && $actual.size > 0 && $actual.all { it?.toString() == ${expected.dq()} }"
+                "!${actual.identifier}.isNullOrEmpty() && ${actual.identifier}.all { it == ${expected.dq()} }"
         }
         write(comparison)
     }
 }
+
+private fun buildSyntheticInputOutputShape(opName: String, input: ShapeId, output: ShapeId): StructureShape =
+    StructureShape.Builder()
+        .id("software.amazon.smithy.kotlin.codegen.rendering.waiters.synthetic#${opName}RequestResponse")
+        .addMember("input", input)
+        .addMember("output", output)
+        .build()