Policy nested rule fix

PiperOrigin-RevId: 905187056

Author: l46kok

policy/src/main/java/dev/cel/policy/RuleComposer.java

@@ -18,14 +18,14 @@
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static java.util.stream.Collectors.toCollection;
 
-import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Lists;
 import dev.cel.bundle.Cel;
 import dev.cel.common.CelAbstractSyntaxTree;
 import dev.cel.common.CelMutableAst;
 import dev.cel.common.CelValidationException;
 import dev.cel.common.Operator;
+import dev.cel.common.ast.CelExpr;
 import dev.cel.common.ast.CelExpr.ExprKind.Kind;
 import dev.cel.common.formats.ValueString;
 import dev.cel.common.navigation.CelNavigableMutableAst;
@@ -48,22 +48,11 @@ final class RuleComposer implements CelAstOptimizer {
 
   @Override
   public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel) {
-    RuleOptimizationResult result = optimizeRule(cel, compiledRule);
-    return OptimizationResult.create(result.ast().toParsedAst());
+    Step result = optimizeRule(cel, compiledRule);
+    return OptimizationResult.create(result.expr.toParsedAst());
   }
 
-  @AutoValue
-  abstract static class RuleOptimizationResult {
-    abstract CelMutableAst ast();
-
-    abstract boolean isOptionalResult();
-
-    static RuleOptimizationResult create(CelMutableAst ast, boolean isOptionalResult) {
-      return new AutoValue_RuleComposer_RuleOptimizationResult(ast, isOptionalResult);
-    }
-  }
-
-  private RuleOptimizationResult optimizeRule(Cel cel, CelCompiledRule compiledRule) {
+  private Step optimizeRule(Cel cel, CelCompiledRule compiledRule) {
     cel =
         cel.toCelBuilder()
             .addVarDeclarations(
@@ -72,81 +61,52 @@ private RuleOptimizationResult optimizeRule(Cel cel, CelCompiledRule compiledRul
                     .collect(toImmutableList()))
             .build();
 
-    CelMutableAst matchAst = astMutator.newGlobalCall(Function.OPTIONAL_NONE.getFunction());
-    boolean isOptionalResult = true;
+    Step output = null;
+    if (compiledRule.hasOptionalOutput()) {
+      output =
+          new Step(
+              true,
+              false,
+              newTrueAst(cel),
+              astMutator.newGlobalCall(Function.OPTIONAL_NONE.getFunction()));
+    }
+
     // Keep track of the last output ID that might cause type-check failure while attempting to
     // compose the subgraphs.
     long lastOutputId = 0;
+
     for (CelCompiledMatch match : Lists.reverse(compiledRule.matches())) {
       CelAbstractSyntaxTree conditionAst = match.condition();
-      // If the condition is trivially true, none of the matches in the rule causes the result
-      // to become optional, and the rule is not the last match, then this will introduce
-      // unreachable outputs or rules.
       boolean isTriviallyTrue = match.isConditionTriviallyTrue();
+      CelMutableAst condAst = CelMutableAst.fromCelAst(conditionAst);
 
       switch (match.result().kind()) {
-        // For the match's output, determine whether the output should be wrapped
-        // into an optional value, a conditional, or both.
         case OUTPUT:
           OutputValue matchOutput = match.result().output();
           CelMutableAst outAst = CelMutableAst.fromCelAst(matchOutput.ast());
-          if (isTriviallyTrue) {
-            matchAst = outAst;
-            isOptionalResult = false;
-            lastOutputId = matchOutput.sourceId();
-            continue;
-          }
-          if (isOptionalResult) {
-            outAst = astMutator.newGlobalCall(Function.OPTIONAL_OF.getFunction(), outAst);
-          }
+          Step step = new Step(false, !isTriviallyTrue, condAst, outAst);
+          output = combine(cel, astMutator, step, output);
 
-          matchAst =
-              astMutator.newGlobalCall(
-                  Operator.CONDITIONAL.getFunction(),
-                  CelMutableAst.fromCelAst(conditionAst),
-                  outAst,
-                  matchAst);
           assertComposedAstIsValid(
               cel,
-              matchAst,
+              output.expr,
               "conflicting output types found.",
               matchOutput.sourceId(),
               lastOutputId);
+
           lastOutputId = matchOutput.sourceId();
           continue;
         case RULE:
-          // If the match has a nested rule, then compute the rule and whether it has
-          // an optional return value.
           CelCompiledRule matchNestedRule = match.result().rule();
-          RuleOptimizationResult nestedRule = optimizeRule(cel, matchNestedRule);
+          Step nestedRule = optimizeRule(cel, matchNestedRule);
           boolean nestedHasOptional = matchNestedRule.hasOptionalOutput();
-          CelMutableAst nestedRuleAst = nestedRule.ast();
-          if (isOptionalResult && !nestedHasOptional) {
-            nestedRuleAst =
-                astMutator.newGlobalCall(Function.OPTIONAL_OF.getFunction(), nestedRuleAst);
-          }
-          if (!isOptionalResult && nestedHasOptional) {
-            matchAst = astMutator.newGlobalCall(Function.OPTIONAL_OF.getFunction(), matchAst);
-            isOptionalResult = true;
-          }
-          // If either the nested rule or current condition output are optional then
-          // use optional.or() to specify the combination of the first and second results
-          // Note, the argument order is reversed due to the traversal of matches in
-          // reverse order.
-          if (isOptionalResult && isTriviallyTrue) {
-            matchAst = astMutator.newMemberCall(nestedRuleAst, Function.OR.getFunction(), matchAst);
-          } else {
-            matchAst =
-                astMutator.newGlobalCall(
-                    Operator.CONDITIONAL.getFunction(),
-                    CelMutableAst.fromCelAst(conditionAst),
-                    nestedRuleAst,
-                    matchAst);
-          }
+
+          Step ruleStep = new Step(nestedHasOptional, !isTriviallyTrue, condAst, nestedRule.expr);
+          output = combine(cel, astMutator, ruleStep, output);
 
           assertComposedAstIsValid(
               cel,
-              matchAst,
+              output.expr,
               String.format(
                   "failed composing the subrule '%s' due to conflicting output types.",
                   matchNestedRule.ruleId().map(ValueString::value).orElse("")),
@@ -155,11 +115,106 @@ private RuleOptimizationResult optimizeRule(Cel cel, CelCompiledRule compiledRul
       }
     }
 
-    CelMutableAst result = inlineCompiledVariables(matchAst, compiledRule.variables());
+    CelMutableAst resultExpr = output.expr;
+    resultExpr = inlineCompiledVariables(resultExpr, compiledRule.variables());
+    resultExpr = astMutator.renumberIdsConsecutively(resultExpr);
 
-    result = astMutator.renumberIdsConsecutively(result);
+    return new Step(output.isOptional, false, newTrueAst(cel), resultExpr);
+  }
 
-    return RuleOptimizationResult.create(result, isOptionalResult);
+  static RuleComposer newInstance(
+      CelCompiledRule compiledRule, String variablePrefix, int iterationLimit) {
+    return new RuleComposer(compiledRule, variablePrefix, iterationLimit);
+  }
+
+  // For the match's output, determine whether the output should be wrapped
+  // into an optional value, a conditional, or both.
+  //
+  // Assembles two output expressions into a single output step.
+  private Step combine(Cel cel, AstMutator astMutator, Step s, Step step) {
+    if (step == null) {
+      return s;
+    }
+    CelMutableAst trueCondition = newTrueAst(cel);
+
+    if (s.isOptional) {
+      if (step.isOptional) {
+        if (s.isConditional) {
+          return new Step(
+              true,
+              false,
+              trueCondition,
+              astMutator.newGlobalCall(
+                  Operator.CONDITIONAL.getFunction(), s.cond, s.expr, step.expr));
+        } else {
+          if (!isOptionalNone(step.expr)) {
+            // If either the nested rule or current condition output are optional then
+            // use optional.or() to specify the combination of the first and second results
+            // Note, the argument order is reversed due to the traversal of matches in
+            // reverse order.
+            return new Step(
+                true, false, trueCondition, astMutator.newMemberCall(s.expr, "or", step.expr));
+          }
+          return s;
+        }
+      } else { // step is non-optional
+        if (s.isConditional) {
+          return new Step(
+              true,
+              false,
+              trueCondition,
+              astMutator.newGlobalCall(
+                  Operator.CONDITIONAL.getFunction(),
+                  s.cond,
+                  s.expr,
+                  astMutator.newGlobalCall(Function.OPTIONAL_OF.getFunction(), step.expr)));
+        } else {
+          return new Step(
+              false, false, trueCondition, astMutator.newMemberCall(s.expr, "orValue", step.expr));
+        }
+      }
+    } else { // s is non-optional
+      if (step.isOptional) {
+        if (s.isConditional) {
+          return new Step(
+              true,
+              false,
+              trueCondition,
+              astMutator.newGlobalCall(
+                  Operator.CONDITIONAL.getFunction(),
+                  s.cond,
+                  astMutator.newGlobalCall(Function.OPTIONAL_OF.getFunction(), s.expr),
+                  step.expr));
+        } else {
+          // If the condition is trivially true, none of the matches in the rule causes the result
+          // to become optional, and the rule is not the last match, then this will introduce
+          // unreachable outputs or rules (pruning away 'step').
+          return s;
+        }
+      } else { // step is non-optional
+        return new Step(
+            false,
+            false,
+            trueCondition,
+            astMutator.newGlobalCall(
+                Operator.CONDITIONAL.getFunction(), s.cond, s.expr, step.expr));
+      }
+    }
+  }
+
+  private static boolean isOptionalNone(CelMutableAst ast) {
+    CelExpr expr = ast.toParsedAst().getExpr();
+    return expr.getKind().equals(Kind.CALL)
+        && expr.call().function().equals("optional.none")
+        && expr.call().args().isEmpty();
+  }
+
+  private static CelMutableAst newTrueAst(Cel cel) {
+    try {
+      return CelMutableAst.fromCelAst(cel.compile("true").getAst());
+    } catch (CelValidationException e) {
+      throw new IllegalStateException("Failed to compile 'true'", e);
+    }
   }
 
   private CelMutableAst inlineCompiledVariables(
@@ -186,11 +241,6 @@ private CelMutableAst inlineCompiledVariables(
     return mutatedAst;
   }
 
-  static RuleComposer newInstance(
-      CelCompiledRule compiledRule, String variablePrefix, int iterationLimit) {
-    return new RuleComposer(compiledRule, variablePrefix, iterationLimit);
-  }
-
   private void assertComposedAstIsValid(
       Cel cel, CelMutableAst composedAst, String failureMessage, Long... ids) {
     assertComposedAstIsValid(cel, composedAst, failureMessage, Arrays.asList(ids));
@@ -206,10 +256,35 @@ private void assertComposedAstIsValid(
     }
   }
 
-  private RuleComposer(CelCompiledRule compiledRule, String variablePrefix, int iterationLimit) {
-    this.compiledRule = checkNotNull(compiledRule);
-    this.variablePrefix = variablePrefix;
-    this.astMutator = AstMutator.newInstance(iterationLimit);
+  // Step represents an intermediate stage of rule and match expression composition.
+  //
+  // The CelCompiledRule and CelCompiledMatch types are meant to represent standalone tuples of
+  // condition and output expressions, and have no notion of how the order of combination would
+  // impact composition since composition rules may vary based on the policy execution semantic,
+  // e.g. first-match versus logical-or, logical-and, or accumulation.
+  private static class Step {
+    /**
+     * Indicates whether the output step has an optional result. Individual conditional attributes
+     * are not optional; however, rules and subrules can have optional output.
+     */
+    private final boolean isOptional;
+
+    /** True if the condition expression is not trivially true. */
+    private final boolean isConditional;
+
+    /** The condition associated with the output. */
+    private final CelMutableAst cond;
+
+    /** The output expression for the step. */
+    private final CelMutableAst expr;
+
+    private Step(
+        boolean isOptional, boolean isConditional, CelMutableAst cond, CelMutableAst expr) {
+      this.isOptional = isOptional;
+      this.isConditional = isConditional;
+      this.cond = cond;
+      this.expr = expr;
+    }
   }
 
   static final class RuleCompositionException extends RuntimeException {
@@ -225,4 +300,10 @@ private RuleCompositionException(
       this.compileException = e;
     }
   }
+
+  private RuleComposer(CelCompiledRule compiledRule, String variablePrefix, int iterationLimit) {
+    this.compiledRule = checkNotNull(compiledRule);
+    this.variablePrefix = variablePrefix;
+    this.astMutator = AstMutator.newInstance(iterationLimit);
+  }
 }

policy/src/test/java/dev/cel/policy/CelPolicyCompilerImplTest.java

@@ -266,8 +266,8 @@ public void evaluateYamlPolicy_nestedRuleProducesOptionalOutput() throws Excepti
         CelPolicyCompilerFactory.newPolicyCompiler(cel).build().compile(policy);
     Optional<Object> evalResult = (Optional<Object>) cel.createProgram(compiledPolicyAst).eval();
 
-    // Result is Optional<Optional<Object>>
-    assertThat(evalResult).hasValue(Optional.of(true));
+    // Result is Optional<Object> containing true
+    assertThat(evalResult).hasValue(true);
   }
 
   @Test

policy/src/test/java/dev/cel/policy/PolicyTestHelper.java

@@ -40,11 +40,11 @@ final class PolicyTestHelper {
   enum TestYamlPolicy {
     NESTED_RULE(
         "nested_rule",
-        true,
+        false,
         "cel.@block([resource.origin, @index0 in [\"us\", \"uk\", \"es\"], {\"banned\": true}],"
             + " ((@index0 in {\"us\": false, \"ru\": false, \"ir\": false} && !@index1) ?"
-            + " optional.of(@index2) : optional.none()).or(optional.of(@index1 ? {\"banned\":"
-            + " false} : @index2)))"),
+            + " optional.of(@index2) : optional.none()).orValue(@index1 ? {\"banned\":"
+            + " false} : @index2))"),
     NESTED_RULE2(
         "nested_rule2",
         false,
@@ -61,6 +61,22 @@ enum TestYamlPolicy {
             + " false, \"ru\": false, \"ir\": false} && @index1) ? {\"banned\":"
             + " \"restricted_region\"} : {\"banned\": \"bad_actor\"}) : (@index1 ?"
             + " optional.of({\"banned\": \"unconfigured_region\"}) : optional.none()))"),
+    NESTED_RULE4("nested_rule4", false, "(x > 0) ? true : false"),
+    NESTED_RULE5(
+        "nested_rule5",
+        true,
+        "cel.@block([optional.of(true), optional.none()], (x > 0) ? ((x > 2) ? @index0 : @index1) :"
+            + " ((x > 1) ? ((x >= 2) ? @index0 : @index1) : optional.of(false)))"),
+    NESTED_RULE6(
+        "nested_rule6",
+        false,
+        "cel.@block([optional.of(true), optional.none()], ((x > 2) ? @index0 : @index1).orValue(((x"
+            + " > 3) ? @index0 : @index1).orValue(false)))"),
+    NESTED_RULE7(
+        "nested_rule7",
+        true,
+        "cel.@block([optional.of(true), optional.none()], ((x > 2) ? @index0 : @index1).or(((x > 3)"
+            + " ? @index0 : @index1).or((x > 1) ? optional.of(false) : @index1)))"),
     REQUIRED_LABELS(
         "required_labels",
         true,

testing/src/test/resources/policy/nested_rule4/config.yaml

@@ -0,0 +1,19 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+name: "nested_rule4"
+variables:
+  - name: x
+    type:
+      type_name: int