Skip to content

Commit 6c8575f

Browse files
adwsinghadwsingh
committed
Rules engine optimizations
1 parent 4545d60 commit 6c8575f

11 files changed

Lines changed: 525 additions & 78 deletions

File tree

‎rulesengine/src/main/java/software/amazon/smithy/java/rulesengine/ArrayPropertyGetter.java‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package software.amazon.smithy.java.rulesengine;
7+
8+
/**
9+
* Lightweight PropertyGetter backed by parallel key/value arrays.
10+
* More efficient than Map for small fixed-key lookups (linear scan beats hashing for ~4 entries).
11+
*/
12+
record ArrayPropertyGetter(String[] keys, Object[] values) implements PropertyGetter {
13+
@Override
14+
public Object getProperty(String name) {
15+
for (int i = 0; i < keys.length; i++) {
16+
if (name.equals(keys[i])) {
17+
return values[i];
18+
}
19+
}
20+
return null;
21+
}
22+
}

‎rulesengine/src/main/java/software/amazon/smithy/java/rulesengine/Bytecode.java‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -178,6 +178,17 @@ public final class Bytecode {
178178
private final int[] hardRequiredIndices;
179179
private final Map<String, Integer> inputRegisterMap;
180180

181+
// Inline condition types for fast BDD evaluation.
182+
static final byte COND_ISSET = 1;
183+
static final byte COND_IS_TRUE = 2;
184+
static final byte COND_IS_FALSE = 3;
185+
static final byte COND_NOT_SET = 4;
186+
static final byte COND_STRING_EQ_REG_CONST = 5;
187+
188+
// Condition classification arrays for inline BDD evaluation
189+
final byte[] conditionTypes;
190+
final int[] conditionOperands;
191+
181192
private Bdd bdd;
182193

183194
Bytecode(
@@ -230,6 +241,67 @@ public final class Bytecode {
230241
this.hardRequiredIndices = findRequiredIndicesWithoutDefaultsOrBuiltins(registerDefinitions);
231242
this.inputRegisterMap = createInputRegisterMap(registerDefinitions);
232243
this.version = version;
244+
245+
// Classify conditions for inline BDD evaluation
246+
this.conditionTypes = new byte[conditionOffsets.length];
247+
this.conditionOperands = new int[conditionOffsets.length];
248+
classifyConditions();
249+
}
250+
251+
private void classifyConditions() {
252+
int len = bytecode.length;
253+
for (int i = 0; i < conditionOffsets.length; i++) {
254+
int offset = conditionOffsets[i];
255+
256+
// 2-byte opcode patterns: <opcode> <register> <return> (need offset+2 in bounds)
257+
if (offset + 2 < len) {
258+
int firstOpcode = bytecode[offset] & 0xFF;
259+
int reg = bytecode[offset + 1] & 0xFF;
260+
int next = bytecode[offset + 2] & 0xFF;
261+
// Only inline conditions that end with RETURN_VALUE (no binding).
262+
// Conditions with SET_REG_RETURN have a side effect (register write) that
263+
// the inline path cannot replicate — they must go through full bytecode eval.
264+
boolean isReturn = next == (Opcodes.RETURN_VALUE & 0xFF);
265+
266+
if (isReturn) {
267+
switch (firstOpcode) {
268+
case Opcodes.TEST_REGISTER_ISSET & 0xFF -> {
269+
conditionTypes[i] = COND_ISSET;
270+
conditionOperands[i] = reg;
271+
continue;
272+
}
273+
case Opcodes.TEST_REGISTER_IS_TRUE & 0xFF -> {
274+
conditionTypes[i] = COND_IS_TRUE;
275+
conditionOperands[i] = reg;
276+
continue;
277+
}
278+
case Opcodes.TEST_REGISTER_IS_FALSE & 0xFF -> {
279+
conditionTypes[i] = COND_IS_FALSE;
280+
conditionOperands[i] = reg;
281+
continue;
282+
}
283+
case Opcodes.TEST_REGISTER_NOT_SET & 0xFF -> {
284+
conditionTypes[i] = COND_NOT_SET;
285+
conditionOperands[i] = reg;
286+
continue;
287+
}
288+
default -> {
289+
}
290+
}
291+
}
292+
}
293+
294+
// 4-byte opcode pattern: STRING_EQUALS_REG_CONST [reg:1] [const:2] <return>
295+
if (offset + 4 < len && (bytecode[offset] & 0xFF) == (Opcodes.STRING_EQUALS_REG_CONST & 0xFF)) {
296+
int reg = bytecode[offset + 1] & 0xFF;
297+
int constIdx = ((bytecode[offset + 2] & 0xFF) << 8) | (bytecode[offset + 3] & 0xFF);
298+
int next = bytecode[offset + 4] & 0xFF;
299+
if (next == (Opcodes.RETURN_VALUE & 0xFF)) {
300+
conditionTypes[i] = COND_STRING_EQ_REG_CONST;
301+
conditionOperands[i] = (constIdx << 8) | reg;
302+
}
303+
}
304+
}
233305
}
234306

235307
/**

‎rulesengine/src/main/java/software/amazon/smithy/java/rulesengine/BytecodeCompiler.java‎

Lines changed: 119 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -113,12 +113,14 @@ Bytecode compile() {
113113

114114
private void compileCondition(Condition condition) {
115115
compileExpression(condition.getFunction());
116-
condition.getResult().ifPresent(result -> {
117-
byte register = registerAllocator.getOrAllocateRegister(result.toString());
118-
writer.writeByte(Opcodes.SET_REGISTER);
116+
var result = condition.getResult();
117+
if (result.isPresent()) {
118+
byte register = registerAllocator.getOrAllocateRegister(result.get().toString());
119+
writer.writeByte(Opcodes.SET_REG_RETURN);
119120
writer.writeByte(register);
120-
});
121-
writer.writeByte(Opcodes.RETURN_VALUE);
121+
} else {
122+
writer.writeByte(Opcodes.RETURN_VALUE);
123+
}
122124
}
123125

124126
private void compileEndpointRule(EndpointRule rule) {
@@ -150,7 +152,7 @@ private void compileEndpointRule(EndpointRule rule) {
150152
compileMapCreation(e.getProperties().size());
151153
}
152154

153-
compileExpression(e.getUrl());
155+
compileEndpointUrl(e.getUrl());
154156

155157
// Add the return endpoint instruction
156158
writer.writeByte(Opcodes.RETURN_ENDPOINT);
@@ -164,6 +166,95 @@ private void compileEndpointRule(EndpointRule rule) {
164166
writer.writeByte(packed);
165167
}
166168

169+
/**
170+
* Compile the URL expression for an endpoint, attempting to decompose it into scheme/host/path
171+
* so that BUILD_URI can construct a SmithyUri directly without string-to-URI parsing.
172+
*
173+
* <p>Falls back to compileExpression (producing a String) if the URL can't be decomposed.
174+
*/
175+
private void compileEndpointUrl(Expression urlExpression) {
176+
// Only optimize StringLiteral templates with multiple parts where the first part
177+
// is a literal starting with a scheme (e.g., "https://...")
178+
if (urlExpression instanceof StringLiteral sl) {
179+
var template = sl.value();
180+
var parts = template.getParts();
181+
if (parts.size() > 1 && parts.getFirst() instanceof Template.Literal firstLit) {
182+
String firstStr = firstLit.toString();
183+
int schemeEnd = firstStr.indexOf("://");
184+
if (schemeEnd > 0) {
185+
String scheme = firstStr.substring(0, schemeEnd);
186+
String afterScheme = firstStr.substring(schemeEnd + 3);
187+
188+
// Find where path starts: look for a part that is "/" or starts with "/"
189+
// The host is everything between "://" and the first "/" separator
190+
// In most S3 templates, there's no explicit "/" — the path is empty
191+
// In some, there's a "/" part followed by bucket + url.path
192+
int pathPartIndex = -1;
193+
for (int i = 0; i < parts.size(); i++) {
194+
if (parts.get(i) instanceof Template.Literal lit && lit.toString().startsWith("/")
195+
&& i > 0) {
196+
pathPartIndex = i;
197+
break;
198+
}
199+
}
200+
201+
// Compile host parts: afterScheme + parts[1..pathPartIndex)
202+
int hostPartCount = 0;
203+
if (!afterScheme.isEmpty()) {
204+
addLoadConst(afterScheme);
205+
hostPartCount++;
206+
}
207+
int hostEnd = pathPartIndex > 0 ? pathPartIndex : parts.size();
208+
for (int i = 1; i < hostEnd; i++) {
209+
var part = parts.get(i);
210+
if (part instanceof Template.Dynamic d) {
211+
compileExpression(d.toExpression());
212+
} else {
213+
addLoadConst(part.toString());
214+
}
215+
hostPartCount++;
216+
}
217+
// Resolve host template to a single string
218+
if (hostPartCount == 1) {
219+
// Already a single value on stack
220+
} else if (hostPartCount > 1) {
221+
writer.writeByte(Opcodes.RESOLVE_TEMPLATE);
222+
writer.writeByte(hostPartCount);
223+
} else {
224+
addLoadConst("");
225+
}
226+
227+
// Compile path parts
228+
if (pathPartIndex > 0) {
229+
int pathPartCount = 0;
230+
for (int i = pathPartIndex; i < parts.size(); i++) {
231+
var part = parts.get(i);
232+
if (part instanceof Template.Dynamic d) {
233+
compileExpression(d.toExpression());
234+
} else {
235+
addLoadConst(part.toString());
236+
}
237+
pathPartCount++;
238+
}
239+
if (pathPartCount > 1) {
240+
writer.writeByte(Opcodes.RESOLVE_TEMPLATE);
241+
writer.writeByte(pathPartCount);
242+
}
243+
} else {
244+
addLoadConst("");
245+
}
246+
247+
// BUILD_URI pops host and path, pushes SmithyUri
248+
writer.writeByte(Opcodes.BUILD_URI);
249+
writer.writeShort(writer.getConstantIndex(scheme));
250+
return;
251+
}
252+
}
253+
}
254+
// Fallback: compile as regular string expression
255+
compileExpression(urlExpression);
256+
}
257+
167258
private void compileErrorRule(ErrorRule rule) {
168259
compileExpression(rule.getError());
169260
writer.writeByte(Opcodes.RETURN_ERROR);
@@ -288,6 +379,20 @@ public Void visitStringEquals(Expression left, Expression right) {
288379
}
289380
}
290381
}
382+
if (left instanceof Reference ref && right instanceof StringLiteral sl
383+
&& sl.value().getParts().size() == 1) {
384+
writer.writeByte(Opcodes.STRING_EQUALS_REG_CONST);
385+
writer.writeByte(registerAllocator.getRegister(ref.getName().toString()));
386+
writer.writeShort(writer.getConstantIndex(sl.value().getParts().get(0).toString()));
387+
return null;
388+
}
389+
if (right instanceof Reference ref && left instanceof StringLiteral sl
390+
&& sl.value().getParts().size() == 1) {
391+
writer.writeByte(Opcodes.STRING_EQUALS_REG_CONST);
392+
writer.writeByte(registerAllocator.getRegister(ref.getName().toString()));
393+
writer.writeShort(writer.getConstantIndex(sl.value().getParts().get(0).toString()));
394+
return null;
395+
}
291396
compileExpression(left);
292397
compileExpression(right);
293398
writer.writeByte(Opcodes.STRING_EQUALS);
@@ -532,7 +637,14 @@ private void compileLiteral(Literal literal) {
532637
compileLiteral(e.getValue()); // value then key to make popping ordered
533638
addLoadConst(e.getKey().toString());
534639
}
535-
compileMapCreation(r.members().size());
640+
int size = r.members().size();
641+
if (size <= 8) {
642+
// Small records: PropertyGetter with linear scan beats Map hashing
643+
writer.writeByte(Opcodes.STRUCTN);
644+
writer.writeByte(size);
645+
} else {
646+
compileMapCreation(size);
647+
}
536648
}
537649
case BooleanLiteral b -> addLoadConst(b.value().getValue());
538650
case IntegerLiteral i -> addLoadConst(i.toNode().expectNumberNode().getValue());

‎rulesengine/src/main/java/software/amazon/smithy/java/rulesengine/BytecodeDisassembler.java‎

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,12 @@ final class BytecodeDisassembler {
8585
Map.entry(Opcodes.JUMP, new InstructionDef("JUMP", Show.JUMP_OFFSET)),
8686
Map.entry(Opcodes.SUBSTRING_EQ, new InstructionDef("SUBSTRING_EQ", Show.SUBSTRING_EQ)),
8787
Map.entry(Opcodes.SPLIT_GET, new InstructionDef("SPLIT_GET", Show.SPLIT_GET)),
88-
Map.entry(Opcodes.SELECT_BOOL_REG, new InstructionDef("SELECT_BOOL_REG", Show.SELECT_BOOL)));
88+
Map.entry(Opcodes.SELECT_BOOL_REG, new InstructionDef("SELECT_BOOL_REG", Show.SELECT_BOOL))),
89+
Map.entry(Opcodes.SPLIT_GET, new InstructionDef("SPLIT_GET", Show.SPLIT_GET)),
90+
Map.entry(Opcodes.STRING_EQUALS_REG_CONST, new InstructionDef("STRING_EQUALS_REG_CONST", Show.REG_CONST)),
91+
Map.entry(Opcodes.SET_REG_RETURN, new InstructionDef("SET_REG_RETURN", Show.REGISTER)),
92+
Map.entry(Opcodes.BUILD_URI, new InstructionDef("BUILD_URI", Show.CONST)),
93+
Map.entry(Opcodes.STRUCTN, new InstructionDef("STRUCTN", Show.NUMBER)));
8994
9095
private enum Show {
9196
CONST,
@@ -101,7 +106,9 @@ private enum Show {
101106
ARG_COUNT,
102107
SUBSTRING_EQ,
103108
SPLIT_GET,
104-
SELECT_BOOL
109+
SELECT_BOOL,
110+
SPLIT_GET,
111+
REG_CONST
105112
}
106113
107114
private record InstructionDef(String name, Show show) {
@@ -263,8 +270,7 @@ private void writeInstruction(StringBuilder s, BytecodeWalker walker, String ind
263270
opcode == Opcodes.JNN_OR_POP
264271
|| (opcode == Opcodes.GET_PROPERTY_REG && i == 1)
265272
|| (opcode == Opcodes.SELECT_BOOL_REG && i >= 1)
266-
||
267-
(opcode == Opcodes.RESOLVE_TEMPLATE && i == 1)) {
273+
|| (opcode == Opcodes.STRING_EQUALS_REG_CONST && i == 1)) {
268274
s.append(String.format("%5d", value));
269275
} else {
270276
s.append(String.format("%3d", value));
@@ -399,6 +405,17 @@ private void appendSymbolicInfo(StringBuilder s, BytecodeWalker walker, Show sho
399405
s.append(formatConstant(bytecode.getConstant(falseIdx)));
400406
}
401407
}
408+
case REG_CONST -> {
409+
int regIndex = walker.getOperand(0);
410+
int constIdx = walker.getOperand(1);
411+
if (regIndex >= 0 && regIndex < bytecode.getRegisterDefinitions().length) {
412+
s.append(bytecode.getRegisterDefinitions()[regIndex].name());
413+
}
414+
s.append(" == ");
415+
if (constIdx >= 0 && constIdx < bytecode.getConstantPoolCount()) {
416+
s.append(formatConstant(bytecode.getConstant(constIdx)));
417+
}
418+
}
402419
}
403420
}
404421

‎rulesengine/src/main/java/software/amazon/smithy/java/rulesengine/BytecodeEndpointResolver.java‎

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313
import software.amazon.smithy.java.endpoints.EndpointResolver;
1414
import software.amazon.smithy.java.endpoints.EndpointResolverParams;
1515
import software.amazon.smithy.java.logging.InternalLogger;
16-
import software.amazon.smithy.rulesengine.logic.bdd.Bdd;
1716

1817
/**
1918
* Endpoint resolver that uses a compiled endpoint rules program from a BDD.
@@ -23,7 +22,6 @@ public final class BytecodeEndpointResolver implements EndpointResolver {
2322
private static final InternalLogger LOGGER = InternalLogger.getLogger(BytecodeEndpointResolver.class);
2423

2524
private final Bytecode bytecode;
26-
private final Bdd bdd;
2725
private final RulesExtension[] extensions;
2826
private final RegisterFiller registerFiller;
2927
private final ContextProvider ctxProvider = new ContextProvider.OrchestratingProvider();
@@ -36,7 +34,6 @@ public BytecodeEndpointResolver(
3634
) {
3735
this.bytecode = bytecode;
3836
this.extensions = extensions.toArray(new RulesExtension[0]);
39-
this.bdd = bytecode.getBdd();
4037

4138
// Create and reuse this register filler across thread local evaluators.
4239
this.registerFiller = RegisterFiller.of(bytecode, builtinProviders);
@@ -64,10 +61,6 @@ public Endpoint resolveEndpoint(EndpointResolverParams params) {
6461

6562
LOGGER.debug("Resolving endpoint of {} using VM", operation);
6663

67-
var resultIndex = bdd.evaluate(evaluator);
68-
if (resultIndex < 0) {
69-
return null;
70-
}
71-
return evaluator.resolveResult(resultIndex);
64+
return evaluator.evaluateBdd();
7265
}
7366
}

0 commit comments

Comments
 (0)