Make various improvements to the VM

Author: mtdowling

aws/client/aws-client-rulesengine/src/jmh/java/software/amazon/smithy/java/aws/client/rulesengine/Bench.java

@@ -40,18 +40,21 @@ public class Bench {
     private EndpointResolver endpointResolver;
     private EndpointResolverParams endpointParams;
 
+    private Model model;
+    private ServiceShape service;
+    private final RulesEngineBuilder engine = new RulesEngineBuilder();
+
     @Setup
     public void setup() {
-        var model = Model.assembler()
+        model = Model.assembler()
                 .discoverModels()
                 .addImport(ResolverTest.class.getResource("s3.json"))
                 .putProperty(ModelAssembler.ALLOW_UNKNOWN_TRAITS, true)
                 .assemble()
                 .unwrap();
         model = customizeS3Model(model);
-        var service = model.expectShape(ShapeId.from("com.amazonaws.s3#AmazonS3"), ServiceShape.class);
+        service = model.expectShape(ShapeId.from("com.amazonaws.s3#AmazonS3"), ServiceShape.class);
 
-        var engine = new RulesEngineBuilder();
         var plugin = EndpointRulesPlugin.create(engine);
 
         client = DynamicClient.builder()
@@ -61,17 +64,18 @@ public void setup() {
                 .addPlugin(plugin)
                 .build();
         endpointResolver = client.config().endpointResolver();
-        var ctx = Context.create();
-        ctx.put(EndpointSettings.REGION, "us-east-1");
-//        ctx.put(EndpointRulesPlugin.ADDITIONAL_ENDPOINT_PARAMS, Map.of("ForcePathStyle", true));
 
         var inputValue = client.createStruct(ShapeId.from("com.amazonaws.s3#GetObjectRequest"),
                 Document.of(Map.of(
                         "Bucket",
                         Document.of("miked"),
-                        //, Document.of("foo"),
                         "Key",
                         Document.of("bar"))));
+
+        var ctx = Context.create();
+        ctx.put(EndpointSettings.REGION, "us-east-1");
+        //        ctx.put(EndpointRulesPlugin.ADDITIONAL_ENDPOINT_PARAMS, Map.of("UseFIPS", true, "UseDualStack", true));
+
         endpointParams = EndpointResolverParams.builder()
                 .context(ctx)
                 .inputValue(inputValue)
@@ -101,4 +105,17 @@ private static Model customizeS3Model(Model m) {
     public Object evaluate() throws Exception {
         return endpointResolver.resolveEndpoint(endpointParams).get();
     }
+
+    //    public static final TraitKey<EndpointRuleSetTrait> ENDPOINT_RULESET_TRAIT =
+    //            TraitKey.get(EndpointRuleSetTrait.class);
+    //
+    //    @Benchmark
+    //    public Object compileBdd() {
+    //        var ruleset = client.config().service().schema().getTrait(ENDPOINT_RULESET_TRAIT);
+    //        var cfg = Cfg.from(ruleset.getEndpointRuleSet());
+    //        var bdd = BddTrait.from(cfg).transform(SiftingOptimization.builder().cfg(cfg).build()).transform(new NodeReversal());
+    //        var bytecode = engine.compile(bdd);
+    //        var providers = engine.getBuiltinProviders();
+    //        return new BytecodeEndpointResolver(bytecode, engine.getExtensions(), providers);
+    //    }
 }

client/client-rulesengine/src/main/java/software/amazon/smithy/java/client/rulesengine/Bytecode.java

@@ -51,29 +51,25 @@
  * <ol>
  *   <li><b>Condition Table</b> - Array of 4-byte offsets pointing to each condition's bytecode</li>
  *   <li><b>Result Table</b> - Array of 4-byte offsets pointing to each result's bytecode</li>
+ *   <li><b>Register Definitions</b> - Array of parameter/register metadata (immediately after result table)</li>
  *   <li><b>Function Table</b> - Array of function names</li>
- *   <li><b>Register Definitions</b> - Array of parameter/register metadata</li>
- *   <li><b>BDD Table</b> - Array of BDD nodes</li>
+ *   <li><b>BDD Table</b> - Array of BDD nodes (3 ints per node)</li>
  *   <li><b>Bytecode Section</b> - Compiled instructions for conditions and results</li>
  *   <li><b>Constant Pool</b> - All constants referenced by the bytecode</li>
  * </ol>
  *
  * <h3>Condition Table</h3>
  * <p>Array of 4-byte offsets pointing to the start of each condition's bytecode.
- * Each offset is absolute from the start of the file.
+ * Each offset is absolute from the start of the file. When loaded, these are
+ * adjusted to be relative to the bytecode section start for efficient access.
  *
  * <h3>Result Table</h3>
  * <p>Array of 4-byte offsets pointing to the start of each result's bytecode.
- * Each offset is absolute from the start of the file.
- *
- * <h3>Function Table</h3>
- * <p>Array of function names used in the bytecode. Each function name is encoded as:
- * <pre>
- * [length:2][UTF-8 bytes]
- * </pre>
+ * Each offset is absolute from the start of the file. When loaded, these are
+ * adjusted to be relative to the bytecode section start for efficient access.
  *
  * <h3>Register Definitions</h3>
- * <p>Immediately follows the function table. Each register is encoded as:
+ * <p>Immediately follows the result table. Each register is encoded as:
  * <pre>
  * [nameLen:2][name:UTF-8][required:1][temp:1][hasDefault:1][default:?][hasBuiltin:1][builtin:?]
  * </pre>
@@ -89,16 +85,38 @@
  *   <li>builtin: UTF-8 encoded builtin name (only present if hasBuiltin=1)</li>
  * </ul>
  *
+ * <h3>Function Table</h3>
+ * <p>Array of function names used in the bytecode. Each function name is encoded as:
+ * <pre>
+ * [length:2][UTF-8 bytes]
+ * </pre>
+ *
  * <h3>BDD Table</h3>
  * <p>Array of BDD nodes for efficient condition evaluation. Each node is encoded as 12 bytes:
  * <pre>
  * [varIdx:4][highRef:4][lowRef:4]
  * </pre>
+ * Where:
+ * <ul>
+ *   <li>varIdx: Index of the condition to test (0-based)</li>
+ *   <li>highRef: Reference to follow if condition is true</li>
+ *   <li>lowRef: Reference to follow if condition is false</li>
+ * </ul>
+ *
+ * <p>Reference encoding follows the BDD conventions:
+ * <ul>
+ *   <li>1: TRUE terminal</li>
+ *   <li>-1: FALSE terminal</li>
+ *   <li>2, 3, ...: Node references (node at index ref-1)</li>
+ *   <li>-2, -3, ...: Complement node references (logical NOT)</li>
+ *   <li>100_000_000+: Result terminals (100_000_000 + resultIndex)</li>
+ * </ul>
  *
  * <h3>Bytecode Section</h3>
  * <p>Contains the compiled bytecode instructions for all conditions and results.
- * The condition/result tables point to offsets within this section. Instructions
- * use a stack-based virtual machine with opcodes defined in {@link Opcodes}.
+ * The condition/result tables contain absolute offsets from the start of the file that point into this section.
+ * Instructions use a stack-based virtual machine with opcodes defined in {@link Opcodes}. This section may include
+ * control flow instructions like JT_OR_POP for short-circuit evaluation.
  *
  * <h3>Constant Pool</h3>
  * <p>Contains all constants referenced by the bytecode. Each constant is prefixed
@@ -115,7 +133,7 @@
  * 5     MAP     [count:2]([keyLen:2][key:UTF-8][value:?])...
  * </pre>
  *
- * <p>Lists and maps can contain nested values of any supported type.
+ * <p>Lists and maps can contain nested values of any supported type, up to a maximum nesting depth of 100 levels.
  *
  * <h2>Usage</h2>
  *
@@ -171,6 +189,10 @@ public final class Bytecode {
             int[] bddNodes,
             int bddRootRef
     ) {
+        if (bddNodes.length % 3 != 0) {
+            throw new IllegalArgumentException("BDD nodes length must be multiple of 3, got: " + bddNodes.length);
+        }
+
         this.bytecode = Objects.requireNonNull(bytecode);
         this.conditionOffsets = Objects.requireNonNull(conditionOffsets);
         this.resultOffsets = Objects.requireNonNull(resultOffsets);

client/client-rulesengine/src/main/java/software/amazon/smithy/java/client/rulesengine/BytecodeCompiler.java

@@ -55,7 +55,7 @@ final class BytecodeCompiler {
         this.builtinProviders = builtinProviders;
         this.availableFunctions = functions;
         this.bdd = bdd;
-        this.registerAllocator = new RegisterAllocator.FlatAllocator();
+        this.registerAllocator = new RegisterAllocator();
 
         // Add parameters as registry values
         for (var param : bdd.getParameters()) {
@@ -110,7 +110,7 @@ Bytecode compile() {
     private void compileCondition(Condition condition) {
         compileExpression(condition.getFunction());
         condition.getResult().ifPresent(result -> {
-            byte register = registerAllocator.shadow(result.toString());
+            byte register = registerAllocator.getOrAllocateRegister(result.toString());
             writer.writeByte(Opcodes.SET_REGISTER);
             writer.writeByte(register);
         });
@@ -370,7 +370,7 @@ private void compileLiteral(Literal literal) {
                 // Simple string with no interpolation
                 addLoadConst(parts.get(0).toString());
             } else if (parts.size() == 1 && parts.get(0) instanceof Template.Dynamic dynamic) {
-                // Single dynamic expression - just evaluate it
+                // Single dynamic expression, so just evaluate it
                 compileExpression(dynamic.toExpression());
             } else {
                 // Multiple parts - need to concatenate
@@ -447,7 +447,6 @@ private Bytecode buildProgram() {
         var registerDefs = registerAllocator.getRegistry().toArray(new RegisterDefinition[0]);
         var fns = usedFunctions.toArray(new RulesFunction[0]);
 
-        // Extract BDD nodes to an array using the consumer approach
         Bdd bddData = bdd.getBdd();
         int nodeCount = bddData.getNodeCount();
         int[] bddNodes = new int[nodeCount * 3];