Make coalesce variadic

This commit adds support for variadic functions in the rules engine
and makes coalesce variadic. This makes things like phi functions
a shallow list of expressions rather than a massive list of nested
binary expressions. This also makes it easier to optimize phi nodes
without requiring peephole checks in compilers for optimization.

Author: mtdowling

docs/source-2.0/additional-specs/rules-engine/standard-library.rst

@@ -44,24 +44,28 @@ parameter is equal to the value ``false``:
 =====================
 
 Summary
-    Evaluates the first argument and returns the result if it is present, otherwise evaluates and returns the result
-    of the second argument.
+    Evaluates arguments in order and returns the first non-empty result, otherwise returns the result of the last
+    argument.
 Argument types
-    * value1: ``T`` or ``option<T>``
-    * value2: ``T`` or ``option<T>``
+    * This function is variadic and requires two or more arguments, each of type ``T`` or ``option<T>``
+    * All arguments must have the same inner type ``T``
 Return type
-    * ``coalesce(T, T)`` → ``T``
-    * ``coalesce(option<T>, T)`` → ``T``
-    * ``coalesce(T, option<T>)`` → ``T``
-    * ``coalesce(option<T>, option<T>)`` → ``option<T>``
+    * ``coalesce(T, T, ...)`` → ``T``
+    * ``coalesce(option<T>, T, ...)`` → ``T`` (if any argument is non-optional)
+    * ``coalesce(T, option<T>, ...)`` → ``T`` (if any argument is non-optional)
+    * ``coalesce(option<T>, option<T>, ...)`` → ``option<T>`` (if all arguments are optional)
 Since
     1.1
 
-The ``coalesce`` function provides null-safe chaining by returning the result of the first argument if it returns a
-value, otherwise returns the result of the second argument. This is particularly useful for providing default values
-for optional parameters, chaining multiple optional values together, and related optimizations.
+The ``coalesce`` function provides null-safe chaining by evaluating arguments in order and returning the first
+non-empty result. If all arguments evaluate to empty, it returns the result of the last argument. This is
+particularly useful for providing default values for optional parameters, chaining multiple optional values
+together, and related optimizations.
 
-The following example demonstrates chaining multiple ``coalesce`` calls to try several optional values
+The function accepts two or more arguments, all of which must have the same inner type after unwrapping any
+optionals. The return type is ``option<T>`` only if all arguments are ``option<T>``; otherwise it returns ``T``.
+
+The following example demonstrates using ``coalesce`` with multiple arguments to try several optional values
 in sequence:
 
 .. code-block:: json
@@ -70,20 +74,15 @@ in sequence:
         "fn": "coalesce",
         "argv": [
             {"ref": "customEndpoint"},
-            {
-                "fn": "coalesce",
-                "argv": [
-                    {"ref": "regionalEndpoint"},
-                    {"ref": "defaultEndpoint"}
-                ]
-            }
+            {"ref": "regionalEndpoint"},
+            {"ref": "defaultEndpoint"}
         ]
     }
 
 .. important::
-    Both arguments must be of the same type after unwrapping any optionals (types are known at compile time and do not
-    need to be validated at runtime). Note that the first result is returned even if it's ``false`` (coalesce is
-    looking for a *non-empty* value).
+    All arguments must have the same type after unwrapping any optionals (types are known at compile time and do not
+    need to be validated at runtime). Note that the first non-empty result is returned even if it's ``false``
+    (coalesce is looking for a *non-empty* value, not a truthy value).
 
 
 .. _rules-engine-standard-library-getAttr:

smithy-rules-engine/src/main/java/software/amazon/smithy/rulesengine/language/evaluation/RuleEvaluator.java

@@ -184,9 +184,15 @@ public Value visitIsSet(Expression fn) {
     }
 
     @Override
-    public Value visitCoalesce(Expression left, Expression right) {
-        Value leftValue = left.accept(this);
-        return leftValue.isEmpty() ? right.accept(this) : leftValue;
+    public Value visitCoalesce(List<Expression> expressions) {
+        Value result = Value.emptyValue();
+        for (Expression exp : expressions) {
+            result = exp.accept(this);
+            if (!result.isEmpty()) {
+                return result;
+            }
+        }
+        return result;
     }
 
     @Override

smithy-rules-engine/src/main/java/software/amazon/smithy/rulesengine/language/syntax/expressions/ExpressionVisitor.java

@@ -9,7 +9,6 @@
 import software.amazon.smithy.rulesengine.language.syntax.expressions.functions.FunctionDefinition;
 import software.amazon.smithy.rulesengine.language.syntax.expressions.functions.GetAttr;
 import software.amazon.smithy.rulesengine.language.syntax.expressions.literal.Literal;
-import software.amazon.smithy.utils.ListUtils;
 import software.amazon.smithy.utils.SmithyUnstableApi;
 
 /**
@@ -54,12 +53,11 @@ public interface ExpressionVisitor<R> {
     /**
      * Visits a coalesce function.
      *
-     * @param left the first value to check.
-     * @param right the second value to check.
+     * @param expressions The coalesce expressions to check.
      * @return the value from the visitor.
      */
-    default R visitCoalesce(Expression left, Expression right) {
-        return visitLibraryFunction(Coalesce.getDefinition(), ListUtils.of(left, right));
+    default R visitCoalesce(List<Expression> expressions) {
+        return visitLibraryFunction(Coalesce.getDefinition(), expressions);
     }
 
     /**
@@ -121,7 +119,7 @@ public R visitIsSet(Expression fn) {
         }
 
         @Override
-        public R visitCoalesce(Expression left, Expression right) {
+        public R visitCoalesce(List<Expression> expressions) {
             return getDefault();
         }
 

smithy-rules-engine/src/main/java/software/amazon/smithy/rulesengine/language/syntax/expressions/functions/Coalesce.java

@@ -4,8 +4,9 @@
  */
 package software.amazon.smithy.rulesengine.language.syntax.expressions.functions;
 
-import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
+import java.util.Optional;
 import software.amazon.smithy.rulesengine.language.RulesVersion;
 import software.amazon.smithy.rulesengine.language.evaluation.Scope;
 import software.amazon.smithy.rulesengine.language.evaluation.type.OptionalType;
@@ -18,19 +19,17 @@
 
 /**
  * A coalesce function that returns the first non-empty value.
- * At runtime, returns the left value unless it's EmptyValue, in which case returns the right value.
+ *
+ * <p>This variadic function requires two or more arguments. At runtime, returns the first arguments that returns a
+ * non-EmptyValue, otherwise returns the result of the last argument.
  *
  * <p>Type checking rules:
  * <ul>
- * <li>{@code coalesce(T, T) => T} (same types)</li>
- * <li>{@code coalesce(Optional<T>, T) => T} (unwraps optional)</li>
- * <li>{@code coalesce(T, Optional<T>) => T} (unwraps optional)</li>
- * <li>{@code coalesce(Optional<T>, Optional<T>) => Optional<T>}</li>
+ * <li>{@code coalesce(T, T, T) => T} (same types)</li>
+ * <li>{@code coalesce(Optional<T>, T, T) => T} (any non-optional makes result non-optional)</li>
+ * <li>{@code coalesce(Optional<T>, Optional<T>, Optional<T>) => Optional<T>} (all optional)</li>
  * </ul>
  *
- * <p>Supports chaining:
- * {@code coalesce(opt1, coalesce(opt2, coalesce(opt3, default)))}
- *
  * <p>Available since: rules engine 1.1.
  */
 @SmithyUnstableApi
@@ -52,14 +51,23 @@ public static Definition getDefinition() {
     }
 
     /**
-     * Creates a {@link Coalesce} function from the given expressions.
+     * Creates a {@link Coalesce} function from variadic expressions.
+     *
+     * @param args the expressions to coalesce
+     * @return The resulting {@link Coalesce} function.
+     */
+    public static Coalesce ofExpressions(ToExpression... args) {
+        return DEFINITION.createFunction(FunctionNode.ofExpressions(ID, args));
+    }
+
+    /**
+     * Creates a {@link Coalesce} function from a list of expressions.
      *
-     * @param arg1 the first expression, typically optional.
-     * @param arg2 the second expression, used as fallback.
+     * @param args the expressions to coalesce
      * @return The resulting {@link Coalesce} function.
      */
-    public static Coalesce ofExpressions(ToExpression arg1, ToExpression arg2) {
-        return DEFINITION.createFunction(FunctionNode.ofExpressions(ID, arg1, arg2));
+    public static Coalesce ofExpressions(List<? extends ToExpression> args) {
+        return ofExpressions(args.toArray(new ToExpression[0]));
     }
 
     @Override
@@ -69,37 +77,38 @@ public RulesVersion availableSince() {
 
     @Override
     public <R> R accept(ExpressionVisitor<R> visitor) {
-        List<Expression> args = getArguments();
-        return visitor.visitCoalesce(args.get(0), args.get(1));
+        return visitor.visitCoalesce(getArguments());
     }
 
     @Override
     public Type typeCheck(Scope<Type> scope) {
         List<Expression> args = getArguments();
-
-        if (args.size() != 2) {
-            throw new IllegalArgumentException("Coalesce requires exactly 2 arguments, got " + args.size());
-        }
-
-        Type leftType = args.get(0).typeCheck(scope);
-        Type rightType = args.get(1).typeCheck(scope);
-        Type leftInner = getInnerType(leftType);
-        Type rightInner = getInnerType(rightType);
-
-        // Both must be the same type (after unwrapping optionals)
-        if (!leftInner.equals(rightInner)) {
-            throw new IllegalArgumentException(String.format(
-                    "Type mismatch in coalesce: %s and %s must be the same type",
-                    leftType,
-                    rightType));
+        if (args.size() < 2) {
+            throw new IllegalArgumentException("Coalesce requires at least 2 arguments, got " + args.size());
         }
 
-        // Only return Optional if both sides are optional
-        if (leftType instanceof OptionalType && rightType instanceof OptionalType) {
-            return Type.optionalType(leftInner);
+        // Get the first argument's type as the baseline
+        Type firstType = args.get(0).typeCheck(scope);
+        Type baseInnerType = getInnerType(firstType);
+        boolean hasNonOptional = !(firstType instanceof OptionalType);
+
+        // Check all other arguments match the base type
+        for (int i = 1; i < args.size(); i++) {
+            Type argType = args.get(i).typeCheck(scope);
+            Type innerType = getInnerType(argType);
+
+            if (!innerType.equals(baseInnerType)) {
+                throw new IllegalArgumentException(String.format(
+                        "Type mismatch in coalesce at argument %d: expected %s but got %s",
+                        i + 1,
+                        baseInnerType,
+                        innerType));
+            }
+
+            hasNonOptional = hasNonOptional || !(argType instanceof OptionalType);
         }
 
-        return leftInner;
+        return hasNonOptional ? baseInnerType : Type.optionalType(baseInnerType);
     }
 
     private static Type getInnerType(Type t) {
@@ -119,7 +128,12 @@ public String getId() {
 
         @Override
         public List<Type> getArguments() {
-            return Arrays.asList(Type.anyType(), Type.anyType());
+            return Collections.emptyList();
+        }
+
+        @Override
+        public Optional<Type> getVariadicArguments() {
+            return Optional.of(Type.anyType());
         }
 
         @Override

smithy-rules-engine/src/main/java/software/amazon/smithy/rulesengine/language/syntax/expressions/functions/FunctionDefinition.java

@@ -5,6 +5,7 @@
 package software.amazon.smithy.rulesengine.language.syntax.expressions.functions;
 
 import java.util.List;
+import java.util.Optional;
 import software.amazon.smithy.rulesengine.language.evaluation.type.Type;
 import software.amazon.smithy.rulesengine.language.evaluation.value.Value;
 import software.amazon.smithy.utils.SmithyUnstableApi;
@@ -26,6 +27,18 @@ public interface FunctionDefinition {
      */
     List<Type> getArguments();
 
+    /**
+     * Gets the type of variadic arguments if this function accepts them.
+     *
+     * <p>When present, the function accepts any number of additional arguments of this type after the fixed arguments
+     * from getArguments().
+     *
+     * @return the variadic argument type, or empty if not variadic
+     */
+    default Optional<Type> getVariadicArguments() {
+        return Optional.empty();
+    }
+
     /**
      * The return type of this function definition.
      * @return The function return type

smithy-rules-engine/src/main/java/software/amazon/smithy/rulesengine/language/syntax/expressions/functions/LibraryFunction.java

@@ -8,6 +8,7 @@
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
 import software.amazon.smithy.model.SourceException;
 import software.amazon.smithy.model.SourceLocation;
@@ -105,45 +106,72 @@ public SourceLocation getSourceLocation() {
     protected Type typeCheckLocal(Scope<Type> scope) {
         RuleError.context(String.format("while typechecking the invocation of %s", definition.getId()), this, () -> {
             try {
-                checkTypeSignature(definition.getArguments(), functionNode.getArguments(), scope);
+                checkTypeSignature(scope);
             } catch (InnerParseError e) {
                 throw new RuntimeException(e.getMessage());
             }
         });
         return definition.getReturnType();
     }
 
-    private void checkTypeSignature(List<Type> expectedArgs, List<Expression> actualArguments, Scope<Type> scope)
+    private void checkTypeSignature(Scope<Type> scope) throws InnerParseError {
+        List<Type> expectedArgs = definition.getArguments();
+        Optional<Type> variadicType = definition.getVariadicArguments();
+        List<Expression> actualArguments = functionNode.getArguments();
+
+        if (variadicType.isPresent()) {
+            // check we have at least the fixed arguments
+            if (actualArguments.size() < expectedArgs.size()) {
+                throw new InnerParseError(String.format("Expected at least %s arguments but found %s",
+                        expectedArgs.size(),
+                        actualArguments.size()));
+            }
+            // check fixed arguments
+            for (int i = 0; i < expectedArgs.size(); i++) {
+                checkArgument(i, expectedArgs.get(i), actualArguments.get(i), scope);
+            }
+            // check variadic arguments
+            Type varType = variadicType.get();
+            for (int i = expectedArgs.size(); i < actualArguments.size(); i++) {
+                checkArgument(i, varType, actualArguments.get(i), scope);
+            }
+        } else {
+            // Non-variadic, so exact count required
+            if (expectedArgs.size() != actualArguments.size()) {
+                throw new InnerParseError(String.format("Expected %s arguments but found %s",
+                        expectedArgs.size(),
+                        actualArguments.size()));
+            }
+            // check all positional arguments
+            for (int i = 0; i < expectedArgs.size(); i++) {
+                checkArgument(i, expectedArgs.get(i), actualArguments.get(i), scope);
+            }
+        }
+    }
+
+    private void checkArgument(int index, Type expected, Expression actual, Scope<Type> scope)
             throws InnerParseError {
-        if (expectedArgs.size() != actualArguments.size()) {
-            throw new InnerParseError(
-                    String.format(
-                            "Expected %s arguments but found %s",
-                            expectedArgs.size(),
-                            actualArguments));
+        Type actualType = actual.typeCheck(scope);
+        if (expected.isA(actualType)) {
+            return;
         }
-        for (int i = 0; i < expectedArgs.size(); i++) {
-            Type expected = expectedArgs.get(i);
-            Type actual = actualArguments.get(i).typeCheck(scope);
-            if (!expected.isA(actual)) {
-                Type optAny = Type.optionalType(Type.anyType());
-                String hint = "";
-                if (actual.isA(optAny) && !expected.isA(optAny)
-                        && actual.expectOptionalType().inner().equals(expected)) {
-                    hint = String.format(
-                            "hint: use `assign` in a condition or `isSet(%s)` to prove that this value is non-null",
-                            actualArguments.get(i));
-                    hint = StringUtils.indent(hint, 2);
-                }
-                throw new InnerParseError(
-                        String.format(
-                                "Unexpected type in the %s argument: Expected %s but found %s%n%s",
-                                ordinal(i + 1),
-                                expected,
-                                actual,
-                                hint));
-            }
+
+        Type optAny = Type.optionalType(Type.anyType());
+        String hint = "";
+        if (actualType.isA(optAny)
+                && !expected.isA(optAny)
+                && actualType.expectOptionalType().inner().equals(expected)) {
+            hint = String.format(
+                    "hint: use `assign` in a condition or `isSet(%s)` to prove that this value is non-null",
+                    actual);
+            hint = StringUtils.indent(hint, 2);
         }
+
+        throw new InnerParseError(String.format("Unexpected type in the %s argument: Expected %s but found %s%n%s",
+                ordinal(index + 1),
+                expected,
+                actualType,
+                hint));
     }
 
     private static String ordinal(int arg) {