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,58 @@ private RuleOptimizationResult optimizeRule(Cel cel, CelCompiledRule compiledRul
                     .collect(toImmutableList()))
             .build();
 
-    CelMutableAst matchAst = astMutator.newGlobalCall(Function.OPTIONAL_NONE.getFunction());
-    boolean isOptionalResult = true;
-    // Keep track of the last output ID that might cause type-check failure while attempting to
-    // compose the subgraphs.
+    Step output = null;
+    // If the rule has an optional output, the last result in the ternary should return
+    // `optional.none`. This output is implicit and created here to reflect the desired
+    // last possible output of this type of rule.
+    if (compiledRule.hasOptionalOutput()) {
+      output =
+          Step.newUnconditionalOptionalStep(
+              newTrueAst(cel), astMutator.newGlobalCall(Function.OPTIONAL_NONE.getFunction()));
+    }
+
     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:
+          // If the output is non-nil, then it is considered a non-optional output since
+          // it is explicitly stated. If the rule itself is optional, then the base case value
+          // of output being optional.none() will convert the non-optional value to an optional
+          // one.
           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);
-          }
-
-          matchAst =
-              astMutator.newGlobalCall(
-                  Operator.CONDITIONAL.getFunction(),
-                  CelMutableAst.fromCelAst(conditionAst),
-                  outAst,
-                  matchAst);
+          Step step = Step.newNonOptionalStep(!isTriviallyTrue, condAst, outAst);
+          output = combine(cel, astMutator, step, output);
+
           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 =
+              nestedHasOptional
+                  ? Step.newOptionalStep(!isTriviallyTrue, condAst, nestedRule.expr)
+                  : Step.newNonOptionalStep(!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 +121,116 @@ 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);
+
+    return output.isOptional
+        ? Step.newUnconditionalOptionalStep(newTrueAst(cel), resultExpr)
+        : Step.newUnconditionalNonOptionalStep(newTrueAst(cel), resultExpr);
+  }
+
+  static RuleComposer newInstance(
+      CelCompiledRule compiledRule, String variablePrefix, int iterationLimit) {
+    return new RuleComposer(compiledRule, variablePrefix, iterationLimit);
+  }
+
+  // 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) {
+      return combineWhenCurrentIsOptional(s, step, astMutator, trueCondition);
+    } else {
+      return combineWhenCurrentIsNonOptional(s, step, astMutator, trueCondition);
+    }
+  }
 
-    result = astMutator.renumberIdsConsecutively(result);
+  private Step combineWhenCurrentIsOptional(
+      Step s, Step step, AstMutator astMutator, CelMutableAst trueCondition) {
+    // optional.combine(optional) // optional
+    // (optional && conditional).combine(non-optional) // optional
+    // (optional && unconditional).combine(non-optional) // non-optional
+    if (step.isOptional) {
+      if (s.isConditional) {
+        return Step.newUnconditionalOptionalStep(
+            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 Step.newUnconditionalOptionalStep(
+              trueCondition, astMutator.newMemberCall(s.expr, "or", step.expr));
+        }
+        return s;
+      }
+    } else { // step is non-optional
+      if (s.isConditional) {
+        return Step.newUnconditionalOptionalStep(
+            trueCondition,
+            astMutator.newGlobalCall(
+                Operator.CONDITIONAL.getFunction(),
+                s.cond,
+                s.expr,
+                astMutator.newGlobalCall(Function.OPTIONAL_OF.getFunction(), step.expr)));
+      } else {
+        return Step.newUnconditionalNonOptionalStep(
+            trueCondition, astMutator.newMemberCall(s.expr, "orValue", step.expr));
+      }
+    }
+  }
 
-    return RuleOptimizationResult.create(result, isOptionalResult);
+  private Step combineWhenCurrentIsNonOptional(
+      Step s, Step step, AstMutator astMutator, CelMutableAst trueCondition) {
+    // non-optional.combine(non-optional) // non-optional
+    // (non-optional && conditional).combine(optional) // optional
+    // (non-optional && unconditional).combine(optional) // non-optional
+    //
+    // The last combination case is unusual, but effectively it means that the non-optional value
+    // prunes away
+    // the potential optional output.
+    if (step.isOptional) {
+      if (s.isConditional) {
+        return Step.newUnconditionalOptionalStep(
+            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 Step.newUnconditionalNonOptionalStep(
+          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 +257,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 +272,55 @@ 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;
+    }
+
+    private static Step newOptionalStep(
+        boolean isConditional, CelMutableAst cond, CelMutableAst expr) {
+      return new Step(/* isOptional= */ true, isConditional, cond, expr);
+    }
+
+    private static Step newNonOptionalStep(
+        boolean isConditional, CelMutableAst cond, CelMutableAst expr) {
+      return new Step(/* isOptional= */ false, isConditional, cond, expr);
+    }
+
+    private static Step newUnconditionalOptionalStep(
+        CelMutableAst trueCondition, CelMutableAst expr) {
+      return new Step(/* isOptional= */ true, /* isConditional= */ false, trueCondition, expr);
+    }
+
+    private static Step newUnconditionalNonOptionalStep(
+        CelMutableAst trueCondition, CelMutableAst expr) {
+      return new Step(/* isOptional= */ false, /* isConditional= */ false, trueCondition, expr);
+    }
   }
 
   static final class RuleCompositionException extends RuntimeException {
@@ -225,4 +336,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