1// Copyright 2016 The Go Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5// Package cfg constructs a simple control-flow graph (CFG) of the
6// statements and expressions within a single function.
7//
8// Use cfg.New to construct the CFG for a function body.
9//
10// The blocks of the CFG contain all the function's non-control
11// statements.  The CFG does not contain control statements such as If,
12// Switch, Select, and Branch, but does contain their subexpressions.
13// For example, this source code:
14//
15//	if x := f(); x != nil {
16//		T()
17//	} else {
18//		F()
19//	}
20//
21// produces this CFG:
22//
23//    1:  x := f()
24//        x != nil
25//        succs: 2, 3
26//    2:  T()
27//        succs: 4
28//    3:  F()
29//        succs: 4
30//    4:
31//
32// The CFG does contain Return statements; even implicit returns are
33// materialized (at the position of the function's closing brace).
34//
35// The CFG does not record conditions associated with conditional branch
36// edges, nor the short-circuit semantics of the && and || operators,
37// nor abnormal control flow caused by panic.  If you need this
38// information, use golang.org/x/tools/go/ssa instead.
39//
40package cfg
41
42import (
43	"bytes"
44	"fmt"
45	"go/ast"
46	"go/format"
47	"go/token"
48)
49
50// A CFG represents the control-flow graph of a single function.
51//
52// The entry point is Blocks[0]; there may be multiple return blocks.
53type CFG struct {
54	Blocks []*Block // block[0] is entry; order otherwise undefined
55}
56
57// A Block represents a basic block: a list of statements and
58// expressions that are always evaluated sequentially.
59//
60// A block may have 0-2 successors: zero for a return block or a block
61// that calls a function such as panic that never returns; one for a
62// normal (jump) block; and two for a conditional (if) block.
63type Block struct {
64	Nodes []ast.Node // statements, expressions, and ValueSpecs
65	Succs []*Block   // successor nodes in the graph
66	Index int32      // index within CFG.Blocks
67	Live  bool       // block is reachable from entry
68
69	comment string    // for debugging
70	succs2  [2]*Block // underlying array for Succs
71}
72
73// New returns a new control-flow graph for the specified function body,
74// which must be non-nil.
75//
76// The CFG builder calls mayReturn to determine whether a given function
77// call may return.  For example, calls to panic, os.Exit, and log.Fatal
78// do not return, so the builder can remove infeasible graph edges
79// following such calls.  The builder calls mayReturn only for a
80// CallExpr beneath an ExprStmt.
81func New(body *ast.BlockStmt, mayReturn func(*ast.CallExpr) bool) *CFG {
82	b := builder{
83		mayReturn: mayReturn,
84		cfg:       new(CFG),
85	}
86	b.current = b.newBlock("entry")
87	b.stmt(body)
88
89	// Compute liveness (reachability from entry point), breadth-first.
90	q := make([]*Block, 0, len(b.cfg.Blocks))
91	q = append(q, b.cfg.Blocks[0]) // entry point
92	for len(q) > 0 {
93		b := q[len(q)-1]
94		q = q[:len(q)-1]
95
96		if !b.Live {
97			b.Live = true
98			q = append(q, b.Succs...)
99		}
100	}
101
102	// Does control fall off the end of the function's body?
103	// Make implicit return explicit.
104	if b.current != nil && b.current.Live {
105		b.add(&ast.ReturnStmt{
106			Return: body.End() - 1,
107		})
108	}
109
110	return b.cfg
111}
112
113func (b *Block) String() string {
114	return fmt.Sprintf("block %d (%s)", b.Index, b.comment)
115}
116
117// Return returns the return statement at the end of this block if present, nil otherwise.
118func (b *Block) Return() (ret *ast.ReturnStmt) {
119	if len(b.Nodes) > 0 {
120		ret, _ = b.Nodes[len(b.Nodes)-1].(*ast.ReturnStmt)
121	}
122	return
123}
124
125// Format formats the control-flow graph for ease of debugging.
126func (g *CFG) Format(fset *token.FileSet) string {
127	var buf bytes.Buffer
128	for _, b := range g.Blocks {
129		fmt.Fprintf(&buf, ".%d: # %s\n", b.Index, b.comment)
130		for _, n := range b.Nodes {
131			fmt.Fprintf(&buf, "\t%s\n", formatNode(fset, n))
132		}
133		if len(b.Succs) > 0 {
134			fmt.Fprintf(&buf, "\tsuccs:")
135			for _, succ := range b.Succs {
136				fmt.Fprintf(&buf, " %d", succ.Index)
137			}
138			buf.WriteByte('\n')
139		}
140		buf.WriteByte('\n')
141	}
142	return buf.String()
143}
144
145func formatNode(fset *token.FileSet, n ast.Node) string {
146	var buf bytes.Buffer
147	format.Node(&buf, fset, n)
148	// Indent secondary lines by a tab.
149	return string(bytes.Replace(buf.Bytes(), []byte("\n"), []byte("\n\t"), -1))
150}
151