Skip to content

Commit bf5b4f4

Browse files
mtdowlingmtdowling
committed
Fix bytecode loading and creation
Bytecode was meant to hold the entire program not just instructions. This commit ensures it does that, and correctly.
1 parent 9c12fe7 commit bf5b4f4

File tree

9 files changed

+174
-123
lines changed

9 files changed

+174
-123
lines changed

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

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55

66
package software.amazon.smithy.java.client.rulesengine;
77

8+
import java.nio.ByteBuffer;
89
import java.util.ArrayList;
910
import java.util.Arrays;
1011
import java.util.HashMap;
@@ -54,7 +55,7 @@
5455
* <li><b>Register Definitions</b> - Array of parameter/register metadata (immediately after result table)</li>
5556
* <li><b>Function Table</b> - Array of function names</li>
5657
* <li><b>BDD Table</b> - Array of BDD nodes (3 ints per node)</li>
57-
* <li><b>Bytecode Section</b> - Compiled instructions for conditions and results</li>
58+
* <li><b>Instruction Section</b> - Compiled instructions for conditions and results</li>
5859
* <li><b>Constant Pool</b> - All constants referenced by the bytecode</li>
5960
* </ol>
6061
*
@@ -112,7 +113,7 @@
112113
* <li>100_000_000+: Result terminals (100_000_000 + resultIndex)</li>
113114
* </ul>
114115
*
115-
* <h3>Bytecode Section</h3>
116+
* <h3>Instruction Section</h3>
116117
* <p>Contains the compiled bytecode instructions for all conditions and results.
117118
* The condition/result tables contain absolute offsets from the start of the file that point into this section.
118119
* Instructions use a stack-based virtual machine with opcodes defined in {@link Opcodes}. This section may include
@@ -316,14 +317,29 @@ public RulesFunction[] getFunctions() {
316317
}
317318

318319
/**
319-
* Get the raw bytecode.
320+
* Get the entire raw bytecode, including headers and the instructions.
320321
*
321322
* @return the bytecode.
322323
*/
323324
public byte[] getBytecode() {
324325
return bytecode;
325326
}
326327

328+
/**
329+
* Gets a read-only view of the instruction section from the bytecode.
330+
*
331+
* <p>This section contains the compiled opcodes for all conditions and results.
332+
*
333+
* @return a ByteBuffer view of the instruction bytes
334+
*/
335+
public ByteBuffer getInstructions() {
336+
BytecodeReader reader = new BytecodeReader(bytecode, 0);
337+
reader.skipToBytecodeSection();
338+
int start = reader.offset;
339+
int end = reader.getBytecodeEndPosition();
340+
return ByteBuffer.wrap(bytecode, start, end - start).asReadOnlyBuffer();
341+
}
342+
327343
/**
328344
* Get the register definitions for the bytecode (both input parameters and temp registers).
329345
*

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

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -331,7 +331,6 @@ private void appendSymbolicInfo(StringBuilder s, BytecodeWalker walker, Show sho
331331
}
332332
}
333333
case JUMP_OFFSET -> {
334-
int offset = walker.getOperand(0);
335334
s.append("-> ").append(walker.getJumpTarget());
336335
}
337336
}

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

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,24 @@ private Object readConstant(int depth) {
9999
};
100100
}
101101

102+
void skipToBytecodeSection() {
103+
int bddNodeCount = readIntAtSpecificPosition(16);
104+
int bddTableOffset = readIntAtSpecificPosition(40);
105+
this.offset = bddTableOffset + (bddNodeCount * 12);
106+
}
107+
108+
private int readIntAtSpecificPosition(int position) {
109+
int savedOffset = this.offset;
110+
this.offset = position;
111+
int value = readInt();
112+
this.offset = savedOffset;
113+
return value;
114+
}
115+
116+
int getBytecodeEndPosition() {
117+
return readIntAtSpecificPosition(36); // constantPoolOffset
118+
}
119+
102120
RegisterDefinition[] readRegisterDefinitions(int count) {
103121
RegisterDefinition[] registers = new RegisterDefinition[count];
104122

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

Lines changed: 25 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -5,25 +5,35 @@
55

66
package software.amazon.smithy.java.client.rulesengine;
77

8+
import java.nio.ByteBuffer;
9+
810
/**
911
* Utility for walking through bytecode instructions, understanding instruction boundaries
1012
* and operand sizes. This prevents misinterpreting operands as opcodes.
1113
*/
1214
final class BytecodeWalker {
13-
private final byte[] code;
15+
private final ByteBuffer code;
1416
private int pc;
1517

1618
public BytecodeWalker(byte[] code) {
17-
this(code, 0);
19+
this(ByteBuffer.wrap(code));
1820
}
1921

2022
public BytecodeWalker(byte[] code, int startOffset) {
21-
this.code = code;
23+
this(ByteBuffer.wrap(code), startOffset);
24+
}
25+
26+
public BytecodeWalker(ByteBuffer code) {
27+
this(code, 0);
28+
}
29+
30+
public BytecodeWalker(ByteBuffer code, int startOffset) {
31+
this.code = code.duplicate();
2232
this.pc = startOffset;
2333
}
2434

2535
public boolean hasNext() {
26-
return pc < code.length;
36+
return pc < code.limit();
2737
}
2838

2939
public int getPosition() {
@@ -34,7 +44,7 @@ public byte currentOpcode() {
3444
if (!hasNext()) {
3545
throw new IllegalStateException("No more instructions");
3646
}
37-
return code[pc];
47+
return code.get(pc);
3848
}
3949

4050
public boolean advance() {
@@ -53,7 +63,7 @@ public int getInstructionLength() {
5363
if (!hasNext()) {
5464
return -1;
5565
}
56-
return getInstructionLength(code[pc]);
66+
return getInstructionLength(code.get(pc));
5767
}
5868

5969
public int getOperandCount() {
@@ -98,7 +108,7 @@ public int getOperand(int index) {
98108
case Opcodes.TEST_REGISTER_IS_FALSE:
99109
case Opcodes.RETURN_ENDPOINT:
100110
if (index == 0) {
101-
return code[pc + 1] & 0xFF;
111+
return code.get(pc + 1) & 0xFF;
102112
}
103113
break;
104114

@@ -107,38 +117,38 @@ public int getOperand(int index) {
107117
case Opcodes.GET_PROPERTY:
108118
case Opcodes.JNN_OR_POP:
109119
if (index == 0) {
110-
return ((code[pc + 1] & 0xFF) << 8) | (code[pc + 2] & 0xFF);
120+
return ((code.get(pc + 1) & 0xFF) << 8) | (code.get(pc + 2) & 0xFF);
111121
}
112122
break;
113123

114124
// Mixed operand instructions
115125
case Opcodes.GET_PROPERTY_REG:
116126
if (index == 0) {
117-
return code[pc + 1] & 0xFF; // register
127+
return code.get(pc + 1) & 0xFF; // register
118128
} else if (index == 1) {
119-
return ((code[pc + 2] & 0xFF) << 8) | (code[pc + 3] & 0xFF); // property index
129+
return ((code.get(pc + 2) & 0xFF) << 8) | (code.get(pc + 3) & 0xFF); // property index
120130
}
121131
break;
122132

123133
case Opcodes.GET_INDEX_REG:
124134
if (index == 0) {
125-
return code[pc + 1] & 0xFF; // register
135+
return code.get(pc + 1) & 0xFF; // register
126136
} else if (index == 1) {
127-
return code[pc + 2] & 0xFF; // index
137+
return code.get(pc + 2) & 0xFF; // index
128138
}
129139
break;
130140

131141
case Opcodes.RESOLVE_TEMPLATE:
132142
if (index == 0) {
133-
return code[pc + 1] & 0xFF; // arg count
143+
return code.get(pc + 1) & 0xFF; // arg count
134144
} else if (index == 1) {
135-
return ((code[pc + 2] & 0xFF) << 8) | (code[pc + 3] & 0xFF); // template index
145+
return ((code.get(pc + 2) & 0xFF) << 8) | (code.get(pc + 3) & 0xFF); // template index
136146
}
137147
break;
138148

139149
case Opcodes.SUBSTRING:
140150
if (index >= 0 && index < 3) {
141-
return code[pc + 1 + index] & 0xFF;
151+
return code.get(pc + 1 + index) & 0xFF;
142152
}
143153
break;
144154
}

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

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -139,16 +139,17 @@ Bytecode build(
139139
writeBddTable(dos, bddNodes);
140140

141141
// Apply jump patches to get the final bytecode
142-
byte[] originalBytecode = bytecodeStream.toByteArray();
143-
byte[] bytecode = originalBytecode;
144-
if (!jumpPatches.isEmpty()) {
142+
if (jumpPatches.isEmpty()) {
143+
// If no jumps, the just copy directly
144+
bytecodeStream.writeTo(dos);
145+
} else {
146+
byte[] originalBytecode = bytecodeStream.toByteArray();
145147
ByteArrayOutputStream patchedStream = new ByteArrayOutputStream();
146148
DataOutputStream patchedDos = new DataOutputStream(patchedStream);
147149
writePatchedBytecode(patchedDos, originalBytecode);
148150
patchedDos.flush();
149-
bytecode = patchedStream.toByteArray();
151+
dos.write(patchedStream.toByteArray());
150152
}
151-
dos.write(bytecode);
152153

153154
int constantPoolOffset = complete.size();
154155
writeConstantPool(dos);
@@ -162,9 +163,14 @@ Bytecode build(
162163
patchInt(data, 40, bddTableOffset);
163164

164165
return new Bytecode(
165-
bytecode,
166-
conditionOffsets.stream().mapToInt(Integer::intValue).toArray(),
167-
resultOffsets.stream().mapToInt(Integer::intValue).toArray(),
166+
data,
167+
// Convert to absolute offsets for the Bytecode object
168+
conditionOffsets.stream()
169+
.mapToInt(offset -> bytecodeOffset + offset)
170+
.toArray(),
171+
resultOffsets.stream()
172+
.mapToInt(offset -> bytecodeOffset + offset)
173+
.toArray(),
168174
registerDefinitions,
169175
constants.toArray(),
170176
functions,

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

Lines changed: 4 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -200,11 +200,11 @@ public Bytecode load(byte[] data) {
200200
}
201201

202202
// Load function names and resolve them using this builder's functions
203+
// This advances reader.offset to the end of the function table
203204
reader.offset = functionTableOffset;
204205
RulesFunction[] resolvedFunctions = loadFunctions(reader, functionCount);
205206

206-
// Load register definitions (after result table)
207-
reader.offset = resultTableOffset + (resultCount * 4);
207+
// Load register definitions
208208
RegisterDefinition[] registers = reader.readRegisterDefinitions(registerCount);
209209

210210
// Load BDD nodes as flat array
@@ -217,36 +217,10 @@ public Bytecode load(byte[] data) {
217217
bddNodes[baseIdx + 2] = reader.readInt(); // low
218218
}
219219

220-
// Find bytecode start and length
221-
int bytecodeStart = bddTableOffset + (bddNodeCount * 12);
222-
int bytecodeLength = constantPoolOffset - bytecodeStart;
223-
224-
if (bytecodeLength < 0) {
225-
throw new IllegalArgumentException("Invalid bytecode section length");
226-
}
227-
228-
// Extract bytecode section (with relative offsets)
229-
byte[] bytecode = new byte[bytecodeLength];
230-
System.arraycopy(data, bytecodeStart, bytecode, 0, bytecodeLength);
231-
232-
// Adjust offsets to be relative to bytecode start
233-
for (int i = 0; i < conditionCount; i++) {
234-
conditionOffsets[i] -= bytecodeStart;
235-
if (conditionOffsets[i] < 0 || conditionOffsets[i] >= bytecodeLength) {
236-
throw new IllegalArgumentException("Invalid condition offset at index " + i);
237-
}
238-
}
239-
240-
for (int i = 0; i < resultCount; i++) {
241-
resultOffsets[i] -= bytecodeStart;
242-
if (resultOffsets[i] < 0 || resultOffsets[i] >= bytecodeLength) {
243-
throw new IllegalArgumentException("Invalid result offset at index " + i);
244-
}
245-
}
246-
247220
Object[] constantPool = loadConstantPool(data, constantPoolOffset, constantCount);
221+
248222
return new Bytecode(
249-
bytecode,
223+
data,
250224
conditionOffsets,
251225
resultOffsets,
252226
registers,

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

Lines changed: 33 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import static org.junit.jupiter.api.Assertions.assertNotNull;
1010
import static org.junit.jupiter.api.Assertions.assertTrue;
1111

12+
import java.nio.ByteBuffer;
1213
import java.util.ArrayList;
1314
import java.util.HashMap;
1415
import java.util.List;
@@ -507,7 +508,7 @@ void testCompileLargeList() {
507508
Bytecode bytecode = compiler.compile();
508509

509510
OpcodeWithValue result = findOpcodeWithValue(bytecode, Opcodes.LISTN);
510-
assertTrue(result.found());
511+
assertTrue(result.found(), "LISTN opcode not found in bytecode");
511512
assertEquals(5, result.value());
512513
}
513514

@@ -527,18 +528,17 @@ void testCompileLargeMap() {
527528
assertOpcodePresent(bytecode, Opcodes.MAPN);
528529
}
529530

530-
private void assertOpcodePresent(Bytecode bytecode, byte expectedOpcode) {
531-
BytecodeWalker walker = new BytecodeWalker(bytecode.getBytecode());
532-
boolean found = false;
531+
private void assertOpcodePresent(Bytecode bytecode, int expectedOpcode) {
532+
ByteBuffer instructions = bytecode.getInstructions();
533+
byte[] instructionBytes = new byte[instructions.remaining()];
534+
instructions.get(instructionBytes);
533535

534-
while (walker.hasNext()) {
535-
if (walker.currentOpcode() == expectedOpcode) {
536+
boolean found = false;
537+
for (byte b : instructionBytes) {
538+
if ((b & 0xFF) == expectedOpcode) {
536539
found = true;
537540
break;
538541
}
539-
if (!walker.advance()) {
540-
break; // Unknown opcode encountered
541-
}
542542
}
543543

544544
assertTrue(found, "Expected opcode " + expectedOpcode + " to be present in bytecode");
@@ -556,22 +556,42 @@ private void assertConstantPresent(Bytecode bytecode, Object expectedConstant) {
556556
}
557557

558558
private OpcodeWithValue findOpcodeWithValue(Bytecode bytecode, byte expectedOpcode) {
559-
BytecodeWalker walker = new BytecodeWalker(bytecode.getBytecode());
559+
// Search conditions
560+
for (int i = 0; i < bytecode.getConditionCount(); i++) {
561+
OpcodeWithValue result = findInSection(bytecode, bytecode.getConditionStartOffset(i), expectedOpcode);
562+
if (result.found()) {
563+
return result;
564+
}
565+
}
566+
567+
// Search results
568+
for (int i = 0; i < bytecode.getResultCount(); i++) {
569+
OpcodeWithValue result = findInSection(bytecode, bytecode.getResultOffset(i), expectedOpcode);
570+
if (result.found()) {
571+
return result;
572+
}
573+
}
574+
575+
return new OpcodeWithValue(false, 0);
576+
}
577+
578+
private OpcodeWithValue findInSection(Bytecode bytecode, int offset, byte expectedOpcode) {
579+
BytecodeWalker walker = new BytecodeWalker(bytecode.getInstructions(), offset);
560580

561581
while (walker.hasNext()) {
562582
if (walker.currentOpcode() == expectedOpcode) {
563583
try {
564-
int value = walker.getOperand(0);
565-
return new OpcodeWithValue(true, value);
584+
return new OpcodeWithValue(true, walker.getOperand(0));
566585
} catch (IllegalArgumentException e) {
567586
// Opcode doesn't have operands
568587
return new OpcodeWithValue(true, 0);
569588
}
570589
}
571-
if (!walker.advance()) {
590+
if (walker.isReturnOpcode() || !walker.advance()) {
572591
break;
573592
}
574593
}
594+
575595
return new OpcodeWithValue(false, 0);
576596
}
577597

0 commit comments

Comments
 (0)