go.tools/ssa/ssautil: Switches() utility infers multiway branches.

(Note new subpackage ssa/ssautil.)

R=gri, axwalk, gri
CC=golang-dev
https://golang.org/cl/36710043

ssa/ssautil/switch.go

@@ -0,0 +1,234 @@
+// Copyright 2013 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package ssautil
+
+// This file implements discovery of switch and type-switch constructs
+// from low-level control flow.
+//
+// Many techniques exist for compiling a high-level switch with
+// constant cases to efficient machine code.  The optimal choice will
+// depend on the data type, the specific case values, the code in the
+// body of each case, and the hardware.
+// Some examples:
+// - a lookup table (for a switch that maps constants to constants)
+// - a computed goto
+// - a binary tree
+// - a perfect hash
+// - a two-level switch (to partition constant strings by their first byte).
+
+import (
+	"bytes"
+	"fmt"
+	"go/token"
+
+	"code.google.com/p/go.tools/go/types"
+	"code.google.com/p/go.tools/ssa"
+)
+
+// A ConstCase represents a single constant comparison.
+// It is part of a Switch.
+type ConstCase struct {
+	Block *ssa.BasicBlock // block performing the comparison
+	Body  *ssa.BasicBlock // body of the case
+	Value *ssa.Const      // case comparand
+}
+
+// A TypeCase represents a single type assertion.
+// It is part of a Switch.
+type TypeCase struct {
+	Block   *ssa.BasicBlock // block performing the type assert
+	Body    *ssa.BasicBlock // body of the case
+	Type    types.Type      // case type
+	Binding ssa.Value       // value bound by this case
+}
+
+// A Switch is a logical high-level control flow operation
+// (a multiway branch) discovered by analysis of a CFG containing
+// only if/else chains.  It is not part of the ssa.Instruction set.
+//
+// One of ConstCases and TypeCases has length >= 2;
+// the other is nil.
+//
+// In a value switch, the list of cases may contain duplicate constants.
+// A type switch may contain duplicate types, or types assignable
+// to an interface type also in the list.
+// TODO(adonovan): eliminate such duplicates.
+//
+type Switch struct {
+	Start      *ssa.BasicBlock // block containing start of if/else chain
+	X          ssa.Value       // the switch operand
+	ConstCases []ConstCase     // ordered list of constant comparisons
+	TypeCases  []TypeCase      // ordered list of type assertions
+	Default    *ssa.BasicBlock // successor if all comparisons fail
+}
+
+func (sw *Switch) String() string {
+	// We represent each block by the String() of its
+	// first Instruction, e.g. "print(42:int)".
+	var buf bytes.Buffer
+	if sw.ConstCases != nil {
+		fmt.Fprintf(&buf, "switch %s {\n", sw.X.Name())
+		for _, c := range sw.ConstCases {
+			fmt.Fprintf(&buf, "case %s: %s\n", c.Value, c.Body.Instrs[0])
+		}
+	} else {
+		fmt.Fprintf(&buf, "switch %s.(type) {\n", sw.X.Name())
+		for _, c := range sw.TypeCases {
+			fmt.Fprintf(&buf, "case %s %s: %s\n",
+				c.Binding.Name(), c.Type, c.Body.Instrs[0])
+		}
+	}
+	if sw.Default != nil {
+		fmt.Fprintf(&buf, "default: %s\n", sw.Default.Instrs[0])
+	}
+	fmt.Fprintf(&buf, "}")
+	return buf.String()
+}
+
+// Switches examines the control-flow graph of fn and returns the
+// set of inferred value and type switches.  A value switch tests an
+// ssa.Value for equality against two or more compile-time constant
+// values.  Switches involving link-time constants (addresses) are
+// ignored.  A type switch type-asserts an ssa.Value against two or
+// more types.
+//
+// The switches are returned in dominance order.
+//
+// The resulting switches do not necessarily correspond to uses of the
+// 'switch' keyword in the source: for example, a single source-level
+// switch statement with non-constant cases may result in zero, one or
+// many Switches, one per plural sequence of constant cases.
+// Switches may even be inferred from if/else- or goto-based control flow.
+// (In general, the control flow constructs of the source program
+// cannot be faithfully reproduced from the SSA representation.)
+//
+func Switches(fn *ssa.Function) []Switch {
+	// Traverse the CFG in dominance order, so we don't
+	// enter an if/else-chain in the middle.
+	var switches []Switch
+	seen := make(map[*ssa.BasicBlock]bool) // TODO(adonovan): opt: use ssa.blockSet
+	for _, b := range fn.DomPreorder() {
+		if x, k := isComparisonBlock(b); x != nil {
+			// Block b starts a switch.
+			sw := Switch{Start: b, X: x}
+			valueSwitch(&sw, k, seen)
+			if len(sw.ConstCases) > 1 {
+				switches = append(switches, sw)
+			}
+		}
+
+		if y, x, T := isTypeAssertBlock(b); y != nil {
+			// Block b starts a type switch.
+			sw := Switch{Start: b, X: x}
+			typeSwitch(&sw, y, T, seen)
+			if len(sw.TypeCases) > 1 {
+				switches = append(switches, sw)
+			}
+		}
+	}
+	return switches
+}
+
+func valueSwitch(sw *Switch, k *ssa.Const, seen map[*ssa.BasicBlock]bool) {
+	b := sw.Start
+	x := sw.X
+	for x == sw.X {
+		if seen[b] {
+			break
+		}
+		seen[b] = true
+
+		sw.ConstCases = append(sw.ConstCases, ConstCase{
+			Block: b,
+			Body:  b.Succs[0],
+			Value: k,
+		})
+		b = b.Succs[1]
+		if len(b.Instrs) > 2 {
+			// Block b contains not just 'if x == k',
+			// so it may have side effects that
+			// make it unsafe to elide.
+			break
+		}
+		if len(b.Preds) != 1 {
+			// Block b has multiple predecessors,
+			// so it cannot be treated as a case.
+			break
+		}
+		x, k = isComparisonBlock(b)
+	}
+	sw.Default = b
+}
+
+func typeSwitch(sw *Switch, y ssa.Value, T types.Type, seen map[*ssa.BasicBlock]bool) {
+	b := sw.Start
+	x := sw.X
+	for x == sw.X {
+		if seen[b] {
+			break
+		}
+		seen[b] = true
+
+		sw.TypeCases = append(sw.TypeCases, TypeCase{
+			Block:   b,
+			Body:    b.Succs[0],
+			Type:    T,
+			Binding: y,
+		})
+		b = b.Succs[1]
+		if len(b.Instrs) > 4 {
+			// Block b contains not just
+			//  {TypeAssert; Extract #0; Extract #1; If}
+			// so it may have side effects that
+			// make it unsafe to elide.
+			break
+		}
+		if len(b.Preds) != 1 {
+			// Block b has multiple predecessors,
+			// so it cannot be treated as a case.
+			break
+		}
+		y, x, T = isTypeAssertBlock(b)
+	}
+	sw.Default = b
+}
+
+// isComparisonBlock returns the operands (v, k) if a block ends with
+// a comparison v==k, where k is a compile-time constant.
+//
+func isComparisonBlock(b *ssa.BasicBlock) (v ssa.Value, k *ssa.Const) {
+	if n := len(b.Instrs); n >= 2 {
+		if i, ok := b.Instrs[n-1].(*ssa.If); ok {
+			if binop, ok := i.Cond.(*ssa.BinOp); ok && binop.Block() == b && binop.Op == token.EQL {
+				if k, ok := binop.Y.(*ssa.Const); ok {
+					return binop.X, k
+				}
+				if k, ok := binop.X.(*ssa.Const); ok {
+					return binop.Y, k
+				}
+			}
+		}
+	}
+	return
+}
+
+// isTypeAssertBlock returns the operands (y, x, T) if a block ends with
+// a type assertion "if y, ok := x.(T); ok {".
+//
+func isTypeAssertBlock(b *ssa.BasicBlock) (y, x ssa.Value, T types.Type) {
+	if n := len(b.Instrs); n >= 4 {
+		if i, ok := b.Instrs[n-1].(*ssa.If); ok {
+			if ext1, ok := i.Cond.(*ssa.Extract); ok && ext1.Block() == b && ext1.Index == 1 {
+				if ta, ok := ext1.Tuple.(*ssa.TypeAssert); ok && ta.Block() == b {
+					// hack: relies upon instruction ordering.
+					if ext0, ok := b.Instrs[n-3].(*ssa.Extract); ok {
+						return ext0, ta.X, ta.AssertedType
+					}
+				}
+			}
+		}
+	}
+	return
+}

ssa/ssautil/switch_test.go

@@ -0,0 +1,69 @@
+// Copyright 2013 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package ssautil_test
+
+import (
+	"go/parser"
+	"strings"
+	"testing"
+
+	"code.google.com/p/go.tools/importer"
+	"code.google.com/p/go.tools/ssa"
+	"code.google.com/p/go.tools/ssa/ssautil"
+)
+
+func TestSwitches(t *testing.T) {
+	imp := importer.New(new(importer.Config)) // (uses GCImporter)
+	f, err := parser.ParseFile(imp.Fset, "testdata/switches.go", nil, parser.ParseComments)
+	if err != nil {
+		t.Error(err)
+		return
+	}
+
+	mainInfo := imp.CreatePackage("main", f)
+
+	prog := ssa.NewProgram(imp.Fset, 0)
+	if err := prog.CreatePackages(imp); err != nil {
+		t.Error(err)
+		return
+	}
+	mainPkg := prog.Package(mainInfo.Pkg)
+	mainPkg.Build()
+
+	for _, mem := range mainPkg.Members {
+		if fn, ok := mem.(*ssa.Function); ok {
+			if fn.Synthetic != "" {
+				continue // e.g. init()
+			}
+			// Each (multi-line) "switch" comment within
+			// this function must match the printed form
+			// of a ConstSwitch.
+			var wantSwitches []string
+			for _, c := range f.Comments {
+				if fn.Syntax().Pos() <= c.Pos() && c.Pos() < fn.Syntax().End() {
+					text := strings.TrimSpace(c.Text())
+					if strings.HasPrefix(text, "switch ") {
+						wantSwitches = append(wantSwitches, text)
+					}
+				}
+			}
+
+			switches := ssautil.Switches(fn)
+			if len(switches) != len(wantSwitches) {
+				t.Errorf("in %s, found %d switches, want %d", fn, len(switches), len(wantSwitches))
+			}
+			for i, sw := range switches {
+				got := sw.String()
+				if i >= len(wantSwitches) {
+					continue
+				}
+				want := wantSwitches[i]
+				if got != want {
+					t.Errorf("in %s, found switch %d: got <<%s>>, want <<%s>>", fn, i, got, want)
+				}
+			}
+		}
+	}
+}

ssa/ssautil/testdata/switches.go

@@ -0,0 +1,339 @@
+// +build ignore
+
+package main
+
+// This file is the input to TestFindSwitches in switch_test.go.
+// Each multiway conditional with constant or type cases (Switch)
+// discovered by FindSwitches is printed, and compared with the
+// comments.
+//
+// The body of each case is printed as the value of its first
+// instruction.
+
+// -------- Value switches --------
+
+func SimpleSwitch(x, y int) {
+	// switch x {
+	// case 1:int: print(1:int)
+	// case 2:int: print(23:int)
+	// case 3:int: print(23:int)
+	// case 4:int: print(3:int)
+	// default: x == y
+	// }
+	switch x {
+	case 1:
+		print(1)
+	case 2, 3:
+		print(23)
+		fallthrough
+	case 4:
+		print(3)
+	default:
+		print(4)
+	case y:
+		print(5)
+	}
+	print(6)
+}
+
+func four() int { return 4 }
+
+// A non-constant case makes a switch "impure", but its pure
+// cases form two separate switches.
+func SwitchWithNonConstantCase(x int) {
+	// switch x {
+	// case 1:int: print(1:int)
+	// case 2:int: print(23:int)
+	// case 3:int: print(23:int)
+	// default: four()
+	// }
+
+	// switch x {
+	// case 5:int: print(5:int)
+	// case 6:int: print(6:int)
+	// default: print("done":string)
+	// }
+	switch x {
+	case 1:
+		print(1)
+	case 2, 3:
+		print(23)
+	case four():
+		print(3)
+	case 5:
+		print(5)
+	case 6:
+		print(6)
+	}
+	print("done")
+}
+
+// Switches may be found even where the source
+// program doesn't have a switch statement.
+
+func ImplicitSwitches(x, y int) {
+	// switch x {
+	// case 1:int: print(12:int)
+	// case 2:int: print(12:int)
+	// default: x < 5:int
+	// }
+	if x == 1 || 2 == x || x < 5 {
+		print(12)
+	}
+
+	// switch x {
+	// case 3:int: print(34:int)
+	// case 4:int: print(34:int)
+	// default: x == y
+	// }
+	if x == 3 || 4 == x || x == y {
+		print(34)
+	}
+
+	// Not a switch: no consistent variable.
+	if x == 5 || y == 6 {
+		print(56)
+	}
+
+	// Not a switch: only one constant comparison.
+	if x == 7 || x == y {
+		print(78)
+	}
+}
+
+func IfElseBasedSwitch(x int) {
+	// switch x {
+	// case 1:int: print(1:int)
+	// case 2:int: print(2:int)
+	// default: print("else":string)
+	// }
+	if x == 1 {
+		print(1)
+	} else if x == 2 {
+		print(2)
+	} else {
+		print("else")
+	}
+}
+
+func GotoBasedSwitch(x int) {
+	// switch x {
+	// case 1:int: print(1:int)
+	// case 2:int: print(2:int)
+	// default: print("else":string)
+	// }
+	if x == 1 {
+		goto L1
+	}
+	if x == 2 {
+		goto L2
+	}
+	print("else")
+L1:
+	print(1)
+	goto end
+L2:
+	print(2)
+end:
+}
+
+func SwitchInAForLoop(x int) {
+	// switch x {
+	// case 1:int: print(1:int)
+	// case 2:int: print(2:int)
+	// default: print("head":string)
+	// }
+loop:
+	for {
+		print("head")
+		switch x {
+		case 1:
+			print(1)
+			break loop
+		case 2:
+			print(2)
+			break loop
+		}
+	}
+}
+
+// This case is a switch in a for-loop, both constructed using goto.
+// As before, the default case points back to the block containing the
+// switch, but that's ok.
+func SwitchInAForLoopUsingGoto(x int) {
+	// switch x {
+	// case 1:int: print(1:int)
+	// case 2:int: print(2:int)
+	// default: print("head":string)
+	// }
+loop:
+	print("head")
+	if x == 1 {
+		goto L1
+	}
+	if x == 2 {
+		goto L2
+	}
+	goto loop
+L1:
+	print(1)
+	goto end
+L2:
+	print(2)
+end:
+}
+
+func UnstructuredSwitchInAForLoop(x int) {
+	// switch x {
+	// case 1:int: print(1:int)
+	// case 2:int: x == 1:int
+	// default: print("end":string)
+	// }
+	for {
+		if x == 1 {
+			print(1)
+			return
+		}
+		if x == 2 {
+			continue
+		}
+		break
+	}
+	print("end")
+}
+
+func CaseWithMultiplePreds(x int) {
+	for {
+		if x == 1 {
+			print(1)
+			return
+		}
+	loop:
+		// This block has multiple predecessors,
+		// so can't be treated as a switch case.
+		if x == 2 {
+			goto loop
+		}
+		break
+	}
+	print("end")
+}
+
+func DuplicateConstantsAreNotEliminated(x int) {
+	// switch x {
+	// case 1:int: print(1:int)
+	// case 1:int: print("1a":string)
+	// case 2:int: print(2:int)
+	// default: return
+	// }
+	if x == 1 {
+		print(1)
+	} else if x == 1 { // duplicate => unreachable
+		print("1a")
+	} else if x == 2 {
+		print(2)
+	}
+}
+
+// Interface values (created by comparisons) are not constants,
+// so ConstSwitch.X is never of interface type.
+func MakeInterfaceIsNotAConstant(x interface{}) {
+	if x == "foo" {
+		print("foo")
+	} else if x == 1 {
+		print(1)
+	}
+}
+
+func ZeroInitializedVarsAreConstants(x int) {
+	// switch x {
+	// case 0:int: print(1:int)
+	// case 2:int: print(2:int)
+	// default: print("end":string)
+	// }
+	var zero int // SSA construction replaces zero with 0
+	if x == zero {
+		print(1)
+	} else if x == 2 {
+		print(2)
+	}
+	print("end")
+}
+
+// -------- Select --------
+
+// NB, potentially fragile reliance on register number.
+func SelectDesugarsToSwitch(ch chan int) {
+	// switch t1 {
+	// case 0:int: extract t0 #2
+	// case 1:int: println(0:int)
+	// case 2:int: println(1:int)
+	// default: println("default":string)
+	// }
+	select {
+	case x := <-ch:
+		println(x)
+	case <-ch:
+		println(0)
+	case ch <- 1:
+		println(1)
+	default:
+		println("default")
+	}
+}
+
+// -------- Type switches --------
+
+// NB, reliance on fragile register numbering.
+func SimpleTypeSwitch(x interface{}) {
+	// switch x.(type) {
+	// case t3 int: println(x)
+	// case t7 bool: println(x)
+	// case t10 string: println(t10)
+	// default: println(x)
+	// }
+	switch y := x.(type) {
+	case nil:
+		println(y)
+	case int, bool:
+		println(y)
+	case string:
+		println(y)
+	default:
+		println(y)
+	}
+}
+
+// NB, potentially fragile reliance on register number.
+func DuplicateTypesAreNotEliminated(x interface{}) {
+	// switch x.(type) {
+	// case t1 string: println(1:int)
+	// case t5 interface{}: println(t5)
+	// case t9 int: println(3:int)
+	// default: return
+	// }
+	switch y := x.(type) {
+	case string:
+		println(1)
+	case interface{}:
+		println(y)
+	case int:
+		println(3) // unreachable!
+	}
+}
+
+// NB, potentially fragile reliance on register number.
+func AdHocTypeSwitch(x interface{}) {
+	// switch x.(type) {
+	// case t1 int: println(t1)
+	// case t5 string: println(t5)
+	// default: print("default":string)
+	// }
+	if i, ok := x.(int); ok {
+		println(i)
+	} else if s, ok := x.(string); ok {
+		println(s)
+	} else {
+		print("default")
+	}
+}