Skip to content

Commit 7915900

Browse files
authored
feat: support in memory trace generation (cherry pick) (#1341)
separate format implementation This separates the mechanism for representing columns from trace generation itself, such that the representation can support different implementations and (critically) writing can be performed in memory.
1 parent ba8d346 commit 7915900

File tree

4 files changed

+37
-276
lines changed

4 files changed

+37
-276
lines changed

pkg/cmd/generate/class.go

Lines changed: 18 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -130,14 +130,6 @@ func generateClassContents[F any](className string, super string, mod corset.Sou
130130
}
131131
}
132132
//
133-
if mod.Name == "" {
134-
ninputs := getMaxRegisterIndex(schema)
135-
// Write out constructor function.
136-
constructor := strings.ReplaceAll(javaTraceOpen, "{class}", className)
137-
constructor = strings.ReplaceAll(constructor, "{ninputs}", fmt.Sprintf("%d", ninputs))
138-
builder.WriteIndentedString(constructor)
139-
}
140-
//
141133
generateJavaClassFooter(builder)
142134
}
143135

@@ -172,11 +164,11 @@ func generateJavaModuleHeaders[F any](mod corset.SourceModule, schema sc.AnySche
172164
reg := schema.Register(col.Register)
173165
// Check whether this is part of our module
174166
if reg.IsInputOutput() {
175-
byteWidth := fmt.Sprintf("%d", byteWidth(reg.Width))
167+
bitWidth := fmt.Sprintf("%d", reg.Width)
176168
name := fmt.Sprintf("%s.%s", mod.Name, reg.Name)
177169
regStr := fmt.Sprintf("%d", col.Register.Index(schema.Width()))
178170
i1Builder.WriteIndentedString(
179-
"headers.add(new ColumnHeader(\"", name, "\",", regStr, ",", byteWidth, ",length));\n")
171+
"headers.add(new ColumnHeader(\"", name, "\",", regStr, ",", bitWidth, ",length));\n")
180172
//
181173
count++
182174
}
@@ -279,7 +271,7 @@ func generateJavaModuleRegisterFields[F any](mod corset.SourceModule, schema sc.
279271
// Determine suitable name for field
280272
fieldName := toRegisterName(col.Register, reg.Name)
281273
//
282-
builder.WriteIndentedString("private MappedByteBuffer ", fieldName, ";\n")
274+
builder.WriteIndentedString("private Column ", fieldName, ";\n")
283275
// increase count
284276
count++
285277
}
@@ -337,7 +329,7 @@ func generateJavaModuleMetadata(metadata typed.Map, builder indentBuilder) {
337329
builder.WriteIndentedString("}\n\n")
338330
}
339331

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

343335
func generateJavaModuleConstructor(classname string, mod corset.SourceModule, builder indentBuilder) {
@@ -365,7 +357,7 @@ func generateJavaModuleOpen[F any](mod corset.SourceModule, schema sc.AnySchema[
365357
//
366358
innerBuilder := builder.Indent()
367359
//
368-
builder.WriteIndentedString("private void open(MappedByteBuffer[] registers) {\n")
360+
builder.WriteIndentedString("public void open(Column[] registers) {\n")
369361
innerBuilder.WriteIndentedString("// initialise register(s)\n")
370362
// Write register initialisers
371363
for _, col := range mod.Registers(schema.Width()) {
@@ -443,13 +435,11 @@ func generateJavaModuleColumnSetter[F any](className string, methodName string,
443435
//
444436
switch {
445437
case bitwidth == 1:
446-
i1Builder.WriteIndentedString(fieldName, ".put((byte) (val ? 1 : 0));\n")
447-
case bitwidth <= 8:
448-
i1Builder.WriteIndentedString(fieldName, ".put((byte) val);\n")
438+
i1Builder.WriteIndentedString(fieldName, ".write(val);\n")
449439
case bitwidth <= 63:
450-
generateJavaModuleLongPutter(col.Name, fieldName, bitwidth, i1Builder)
440+
i1Builder.WriteIndentedString(fieldName, ".write(val);\n")
451441
default:
452-
generateJavaModuleBytesPutter(col.Name, fieldName, bitwidth, i1Builder)
442+
i1Builder.WriteIndentedString(fieldName, ".write(val.trimLeadingZeros().toArray());\n")
453443
}
454444
//
455445
i1Builder.WriteIndentedString("\n")
@@ -470,41 +460,6 @@ func generateJavaModuleLegacyColumnSetter(className string, methodName string, b
470460
builder.WriteIndentedString("}\n\n")
471461
}
472462

473-
func generateJavaModuleLongPutter(columnName, fieldName string, bitwidth uint, builder indentBuilder) {
474-
n := byteWidth(bitwidth)
475-
i1Builder := builder.Indent()
476-
builder.WriteIndentedString("if(val < 0 || val >= ", maxValueStr(bitwidth), "L) {\n")
477-
i1Builder.WriteIndentedString(
478-
"throw new IllegalArgumentException(\"", columnName+" has invalid value (\" + val + \")\");\n")
479-
builder.WriteIndentedString("}\n")
480-
//
481-
for i := int(n) - 1; i >= 0; i-- {
482-
shift := (i * 8)
483-
if shift == 0 {
484-
builder.WriteIndentedString(fieldName, ".put((byte) val);\n")
485-
} else {
486-
builder.WriteIndentedString(fieldName, ".put((byte) (val >> ", fmt.Sprintf("%d", shift), "));\n")
487-
}
488-
}
489-
}
490-
491-
func generateJavaModuleBytesPutter(columnName, fieldName string, bitwidth uint, builder indentBuilder) {
492-
i1Builder := builder.Indent()
493-
n := byteWidth(bitwidth)
494-
//
495-
builder.WriteIndentedString("// Trim array to size\n")
496-
builder.WriteIndentedString("Bytes bs = val.trimLeadingZeros();\n")
497-
builder.WriteIndentedString("// Sanity check against expected width\n")
498-
builder.WriteIndentedString(fmt.Sprintf("if(bs.bitLength() > %d) {\n", bitwidth))
499-
i1Builder.WriteIndentedString(
500-
fmt.Sprintf("throw new IllegalArgumentException(\"%s has invalid width (\"+bs.bitLength()+\"bits)\");\n", columnName))
501-
builder.WriteIndentedString("}\n")
502-
builder.WriteIndentedString("// Write padding (if necessary)\n")
503-
builder.WriteIndentedString(fmt.Sprintf("for(int i=bs.size(); i<%d; i++) { %s.put((byte) 0); }\n", n, fieldName))
504-
builder.WriteIndentedString("// Write bytes\n")
505-
builder.WriteIndentedString(fmt.Sprintf("for(int i=0; i<bs.size(); i++) { %s.put(bs.get(i)); }\n", fieldName))
506-
}
507-
508463
func generateJavaModuleValidateRow[F any](className string, mod corset.SourceModule, schema sc.AnySchema[F],
509464
builder indentBuilder) {
510465
//
@@ -546,10 +501,18 @@ func generateJavaModuleFillAndValidateRow[F any](className string, mod corset.So
546501
if reg.IsInputOutput() {
547502
name := toRegisterName(col.Register, reg.Name)
548503
regstr := fmt.Sprintf("%d", col.Register.Index(schema.Width()))
549-
byteWidth := fmt.Sprintf("%d", byteWidth(reg.Width))
550504
// Yes, include register
551505
i1Builder.WriteIndentedString("if(!filled.get(", regstr, ")) {\n")
552-
i2Builder.WriteIndentedString(name, ".position(", name, ".position() + ", byteWidth, ");\n")
506+
//
507+
switch {
508+
case reg.Width == 1:
509+
i2Builder.WriteIndentedString(name, ".write(false);\n")
510+
case reg.Width <= 63:
511+
i2Builder.WriteIndentedString(name, ".write(0L);\n")
512+
default:
513+
i2Builder.WriteIndentedString(name, ".write(new byte[0]);\n")
514+
}
515+
//
553516
i1Builder.WriteIndentedString("}\n")
554517
}
555518
}

pkg/cmd/generate/interface.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,7 @@ func generateInterfaceContents(className string, mod corset.SourceModule, builde
106106
//
107107
if mod.Name == "" {
108108
builder.WriteString(javaColumnHeader)
109+
builder.WriteString(javaColumn)
109110
builder.WriteString(javaAddMetadataSignature)
110111
builder.WriteString(javaOpenSignature)
111112
}

pkg/cmd/generate/snippets.go

Lines changed: 18 additions & 187 deletions
Original file line numberDiff line numberDiff line change
@@ -62,205 +62,36 @@ const javaColumnHeader string = `
6262
* ColumnHeader contains information about a given column in the resulting trace file.
6363
*
6464
* @param name Name of the column, as found in the trace file.
65-
* @param bytesPerElement Bytes required for each element in the column.
65+
* @param bitwidth Max number of bits required for any element in the column.
6666
*/
67-
public record ColumnHeader(String name, int register, long bytesPerElement, long length) { }
67+
public record ColumnHeader(String name, int register, int bitwidth, int length) { }
6868
`
6969

7070
// nolint
71-
const javaAddMetadataSignature string = `
72-
/**
73-
* Add an item of metadata to this trace.
74-
*/
75-
public void addMetadata(String key, Object value);
76-
`
77-
78-
// nolint
79-
const javaOpenSignature string = `
80-
/**
81-
* Construct a new trace which will be written to a given file.
82-
*
83-
* @param file File into which the trace will be written. Observe any previous contents of this file will be lost.
84-
* @return Trace object to use for writing column data.
85-
*
86-
* @throws IOException If an I/O error occurs.
87-
*/
88-
public void open(RandomAccessFile file, List<ColumnHeader> rawHeaders) throws IOException;
89-
`
90-
91-
// nolint
92-
const javaTraceOpen string = `
71+
const javaColumn string = `
9372
/**
94-
* Construct a new trace which will be written to a given file.
95-
*
96-
* @param file File into which the trace will be written. Observe any previous contents of this file will be lost.
97-
* @return Trace object to use for writing column data.
73+
* Column provides an interface for writing column data into the resulting trace file.
9874
*
99-
* @throws IOException If an I/O error occurs.
10075
*/
101-
public void open(RandomAccessFile file, List<ColumnHeader> rawHeaders) throws IOException {
102-
// Convert metadata into JSON bytes
103-
byte[] metadataBytes = getMetadataBytes(metadata);
104-
// Construct trace file header bytes
105-
byte[] header = constructTraceFileHeader(metadataBytes);
106-
// Align headers according to register indices.
107-
ColumnHeader[] columnHeaders = alignHeaders(rawHeaders);
108-
// Determine file size
109-
long headerSize = determineColumnHeadersSize(columnHeaders) + header.length;
110-
long dataSize = determineColumnDataSize(columnHeaders);
111-
file.setLength(headerSize + dataSize);
112-
// Write headers
113-
writeHeaders(file,header,columnHeaders,headerSize);
114-
// Initialise buffers
115-
MappedByteBuffer[] buffers = initialiseByteBuffers(file,columnHeaders,headerSize);
116-
// Done
117-
this.open(buffers);
118-
}
119-
120-
/**
121-
* Construct trace file header containing the given metadata bytes.
122-
*
123-
* @param metadata Metadata bytes to be embedded in the trace file.
124-
*
125-
* @return bytes making up the header.
126-
*/
127-
private static byte[] constructTraceFileHeader(byte[] metadata) {
128-
ByteBuffer buffer = ByteBuffer.allocate(16 + metadata.length);
129-
// File identifier
130-
buffer.put(new byte[]{'z','k','t','r','a','c','e','r'});
131-
// Major version
132-
buffer.putShort((short) 1);
133-
// Minor version
134-
buffer.putShort((short) 0);
135-
// Metadata length
136-
buffer.putInt(metadata.length);
137-
// Metadata
138-
buffer.put(metadata);
139-
// Done
140-
return buffer.array();
76+
public interface Column {
77+
public void write(boolean value);
78+
public void write(long value);
79+
public void write(byte[] value);
14180
}
81+
`
14282

83+
// nolint
84+
const javaAddMetadataSignature string = `
14385
/**
144-
* Align headers ensures that the order in which columns are seen matches the order found in the trace schema.
145-
*
146-
* @param headers The headers to be aligned.
147-
* @return The aligned headers.
86+
* Get static metadata stored within this trace during compilation.
14887
*/
149-
private static ColumnHeader[] alignHeaders(List<ColumnHeader> headers) {
150-
ColumnHeader[] alignedHeaders = new ColumnHeader[{ninputs}];
151-
//
152-
for(ColumnHeader header : headers) {
153-
alignedHeaders[header.register()] = header;
154-
}
155-
//
156-
return alignedHeaders;
157-
}
158-
159-
/**
160-
* Precompute the size of the trace file in order to memory map the buffers.
161-
*
162-
* @param headers Set of headers for the columns being written.
163-
* @return Number of bytes requires for the trace file header.
164-
*/
165-
private static long determineColumnHeadersSize(ColumnHeader[] headers) {
166-
long nBytes = 4; // column count
167-
168-
for (ColumnHeader header : headers) {
169-
if(header != null) {
170-
nBytes += 2; // name length
171-
nBytes += header.name().length();
172-
nBytes += 1; // byte per element
173-
nBytes += 4; // element count
174-
}
175-
}
176-
177-
return nBytes;
178-
}
179-
180-
/**
181-
* Precompute the size of the trace file in order to memory map the buffers.
182-
*
183-
* @param headers Set of headers for the columns being written.
184-
* @return Number of bytes required for storing all column data, excluding the header.
185-
*/
186-
private static long determineColumnDataSize(ColumnHeader[] headers) {
187-
long nBytes = 0;
188-
189-
for (ColumnHeader header : headers) {
190-
if(header != null) {
191-
nBytes += header.length() * header.bytesPerElement();
192-
}
193-
}
194-
195-
return nBytes;
196-
}
197-
198-
/**
199-
* Write header information for the trace file.
200-
*
201-
* @param file Trace file being written.
202-
* @param header Trace file header
203-
* @param headers Column headers.
204-
* @param size Overall size of the header.
205-
*/
206-
private static void writeHeaders(RandomAccessFile file, byte[] header, ColumnHeader[] headers, long size) throws IOException {
207-
final var buffer = file.getChannel().map(FileChannel.MapMode.READ_WRITE, 0, size);
208-
// Write trace file header
209-
buffer.put(header);
210-
// Write column count as uint32
211-
buffer.putInt(countHeaders(headers));
212-
// Write column headers one-by-one
213-
for(ColumnHeader h : headers) {
214-
if(h != null) {
215-
buffer.putShort((short) h.name().length());
216-
buffer.put(h.name().getBytes());
217-
buffer.put((byte) h.bytesPerElement());
218-
buffer.putInt((int) h.length());
219-
}
220-
}
221-
}
222-
223-
/**
224-
* Initialise one memory mapped byte buffer for each column to be written in the trace.
225-
* @param headers Set of headers for the columns being written.
226-
* @param headerSize Space required at start of trace file for header.
227-
* @return Buffer array with one entry per header.
228-
*/
229-
private static MappedByteBuffer[] initialiseByteBuffers(RandomAccessFile file, ColumnHeader[] headers,
230-
long headerSize) throws IOException {
231-
MappedByteBuffer[] buffers = new MappedByteBuffer[{ninputs}];
232-
long offset = headerSize;
233-
for(int i=0;i<headers.length;i++) {
234-
if(headers[i] != null) {
235-
// Determine size (in bytes) required to store all elements of this column.
236-
long length = headers[i].length() * headers[i].bytesPerElement();
237-
// Preallocate space for this column.
238-
buffers[i] = file.getChannel().map(FileChannel.MapMode.READ_WRITE, offset, length);
239-
//
240-
offset += length;
241-
}
242-
}
243-
return buffers;
244-
}
245-
246-
/**
247-
* Counter number of active (i.e. non-null) headers. A header can be null if
248-
* it represents a column in a module which is not activated for this trace.
249-
*/
250-
private static int countHeaders(ColumnHeader[] headers) throws IOException {
251-
int count = 0;
252-
for(ColumnHeader h : headers) {
253-
if(h != null) { count++; }
254-
}
255-
return count;
256-
}
88+
public Map<String,Object> getMetaData();
89+
`
25790

91+
// nolint
92+
const javaOpenSignature string = `
25893
/**
259-
* Object writer is used for generating JSON byte strings.
94+
* Open this trace file for the given set of columns.
26095
*/
261-
private static final ObjectWriter objectWriter = new ObjectMapper().writer();
262-
263-
public static byte[] getMetadataBytes(Map<String, Object> metadata) throws IOException {
264-
return objectWriter.writeValueAsBytes(metadata);
265-
}
96+
public void open(Column[] columns);
26697
`

0 commit comments

Comments
 (0)