Add WASM-ready flamegraph JSON export

7 files changed, 196 insertions, 2 deletions

diff --git a/README.md b/README.md
index 386f874..af824ae 100644
--- a/README.md
+++ b/README.md
@@ -77,6 +77,14 @@ This generates an SVG and starts an embedded web server. The terminal prints a U
 Flamegraph available at http://HOSTNAME:PORT/abs/path/to.svg
 ```
 
+For experimental WebAssembly frontends, you can also emit a flamegraph JSON tree:
+
+```sh
+./ior -ior=trace.ior.zst -flamegraphJson
+```
+
+This writes `<trace>.<fields>-by-<count>.json` next to the SVG.
+
 ## Live Flamegraph Mode
 
 Run live mode (requires root privileges):
diff --git a/internal/flags/flags.go b/internal/flags/flags.go
index a732c3a..bc1654b 100644
--- a/internal/flags/flags.go
+++ b/internal/flags/flags.go
@@ -67,6 +67,7 @@ type Flags struct {
 	LiveInterval     time.Duration
 	OpenCommand      string
 	FlamegraphName   string
+	FlamegraphJSON   bool
 	TUIExportEnable  bool
 
 	// To convert ior data into native SVG format
@@ -125,6 +126,7 @@ func parse() error {
 	flag.DurationVar(&singleton.LiveInterval, "live-interval", 200*time.Millisecond, "Live flamegraph refresh interval")
 	flag.StringVar(&singleton.OpenCommand, "open", "", "Command to open live flamegraph URL (used with -live); use {url} placeholder or URL is appended")
 	flag.StringVar(&singleton.FlamegraphName, "name", "default", "Name of the flamegraph, used to generate the SVG file")
+	flag.BoolVar(&singleton.FlamegraphJSON, "flamegraphJson", false, "Also export flamegraph tree as JSON in -ior mode (experimental WASM-ready output)")
 	flag.BoolVar(&singleton.TUIExportEnable, "tuiExport", true, "Enable writing TUI snapshot export files")
 
 	flag.StringVar(&singleton.IorDataFile, "ior", "", "IOR data file to convert into native SVG flamegraph")
diff --git a/internal/flags/flags_test.go b/internal/flags/flags_test.go
index 32af4bd..11414d8 100644
--- a/internal/flags/flags_test.go
+++ b/internal/flags/flags_test.go
@@ -104,6 +104,16 @@ func TestParseOpenFlags(t *testing.T) {
 	}
 }
 
+func TestParseFlamegraphJSONFlag(t *testing.T) {
+	cfg, err := parseForTest(t, "-flamegraphJson")
+	if err != nil {
+		t.Fatalf("parse returned error: %v", err)
+	}
+	if !cfg.FlamegraphJSON {
+		t.Fatalf("expected -flamegraphJson to enable JSON export")
+	}
+}
+
 func TestParseDefaultCollapsedFieldsOrder(t *testing.T) {
 	cfg, err := parseForTest(t)
 	if err != nil {
diff --git a/internal/flamegraph/nativejson.go b/internal/flamegraph/nativejson.go
new file mode 100644
index 0000000..088bcfc
--- /dev/null
+++ b/internal/flamegraph/nativejson.go
@@ -0,0 +1,86 @@
+package flamegraph
+
+import (
+	"encoding/json"
+	"fmt"
+	"io"
+	"iter"
+	"os"
+	"strings"
+)
+
+type jsonNode struct {
+	Name     string     `json:"name"`
+	Value    uint64     `json:"value"`
+	Total    uint64     `json:"total"`
+	Children []jsonNode `json:"children,omitempty"`
+}
+
+type jsonFlamegraph struct {
+	Fields     []string `json:"fields"`
+	CountField string   `json:"countField"`
+	Root       jsonNode `json:"root"`
+}
+
+func (n NativeSVG) WriteJSONFromFile(iorDataFile string) (outFile string, err error) {
+	outFile = fmt.Sprintf("%s.%s-by-%s.json",
+		strings.TrimSuffix(iorDataFile, ".ior.zst"),
+		strings.Join(n.fields, ":"),
+		n.countField,
+	)
+	defer func() {
+		if err != nil {
+			_ = os.Remove(outFile)
+		}
+	}()
+
+	iod, err := newIorDataFromFile(iorDataFile)
+	if err != nil {
+		return outFile, fmt.Errorf("read ior data: %w", err)
+	}
+
+	fd, err := os.Create(outFile)
+	if err != nil {
+		return outFile, fmt.Errorf("create output %s: %w", outFile, err)
+	}
+	defer fd.Close()
+
+	if err := n.WriteJSONFromIter(iod.iter(), fd); err != nil {
+		return outFile, err
+	}
+	return outFile, nil
+}
+
+func (n NativeSVG) WriteJSONFromIter(records iter.Seq[IterRecord], w io.Writer) error {
+	tr, err := n.buildTrieFromIter(records)
+	if err != nil {
+		return err
+	}
+
+	payload := jsonFlamegraph{
+		Fields:     append([]string(nil), n.fields...),
+		CountField: n.countField,
+		Root:       jsonNodeFromTrieNode(tr.root, "root"),
+	}
+
+	enc := json.NewEncoder(w)
+	enc.SetIndent("", "  ")
+	return enc.Encode(payload)
+}
+
+func jsonNodeFromTrieNode(node *trieNode, name string) jsonNode {
+	out := jsonNode{
+		Name:  name,
+		Value: node.value,
+		Total: node.total,
+	}
+	if len(node.children) == 0 {
+		return out
+	}
+
+	out.Children = make([]jsonNode, 0, len(node.children))
+	for _, child := range node.children {
+		out.Children = append(out.Children, jsonNodeFromTrieNode(child, child.name))
+	}
+	return out
+}
diff --git a/internal/flamegraph/nativejson_test.go b/internal/flamegraph/nativejson_test.go
new file mode 100644
index 0000000..c76d327
--- /dev/null
+++ b/internal/flamegraph/nativejson_test.go
@@ -0,0 +1,75 @@
+package flamegraph
+
+import (
+	"encoding/json"
+	"os"
+	"testing"
+)
+
+type jsonNodeForTest struct {
+	Name     string            `json:"name"`
+	Value    uint64            `json:"value"`
+	Total    uint64            `json:"total"`
+	Children []jsonNodeForTest `json:"children"`
+}
+
+type jsonFlamegraphForTest struct {
+	Fields     []string        `json:"fields"`
+	CountField string          `json:"countField"`
+	Root       jsonNodeForTest `json:"root"`
+}
+
+func TestWriteJSONFromFileContainsFlamegraphTree(t *testing.T) {
+	dir := t.TempDir()
+	iorFile := writeTestIorZst(t, dir)
+
+	n := NewNativeSVG([]string{"comm", "path", "tracepoint"}, "count")
+	outFile, err := n.WriteJSONFromFile(iorFile)
+	if err != nil {
+		t.Fatalf("WriteJSONFromFile returned error: %v", err)
+	}
+
+	data, err := os.ReadFile(outFile)
+	if err != nil {
+		t.Fatalf("read output json: %v", err)
+	}
+
+	var payload jsonFlamegraphForTest
+	if err := json.Unmarshal(data, &payload); err != nil {
+		t.Fatalf("unmarshal output json: %v", err)
+	}
+
+	if payload.CountField != "count" {
+		t.Fatalf("count field = %q, want %q", payload.CountField, "count")
+	}
+	if len(payload.Fields) != 3 {
+		t.Fatalf("fields len = %d, want 3", len(payload.Fields))
+	}
+	if payload.Root.Name != "root" {
+		t.Fatalf("root name = %q, want %q", payload.Root.Name, "root")
+	}
+	if payload.Root.Total != 1 {
+		t.Fatalf("root total = %d, want 1", payload.Root.Total)
+	}
+	if len(payload.Root.Children) != 1 {
+		t.Fatalf("root children len = %d, want 1", len(payload.Root.Children))
+	}
+	if payload.Root.Children[0].Name != "tester" {
+		t.Fatalf("root child name = %q, want %q", payload.Root.Children[0].Name, "tester")
+	}
+}
+
+func TestWriteJSONFromFileCleansUpPartialOutputOnError(t *testing.T) {
+	dir := t.TempDir()
+	iorFile := writeTestIorZst(t, dir)
+
+	n := NewNativeSVG([]string{"invalidField"}, "count")
+	outFile, err := n.WriteJSONFromFile(iorFile)
+	if err == nil {
+		t.Fatal("expected error for invalid field, got nil")
+	}
+
+	if _, statErr := os.Stat(outFile); !os.IsNotExist(statErr) {
+		t.Fatalf("expected partial output to be removed, stat err=%v", statErr)
+	}
+}
diff --git a/internal/flamegraph/nativesvg.go b/internal/flamegraph/nativesvg.go
index 8a2bcd5..80061b4 100644
--- a/internal/flamegraph/nativesvg.go
+++ b/internal/flamegraph/nativesvg.go
@@ -58,18 +58,26 @@ func (n NativeSVG) WriteSVGFromFile(iorDataFile string) (outFile string, err err
 }
 
 func (n NativeSVG) WriteSVGFromIter(records iter.Seq[IterRecord], w io.Writer) error {
+	tr, err := n.buildTrieFromIter(records)
+	if err != nil {
+		return err
+	}
+	return WriteSVG(w, tr, n.config)
+}
+
+func (n NativeSVG) buildTrieFromIter(records iter.Seq[IterRecord]) (*trie, error) {
 	tr := newTrie()
 	var framesBuf []string
 	for record := range records {
 		frames, err := n.recordFrames(record, framesBuf)
 		if err != nil {
-			return err
+			return nil, err
 		}
 		framesBuf = frames
 		tr.add(frames, record.Cnt.ValueByName(n.countField))
 	}
 	tr.computeTotals()
-	return WriteSVG(w, tr, n.config)
+	return tr, nil
 }
 
 func (n NativeSVG) recordFrames(record IterRecord, framesBuf []string) ([]string, error) {
diff --git a/internal/ior.go b/internal/ior.go
index b198d78..865e51d 100644
--- a/internal/ior.go
+++ b/internal/ior.go
@@ -127,6 +127,11 @@ func Run() error {
 		if err != nil {
 			return err
 		}
+		if cfg.FlamegraphJSON {
+			if _, err := native.WriteJSONFromFile(iorFile); err != nil {
+				return err
+			}
+		}
 
 		if err := flamegraph.ServeSVG(svgFile); err != nil {
 			return err