Skip to content

Commit 7aec981

Browse files
mtdowlingmtdowling
committed
Add various optimizations to rules VM
1 parent 4998b35 commit 7aec981

File tree

4 files changed

+112
-82
lines changed

4 files changed

+112
-82
lines changed

‎client/client-rulesengine/src/main/java/software/amazon/smithy/java/client/rulesengine/RulesCompiler.java

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -259,8 +259,8 @@ private void addLiteralOpcodes(Literal literal) {
259259
add_CREATE_LIST((short) t.members().size());
260260
} else if (literal instanceof RecordLiteral r) {
261261
for (var e : r.members().entrySet()) {
262+
addLiteralOpcodes(e.getValue()); // value then key to make popping ordered
262263
add_LOAD_CONST(e.getKey().toString());
263-
addLiteralOpcodes(e.getValue());
264264
}
265265
add_CREATE_MAP((short) r.members().size());
266266
} else if (literal instanceof BooleanLiteral b) {
@@ -389,14 +389,14 @@ private void compileEndpointRule(EndpointRule rule) {
389389
// Add endpoint header instructions.
390390
if (!e.getHeaders().isEmpty()) {
391391
for (var entry : e.getHeaders().entrySet()) {
392-
// Push the header name first.
393-
add_LOAD_CONST(entry.getKey());
394-
// Push the instructions for creating the headers.
392+
// Header values. Then header name.
395393
for (var h : entry.getValue()) {
396394
compileExpression(h);
397395
}
398396
// Process the N header values that are on the stack.
399397
add_CREATE_LIST((short) entry.getValue().size());
398+
// Now the header name.
399+
add_LOAD_CONST(entry.getKey());
400400
}
401401
// Combine the N headers that are on the stack in the form of String followed by List<String>.
402402
add_CREATE_MAP((short) e.getHeaders().size());
@@ -405,8 +405,8 @@ private void compileEndpointRule(EndpointRule rule) {
405405
// Add property instructions.
406406
if (!e.getProperties().isEmpty()) {
407407
for (var entry : e.getProperties().entrySet()) {
408-
add_LOAD_CONST(entry.getKey().toString());
409408
compileExpression(entry.getValue());
409+
add_LOAD_CONST(entry.getKey().toString());
410410
}
411411
add_CREATE_MAP((short) e.getProperties().size());
412412
}

‎client/client-rulesengine/src/main/java/software/amazon/smithy/java/client/rulesengine/RulesVm.java

Lines changed: 69 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import java.net.URI;
99
import java.net.URISyntaxException;
1010
import java.util.Arrays;
11+
import java.util.Collections;
1112
import java.util.HashMap;
1213
import java.util.LinkedHashMap;
1314
import java.util.List;
@@ -36,8 +37,8 @@ protected boolean removeEldestEntry(Map.Entry<String, URI> eldest) {
3637
private static final int MIN_TEMP_ARRAY_SIZE = 8;
3738

3839
// Temp array used during evaluation.
39-
private Object[] tempArray;
40-
private int tempArraySize = 0;
40+
private Object[] tempArray = new Object[8];
41+
private int tempArraySize = 8;
4142

4243
private final Context context;
4344
private final RulesProgram program;
@@ -109,14 +110,18 @@ void initializeRegister(Context context, int index, ParamDefinition definition)
109110

110111
private void push(Object value) {
111112
if (stackPosition == stack.length) {
112-
int newCapacity = stack.length + (stack.length >> 1);
113-
Object[] newStack = new Object[newCapacity];
114-
System.arraycopy(stack, 0, newStack, 0, stack.length);
115-
stack = newStack;
113+
resizeStack();
116114
}
117115
stack[stackPosition++] = value;
118116
}
119117

118+
private void resizeStack() {
119+
int newCapacity = stack.length + (stack.length >> 1);
120+
Object[] newStack = new Object[newCapacity];
121+
System.arraycopy(stack, 0, newStack, 0, stack.length);
122+
stack = newStack;
123+
}
124+
120125
private Object pop() {
121126
return stack[--stackPosition]; // no need to clear out the memory since it's tied to lifetime of the VM.
122127
}
@@ -139,22 +144,13 @@ private Object run() {
139144
// Skip version, params, and register bytes.
140145
for (pc = program.instructionOffset + 3; pc < instructionSize; pc++) {
141146
switch (instructions[pc]) {
142-
case RulesProgram.LOAD_CONST -> {
143-
int constant = instructions[++pc] & 0xFF; // read unsigned byte
144-
push(constantPool[constant]);
145-
}
147+
case RulesProgram.LOAD_CONST -> push(constantPool[instructions[++pc] & 0xFF]); // read unsigned byte
146148
case RulesProgram.LOAD_CONST_W -> {
147149
push(constantPool[readUnsignedShort(pc + 1)]); // read unsigned short
148150
pc += 2;
149151
}
150-
case RulesProgram.SET_REGISTER -> {
151-
int register = instructions[++pc] & 0xFF; // read unsigned byte
152-
registers[register] = peek();
153-
}
154-
case RulesProgram.LOAD_REGISTER -> {
155-
int register = instructions[++pc] & 0xFF; // read unsigned byte
156-
push(registers[register]);
157-
}
152+
case RulesProgram.SET_REGISTER -> registers[instructions[++pc] & 0xFF] = peek(); // read unsigned byte
153+
case RulesProgram.LOAD_REGISTER -> push(registers[instructions[++pc] & 0xFF]); // read unsigned byte
158154
case RulesProgram.JUMP_IF_FALSEY -> {
159155
Object value = pop();
160156
if (value == null || Boolean.FALSE.equals(value)) {
@@ -170,34 +166,19 @@ private Object run() {
170166
push(value != null && !Boolean.FALSE.equals(value));
171167
}
172168
case RulesProgram.TEST_REGISTER_ISSET -> {
173-
int register = instructions[++pc] & 0xFF; // read unsigned byte
174-
var value = registers[register];
169+
var value = registers[instructions[++pc] & 0xFF]; // read unsigned byte
175170
push(value != null && !Boolean.FALSE.equals(value));
176171
}
177172
case RulesProgram.RETURN_ERROR -> {
178173
throw new RulesEvaluationError((String) pop());
179174
}
180175
case RulesProgram.RETURN_ENDPOINT -> {
181-
short packed = instructions[++pc];
182-
boolean hasHeaders = (packed & 1) != 0;
183-
boolean hasProperties = (packed & 2) != 0;
184-
return setEndpoint(hasProperties, hasHeaders);
185-
}
186-
case RulesProgram.CREATE_LIST -> {
187-
int size = instructions[++pc] & 0xFF; // read unsigned byte
188-
var values = new Object[size];
189-
for (var i = size - 1; i >= 0; i--) {
190-
values[i] = pop();
191-
}
192-
push(Arrays.asList(values));
193-
}
194-
case RulesProgram.CREATE_MAP -> {
195-
int size = instructions[++pc] & 0xFF; // read unsigned byte
196-
createMap(size);
176+
return setEndpoint(instructions[++pc]);
197177
}
178+
case RulesProgram.CREATE_LIST -> createList(instructions[++pc] & 0xFF); // read unsigned byte
179+
case RulesProgram.CREATE_MAP -> createMap(instructions[++pc] & 0xFF); // read unsigned byte
198180
case RulesProgram.RESOLVE_TEMPLATE -> {
199-
var constant = readUnsignedShort(pc + 1);
200-
resolveTemplate((StringTemplate) constantPool[constant]);
181+
resolveTemplate((StringTemplate) constantPool[readUnsignedShort(pc + 1)]);
201182
pc += 2;
202183
}
203184
case RulesProgram.FN -> {
@@ -245,17 +226,48 @@ private Object run() {
245226
}
246227

247228
private void createMap(int size) {
248-
Map<String, Object> headers = new HashMap<>(size);
249-
for (var i = 0; i < size; i++) {
250-
var value = pop();
251-
var key = pop();
252-
headers.put((String) key, value);
229+
push(switch (size) {
230+
case 0 -> Map.of();
231+
case 1 -> Map.of((String) pop(), pop());
232+
case 2 -> Map.of((String) pop(), pop(), (String) pop(), pop());
233+
case 3 -> Map.of((String) pop(), pop(), (String) pop(), pop(), (String) pop(), pop());
234+
default -> {
235+
Map<String, Object> map = new HashMap<>((int) (size / 0.75f) + 1); // Avoid rehashing
236+
for (var i = 0; i < size; i++) {
237+
map.put((String) pop(), pop());
238+
}
239+
yield map;
240+
}
241+
});
242+
}
243+
244+
private void createList(int size) {
245+
push(switch (size) {
246+
case 0 -> List.of();
247+
case 1 -> Collections.singletonList(pop());
248+
default -> {
249+
var values = new Object[size];
250+
for (var i = size - 1; i >= 0; i--) {
251+
values[i] = pop();
252+
}
253+
yield Arrays.asList(values);
254+
}
255+
});
256+
}
257+
258+
private void resolveTemplate(StringTemplate template) {
259+
var expressionCount = template.expressionCount();
260+
var temp = getTempArray(expressionCount);
261+
for (var i = 0; i < expressionCount; i++) {
262+
temp[i] = pop();
253263
}
254-
push(headers);
264+
push(template.resolve(expressionCount, temp));
255265
}
256266

257267
@SuppressWarnings("unchecked")
258-
private Endpoint setEndpoint(boolean hasProperties, boolean hasHeaders) {
268+
private Endpoint setEndpoint(byte packed) {
269+
boolean hasHeaders = (packed & 1) != 0;
270+
boolean hasProperties = (packed & 2) != 0;
259271
var urlString = (String) pop();
260272
var properties = (Map<String, Object>) (hasProperties ? pop() : Map.of());
261273
var headers = (Map<String, List<String>>) (hasHeaders ? pop() : Map.of());
@@ -286,27 +298,21 @@ public static URI createUri(String uriStr) {
286298
return uri;
287299
}
288300

289-
private void resolveTemplate(StringTemplate template) {
290-
if (template.expressionCount() == 0) {
291-
push(template.resolve(0, tempArray));
292-
} else {
293-
var temp = getTempArray(template.expressionCount());
294-
for (var i = 0; i < template.expressionCount(); i++) {
295-
temp[i] = pop();
296-
}
297-
push(template.resolve(template.expressionCount(), temp));
301+
private Object[] getTempArray(int requiredSize) {
302+
if (tempArraySize < requiredSize) {
303+
resizeTempArray(requiredSize);
298304
}
305+
return tempArray;
299306
}
300307

301-
private Object[] getTempArray(int requiredSize) {
302-
var size = Math.max(MIN_TEMP_ARRAY_SIZE, requiredSize);
303-
if (tempArraySize < size) {
304-
tempArray = new Object[size];
305-
tempArraySize = size;
306-
} else {
307-
// Clear out any previously set values.
308-
Arrays.fill(tempArray, null);
308+
private void resizeTempArray(int requiredSize) {
309+
// Resize to a power of two.
310+
int newSize = MIN_TEMP_ARRAY_SIZE;
311+
while (newSize < requiredSize) {
312+
newSize <<= 1;
309313
}
310-
return tempArray;
314+
315+
tempArray = new Object[newSize];
316+
tempArraySize = newSize;
311317
}
312318
}

‎client/client-rulesengine/src/main/java/software/amazon/smithy/java/client/rulesengine/StringTemplate.java

Lines changed: 36 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -14,13 +14,32 @@
1414

1515
/**
1616
* Similar to {@link Template}, but built around Object instead of {@link Value}.
17-
*
18-
* @param template The original template string.
19-
* @param parts The template parts.
20-
* @param expressionCount The number of expression in the template.
21-
* @param singularExpression A non-null expression value if the template contains only an expression.
2217
*/
23-
record StringTemplate(String template, Object[] parts, int expressionCount, Expression singularExpression) {
18+
final class StringTemplate {
19+
20+
private static final ThreadLocal<StringBuilder> STRING_BUILDER = ThreadLocal.withInitial(
21+
() -> new StringBuilder(64));
22+
23+
private final String template;
24+
private final Object[] parts;
25+
private final int expressionCount;
26+
private final Expression singularExpression;
27+
28+
StringTemplate(String template, Object[] parts, int expressionCount, Expression singularExpression) {
29+
this.template = template;
30+
this.parts = parts;
31+
this.expressionCount = expressionCount;
32+
this.singularExpression = singularExpression;
33+
}
34+
35+
int expressionCount() {
36+
return expressionCount;
37+
}
38+
39+
Expression singularExpression() {
40+
return singularExpression;
41+
}
42+
2443
/**
2544
* Calls a consumer for every expression in the template.
2645
*
@@ -36,19 +55,24 @@ void forEachExpression(Consumer<Expression> consumer) {
3655
}
3756

3857
String resolve(int arraySize, Object[] strings) {
39-
if (arraySize != expressionCount()) {
58+
if (arraySize != expressionCount) {
4059
throw new RulesEvaluationError("Missing template parameters for a string template `"
4160
+ template + "`. Given: [" + Arrays.asList(strings) + ']');
4261
}
43-
StringBuilder result = new StringBuilder();
62+
63+
var result = STRING_BUILDER.get();
64+
result.setLength(0);
4465
int paramIndex = 0;
4566
for (var part : parts) {
46-
if (part instanceof Expression) {
47-
result.append(strings[paramIndex++]);
67+
if (part == null) {
68+
throw new RulesEvaluationError("Missing part of template " + template + " at part " + paramIndex);
69+
} else if (part.getClass() == String.class) {
70+
result.append((String) part); // we know parts are either strings or Expressions.
4871
} else {
49-
result.append(part);
72+
result.append(strings[paramIndex++]);
5073
}
5174
}
75+
5276
return result.toString();
5377
}
5478

@@ -62,7 +86,7 @@ static StringTemplate from(Template template) {
6286
expressionCount++;
6387
parts[i] = d.toExpression();
6488
} else {
65-
parts[i] = part;
89+
parts[i] = part.toString();
6690
}
6791
}
6892
var singularExpression = (expressionCount == 1 && parts.length == 1) ? (Expression) parts[0] : null;

‎client/client-rulesengine/src/test/java/software/amazon/smithy/java/client/rulesengine/RulesVmTest.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -248,12 +248,12 @@ public void createsMapForEndpointHeaders() {
248248
RulesProgram.VERSION,
249249
(byte) 0, // params
250250
(byte) 0, // registers
251-
RulesProgram.LOAD_CONST, // push map key "abc"
252-
1,
253251
RulesProgram.LOAD_CONST, // push list value 0, "def"
254252
2,
255253
RulesProgram.CREATE_LIST, // push list with one value, ["def"].
256254
1,
255+
RulesProgram.LOAD_CONST, // push map key "abc"
256+
1,
257257
RulesProgram.CREATE_MAP, // push with one KVP: {"abc": ["def"]} (the endpoint headers)
258258
1,
259259
RulesProgram.RESOLVE_TEMPLATE, // push resolved string template at constant 0 (2 byte constant)

0 commit comments

Comments
 (0)