Enable writing to lt files

This now enables writing out filter traces back to `lt` files.

pkg/cmd/trace.go

@@ -3,6 +3,7 @@ package cmd
 import (
 	"fmt"
 	"os"
+	"strings"
 
 	"github.com/consensys/go-corset/pkg/trace"
 	"github.com/consensys/go-corset/pkg/util"
@@ -24,8 +25,12 @@ var traceCmd = &cobra.Command{
 		// Parse trace
 		trace := readTraceFile(args[0])
 		list := getFlag(cmd, "list")
+		filter := getString(cmd, "filter")
 		output := getString(cmd, "out")
 		//
+		if filter != "" {
+			trace = filterColumns(trace, filter)
+		}
 		if list {
 			listColumns(trace)
 		}
@@ -36,6 +41,21 @@ var traceCmd = &cobra.Command{
 	},
 }
 
+// Construct a new trace containing only those columns from the original who
+// name begins with the given prefix.
+func filterColumns(tr trace.Trace, prefix string) trace.Trace {
+	ntr := trace.EmptyArrayTrace()
+	//
+	for i := uint(0); i < tr.Width(); i++ {
+		ith := tr.ColumnByIndex(i)
+		if strings.HasPrefix(ith.Name(), prefix) {
+			ntr.Add(ith)
+		}
+	}
+	// Done
+	return ntr
+}
+
 func listColumns(tr trace.Trace) {
 	tbl := util.NewTablePrinter(3, tr.Width())
 
@@ -55,4 +75,5 @@ func init() {
 	rootCmd.AddCommand(traceCmd)
 	traceCmd.Flags().BoolP("list", "l", false, "detail the columns in the trace file")
 	traceCmd.Flags().StringP("out", "o", "", "Specify output file to write trace")
+	traceCmd.Flags().StringP("filter", "f", "", "Filter columns beginning with prefix")
 }

pkg/cmd/util.go

@@ -62,18 +62,26 @@ func getString(cmd *cobra.Command, flag string) string {
 // Write a given trace file to disk
 func writeTraceFile(filename string, tr trace.Trace) {
 	var err error
+
+	var bytes []byte
 	// Check file extension
 	ext := path.Ext(filename)
 	//
 	switch ext {
 	case ".json":
 		js := json.ToJsonString(tr)
 		//
-		if err := os.WriteFile(filename, []byte(js), 0644); err == nil {
+		if err = os.WriteFile(filename, []byte(js), 0644); err == nil {
 			return
 		}
 	case ".lt":
-
+		bytes, err = lt.ToBytes(tr)
+		//
+		if err == nil {
+			if err = os.WriteFile(filename, bytes, 0644); err == nil {
+				return
+			}
+		}
 	default:
 		err = fmt.Errorf("Unknown trace file format: %s", ext)
 	}

pkg/trace/array_trace.go

@@ -104,6 +104,20 @@ func (p *ArrayTrace) Clone() Trace {
 	return clone
 }
 
+// Add adds a new column of data to this trace.
+func (p *ArrayTrace) Add(column Column) {
+	// Sanity check the column does not already exist.
+	if p.HasColumn(column.Name()) {
+		panic("column already exists")
+	}
+	// Append it
+	p.columns = append(p.columns, column)
+	// Update maximum height
+	if column.Height() > p.height {
+		p.height = column.Height()
+	}
+}
+
 // AddColumn adds a new column of data to this trace.
 func (p *ArrayTrace) AddColumn(name string, data []*fr.Element, padding *fr.Element) {
 	// Sanity check the column does not already exist.

pkg/trace/lt/writer.go

@@ -10,8 +10,8 @@ import (
 )
 
 // ToBytes writes a given trace file as an array of bytes.
-func ToBytes(columns []trace.Column) ([]byte, error) {
-	buf, err := ToBytesBuffer(columns)
+func ToBytes(tr trace.Trace) ([]byte, error) {
+	buf, err := ToBytesBuffer(tr)
 	if err != nil {
 		return nil, err
 	}
@@ -20,24 +20,25 @@ func ToBytes(columns []trace.Column) ([]byte, error) {
 }
 
 // ToBytesBuffer writes a given trace file into a byte buffer.
-func ToBytesBuffer(columns []trace.Column) (*bytes.Buffer, error) {
+func ToBytesBuffer(tr trace.Trace) (*bytes.Buffer, error) {
 	var buf bytes.Buffer
-	if err := WriteBytes(columns, &buf); err != nil {
+	if err := WriteBytes(tr, &buf); err != nil {
 		return nil, err
 	}
 
 	return &buf, nil
 }
 
 // WriteBytes a given trace file to an io.Writer.
-func WriteBytes(cols []trace.Column, buf io.Writer) error {
-	ncols := uint32(len(cols))
+func WriteBytes(tr trace.Trace, buf io.Writer) error {
+	ncols := tr.Width()
 	// Write column count
-	if err := binary.Write(buf, binary.BigEndian, ncols); err != nil {
+	if err := binary.Write(buf, binary.BigEndian, uint32(ncols)); err != nil {
 		return err
 	}
 	// Write header information
-	for _, col := range cols {
+	for i := uint(0); i < ncols; i++ {
+		col := tr.ColumnByIndex(i)
 		// Write name length
 		nameBytes := []byte(col.Name())
 		nameLen := uint16(len(nameBytes))
@@ -60,7 +61,8 @@ func WriteBytes(cols []trace.Column, buf io.Writer) error {
 		}
 	}
 	// Write column data information
-	for _, col := range cols {
+	for i := uint(0); i < ncols; i++ {
+		col := tr.ColumnByIndex(i)
 		if err := col.Write(buf); err != nil {
 			return err
 		}