go-type-chains/cmd/root.go

/*
Copyright © 2025 filifa

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program. If not, see <http://www.gnu.org/licenses/>.
*/
package cmd

import (
	"errors"
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"log"
	"os"
	"path/filepath"

	"github.com/spf13/cobra"
	"gonum.org/v1/gonum/graph/encoding/dot"
	"scm.dairydemon.net/filifa/go-type-chains/cmd/internal/depgraph"
)

var pkg string
var ignoredPkgs []string

type filename string
type specname string

func inspect(node ast.Node, fname filename, typesUsed map[filename][]specname, declaredAt map[specname]filename) {
	switch v := node.(type) {
	case *ast.File:
		for _, decl := range v.Decls {
			inspect(decl, fname, typesUsed, declaredAt)
		}
	case *ast.GenDecl:
		for _, spec := range v.Specs {
			inspect(spec, fname, typesUsed, declaredAt)
		}
	case *ast.TypeSpec:
		inspect(v.Type, fname, typesUsed, declaredAt)
		name := specname(v.Name.Name)
		declaredAt[name] = fname
	case *ast.ValueSpec:
		inspect(v.Type, fname, typesUsed, declaredAt)
		for _, name := range v.Names {
			declaredAt[specname(name.Name)] = fname
		}
	case *ast.Ident:
		typesUsed[fname] = append(typesUsed[fname], specname(v.Name))
	case *ast.ParenExpr:
		inspect(v.X, fname, typesUsed, declaredAt)
	case *ast.SelectorExpr:
		inspect(v.X, fname, typesUsed, declaredAt)
	case *ast.StarExpr:
		inspect(v.X, fname, typesUsed, declaredAt)
	case *ast.ArrayType:
		inspect(v.Elt, fname, typesUsed, declaredAt)
	case *ast.ChanType:
		inspect(v.Value, fname, typesUsed, declaredAt)
	case *ast.FuncType:
		for _, field := range v.TypeParams.List {
			inspect(field.Type, fname, typesUsed, declaredAt)
		}
		for _, field := range v.Params.List {
			inspect(field.Type, fname, typesUsed, declaredAt)
		}
		for _, field := range v.Results.List {
			inspect(field.Type, fname, typesUsed, declaredAt)
		}
	case *ast.InterfaceType:
		for _, field := range v.Methods.List {
			inspect(field.Type, fname, typesUsed, declaredAt)
		}
	case *ast.MapType:
		inspect(v.Key, fname, typesUsed, declaredAt)
		inspect(v.Value, fname, typesUsed, declaredAt)
	case *ast.StructType:
		for _, field := range v.Fields.List {
			inspect(field.Type, fname, typesUsed, declaredAt)
		}
	}
}

func constructDependencyGraph(typesUsed map[filename][]specname, declaredAt map[specname]filename, ignores map[specname]bool) (*depgraph.DependencyGraph, error) {
	depGraph := depgraph.NewDependencyGraph()
	for fname, types := range typesUsed {
		from := depgraph.NewSourcefile(string(fname))
		if depGraph.Node(from.ID()) == nil {
			depGraph.AddNode(from)
		}

		for _, t := range types {
			if ignores[t] {
				continue
			}

			declFile, ok := declaredAt[t]
			if !ok {
				return nil, errors.New("did not find declaration for " + string(t))
			}

			if fname == declFile {
				continue
			}

			to := depgraph.NewSourcefile(string(declFile))
			if depGraph.Node(to.ID()) == nil {
				depGraph.AddNode(to)
			}

			e := depGraph.NewEdge(from, to)
			depGraph.SetEdge(e)
		}
	}

	return depGraph, nil
}

func parse(cmd *cobra.Command, args []string) {
	fset := token.NewFileSet()
	pkgs, err := parser.ParseDir(fset, pkg, nil, 0)
	if err != nil {
		log.Fatal(err)
	}

	typesUsed := make(map[filename][]specname)
	declaredAt := make(map[specname]filename)

	tree := pkgs["models"]
	for fpath, file := range tree.Files {
		fname := filename(filepath.Base(fpath))
		typesUsed[fname] = make([]specname, 0)
		inspect(file, fname, typesUsed, declaredAt)
	}

	ignores := map[specname]bool{
		"any":         true,
		"bool":        true,
		"byte":        true,
		"comparable":  true,
		"complext128": true,
		"complex64":   true,
		"error":       true,
		"float32":     true,
		"float64":     true,
		"int":         true,
		"int16":       true,
		"int32":       true,
		"int64":       true,
		"int8":        true,
		"rune":        true,
		"string":      true,
		"uint":        true,
		"uint16":      true,
		"uint32":      true,
		"uint64":      true,
		"uint8":       true,
		"uintptr":     true,
	}

	for _, ignoredPkg := range ignoredPkgs {
		ignores[specname(ignoredPkg)] = true
	}

	depGraph, err := constructDependencyGraph(typesUsed, declaredAt, ignores)
	if err != nil {
		log.Fatal(err)
	}

	viz, err := dot.Marshal(depGraph, "deps", "", "\t")
	if err != nil {
		log.Fatal(err)
	}

	fmt.Println(string(viz))
}

// rootCmd represents the base command when called without any subcommands
var rootCmd = &cobra.Command{
	Use:   "go-type-chains",
	Short: "Output Graphviz encoding of dependencies between files in a Go package",
	Long: `Output Graphviz encoding of dependencies between files in a Go package.

The outputted digraph will create an edge from source file A to source file B if A uses a type or value that is declared in B.`,
	// Uncomment the following line if your bare application
	// has an action associated with it:
	Run: parse,
}

// Execute adds all child commands to the root command and sets flags appropriately.
// This is called by main.main(). It only needs to happen once to the rootCmd.
func Execute() {
	err := rootCmd.Execute()
	if err != nil {
		os.Exit(1)
	}
}

func init() {
	// Here you will define your flags and configuration settings.
	// Cobra supports persistent flags, which, if defined here,
	// will be global for your application.

	// rootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is $HOME/.go-type-chains.yaml)")

	// Cobra also supports local flags, which will only run
	// when this action is called directly.

	rootCmd.Flags().StringVarP(&pkg, "package", "p", "", "package to parse")
	rootCmd.MarkFlagRequired("package")

	rootCmd.Flags().StringArrayVarP(&ignoredPkgs, "ignore", "i", []string{}, "package to ignore")
}