feat: add functions (sort, sort_by, max_by, min_by, map) to JMESPath visitor (#969)

Author: 0marperez

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

@@ -26,8 +26,6 @@ import software.amazon.smithy.model.shapes.ListShape
 import software.amazon.smithy.model.shapes.MapShape
 import software.amazon.smithy.model.shapes.MemberShape
 import software.amazon.smithy.model.shapes.Shape
-import kotlin.contracts.ExperimentalContracts
-import kotlin.contracts.contract
 
 private val suffixSequence = sequenceOf("") + generateSequence(2) { it + 1 }.map(Int::toString) // "", "2", "3", etc.
 
@@ -76,14 +74,6 @@ class KotlinJmespathExpressionVisitor(
     private fun bestTempVarName(preferredName: String): String =
         suffixSequence.map { "$preferredName$it" }.first(tempVars::add)
 
-    @OptIn(ExperimentalContracts::class)
-    private fun codegenReq(condition: Boolean, lazyMessage: () -> String) {
-        contract {
-            returns() implies condition
-        }
-        if (!condition) throw CodegenException(lazyMessage())
-    }
-
     private fun flatMappingBlock(right: JmespathExpression, leftName: String, leftShape: Shape, innerShape: Shape?): VisitedExpression {
         if (right is CurrentExpression) return VisitedExpression(leftName, leftShape) // nothing to map
 
@@ -105,7 +95,9 @@ class KotlinJmespathExpressionVisitor(
         return VisitedExpression(outerName, leftShape, innerResult.shape)
     }
 
-    private fun subfield(expression: FieldExpression, parentName: String, isObject: Boolean = false): VisitedExpression {
+    private data class SubFieldData(val name: String, val codegen: String, val member: Shape?)
+
+    private fun subfieldLogic(expression: FieldExpression, parentName: String, isObject: Boolean = false): SubFieldData {
         val member = currentShape.targetOrSelf(ctx.model).getMember(expression.name).getOrNull()
 
         val name = expression.name.toCamelCase()
@@ -130,9 +122,17 @@ class KotlinJmespathExpressionVisitor(
         }
 
         member?.let { shapeCursor.addLast(it) }
-        return VisitedExpression(addTempVar(name, codegen), member)
+        return SubFieldData(name, codegen, member)
+    }
+
+    private fun subfield(expression: FieldExpression, parentName: String, isObject: Boolean = false): VisitedExpression {
+        val (name, codegen, member) = subfieldLogic(expression, parentName, isObject)
+        return VisitedExpression(addTempVar(name, codegen), member, nullable = currentShape.isNullable)
     }
 
+    private fun subfieldCodegen(expression: FieldExpression, parentName: String, isObject: Boolean = false): String =
+        subfieldLogic(expression, parentName, isObject).codegen
+
     override fun visitAnd(expression: AndExpression): VisitedExpression {
         writer.addImport(RuntimeTypes.Core.Utils.truthiness)
 
@@ -151,22 +151,21 @@ class KotlinJmespathExpressionVisitor(
 
         val codegen = buildString {
             val nullables = buildList {
-                if (left.shape?.isNullable == true) add("${left.identifier} == null")
-                if (right.shape?.isNullable == true) add("${right.identifier} == null")
+                if (left.shape?.isNullable == true || left.nullable) add("${left.identifier} == null")
+                if (right.shape?.isNullable == true || right.nullable) add("${right.identifier} == null")
             }
+
             if (nullables.isNotEmpty()) {
                 val isNullExpr = nullables.joinToString(" || ")
                 append("if ($isNullExpr) null else ")
             }
 
-            val unSafeComparatorExpr = "compareTo(${right.identifier}) ${expression.comparator} 0"
-            val comparatorExpr = if (left.nullable) "?.$unSafeComparatorExpr" else ".$unSafeComparatorExpr"
-
+            val comparatorExpr = ".compareTo(${right.identifier}) ${expression.comparator} 0"
             append("${left.identifier}$comparatorExpr")
         }
 
-        val ident = addTempVar("comparison", codegen)
-        return VisitedExpression(ident)
+        val identifier = addTempVar("comparison", codegen)
+        return VisitedExpression(identifier)
     }
 
     override fun visitCurrentNode(expression: CurrentExpression): VisitedExpression {
@@ -207,14 +206,11 @@ class KotlinJmespathExpressionVisitor(
         return VisitedExpression(ident, currentShape, inner.projected)
     }
 
-    private fun FunctionExpression.singleArg(): VisitedExpression {
-        codegenReq(arguments.size == 1) { "Unexpected number of arguments to $this" }
-        return acceptSubexpression(this.arguments[0])
-    }
-    private fun FunctionExpression.twoArgs(): Pair<VisitedExpression, VisitedExpression> {
-        codegenReq(arguments.size == 2) { "Unexpected number of arguments to $this" }
-        return acceptSubexpression(this.arguments[0]) to acceptSubexpression(this.arguments[1])
-    }
+    private fun FunctionExpression.singleArg(): VisitedExpression =
+        acceptSubexpression(this.arguments[0])
+
+    private fun FunctionExpression.twoArgs(): Pair<VisitedExpression, VisitedExpression> =
+        acceptSubexpression(this.arguments[0]) to acceptSubexpression(this.arguments[1])
 
     private fun FunctionExpression.args(): List<VisitedExpression> =
         this.arguments.map { acceptSubexpression(it) }
@@ -257,9 +253,7 @@ class KotlinJmespathExpressionVisitor(
 
         "avg" -> {
             val numbers = expression.singleArg()
-            val average = numbers.dotFunction(expression, "average()").identifier
-            val isNaN = ensureNullGuard(numbers.shape, "isNaN() == true")
-            VisitedExpression(addTempVar("averageOrNull", "if($average$isNaN) null else $average"), nullable = true)
+            numbers.dotFunction(expression, "average()")
         }
 
         "join" -> {
@@ -334,6 +328,35 @@ class KotlinJmespathExpressionVisitor(
             arg.dotFunction(expression, "type()", ensureNullGuard = false)
         }
 
+        "sort" -> {
+            val arg = expression.singleArg()
+            arg.dotFunction(expression, "sorted()")
+        }
+
+        "sort_by" -> {
+            val list = expression.arguments[0].accept(this)
+            val expressionValue = expression.arguments[1]
+            list.applyFunction(expression.name.toCamelCase(), "sortedBy", expressionValue)
+        }
+
+        "max_by" -> {
+            val list = expression.arguments[0].accept(this)
+            val expressionValue = expression.arguments[1]
+            list.applyFunction(expression.name.toCamelCase(), "maxBy", expressionValue)
+        }
+
+        "min_by" -> {
+            val list = expression.arguments[0].accept(this)
+            val expressionValue = expression.arguments[1]
+            list.applyFunction(expression.name.toCamelCase(), "minBy", expressionValue)
+        }
+
+        "map" -> {
+            val list = expression.arguments[1].accept(this)
+            val expressionValue = expression.arguments[0]
+            list.applyFunction(expression.name.toCamelCase(), "map", expressionValue)
+        }
+
         else -> throw CodegenException("Unknown function type in $expression")
     }
 
@@ -548,6 +571,21 @@ class KotlinJmespathExpressionVisitor(
         return notNull
     }
 
+    private fun VisitedExpression.applyFunction(
+        name: String,
+        operation: String,
+        expression: JmespathExpression,
+    ): VisitedExpression {
+        val result = bestTempVarName(name)
+
+        writer.withBlock("val $result = ${this.identifier}?.$operation {", "}") {
+            val expressionValue = subfieldCodegen((expression as ExpressionTypeExpression).expression as FieldExpression, "it")
+            write("$expressionValue!!")
+        }
+
+        return VisitedExpression(result)
+    }
+
     private val Shape.isNullable: Boolean
         get() = this is MemberShape &&
             ctx.model.expectShape(target).let { !it.hasTrait<OperationInput>() && !it.hasTrait<OperationOutput>() }

tests/codegen/waiter-tests/model/function-map.smithy

@@ -0,0 +1,29 @@
+$version: "2"
+namespace com.test
+
+use smithy.waiters#waitable
+
+@suppress(["WaitableTraitJmespathProblem"])
+@waitable(
+    MapStructEquals: {
+        acceptors: [
+            {
+                state: "success",
+                matcher: {
+                    output: {
+                        path: "map(&string, lists.structs)",
+                        expected: "foo",
+                        comparator: "allStringEquals"
+                    }
+                }
+            }
+        ]
+    },
+)
+@readonly
+@http(method: "GET", uri: "/map/{name}", code: 200)
+operation GetFunctionMapEquals {
+    input: GetEntityRequest,
+    output: GetEntityResponse,
+    errors: [NotFound],
+}

tests/codegen/waiter-tests/model/function-max-by.smithy

@@ -0,0 +1,43 @@
+$version: "2"
+namespace com.test
+
+use smithy.waiters#waitable
+
+@suppress(["WaitableTraitJmespathProblem"])
+@waitable(
+    MaxByNumberEquals: {
+        acceptors: [
+            {
+                state: "success",
+                matcher: {
+                    output: {
+                        path: "max_by(lists.structs, &integer).integer == `100`",
+                        expected: "true",
+                        comparator: "booleanEquals"
+                    }
+                }
+            }
+        ]
+    },
+    MaxByStringEquals: {
+        acceptors: [
+            {
+                state: "success",
+                matcher: {
+                    output: {
+                        path: "max_by(lists.structs, &string).string",
+                        expected: "foo",
+                        comparator: "stringEquals"
+                    }
+                }
+            }
+        ]
+    },
+)
+@readonly
+@http(method: "GET", uri: "/maxBy/{name}", code: 200)
+operation GetFunctionMaxByEquals {
+    input: GetEntityRequest,
+    output: GetEntityResponse,
+    errors: [NotFound],
+}

tests/codegen/waiter-tests/model/function-min-by.smithy

@@ -0,0 +1,43 @@
+$version: "2"
+namespace com.test
+
+use smithy.waiters#waitable
+
+@suppress(["WaitableTraitJmespathProblem"])
+@waitable(
+    MinByNumberEquals: {
+        acceptors: [
+            {
+                state: "success",
+                matcher: {
+                    output: {
+                        path: "min_by(lists.structs, &integer).integer == `100`",
+                        expected: "true",
+                        comparator: "booleanEquals"
+                    }
+                }
+            }
+        ]
+    },
+    MinByStringEquals: {
+        acceptors: [
+            {
+                state: "success",
+                matcher: {
+                    output: {
+                        path: "min_by(lists.structs, &string).string",
+                        expected: "foo",
+                        comparator: "stringEquals"
+                    }
+                }
+            }
+        ]
+    },
+)
+@readonly
+@http(method: "GET", uri: "/minBy/{name}", code: 200)
+operation GetFunctionMinByEquals {
+    input: GetEntityRequest,
+    output: GetEntityResponse,
+    errors: [NotFound],
+}

tests/codegen/waiter-tests/model/function-sort-by.smithy

@@ -0,0 +1,43 @@
+$version: "2"
+namespace com.test
+
+use smithy.waiters#waitable
+
+@suppress(["WaitableTraitJmespathProblem"])
+@waitable(
+    SortByNumberEquals: {
+        acceptors: [
+            {
+                state: "success",
+                matcher: {
+                    output: {
+                        path: "sort_by(lists.structs, &integer)[0].integer == `1`",
+                        expected: "true",
+                        comparator: "booleanEquals"
+                    }
+                }
+            }
+        ]
+    },
+    SortByStringEquals: {
+        acceptors: [
+            {
+                state: "success",
+                matcher: {
+                    output: {
+                        path: "sort_by(lists.structs, &string)[2].string",
+                        expected: "foo",
+                        comparator: "stringEquals"
+                    }
+                }
+            }
+        ]
+    },
+)
+@readonly
+@http(method: "GET", uri: "/sortBy/{name}", code: 200)
+operation GetFunctionSortByEquals {
+    input: GetEntityRequest,
+    output: GetEntityResponse,
+    errors: [NotFound],
+}

tests/codegen/waiter-tests/model/function-sort.smithy

@@ -0,0 +1,43 @@
+$version: "2"
+namespace com.test
+
+use smithy.waiters#waitable
+
+@suppress(["WaitableTraitJmespathProblem"])
+@waitable(
+    SortNumberEquals: {
+        acceptors: [
+            {
+                state: "success",
+                matcher: {
+                    output: {
+                        path: "sort(lists.integers)[2] == `2`",
+                        expected: "true",
+                        comparator: "booleanEquals"
+                    }
+                }
+            }
+        ]
+    },
+    SortStringEquals: {
+        acceptors: [
+            {
+                state: "success",
+                matcher: {
+                    output: {
+                        path: "sort(lists.strings)[2]",
+                        expected: "foo",
+                        comparator: "stringEquals"
+                    }
+                }
+            }
+        ]
+    },
+)
+@readonly
+@http(method: "GET", uri: "/sort/{name}", code: 200)
+operation GetFunctionSortEquals {
+    input: GetEntityRequest,
+    output: GetEntityResponse,
+    errors: [NotFound],
+}

tests/codegen/waiter-tests/model/utils/structures.smithy

@@ -155,6 +155,8 @@ map StructMap {
 structure Struct {
     primitives: EntityPrimitives,
     strings: StringList,
+    integer: Integer,
+    string: String,
     enums: EnumList,
     subStructs: SubStructList,
 }

tests/codegen/waiter-tests/model/utils/waiter-operations.smithy

@@ -31,5 +31,10 @@ service WaitersTestService {
         GetFunctionToStringEquals,
         GetFunctionToNumberEquals,
         GetFunctionTypeEquals,
+        GetFunctionSortByEquals,
+        GetFunctionSortEquals,
+        GetFunctionMaxByEquals,
+        GetFunctionMinByEquals,
+        GetFunctionMapEquals,
     ]
 }