Improve transforms

Rename CfgGuidedOrdering to InitialOrdering

Author: mtdowling

smithy-rules-engine/src/main/java/software/amazon/smithy/rulesengine/logic/bdd/BddEquivalenceChecker.java

@@ -60,7 +60,7 @@ private BddEquivalenceChecker(Cfg cfg, Bdd bdd, List<Condition> conditions, List
         this.bdd = bdd;
         this.conditions = conditions;
         this.results = results;
-        this.parameters = new ArrayList<>(cfg.getRuleSet().getParameters().toList());
+        this.parameters = new ArrayList<>(cfg.getParameters().toList());
 
         for (int i = 0; i < conditions.size(); i++) {
             conditionToIndex.put(conditions.get(i), i);

smithy-rules-engine/src/main/java/software/amazon/smithy/rulesengine/logic/bdd/BddTrait.java

@@ -73,7 +73,7 @@ public static BddTrait from(Cfg cfg) {
         }
 
         return builder()
-                .parameters(cfg.getRuleSet().getParameters())
+                .parameters(cfg.getParameters())
                 .conditions(compiler.getOrderedConditions())
                 .results(compiler.getIndexedResults())
                 .bdd(bdd)

smithy-rules-engine/src/main/java/software/amazon/smithy/rulesengine/logic/bdd/CfgGuidedOrdering.java renamed to smithy-rules-engine/src/main/java/software/amazon/smithy/rulesengine/logic/bdd/InitialOrdering.java

@@ -23,15 +23,15 @@
  * together in the BDD, enabling better node sharing. This ordering implementation flattens the tree structure while
  * respecting data dependencies.
  */
-final class CfgGuidedOrdering implements OrderingStrategy {
-    private static final Logger LOGGER = Logger.getLogger(CfgGuidedOrdering.class.getName());
+final class InitialOrdering implements OrderingStrategy {
+    private static final Logger LOGGER = Logger.getLogger(InitialOrdering.class.getName());
 
     /** How many distinct consumers make an isSet() a "gate". */
     private static final int GATE_SUCCESSOR_THRESHOLD = 2;
 
     private final Cfg cfg;
 
-    CfgGuidedOrdering(Cfg cfg) {
+    InitialOrdering(Cfg cfg) {
         this.cfg = cfg;
     }
 
@@ -51,6 +51,7 @@ public List<Condition> orderConditions(Condition[] conditions) {
 
         long elapsed = System.currentTimeMillis() - startTime;
         LOGGER.info(() -> String.format("Initial ordering: %d conditions in %dms", conditions.length, elapsed));
+        result.forEach(System.out::println);
         return result;
     }
 

smithy-rules-engine/src/main/java/software/amazon/smithy/rulesengine/logic/bdd/OrderingStrategy.java

@@ -28,7 +28,7 @@ interface OrderingStrategy {
      * @return the initial ordering strategy.
      */
     static OrderingStrategy initialOrdering(Cfg cfg) {
-        return new CfgGuidedOrdering(cfg);
+        return new InitialOrdering(cfg);
     }
 
     /**

smithy-rules-engine/src/main/java/software/amazon/smithy/rulesengine/logic/cfg/Cfg.java

@@ -17,6 +17,7 @@
 import java.util.NoSuchElementException;
 import java.util.Set;
 import software.amazon.smithy.rulesengine.language.EndpointRuleSet;
+import software.amazon.smithy.rulesengine.language.syntax.parameters.Parameters;
 import software.amazon.smithy.rulesengine.language.syntax.rule.Condition;
 import software.amazon.smithy.rulesengine.language.syntax.rule.EndpointRule;
 import software.amazon.smithy.rulesengine.language.syntax.rule.ErrorRule;
@@ -38,16 +39,20 @@
  */
 public final class Cfg implements Iterable<CfgNode> {
 
-    private final EndpointRuleSet ruleSet;
+    private final Parameters parameters;
     private final CfgNode root;
 
     // Lazily computed condition data
     private Condition[] conditions;
     private Map<Condition, Integer> conditionToIndex;
 
     Cfg(EndpointRuleSet ruleSet, CfgNode root) {
-        this.ruleSet = ruleSet;
+        this(ruleSet == null ? Parameters.builder().build() : ruleSet.getParameters(), root);
+    }
+
+    Cfg(Parameters parameters, CfgNode root) {
         this.root = SmithyBuilder.requiredState("root", root);
+        this.parameters = parameters;
     }
 
     /**
@@ -125,8 +130,8 @@ private synchronized void extractConditions() {
         this.conditionToIndex = indexMap;
     }
 
-    public EndpointRuleSet getRuleSet() {
-        return ruleSet;
+    public Parameters getParameters() {
+        return parameters;
     }
 
     @Override

smithy-rules-engine/src/main/java/software/amazon/smithy/rulesengine/logic/cfg/CoalesceTransform.java

@@ -45,21 +45,17 @@ static EndpointRuleSet transform(EndpointRuleSet ruleSet) {
         CoalesceTransform transform = new CoalesceTransform();
 
         List<Rule> transformedRules = new ArrayList<>();
-        for (Rule rule : ruleSet.getRules()) {
-            transformedRules.add(transform.transformRule(rule));
+        for (int i = 0; i < ruleSet.getRules().size(); i++) {
+            transformedRules.add(transform.transformRule(ruleSet.getRules().get(i), "root/rule[" + i + "]"));
         }
 
         if (LOGGER.isLoggable(Level.INFO)) {
-            StringBuilder msg = new StringBuilder();
-            msg.append("\n=== Coalescing Transform Complete ===\n");
-            msg.append("Total: ").append(transform.coalesceCount).append(" coalesced, ");
-            msg.append(transform.cacheHits).append(" cache hits, ");
-            msg.append(transform.skippedNoZeroValue).append(" skipped (no zero value), ");
-            msg.append(transform.skippedMultipleUses).append(" skipped (multiple uses)");
-            if (!transform.skippedRecordTypes.isEmpty()) {
-                msg.append("\nSkipped record-returning functions: ").append(transform.skippedRecordTypes);
-            }
-            LOGGER.info(msg.toString());
+            LOGGER.info(String.format(
+                    "Coalescing: %d coalesced, %d cache hits, %d skipped (no zero), %d skipped (multiple uses)",
+                    transform.coalesceCount,
+                    transform.cacheHits,
+                    transform.skippedNoZeroValue,
+                    transform.skippedMultipleUses));
         }
 
         return EndpointRuleSet.builder()
@@ -69,50 +65,31 @@ static EndpointRuleSet transform(EndpointRuleSet ruleSet) {
                 .build();
     }
 
-    private Rule transformRule(Rule rule) {
+    private Rule transformRule(Rule rule, String rulePath) {
         Set<Condition> eliminatedConditions = new HashSet<>();
-        List<Condition> conditions = rule.getConditions();
-        Map<String, Integer> localVarUsage = countLocalVariableUsage(conditions);
-        List<Condition> transformedConditions = transformConditions(conditions, eliminatedConditions, localVarUsage);
 
-        if (rule instanceof TreeRule) {
-            TreeRule treeRule = (TreeRule) rule;
-            List<Rule> transformedNestedRules = new ArrayList<>();
-            boolean nestedChanged = false;
-
-            for (Rule nestedRule : treeRule.getRules()) {
-                Rule transformedNested = transformRule(nestedRule);
-                transformedNestedRules.add(transformedNested);
-                if (transformedNested != nestedRule) {
-                    nestedChanged = true;
-                }
-            }
-
-            if (!transformedConditions.equals(conditions) || nestedChanged) {
-                return TreeRule.builder()
-                        .description(rule.getDocumentation().orElse(null))
-                        .conditions(transformedConditions)
-                        .treeRule(transformedNestedRules);
+        // Count local usage for THIS rule's conditions
+        Map<String, Integer> localVarUsage = new HashMap<>();
+        for (Condition condition : rule.getConditions()) {
+            for (String ref : condition.getFunction().getReferences()) {
+                localVarUsage.merge(ref, 1, Integer::sum);
             }
-        } else if (!transformedConditions.equals(conditions)) {
-            // For other rule types, just update conditions
-            return rule.withConditions(transformedConditions);
         }
 
-        return rule;
-    }
-
-    private Map<String, Integer> countLocalVariableUsage(List<Condition> conditions) {
-        Map<String, Integer> usage = new HashMap<>();
+        List<Condition> transformedConditions = transformConditions(
+                rule.getConditions(),
+                eliminatedConditions,
+                localVarUsage);
 
-        // Count how many times each variable is used within this specific rule
-        for (Condition condition : conditions) {
-            for (String ref : condition.getFunction().getReferences()) {
-                usage.merge(ref, 1, Integer::sum);
-            }
+        if (rule instanceof TreeRule) {
+            TreeRule treeRule = (TreeRule) rule;
+            return TreeRule.builder()
+                    .description(rule.getDocumentation().orElse(null))
+                    .conditions(transformedConditions)
+                    .treeRule(TreeRewriter.transformNestedRules(treeRule, rulePath, this::transformRule));
         }
 
-        return usage;
+        return rule.withConditions(transformedConditions);
     }
 
     private List<Condition> transformConditions(
@@ -128,25 +105,19 @@ private List<Condition> transformConditions(
                 continue;
             }
 
-            // Check if this is a bind that can be coalesced with the next condition
             if (i + 1 < conditions.size() && current.getResult().isPresent()) {
                 String var = current.getResult().get().toString();
                 Condition next = conditions.get(i + 1);
 
                 if (canCoalesce(var, current, next, localVarUsage)) {
-                    // Create coalesced condition
-                    Condition coalesced = createCoalescedCondition(current, next, var);
-                    result.add(coalesced);
-                    // Mark both conditions as eliminated
+                    result.add(createCoalescedCondition(current, next, var));
                     eliminatedConditions.add(current);
                     eliminatedConditions.add(next);
-                    // Skip the next condition
-                    i++;
+                    i++; // Skip next
                     continue;
                 }
             }
 
-            // No coalescing possible, keep the condition as-is
             result.add(current);
         }
 
@@ -155,28 +126,22 @@ private List<Condition> transformConditions(
 
     private boolean canCoalesce(String var, Condition bind, Condition use, Map<String, Integer> localVarUsage) {
         if (!use.getFunction().getReferences().contains(var)) {
-            // The use condition must reference the variable
             return false;
-        } else if (use.getFunction().getFunctionDefinition() == IsSet.getDefinition()) {
-            // Never coalesce into presence checks (isSet)
+        }
+
+        if (use.getFunction().getFunctionDefinition() == IsSet.getDefinition()) {
             return false;
         }
 
-        // Check if variable is only used once in this local rule context (even if it appears multiple times globally)
+        // Check local usage
         Integer localUses = localVarUsage.get(var);
         if (localUses == null || localUses > 1) {
             skippedMultipleUses++;
             return false;
         }
 
-        // Get the actual return type (could be Optional<T> or T)
         Type type = bind.getFunction().getFunctionDefinition().getReturnType();
-
-        // Check if we can get a zero value for this type. For OptionalType, we use the inner type's zero value
-        Type innerType = type;
-        if (type instanceof OptionalType) {
-            innerType = ((OptionalType) type).inner();
-        }
+        Type innerType = type instanceof OptionalType ? ((OptionalType) type).inner() : type;
 
         if (innerType instanceof RecordType) {
             skippedNoZeroValue++;
@@ -196,16 +161,11 @@ private Condition createCoalescedCondition(Condition bind, Condition use, String
         LibraryFunction bindExpr = bind.getFunction();
         LibraryFunction useExpr = use.getFunction();
 
-        // Get the type and its zero value
         Type type = bindExpr.getFunctionDefinition().getReturnType();
-        Type innerType = type;
-        if (type instanceof OptionalType) {
-            innerType = ((OptionalType) type).inner();
-        }
+        Type innerType = type instanceof OptionalType ? ((OptionalType) type).inner() : type;
 
         Literal zero = innerType.getZeroValue().get();
 
-        // Create cache key based on canonical representations
         String bindCanonical = bindExpr.canonicalize().toString();
         String zeroCanonical = zero.toString();
         String useCanonical = useExpr.canonicalize().toString();
@@ -220,10 +180,9 @@ private Condition createCoalescedCondition(Condition bind, Condition use, String
 
         Expression coalesced = Coalesce.ofExpressions(bindExpr, zero);
 
-        // Replace the variable reference in the use expression
         Map<String, Expression> replacements = new HashMap<>();
         replacements.put(var, coalesced);
-        ReferenceRewriter rewriter = ReferenceRewriter.forReplacements(replacements);
+        TreeRewriter rewriter = TreeRewriter.forReplacements(replacements);
         Expression replaced = rewriter.rewrite(useExpr);
         LibraryFunction canonicalized = ((LibraryFunction) replaced).canonicalize();