From 0c5541b647d9c8367f944e15ae775d6859e52a73 Mon Sep 17 00:00:00 2001
From: Kent Gruber <kent.picat.gruber@gmail.com>
Date: Fri, 5 Jan 2024 23:25:10 -0500
Subject: [PATCH] Add `callgraphutil.WriteCosmograph` (#28)

* Add `callgraphutil.WriteCosmograph`
* Update dot_test.go
---
 callgraphutil/cosmograph.go      | 68 ++++++++++++++++++++++++++++++++
 callgraphutil/cosmograph_test.go | 54 +++++++++++++++++++++++++
 callgraphutil/dot_test.go        | 21 ++++++++++
 3 files changed, 143 insertions(+)
 create mode 100644 callgraphutil/cosmograph.go
 create mode 100644 callgraphutil/cosmograph_test.go

diff --git a/callgraphutil/cosmograph.go b/callgraphutil/cosmograph.go
new file mode 100644
index 0000000..602377f
--- /dev/null
+++ b/callgraphutil/cosmograph.go
@@ -0,0 +1,68 @@
+package callgraphutil
+
+import (
+	"encoding/csv"
+	"fmt"
+	"io"
+
+	"golang.org/x/tools/go/callgraph"
+)
+
+// WriteComsmograph writes the given callgraph.Graph to the given io.Writer in CSV
+// format, which can be used to generate a visual representation of the call
+// graph using Comsmograph.
+//
+// https://cosmograph.app/run/
+func WriteCosmograph(graph, metadata io.Writer, g *callgraph.Graph) error {
+	graphWriter := csv.NewWriter(graph)
+	graphWriter.Comma = ','
+	defer graphWriter.Flush()
+
+	metadataWriter := csv.NewWriter(metadata)
+	metadataWriter.Comma = ','
+	defer metadataWriter.Flush()
+
+	// Write header.
+	if err := graphWriter.Write([]string{"source", "target"}); err != nil {
+		return fmt.Errorf("failed to write header: %w", err)
+	}
+
+	// Write metadata header.
+	if err := metadataWriter.Write([]string{"id", "pkg", "func"}); err != nil {
+		return fmt.Errorf("failed to write metadata header: %w", err)
+	}
+
+	// Write edges.
+	for _, n := range g.Nodes {
+		// TODO: fix this so there's not so many "shared" functions?
+		//
+		// It is a bit of a hack, but it works for now.
+		var pkgPath string
+		if n.Func.Pkg != nil {
+			pkgPath = n.Func.Pkg.Pkg.Path()
+		} else {
+			pkgPath = "shared"
+		}
+
+		// Write metadata.
+		if err := metadataWriter.Write([]string{
+			fmt.Sprintf("%d", n.ID),
+			pkgPath,
+			n.Func.String(),
+		}); err != nil {
+			return fmt.Errorf("failed to write metadata: %w", err)
+		}
+
+		for _, e := range n.Out {
+			// Write edge.
+			if err := graphWriter.Write([]string{
+				fmt.Sprintf("%d", n.ID),
+				fmt.Sprintf("%d", e.Callee.ID),
+			}); err != nil {
+				return fmt.Errorf("failed to write edge: %w", err)
+			}
+		}
+	}
+
+	return nil
+}
diff --git a/callgraphutil/cosmograph_test.go b/callgraphutil/cosmograph_test.go
new file mode 100644
index 0000000..0fe84e6
--- /dev/null
+++ b/callgraphutil/cosmograph_test.go
@@ -0,0 +1,54 @@
+package callgraphutil_test
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"testing"
+
+	"github.com/picatz/taint/callgraphutil"
+)
+
+func TestWriteCosmograph(t *testing.T) {
+	var (
+		ownerName = "picatz"
+		repoName  = "taint"
+	)
+
+	repo, _, err := cloneGitHubRepository(context.Background(), ownerName, repoName)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	pkgs, err := loadPackages(context.Background(), repo, "./...")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	mainFn, srcFns, err := loadSSA(context.Background(), pkgs)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	cg, err := loadCallGraph(context.Background(), mainFn, srcFns)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	graphOutput, err := os.Create(fmt.Sprintf("%s.csv", repoName))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer graphOutput.Close()
+
+	metadataOutput, err := os.Create(fmt.Sprintf("%s-metadata.csv", repoName))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer metadataOutput.Close()
+
+	err = callgraphutil.WriteCosmograph(graphOutput, metadataOutput, cg)
+	if err != nil {
+		t.Fatal(err)
+	}
+}
diff --git a/callgraphutil/dot_test.go b/callgraphutil/dot_test.go
index 8cea70c..924a8ef 100644
--- a/callgraphutil/dot_test.go
+++ b/callgraphutil/dot_test.go
@@ -111,12 +111,33 @@ func loadSSA(ctx context.Context, pkgs []*packages.Package) (mainFn *ssa.Functio
 	// Analyze the package.
 	ssaProg, ssaPkgs := ssautil.Packages(pkgs, ssaBuildMode)
 
+	// It's possible that the ssaProg is nil?
+	if ssaProg == nil {
+		err = fmt.Errorf("failed to create new ssa program")
+		return
+	}
+
 	ssaProg.Build()
 
 	for _, pkg := range ssaPkgs {
+		if pkg == nil {
+			continue
+		}
 		pkg.Build()
 	}
 
+	// Remove nil ssaPkgs by iterating over the slice of packages
+	// and for each nil package, we append the slice up to that
+	// index and then append the slice from the next index to the
+	// end of the slice. This effectively removes the nil package
+	// from the slice without having to allocate a new slice.
+	for i := 0; i < len(ssaPkgs); i++ {
+		if ssaPkgs[i] == nil {
+			ssaPkgs = append(ssaPkgs[:i], ssaPkgs[i+1:]...)
+			i--
+		}
+	}
+
 	mainPkgs := ssautil.MainPackages(ssaPkgs)
 
 	mainFn = mainPkgs[0].Members["main"].(*ssa.Function)