more on this

2 files changed, 77 insertions, 27 deletions

diff --git a/internal/flags/flags.go b/internal/flags/flags.go
index 890c18f..ab9a342 100644
--- a/internal/flags/flags.go
+++ b/internal/flags/flags.go
@@ -19,8 +19,12 @@ var validCollapsedFields = []string{
 	"path",
 	"comm",
 	"tracepoint",
-	"pid", "tid",
+	"pid",
+	"tid",
 	"count",
+	"duration",
+	"durationToPrev",
+	"bytes",
 }
 
 func Get() Flags {
@@ -73,7 +77,8 @@ func parse() {
 	flag.StringVar(&singleton.FlamegraphName, "name", "foo", "Name of the flamegraph data output")
 
 	flag.StringVar(&singleton.IorDataFile, "ior", "", "IOR data file to convert into collapsed format")
-	fields := flag.String("fields", "", "Comma separated list of fields to collapse")
+	fields := flag.String("fields", "",
+		fmt.Sprintf("Comma separated list of fields to collapse, valid are: %v", validCollapsedFields))
 	flag.Parse()
 
 	singleton.TracepointsToAttach = extractTracepointFlags(*tracepointsToAttach)
diff --git a/internal/flamegraph/iordata.go b/internal/flamegraph/iordata.go
index 44c4174..6974329 100644
--- a/internal/flamegraph/iordata.go
+++ b/internal/flamegraph/iordata.go
@@ -34,6 +34,14 @@ func newIorData() iorData {
 	return iorData{paths: make(pathMap)}
 }
 
+func newIorDataFromFile(filename string) (iorData, error) {
+	iod := newIorData()
+	if err := iod.loadFromFile(filename); err != nil {
+		return iorData{}, err
+	}
+	return iod, nil
+}
+
 func cloneString(s string) string {
 	// Clone the string by creating a new string with the same content
 	// This is a workaround to avoid using unsafe package
@@ -164,25 +172,74 @@ func (iod iorData) loadFromFile(filename string) error {
 	return iod.deserialize(&buffer)
 }
 
-func (iod iorData) lines() iter.Seq[string] {
-	return func(yield func(string) bool) {
+func (iod iorData) serialize() ([]byte, error) {
+	var buf bytes.Buffer
+	enc := gob.NewEncoder(&buf)
+	err := enc.Encode(iod.paths)
+	return buf.Bytes(), err
+}
+
+func (iod *iorData) deserialize(buf *bytes.Buffer) error {
+	dec := gob.NewDecoder(buf)
+	return dec.Decode(&iod.paths)
+}
+
+// Record returned by the iterator
+type iterRecord struct {
+	path    pathType
+	traceId traceIdType
+	comm    commType
+	pid     pidType
+	tid     tidType
+	flags   flagsType
+	cnt     Counter
+}
+
+func (ir iterRecord) StringByName(name string) string {
+	switch name {
+	case "path":
+		return ir.path
+	case "comm":
+		return ir.comm
+	case "tracepoint":
+		return ir.traceId.String()
+	case "pid":
+		return fmt.Sprint(ir.pid)
+	case "tid":
+		return fmt.Sprint(ir.tid)
+	case "flags":
+		return ir.flags.String()
+	case "count":
+		return fmt.Sprint(ir.cnt.Count)
+	case "duration":
+		return fmt.Sprint(ir.cnt.Duration)
+	case "durationToPrev":
+		return fmt.Sprint(ir.cnt.DurationToPrev)
+	case "bytes":
+		return fmt.Sprint(ir.cnt.Bytes)
+	default:
+		panic(fmt.Sprintln("No", name, "in record"))
+	}
+}
+
+func (iod iorData) iter() iter.Seq[iterRecord] {
+	return func(yield func(iterRecord) bool) {
 		for path, traceIdMap := range iod.paths {
 			for traceId, commMap := range traceIdMap {
 				for comm, pidMap := range commMap {
 					for pid, tidMap := range pidMap {
 						for tid, flagsMap := range tidMap {
 							for flags, cnt := range flagsMap {
-								joinedStr := strings.Join([]string{
-									path,
-									traceId.String(),
-									comm,
-									fmt.Sprint(pid),
-									fmt.Sprint(tid),
-									flags.String(),
-									fmt.Sprintf("%d %d %d %d", cnt.Count, cnt.Duration, cnt.DurationToPrev, cnt.Bytes),
-								},
-									" --- ")
-								if !yield(joinedStr) {
+								record := iterRecord{
+									path:    path,
+									traceId: traceId,
+									comm:    comm,
+									pid:     pid,
+									tid:     tid,
+									flags:   flags,
+									cnt:     cnt,
+								}
+								if !yield(record) {
 									// Stop iteration if yield returns false
 									return
 								}
@@ -194,15 +251,3 @@ func (iod iorData) lines() iter.Seq[string] {
 		}
 	}
 }
-
-func (iod iorData) serialize() ([]byte, error) {
-	var buf bytes.Buffer
-	enc := gob.NewEncoder(&buf)
-	err := enc.Encode(iod.paths)
-	return buf.Bytes(), err
-}
-
-func (iod *iorData) deserialize(buf *bytes.Buffer) error {
-	dec := gob.NewDecoder(buf)
-	return dec.Decode(&iod.paths)
-}