Skip to content

Commit 60925f9

Browse files
gnidangnidan
committed
Fix partial reads for returndata, code, and transient storage
Copy-based reads (returndata, code) returned left-aligned values from MLOAD without right-aligning for partial reads (length < 32). Add shift+mask path matching calldata's partial read handling. Transient storage reads used TLOAD without offset/length handling. Add the same partial-slot extraction as regular storage reads.
1 parent f04ed54 commit 60925f9

File tree

2 files changed

+207
-13
lines changed

2 files changed

+207
-13
lines changed

‎packages/bugc/src/evmgen/generation/instructions/storage.test.ts‎

Lines changed: 92 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,7 @@ describe("generateRead", () => {
102102
});
103103

104104
describe("returndata reads", () => {
105-
it("should use RETURNDATACOPY + MLOAD", () => {
105+
it("should use RETURNDATACOPY + MLOAD for full read", () => {
106106
const mnemonics = mnemonicsFor([
107107
{
108108
kind: "read",
@@ -127,11 +127,40 @@ describe("generateRead", () => {
127127
expect(mnemonics).toContain("MLOAD");
128128
// Should zero scratch memory first
129129
expect(mnemonics).toContain("MSTORE");
130+
// Full read — no shift/mask needed
131+
expect(mnemonics).not.toContain("SHR");
132+
});
133+
134+
it("should shift+mask for partial returndata read", () => {
135+
const mnemonics = mnemonicsFor([
136+
{
137+
kind: "read",
138+
location: "returndata",
139+
offset: {
140+
kind: "const",
141+
value: 0n,
142+
type: Ir.Type.Scalar.uint256,
143+
},
144+
length: {
145+
kind: "const",
146+
value: 20n,
147+
type: Ir.Type.Scalar.uint256,
148+
},
149+
type: Ir.Type.Scalar.address,
150+
dest: "%1",
151+
operationDebug: {},
152+
},
153+
]);
154+
155+
expect(mnemonics).toContain("RETURNDATACOPY");
156+
expect(mnemonics).toContain("MLOAD");
157+
expect(mnemonics).toContain("SHR");
158+
expect(mnemonics).toContain("AND");
130159
});
131160
});
132161

133162
describe("code reads", () => {
134-
it("should use CODECOPY + MLOAD", () => {
163+
it("should use CODECOPY + MLOAD for full read", () => {
135164
const mnemonics = mnemonicsFor([
136165
{
137166
kind: "read",
@@ -154,11 +183,59 @@ describe("generateRead", () => {
154183

155184
expect(mnemonics).toContain("CODECOPY");
156185
expect(mnemonics).toContain("MLOAD");
186+
expect(mnemonics).not.toContain("SHR");
187+
});
188+
189+
it("should shift+mask for partial code read", () => {
190+
const mnemonics = mnemonicsFor([
191+
{
192+
kind: "read",
193+
location: "code",
194+
offset: {
195+
kind: "const",
196+
value: 0n,
197+
type: Ir.Type.Scalar.uint256,
198+
},
199+
length: {
200+
kind: "const",
201+
value: 4n,
202+
type: Ir.Type.Scalar.uint256,
203+
},
204+
type: Ir.Type.Scalar.uint256,
205+
dest: "%1",
206+
operationDebug: {},
207+
},
208+
]);
209+
210+
expect(mnemonics).toContain("CODECOPY");
211+
expect(mnemonics).toContain("MLOAD");
212+
expect(mnemonics).toContain("SHR");
213+
expect(mnemonics).toContain("AND");
157214
});
158215
});
159216

160217
describe("transient storage reads", () => {
161-
it("should use TLOAD", () => {
218+
it("should use TLOAD for full read", () => {
219+
const mnemonics = mnemonicsFor([
220+
{
221+
kind: "read",
222+
location: "transient",
223+
slot: {
224+
kind: "const",
225+
value: 0n,
226+
type: Ir.Type.Scalar.uint256,
227+
},
228+
type: Ir.Type.Scalar.uint256,
229+
dest: "%1",
230+
operationDebug: {},
231+
},
232+
]);
233+
234+
expect(mnemonics).toContain("TLOAD");
235+
expect(mnemonics).not.toContain("SHR");
236+
});
237+
238+
it("should shift+mask for partial transient read", () => {
162239
const mnemonics = mnemonicsFor([
163240
{
164241
kind: "read",
@@ -168,13 +245,25 @@ describe("generateRead", () => {
168245
value: 0n,
169246
type: Ir.Type.Scalar.uint256,
170247
},
248+
offset: {
249+
kind: "const",
250+
value: 0n,
251+
type: Ir.Type.Scalar.uint256,
252+
},
253+
length: {
254+
kind: "const",
255+
value: 1n,
256+
type: Ir.Type.Scalar.uint256,
257+
},
171258
type: Ir.Type.Scalar.uint256,
172259
dest: "%1",
173260
operationDebug: {},
174261
},
175262
]);
176263

177264
expect(mnemonics).toContain("TLOAD");
265+
expect(mnemonics).toContain("SHR");
266+
expect(mnemonics).toContain("AND");
178267
});
179268
});
180269
});

‎packages/bugc/src/evmgen/generation/instructions/storage.ts‎

Lines changed: 115 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -45,11 +45,7 @@ export function generateRead<S extends Stack>(
4545

4646
// Handle transient storage reads
4747
if (inst.location === "transient" && inst.slot) {
48-
return pipe<S>()
49-
.then(loadValue(inst.slot, { debug }), { as: "key" })
50-
.then(TLOAD({ debug }), { as: "value" })
51-
.then(storeValueIfNeeded(inst.dest, { debug }))
52-
.done();
48+
return generateTransientRead(inst, debug);
5349
}
5450

5551
// Handle memory reads
@@ -131,6 +127,57 @@ function generateStorageRead<S extends Stack>(
131127
);
132128
}
133129

130+
/**
131+
* Transient storage read: TLOAD with optional partial extraction.
132+
* Same shift+mask logic as regular storage reads.
133+
*/
134+
function generateTransientRead<S extends Stack>(
135+
inst: Ir.Instruction.Read,
136+
debug: Ir.Instruction.Debug,
137+
): Transition<S, readonly ["value", ...S]> {
138+
const offset = inst.offset?.kind === "const" ? inst.offset.value : 0n;
139+
const length = inst.length?.kind === "const" ? inst.length.value : 32n;
140+
141+
if (offset === 0n && length === 32n) {
142+
return pipe<S>()
143+
.then(loadValue(inst.slot!, { debug }), { as: "key" })
144+
.then(TLOAD({ debug }), { as: "value" })
145+
.then(storeValueIfNeeded(inst.dest, { debug }))
146+
.done();
147+
}
148+
149+
// Partial read - same shift+mask as storage
150+
return (
151+
pipe<S>()
152+
.then(loadValue(inst.slot!, { debug }), { as: "key" })
153+
.then(TLOAD({ debug }), { as: "value" })
154+
155+
.then(PUSHn((32n - BigInt(offset) - BigInt(length)) * 8n, { debug }), {
156+
as: "shift",
157+
})
158+
.then(SHR({ debug }), { as: "shiftedValue" })
159+
.then(PUSHn(1n, { debug }), { as: "b" })
160+
161+
// mask = (1 << (length * 8)) - 1
162+
.then(PUSHn(1n, { debug }), { as: "value" })
163+
.then(PUSHn(BigInt(length) * 8n, { debug }), {
164+
as: "shift",
165+
})
166+
.then(SHL({ debug }), { as: "a" })
167+
.then(SUB({ debug }), { as: "mask" })
168+
.then(
169+
rebrand<"mask", "a", "shiftedValue", "b">({
170+
1: "a",
171+
2: "b",
172+
}),
173+
)
174+
175+
.then(AND({ debug }), { as: "value" })
176+
.then(storeValueIfNeeded(inst.dest, { debug }))
177+
.done()
178+
);
179+
}
180+
134181
/**
135182
* Calldata read: CALLDATALOAD reads 32 bytes at a given offset.
136183
* For partial reads, shift+mask to extract the desired bytes.
@@ -194,7 +241,38 @@ function generateCopyBasedRead<S extends Stack>(
194241
): Transition<S, readonly ["value", ...S]> {
195242
const length = inst.length?.kind === "const" ? inst.length.value : 32n;
196243

197-
// Clear scratch memory first so partial copies are zero-padded
244+
if (length === 32n) {
245+
// Full 32-byte read — copy and load directly
246+
return (
247+
pipe<S>()
248+
// Zero out scratch: MSTORE(0x60, 0)
249+
.then(PUSHn(0n, { debug }), { as: "value" })
250+
.then(PUSHn(SCRATCH_OFFSET, { debug }), { as: "offset" })
251+
.then(MSTORE({ debug }))
252+
253+
// COPY(destOffset=0x60, offset, size=32)
254+
.then(PUSHn(32n, { debug }), { as: "size" })
255+
.then(loadValue(inst.offset!, { debug }), {
256+
as: "offset",
257+
})
258+
.then(PUSHn(SCRATCH_OFFSET, { debug }), {
259+
as: "destOffset",
260+
})
261+
.then(copyOp({ debug }))
262+
263+
// MLOAD from scratch
264+
.then(PUSHn(SCRATCH_OFFSET, { debug }), {
265+
as: "offset",
266+
})
267+
.then(MLOAD({ debug }), { as: "value" })
268+
.then(storeValueIfNeeded(inst.dest, { debug }))
269+
.done()
270+
);
271+
}
272+
273+
// Partial read: copy `length` bytes to scratch, MLOAD
274+
// returns left-aligned data, shift right to right-align,
275+
// then mask.
198276
return (
199277
pipe<S>()
200278
// Zero out scratch: MSTORE(0x60, 0)
@@ -204,13 +282,40 @@ function generateCopyBasedRead<S extends Stack>(
204282

205283
// COPY(destOffset=0x60, offset, size=length)
206284
.then(PUSHn(BigInt(length), { debug }), { as: "size" })
207-
.then(loadValue(inst.offset!, { debug }), { as: "offset" })
208-
.then(PUSHn(SCRATCH_OFFSET, { debug }), { as: "destOffset" })
285+
.then(loadValue(inst.offset!, { debug }), {
286+
as: "offset",
287+
})
288+
.then(PUSHn(SCRATCH_OFFSET, { debug }), {
289+
as: "destOffset",
290+
})
209291
.then(copyOp({ debug }))
210292

211-
// MLOAD from scratch
212-
.then(PUSHn(SCRATCH_OFFSET, { debug }), { as: "offset" })
293+
// MLOAD from scratch — value is left-aligned
294+
.then(PUSHn(SCRATCH_OFFSET, { debug }), {
295+
as: "offset",
296+
})
213297
.then(MLOAD({ debug }), { as: "value" })
298+
299+
// Shift right to right-align
300+
.then(PUSHn((32n - BigInt(length)) * 8n, { debug }), { as: "shift" })
301+
.then(SHR({ debug }), { as: "shiftedValue" })
302+
.then(PUSHn(1n, { debug }), { as: "b" })
303+
304+
// mask = (1 << (length * 8)) - 1
305+
.then(PUSHn(1n, { debug }), { as: "value" })
306+
.then(PUSHn(BigInt(length) * 8n, { debug }), {
307+
as: "shift",
308+
})
309+
.then(SHL({ debug }), { as: "a" })
310+
.then(SUB({ debug }), { as: "mask" })
311+
.then(
312+
rebrand<"mask", "a", "shiftedValue", "b">({
313+
1: "a",
314+
2: "b",
315+
}),
316+
)
317+
318+
.then(AND({ debug }), { as: "value" })
214319
.then(storeValueIfNeeded(inst.dest, { debug }))
215320
.done()
216321
);

0 commit comments

Comments
 (0)