github.com/google/capslock@v0.2.3-0.20240517042941-dac19fc347c0/analyzer/graph.go (about)

     1  // Copyright 2023 Google LLC
     2  //
     3  // Use of this source code is governed by a BSD-style
     4  // license that can be found in the LICENSE file or at
     5  // https://developers.google.com/open-source/licenses/bsd
     6  
     7  package analyzer
     8  
     9  import (
    10  	"bufio"
    11  	"go/types"
    12  	"io"
    13  	"os"
    14  	"strconv"
    15  	"strings"
    16  
    17  	cpb "github.com/google/capslock/proto"
    18  	"golang.org/x/tools/go/callgraph"
    19  	"golang.org/x/tools/go/packages"
    20  )
    21  
    22  func graphOutput(pkgs []*packages.Package, queriedPackages map[*types.Package]struct{}, config *Config) error {
    23  	w := bufio.NewWriterSize(os.Stdout, 1<<20)
    24  	gb := newGraphBuilder(w, func(v interface{}) string {
    25  		switch v := v.(type) {
    26  		case *callgraph.Node:
    27  			if v.Func != nil {
    28  				return v.Func.String()
    29  			}
    30  			return strconv.Itoa(v.ID)
    31  		case cpb.Capability:
    32  			return v.String()
    33  		default:
    34  			panic("unexpected node type")
    35  		}
    36  	})
    37  	callEdge := func(caller, callee *callgraph.Node) {
    38  		gb.Edge(caller, callee)
    39  	}
    40  	capabilityEdge := func(fn *callgraph.Node, c cpb.Capability) {
    41  		gb.Edge(fn, c)
    42  	}
    43  	CapabilityGraph(pkgs, queriedPackages, config, callEdge, capabilityEdge)
    44  	gb.Done()
    45  	return w.Flush()
    46  }
    47  
    48  type graphBuilder struct {
    49  	io.Writer
    50  	nodeNamer func(any) string
    51  	started   bool
    52  	done      bool
    53  }
    54  
    55  func newGraphBuilder(w io.Writer, nodeNamer func(any) string) graphBuilder {
    56  	return graphBuilder{
    57  		Writer:    w,
    58  		nodeNamer: nodeNamer,
    59  	}
    60  }
    61  
    62  func (gb *graphBuilder) Edge(from, to interface{}) {
    63  	if gb.done {
    64  		panic("done")
    65  	}
    66  	if !gb.started {
    67  		gb.Write([]byte("digraph {\n"))
    68  		gb.started = true
    69  	}
    70  	gb.Write([]byte("\t"))
    71  	gb.Write([]byte(`"`))
    72  	gb.Write([]byte(strings.ReplaceAll(gb.nodeNamer(from), `"`, `\"`)))
    73  	gb.Write([]byte(`" -> "`))
    74  	gb.Write([]byte(strings.ReplaceAll(gb.nodeNamer(to), `"`, `\"`)))
    75  	gb.Write([]byte("\"\n"))
    76  }
    77  
    78  func (gb *graphBuilder) Done() {
    79  	if gb.done {
    80  		panic("done")
    81  	}
    82  	gb.Write([]byte("}\n"))
    83  	gb.done = true
    84  }