vta: extends VTA graph construction to handle collections

Collections include channels, slices, ranges, maps, and their related
SSA statements.

Statements involving functions and function calls will be addressed in a
future CL.

Change-Id: I552cbf3ee9d65e125270db69d1fc3c3f6491b121
Reviewed-on: https://go-review.googlesource.com/c/tools/+/322951
Run-TryBot: Zvonimir Pavlinovic <zpavlinovic@google.com>
gopls-CI: kokoro <noreply+kokoro@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
Trust: Zvonimir Pavlinovic <zpavlinovic@google.com>
Reviewed-by: Roland Shoemaker <roland@golang.org>

Author: zpavlinovic

go/callgraph/vta/graph.go

@@ -319,6 +319,20 @@ func (b *builder) instr(instr ssa.Instruction) {
 		b.field(i)
 	case *ssa.FieldAddr:
 		b.fieldAddr(i)
+	case *ssa.Send:
+		b.send(i)
+	case *ssa.Select:
+		b.selekt(i)
+	case *ssa.Index:
+		b.index(i)
+	case *ssa.IndexAddr:
+		b.indexAddr(i)
+	case *ssa.Lookup:
+		b.lookup(i)
+	case *ssa.MapUpdate:
+		b.mapUpdate(i)
+	case *ssa.Next:
+		b.next(i)
 	case *ssa.MakeChan, *ssa.MakeMap, *ssa.MakeSlice, *ssa.BinOp,
 		*ssa.Alloc, *ssa.DebugRef, *ssa.Convert, *ssa.Jump, *ssa.If,
 		*ssa.Slice, *ssa.Range, *ssa.RunDefers:
@@ -336,7 +350,8 @@ func (b *builder) unop(u *ssa.UnOp) {
 		// Multiplication operator * is used here as a dereference operator.
 		b.addInFlowAliasEdges(b.nodeFromVal(u), b.nodeFromVal(u.X))
 	case token.ARROW:
-		// TODO(zpavlinovic): add support for channels.
+		t := u.X.Type().Underlying().(*types.Chan).Elem()
+		b.addInFlowAliasEdges(b.nodeFromVal(u), channelElem{typ: t})
 	default:
 		// There is no interesting type flow otherwise.
 	}
@@ -388,6 +403,91 @@ func (b *builder) fieldAddr(f *ssa.FieldAddr) {
 	b.addInFlowEdge(b.nodeFromVal(f), fnode)
 }
 
+func (b *builder) send(s *ssa.Send) {
+	t := s.Chan.Type().Underlying().(*types.Chan).Elem()
+	b.addInFlowAliasEdges(channelElem{typ: t}, b.nodeFromVal(s.X))
+}
+
+// selekt generates flows for select statement
+//   a = select blocking/nonblocking [c_1 <- t_1, c_2 <- t_2, ..., <- o_1, <- o_2, ...]
+// between receiving channel registers c_i and corresponding input register t_i. Further,
+// flows are generated between o_i and a[2 + i]. Note that a is a tuple register of type
+// <int, bool, r_1, r_2, ...> where the type of r_i is the element type of channel o_i.
+func (b *builder) selekt(s *ssa.Select) {
+	recvIndex := 0
+	for _, state := range s.States {
+		t := state.Chan.Type().Underlying().(*types.Chan).Elem()
+
+		if state.Dir == types.SendOnly {
+			b.addInFlowAliasEdges(channelElem{typ: t}, b.nodeFromVal(state.Send))
+		} else {
+			// state.Dir == RecvOnly by definition of select instructions.
+			tupEntry := indexedLocal{val: s, typ: t, index: 2 + recvIndex}
+			b.addInFlowAliasEdges(tupEntry, channelElem{typ: t})
+			recvIndex++
+		}
+	}
+}
+
+// index instruction a := b[c] on slices creates flows between a and
+// SliceElem(t) flow where t is an interface type of c. Arrays and
+// slice elements are both modeled as SliceElem.
+func (b *builder) index(i *ssa.Index) {
+	et := sliceArrayElem(i.X.Type())
+	b.addInFlowAliasEdges(b.nodeFromVal(i), sliceElem{typ: et})
+}
+
+// indexAddr instruction a := &b[c] fetches address of a index
+// into the field so we create bidirectional flow a <-> SliceElem(t)
+// where t is an interface type of c. Arrays and slice elements are
+// both modeled as SliceElem.
+func (b *builder) indexAddr(i *ssa.IndexAddr) {
+	et := sliceArrayElem(i.X.Type())
+	b.addInFlowEdge(sliceElem{typ: et}, b.nodeFromVal(i))
+	b.addInFlowEdge(b.nodeFromVal(i), sliceElem{typ: et})
+}
+
+// lookup handles map query commands a := m[b] where m is of type
+// map[...]V and V is an interface. It creates flows between `a`
+// and MapValue(V).
+func (b *builder) lookup(l *ssa.Lookup) {
+	t, ok := l.X.Type().Underlying().(*types.Map)
+	if !ok {
+		// No interesting flows for string lookups.
+		return
+	}
+	b.addInFlowAliasEdges(b.nodeFromVal(l), mapValue{typ: t.Elem()})
+}
+
+// mapUpdate handles map update commands m[b] = a where m is of type
+// map[K]V and K and V are interfaces. It creates flows between `a`
+// and MapValue(V) as well as between MapKey(K) and `b`.
+func (b *builder) mapUpdate(u *ssa.MapUpdate) {
+	t, ok := u.Map.Type().Underlying().(*types.Map)
+	if !ok {
+		// No interesting flows for string updates.
+		return
+	}
+
+	b.addInFlowAliasEdges(mapKey{typ: t.Key()}, b.nodeFromVal(u.Key))
+	b.addInFlowAliasEdges(mapValue{typ: t.Elem()}, b.nodeFromVal(u.Value))
+}
+
+// next instruction <ok, key, value> := next r, where r
+// is a range over map or string generates flow between
+// key and MapKey as well value and MapValue nodes.
+func (b *builder) next(n *ssa.Next) {
+	if n.IsString {
+		return
+	}
+	tup := n.Type().Underlying().(*types.Tuple)
+	kt := tup.At(1).Type()
+	vt := tup.At(2).Type()
+
+	b.addInFlowAliasEdges(indexedLocal{val: n, typ: kt, index: 1}, mapKey{typ: kt})
+	b.addInFlowAliasEdges(indexedLocal{val: n, typ: vt, index: 2}, mapValue{typ: vt})
+}
+
 // addInFlowAliasEdges adds an edge r -> l to b.graph if l is a node that can
 // have an inflow, i.e., a node that represents an interface or an unresolved
 // function value. Similarly for the edge l -> r with an additional condition

go/callgraph/vta/graph_test.go

@@ -236,6 +236,11 @@ func TestVTAGraphConstruction(t *testing.T) {
 		"testdata/node_uniqueness.go",
 		"testdata/store_load_alias.go",
 		"testdata/phi_alias.go",
+		"testdata/channels.go",
+		"testdata/select.go",
+		"testdata/stores_arrays.go",
+		"testdata/maps.go",
+		"testdata/ranges.go",
 	} {
 		t.Run(file, func(t *testing.T) {
 			prog, want, err := testProg(file)

go/callgraph/vta/testdata/channels.go

@@ -0,0 +1,36 @@
+// Copyright 2021 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.
+
+// go:build ignore
+
+package testdata
+
+func foo(c chan interface{}, j int) {
+	c <- j + 1
+}
+
+func Baz(i int) {
+	c := make(chan interface{})
+	go foo(c, i)
+	x := <-c
+	print(x)
+}
+
+// Relevant SSA:
+//  func foo(c chan interface{}, j int):
+//  t0 = j + 1:int
+//  t1 = make interface{} <- int (t0)
+//  send c <- t1                        // t1 -> chan {}interface
+//  return
+//
+// func Baz(i int):
+//  t0 = make chan interface{} 0:int
+//  go foo(t0, i)
+//  t1 = <-t0                           // chan {}interface -> t1
+//  t2 = print(t1)
+//  return
+
+// WANT:
+// Channel(chan interface{}) -> Local(t1)
+// Local(t1) -> Channel(chan interface{})

go/callgraph/vta/testdata/maps.go

@@ -0,0 +1,52 @@
+// Copyright 2021 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.
+
+// go:build ignore
+
+package testdata
+
+type I interface {
+	Foo() string
+}
+
+type J interface {
+	Foo() string
+	Bar()
+}
+
+type B struct {
+	p string
+}
+
+func (b B) Foo() string { return b.p }
+func (b B) Bar()        {}
+
+func Baz(m map[I]I, b1, b2 B, n map[string]*J) *J {
+	m[b1] = b2
+
+	return n[b1.Foo()]
+}
+
+// Relevant SSA:
+// func Baz(m map[I]I, b1 B, b2 B, n map[string]*J) *J:
+//   t0 = local B (b1)
+//   *t0 = b1
+//   t1 = local B (b2)
+//   *t1 = b2
+//   t2 = *t0
+//   t3 = make I <- B (t2)
+//   t4 = *t1
+//   t5 = make I <- B (t4)
+//   m[t3] = t5
+//   t6 = *t0
+//   t7 = (B).Foo(t6)
+//   t8 = n[t7]
+//   return t8
+
+// WANT:
+// Local(t4) -> Local(t5)
+// Local(t5) -> MapValue(testdata.I)
+// Local(t3) -> MapKey(testdata.I)
+// Local(t8) -> MapValue(*testdata.J)
+// MapValue(*testdata.J) -> Local(t8)

go/callgraph/vta/testdata/ranges.go

@@ -0,0 +1,55 @@
+// Copyright 2021 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.
+
+// go:build ignore
+
+package testdata
+
+type I interface {
+	Foo() string
+}
+
+type B struct {
+	p string
+}
+
+func (b B) Foo() string { return b.p }
+
+func Baz(m map[I]*I) {
+	for i, v := range m {
+		*v = B{p: i.Foo()}
+	}
+}
+
+// Relevant SSA:
+//  func Baz(m map[I]*I):
+//   0:
+//    t0 = range m
+//         jump 1
+//   1:
+//    t1 = next t0
+//    t2 = extract t1 #0
+//    if t2 goto 2 else 3
+//   2:
+//    t3 = extract t1 #1
+//    t4 = extract t1 #2
+//    t5 = local B (complit)
+//    t6 = &t5.p [#0]
+//    t7 = invoke t3.Foo()
+//    *t6 = t7
+//    t8 = *t5
+//    t9 = make I <- B (t8)
+//    *t4 = t9
+//    jump 1
+//   3:
+//    return
+
+// WANT:
+// MapKey(testdata.I) -> Local(t1[1])
+// Local(t1[1]) -> Local(t3)
+// MapValue(*testdata.I) -> Local(t1[2])
+// Local(t1[2]) -> Local(t4), MapValue(*testdata.I)
+// Local(t8) -> Local(t9)
+// Local(t9) -> Local(t4)
+// Local(t4) -> Local(t1[2])

go/callgraph/vta/testdata/select.go

@@ -0,0 +1,60 @@
+// Copyright 2021 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.
+
+// go:build ignore
+
+package testdata
+
+type I interface {
+	Foo() string
+}
+
+type J interface {
+	I
+}
+
+type B struct {
+	p string
+}
+
+func (b B) Foo() string { return b.p }
+
+func Baz(b1, b2 B, c1 chan I, c2 chan J) {
+	for {
+		select {
+		case c1 <- b1:
+			print("b1")
+		case c2 <- b2:
+			print("b2")
+		case <-c1:
+			print("c1")
+		case k := <-c2:
+			print(k.Foo())
+			return
+		}
+	}
+}
+
+// Relevant SSA:
+// func Baz(b1 B, b2 B, c1 chan I, c2 chan J):
+//   ...
+//   t2 = *t0
+//   t3 = make I <- B (t2)
+//   t4 = *t1
+//   t5 = make J <- B (t4)
+//   t6 = select blocking [c1<-t3, c2<-t5, <-c1, <-c2] (index int, ok bool, I, J)
+//   t7 = extract t6 #0
+//   t8 = t7 == 0:int
+//   if t8 goto 2 else 3
+//         ...
+//  8:
+//   t15 = extract t6 #3
+//   t16 = invoke t15.Foo()
+//   t17 = print(t18)
+
+// WANT:
+// Local(t3) -> Channel(chan testdata.I)
+// Local(t5) -> Channel(chan testdata.J)
+// Channel(chan testdata.I) -> Local(t6[2])
+// Channel(chan testdata.J) -> Local(t6[3])

go/callgraph/vta/utils.go

@@ -84,3 +84,19 @@ func interfaceUnderPtr(t types.Type) types.Type {
 
 	return interfaceUnderPtr(p.Elem())
 }
+
+// sliceArrayElem returns the element type of type `t` that is
+// expected to be a (pointer to) array or slice, consistent with
+// the ssa.Index and ssa.IndexAddr instructions. Panics otherwise.
+func sliceArrayElem(t types.Type) types.Type {
+	u := t.Underlying()
+
+	if p, ok := u.(*types.Pointer); ok {
+		u = p.Elem().Underlying()
+	}
+
+	if a, ok := u.(*types.Array); ok {
+		return a.Elem()
+	}
+	return u.(*types.Slice).Elem()
+}