Skip to content

feat: support in memory trace generation (cherry pick) #1341

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Nov 21, 2025
Merged

feat: support in memory trace generation (cherry pick) #1341

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
73 changes: 18 additions & 55 deletions pkg/cmd/generate/class.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,14 +130,6 @@ func generateClassContents[F any](className string, super string, mod corset.Sou
}
}
//
if mod.Name == "" {
ninputs := getMaxRegisterIndex(schema)
// Write out constructor function.
constructor := strings.ReplaceAll(javaTraceOpen, "{class}", className)
constructor = strings.ReplaceAll(constructor, "{ninputs}", fmt.Sprintf("%d", ninputs))
builder.WriteIndentedString(constructor)
}
//
generateJavaClassFooter(builder)
}

Expand Down Expand Up @@ -172,11 +164,11 @@ func generateJavaModuleHeaders[F any](mod corset.SourceModule, schema sc.AnySche
reg := schema.Register(col.Register)
// Check whether this is part of our module
if reg.IsInputOutput() {
byteWidth := fmt.Sprintf("%d", byteWidth(reg.Width))
bitWidth := fmt.Sprintf("%d", reg.Width)
name := fmt.Sprintf("%s.%s", mod.Name, reg.Name)
regStr := fmt.Sprintf("%d", col.Register.Index(schema.Width()))
i1Builder.WriteIndentedString(
"headers.add(new ColumnHeader(\"", name, "\",", regStr, ",", byteWidth, ",length));\n")
"headers.add(new ColumnHeader(\"", name, "\",", regStr, ",", bitWidth, ",length));\n")
//
count++
}
Expand Down Expand Up @@ -279,7 +271,7 @@ func generateJavaModuleRegisterFields[F any](mod corset.SourceModule, schema sc.
// Determine suitable name for field
fieldName := toRegisterName(col.Register, reg.Name)
//
builder.WriteIndentedString("private MappedByteBuffer ", fieldName, ";\n")
builder.WriteIndentedString("private Column ", fieldName, ";\n")
// increase count
count++
}
Expand Down Expand Up @@ -337,7 +329,7 @@ func generateJavaModuleMetadata(metadata typed.Map, builder indentBuilder) {
builder.WriteIndentedString("}\n\n")
}

builder.WriteIndentedString("public void addMetadata(String key, Object value) { metadata.put(key,value); }\n")
builder.WriteIndentedString("public Map<String,Object> getMetaData() { return metadata; }\n")
}

func generateJavaModuleConstructor(classname string, mod corset.SourceModule, builder indentBuilder) {
Expand Down Expand Up @@ -365,7 +357,7 @@ func generateJavaModuleOpen[F any](mod corset.SourceModule, schema sc.AnySchema[
//
innerBuilder := builder.Indent()
//
builder.WriteIndentedString("private void open(MappedByteBuffer[] registers) {\n")
builder.WriteIndentedString("public void open(Column[] registers) {\n")
innerBuilder.WriteIndentedString("// initialise register(s)\n")
// Write register initialisers
for _, col := range mod.Registers(schema.Width()) {
Expand Down Expand Up @@ -443,13 +435,11 @@ func generateJavaModuleColumnSetter[F any](className string, methodName string,
//
switch {
case bitwidth == 1:
i1Builder.WriteIndentedString(fieldName, ".put((byte) (val ? 1 : 0));\n")
case bitwidth <= 8:
i1Builder.WriteIndentedString(fieldName, ".put((byte) val);\n")
i1Builder.WriteIndentedString(fieldName, ".write(val);\n")
case bitwidth <= 63:
generateJavaModuleLongPutter(col.Name, fieldName, bitwidth, i1Builder)
i1Builder.WriteIndentedString(fieldName, ".write(val);\n")
default:
generateJavaModuleBytesPutter(col.Name, fieldName, bitwidth, i1Builder)
i1Builder.WriteIndentedString(fieldName, ".write(val.trimLeadingZeros().toArray());\n")
}
//
i1Builder.WriteIndentedString("\n")
Expand All @@ -470,41 +460,6 @@ func generateJavaModuleLegacyColumnSetter(className string, methodName string, b
builder.WriteIndentedString("}\n\n")
}

func generateJavaModuleLongPutter(columnName, fieldName string, bitwidth uint, builder indentBuilder) {
n := byteWidth(bitwidth)
i1Builder := builder.Indent()
builder.WriteIndentedString("if(val < 0 || val >= ", maxValueStr(bitwidth), "L) {\n")
i1Builder.WriteIndentedString(
"throw new IllegalArgumentException(\"", columnName+" has invalid value (\" + val + \")\");\n")
builder.WriteIndentedString("}\n")
//
for i := int(n) - 1; i >= 0; i-- {
shift := (i * 8)
if shift == 0 {
builder.WriteIndentedString(fieldName, ".put((byte) val);\n")
} else {
builder.WriteIndentedString(fieldName, ".put((byte) (val >> ", fmt.Sprintf("%d", shift), "));\n")
}
}
}

func generateJavaModuleBytesPutter(columnName, fieldName string, bitwidth uint, builder indentBuilder) {
i1Builder := builder.Indent()
n := byteWidth(bitwidth)
//
builder.WriteIndentedString("// Trim array to size\n")
builder.WriteIndentedString("Bytes bs = val.trimLeadingZeros();\n")
builder.WriteIndentedString("// Sanity check against expected width\n")
builder.WriteIndentedString(fmt.Sprintf("if(bs.bitLength() > %d) {\n", bitwidth))
i1Builder.WriteIndentedString(
fmt.Sprintf("throw new IllegalArgumentException(\"%s has invalid width (\"+bs.bitLength()+\"bits)\");\n", columnName))
builder.WriteIndentedString("}\n")
builder.WriteIndentedString("// Write padding (if necessary)\n")
builder.WriteIndentedString(fmt.Sprintf("for(int i=bs.size(); i<%d; i++) { %s.put((byte) 0); }\n", n, fieldName))
builder.WriteIndentedString("// Write bytes\n")
builder.WriteIndentedString(fmt.Sprintf("for(int i=0; i<bs.size(); i++) { %s.put(bs.get(i)); }\n", fieldName))
}

func generateJavaModuleValidateRow[F any](className string, mod corset.SourceModule, schema sc.AnySchema[F],
builder indentBuilder) {
//
Expand Down Expand Up @@ -546,10 +501,18 @@ func generateJavaModuleFillAndValidateRow[F any](className string, mod corset.So
if reg.IsInputOutput() {
name := toRegisterName(col.Register, reg.Name)
regstr := fmt.Sprintf("%d", col.Register.Index(schema.Width()))
byteWidth := fmt.Sprintf("%d", byteWidth(reg.Width))
// Yes, include register
i1Builder.WriteIndentedString("if(!filled.get(", regstr, ")) {\n")
i2Builder.WriteIndentedString(name, ".position(", name, ".position() + ", byteWidth, ");\n")
//
switch {
case reg.Width == 1:
i2Builder.WriteIndentedString(name, ".write(false);\n")
case reg.Width <= 63:
i2Builder.WriteIndentedString(name, ".write(0L);\n")
default:
i2Builder.WriteIndentedString(name, ".write(new byte[0]);\n")
}
//
i1Builder.WriteIndentedString("}\n")
}
}
Expand Down
1 change: 1 addition & 0 deletions pkg/cmd/generate/interface.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,7 @@ func generateInterfaceContents(className string, mod corset.SourceModule, builde
//
if mod.Name == "" {
builder.WriteString(javaColumnHeader)
builder.WriteString(javaColumn)
builder.WriteString(javaAddMetadataSignature)
builder.WriteString(javaOpenSignature)
}
Expand Down
205 changes: 18 additions & 187 deletions pkg/cmd/generate/snippets.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,205 +62,36 @@ const javaColumnHeader string = `
* ColumnHeader contains information about a given column in the resulting trace file.
*
* @param name Name of the column, as found in the trace file.
* @param bytesPerElement Bytes required for each element in the column.
* @param bitwidth Max number of bits required for any element in the column.
*/
public record ColumnHeader(String name, int register, long bytesPerElement, long length) { }
public record ColumnHeader(String name, int register, int bitwidth, int length) { }
`

// nolint
const javaAddMetadataSignature string = `
/**
* Add an item of metadata to this trace.
*/
public void addMetadata(String key, Object value);
`

// nolint
const javaOpenSignature string = `
/**
* Construct a new trace which will be written to a given file.
*
* @param file File into which the trace will be written. Observe any previous contents of this file will be lost.
* @return Trace object to use for writing column data.
*
* @throws IOException If an I/O error occurs.
*/
public void open(RandomAccessFile file, List<ColumnHeader> rawHeaders) throws IOException;
`

// nolint
const javaTraceOpen string = `
const javaColumn string = `
/**
* Construct a new trace which will be written to a given file.
*
* @param file File into which the trace will be written. Observe any previous contents of this file will be lost.
* @return Trace object to use for writing column data.
* Column provides an interface for writing column data into the resulting trace file.
*
* @throws IOException If an I/O error occurs.
*/
public void open(RandomAccessFile file, List<ColumnHeader> rawHeaders) throws IOException {
// Convert metadata into JSON bytes
byte[] metadataBytes = getMetadataBytes(metadata);
// Construct trace file header bytes
byte[] header = constructTraceFileHeader(metadataBytes);
// Align headers according to register indices.
ColumnHeader[] columnHeaders = alignHeaders(rawHeaders);
// Determine file size
long headerSize = determineColumnHeadersSize(columnHeaders) + header.length;
long dataSize = determineColumnDataSize(columnHeaders);
file.setLength(headerSize + dataSize);
// Write headers
writeHeaders(file,header,columnHeaders,headerSize);
// Initialise buffers
MappedByteBuffer[] buffers = initialiseByteBuffers(file,columnHeaders,headerSize);
// Done
this.open(buffers);
}

/**
* Construct trace file header containing the given metadata bytes.
*
* @param metadata Metadata bytes to be embedded in the trace file.
*
* @return bytes making up the header.
*/
private static byte[] constructTraceFileHeader(byte[] metadata) {
ByteBuffer buffer = ByteBuffer.allocate(16 + metadata.length);
// File identifier
buffer.put(new byte[]{'z','k','t','r','a','c','e','r'});
// Major version
buffer.putShort((short) 1);
// Minor version
buffer.putShort((short) 0);
// Metadata length
buffer.putInt(metadata.length);
// Metadata
buffer.put(metadata);
// Done
return buffer.array();
public interface Column {
public void write(boolean value);
public void write(long value);
public void write(byte[] value);
}
`

// nolint
const javaAddMetadataSignature string = `
/**
* Align headers ensures that the order in which columns are seen matches the order found in the trace schema.
*
* @param headers The headers to be aligned.
* @return The aligned headers.
* Get static metadata stored within this trace during compilation.
*/
private static ColumnHeader[] alignHeaders(List<ColumnHeader> headers) {
ColumnHeader[] alignedHeaders = new ColumnHeader[{ninputs}];
//
for(ColumnHeader header : headers) {
alignedHeaders[header.register()] = header;
}
//
return alignedHeaders;
}

/**
* Precompute the size of the trace file in order to memory map the buffers.
*
* @param headers Set of headers for the columns being written.
* @return Number of bytes requires for the trace file header.
*/
private static long determineColumnHeadersSize(ColumnHeader[] headers) {
long nBytes = 4; // column count

for (ColumnHeader header : headers) {
if(header != null) {
nBytes += 2; // name length
nBytes += header.name().length();
nBytes += 1; // byte per element
nBytes += 4; // element count
}
}

return nBytes;
}

/**
* Precompute the size of the trace file in order to memory map the buffers.
*
* @param headers Set of headers for the columns being written.
* @return Number of bytes required for storing all column data, excluding the header.
*/
private static long determineColumnDataSize(ColumnHeader[] headers) {
long nBytes = 0;

for (ColumnHeader header : headers) {
if(header != null) {
nBytes += header.length() * header.bytesPerElement();
}
}

return nBytes;
}

/**
* Write header information for the trace file.
*
* @param file Trace file being written.
* @param header Trace file header
* @param headers Column headers.
* @param size Overall size of the header.
*/
private static void writeHeaders(RandomAccessFile file, byte[] header, ColumnHeader[] headers, long size) throws IOException {
final var buffer = file.getChannel().map(FileChannel.MapMode.READ_WRITE, 0, size);
// Write trace file header
buffer.put(header);
// Write column count as uint32
buffer.putInt(countHeaders(headers));
// Write column headers one-by-one
for(ColumnHeader h : headers) {
if(h != null) {
buffer.putShort((short) h.name().length());
buffer.put(h.name().getBytes());
buffer.put((byte) h.bytesPerElement());
buffer.putInt((int) h.length());
}
}
}

/**
* Initialise one memory mapped byte buffer for each column to be written in the trace.
* @param headers Set of headers for the columns being written.
* @param headerSize Space required at start of trace file for header.
* @return Buffer array with one entry per header.
*/
private static MappedByteBuffer[] initialiseByteBuffers(RandomAccessFile file, ColumnHeader[] headers,
long headerSize) throws IOException {
MappedByteBuffer[] buffers = new MappedByteBuffer[{ninputs}];
long offset = headerSize;
for(int i=0;i<headers.length;i++) {
if(headers[i] != null) {
// Determine size (in bytes) required to store all elements of this column.
long length = headers[i].length() * headers[i].bytesPerElement();
// Preallocate space for this column.
buffers[i] = file.getChannel().map(FileChannel.MapMode.READ_WRITE, offset, length);
//
offset += length;
}
}
return buffers;
}

/**
* Counter number of active (i.e. non-null) headers. A header can be null if
* it represents a column in a module which is not activated for this trace.
*/
private static int countHeaders(ColumnHeader[] headers) throws IOException {
int count = 0;
for(ColumnHeader h : headers) {
if(h != null) { count++; }
}
return count;
}
public Map<String,Object> getMetaData();
`

// nolint
const javaOpenSignature string = `
/**
* Object writer is used for generating JSON byte strings.
* Open this trace file for the given set of columns.
*/
private static final ObjectWriter objectWriter = new ObjectMapper().writer();

public static byte[] getMetadataBytes(Map<String, Object> metadata) throws IOException {
return objectWriter.writeValueAsBytes(metadata);
}
public void open(Column[] columns);
`
Loading
Loading