Repository: google/go-flow-levee
Pull request #309 (2021-05-05 23:42:42)

EAR pointer analysis: add context sensitivity K>0
Fixes #<issue_number_goes_here>

> It's a good idea to open an issue first for discussion.

- [ ] Running against a large codebase such as [Kubernetes](https://github.com/kubernetes/kubernetes) does not error out. (See [DEVELOPING.md](https://github.com/google/go-flow-levee/blob/master/DEVELOPING.md) for instructions on how to do that.)
- [ ] Appropriate changes to README are included in PR


guodongli-google: Add @timothy-king  as reviewer.

Patch:
diff --git a/internal/pkg/earpointer/analysis.go b/internal/pkg/earpointer/analysis.go
index e426eb9..d193f30 100644
--- a/internal/pkg/earpointer/analysis.go
+++ b/internal/pkg/earpointer/analysis.go
@@ -22,15 +22,25 @@ import (
 	"strconv"
 	"strings"
 
-	"golang.org/x/tools/go/ssa/ssautil"
-
 	"github.com/google/go-flow-levee/internal/pkg/config"
 
 	"golang.org/x/tools/go/analysis"
 	"golang.org/x/tools/go/analysis/passes/buildssa"
+	"golang.org/x/tools/go/callgraph"
+	"golang.org/x/tools/go/callgraph/static"
 	"golang.org/x/tools/go/ssa"
+	"golang.org/x/tools/go/ssa/ssautil"
 )
 
+var (
+	contextK = 0 // the number of call sites in each context
+)
+
+func init() {
+	Analyzer.Flags.IntVar(&contextK, "contextK", 0,
+		`the K value (default=0) in context sensitivity.`)
+}
+
 // Analyzer traverses the packages and constructs an EAR partitions
 // unifying all the IR elements in these packages.
 var Analyzer = &analysis.Analyzer{
@@ -46,17 +56,30 @@ var Analyzer = &analysis.Analyzer{
 // for all reachable contexts for that function. Both the intra-procedural
 // instructions and inter-procedural instructions are handled.
 type visitor struct {
-	state    *state     // mutable state
-	contexts []*Context // for context sensitive analysis
+	state    *state                              // mutable state
+	callees  map[*ssa.CallCommon][]*ssa.Function // callee functions at each callsite
+	contexts map[*ssa.Function][]*Context        // for context sensitive analysis
+	contextK int
 }
 
 func run(pass *analysis.Pass) (interface{}, error) {
 	ssainput := pass.ResultOf[buildssa.Analyzer].(*buildssa.SSA)
-	vis := visitor{state: NewState(), contexts: []*Context{&emptyContext}}
-	vis.initGlobalReferences(ssainput)
+	p := analyze(ssainput.Pkg)
+	return p, nil
+}
+
+// Analyzes an SSA program and build the partition information.
+func analyze(pkg *ssa.Package) *Partitions {
+	prog := pkg.Prog
+	// Use the call graph to initialize the contexts.
+	// TODO: the call graph can be CHA, RTA, VTA, etc.
+	cg := static.CallGraph(prog)
+	vis := visitor{state: NewState(), callees: mapCallees(cg)}
+	vis.initContexts(cg)
+	vis.initGlobalReferences(pkg)
 	// Analyze all the functions and methods in the package,
 	// not just those in ssainput.ScrFuncs.
-	fns := ssautil.AllFunctions(ssainput.Pkg.Prog)
+	fns := ssautil.AllFunctions(prog)
 	for fn := range fns {
 		vis.initFunction(fn)
 	}
@@ -67,13 +90,47 @@ func run(pass *analysis.Pass) (interface{}, error) {
 			}
 		}
 	}
-	return vis.state.ToPartitions(), nil
+	return vis.state.ToPartitions()
+}
+
+// Builds the calling context set for each function.
+func (vis *visitor) initContexts(cg *callgraph.Graph) {
+	vis.contextK = contextK
+	vis.contexts = make(map[*ssa.Function][]*Context)
+	for fn, node := range cg.Nodes {
+		if fn == nil {
+			continue
+		}
+		// Create the contexts for this function.
+		vis.contexts[fn] = collectContext(node, vis.contextK)
+	}
+}
+
+// Collect the contexts from the node backward up to k edges.
+func collectContext(node *callgraph.Node, k int) []*Context {
+	// This implementation can be optimized in two ways:
+	// (1) Reuse the k-1, k-2, ... contexts when calculating k so as to
+	//     avoid duplicate computations.
+	// (2) Make the contexts share memory since they may overlap a lot.
+	// However, we may use only small k values so the performance
+	// and memory consumption issue is not a concern in practice.
+	if k <= 0 || len(node.In) == 0 {
+		return []*Context{&emptyContext}
+	}
+	var kContexts []*Context
+	for _, in := range node.In {
+		for _, prev := range collectContext(in.Caller, k-1) {
+			cur := append(*prev, in.Site)
+			kContexts = append(kContexts, &cur)
+		}
+	}
+	return kContexts
 }
 
 // Insert into the state all the global references.
-func (vis *visitor) initGlobalReferences(ssainput *buildssa.SSA) {
+func (vis *visitor) initGlobalReferences(pkg *ssa.Package) {
 	state := vis.state
-	for _, member := range ssainput.Pkg.Members {
+	for _, member := range pkg.Members {
 		if g, ok := member.(*ssa.Global); ok {
 			if typeMayShareObject(g.Type()) &&
 				// skip some synthetic variables
@@ -88,14 +145,20 @@ func (vis *visitor) initGlobalReferences(ssainput *buildssa.SSA) {
 // Insert into the state all the local references.
 func (vis *visitor) initFunction(fn *ssa.Function) {
 	state := vis.state
+	// A function to insert a local into the state under all related contexts.
+	initLocal := func(v ssa.Value) {
+		for _, c := range vis.getContexts(v) {
+			state.Insert(MakeLocal(c, v))
+		}
+	}
 	for _, param := range fn.Params {
 		if typeMayShareObject(param.Type()) {
-			state.Insert(MakeLocal(&emptyContext, param))
+			initLocal(param)
 		}
 	}
 	for _, fv := range fn.FreeVars {
 		if typeMayShareObject(fv.Type()) {
-			state.Insert(MakeLocal(&emptyContext, fv))
+			initLocal(fv)
 		}
 	}
 	// Usual tuple-type registers are skipped since they will be processed
@@ -115,7 +178,7 @@ func (vis *visitor) initFunction(fn *ssa.Function) {
 		for _, instr := range blk.Instrs {
 			if v, ok := instr.(ssa.Value); ok {
 				if returnTuple(v) || typeMayShareObject(v.Type()) {
-					state.Insert(MakeLocal(&emptyContext, v))
+					initLocal(v)
 				}
 			}
 		}
@@ -202,7 +265,7 @@ func (vis *visitor) visitFieldAddr(addr *ssa.FieldAddr) {
 	// "t1 = t0.x" is implemented by t0[x -> t1] through address unification.
 	obj := addr.X
 	field := makeField(obj.Type(), addr.Field)
-	for _, c := range vis.contexts {
+	for _, c := range vis.getContexts(addr) {
 		vis.unifyFieldAddress(c, obj, field, addr)
 	}
 }
@@ -346,45 +409,46 @@ func (vis *visitor) visitSelect(s *ssa.Select) {
 	// TODO: to be implemented.
 }
 
-func (vis *visitor) visitCall(call *ssa.CallCommon, instr ssa.Instruction) {
+func (vis *visitor) visitCall(call *ssa.CallCommon, callsite ssa.Instruction) {
 	if builtin, ok := call.Value.(*ssa.Builtin); ok {
 		// special handling of builtin functions
-		vis.visitBuiltin(builtin, instr)
+		vis.visitBuiltin(builtin, callsite)
 		return
 	}
 
-	// Collect unification constraints
-	// Here only static calls are supported.
-	// A future extension is to use a more advanced call graph.
-	fn := call.StaticCallee()
-	if fn == nil {
-		return
-	}
-	paramCstrs, retCstrs := vis.collectCalleeConstraints(call, fn, instr)
-	// Unify caller arguments and callee parameters with suitable context.
-	for arg, params := range paramCstrs {
-		for _, param := range params {
-			// The copy-by-value semantics is implemented by the SSA call.
-			// For example:
-			//   type S struct { ... }
-			//   var v S
-			//   g(v)
-			//   func g(a S) { ... }
-			// SSA creates a copy of v and passes it to g. Hence we directly
-			// unify this copy and g's argument "a".
-			if mayShareObject(arg) {
-				vis.unifyLocals(arg, param)
+	// Collect unification constraints using the call graph.
+	for _, fn := range vis.callees[call] {
+		if fn == nil {
+			continue
+		}
+		paramCstrs, retCstrs := vis.collectCalleeConstraints(call, fn, callsite)
+		// Unify caller arguments and callee parameters in matching contexts.
+		for arg, params := range paramCstrs {
+			for _, param := range params {
+				// The copy-by-value semantics is implemented by the SSA call.
+				// For example:
+				//   type S struct { ... }
+				//   var v S
+				//   g(v)
+				//   func g(a S) { ... }
+				// SSA creates a copy of v and passes it to g. Hence we directly
+				// unify this copy and g's argument "a".
+				if mayShareObject(arg) {
+					vis.unifyCallWithContexts(arg, param, callsite)
+				}
 			}
 		}
-	}
-	// Handle return value unification.
-	for callerRet, rets := range retCstrs {
-		for _, calleeRet := range rets {
-			if len(calleeRet) == 1 { // non-tuple case
-				if mayShareObject(calleeRet[0]) {
-					vis.unifyLocals(callerRet, calleeRet[0])
+		// Handle return value unification.
+		for callerRet, rets := range retCstrs {
+			for _, calleeRet := range rets {
+				if len(calleeRet) == 1 { // non-tuple case
+					param := calleeRet[0]
+					if mayShareObject(param) {
+						vis.unifyCallWithContexts(callerRet, param, callsite)
+					}
+					continue
 				}
-			} else { // tuple case
+				// tuple case
 				// For call "return_r = f(...)", if the callee returns a tuple
 				// <r0, r1, ...>, then the i^{th} element in the tuple is
 				// connected to the caller's return_r by making
@@ -396,13 +460,50 @@ func (vis *visitor) visitCall(call *ssa.CallCommon, instr ssa.Instruction) {
 					}
 					// Implement "return_r[i] = r{i}".
 					field := Field{Name: strconv.Itoa(k)}
-					vis.processHeapAccess(retV, callerRet, field)
+					// unify instance field
+					for _, c := range vis.getContexts(callerRet) {
+						callerCxt := append(*c, callsite)
+						for _, calleeCxt := range vis.getContexts(retV) {
+							if vis.contextKEqual(callerCxt, *calleeCxt) {
+								vis.unifyField(c, callerRet, field, retV)
+							}
+						}
+					}
 				}
 			}
 		}
 	}
 }
 
+// Unify the argument in the caller with the parameter in the callee under all contexts
+func (vis *visitor) unifyCallWithContexts(arg ssa.Value, param ssa.Value, callsite ssa.Instruction) {
+	state := vis.state
+	// The caller may have multiple contexts;
+	// perform the unification for each context.
+	for _, c := range vis.getContexts(arg) {
+		// For each context of the callee, match it with the caller context
+		// plus the call site. Unify the caller' arg and the callee's param
+		// for each matched context.
+		// For example, assume K=1,
+		//   func f(x, y *T) {
+		//       g(x)
+		//       g(y)
+		//   }
+		//   func g(a *T) {}
+		// g() is called at contexts [g(x)] and [g(y)], so g.a is unified
+		// with f.x and f.y w.r.t. these contexts, resulting in two
+		// partitions {[g(x)]g.a, f.x} and {[g(y)]g.a, f.y}.
+		argCxt := append(*c, callsite)
+		argRef := MakeReference(c, arg)
+		for _, paramCxt := range vis.getContexts(param) {
+			if vis.contextKEqual(argCxt, *paramCxt) {
+				state.Unify(argRef, MakeReference(paramCxt, param))
+				break
+			}
+		}
+	}
+}
+
 // Handle calls to builtin functions: https://golang.org/pkg/builtin/.
 func (vis *visitor) visitBuiltin(builtin *ssa.Builtin, instr ssa.Instruction) {
 	switch builtin.Name() {
@@ -423,7 +524,7 @@ func (vis *visitor) visitBuiltin(builtin *ssa.Builtin, instr ssa.Instruction) {
 			// ops[0] is the "copy" function itself.
 			dst, src := *ops[1], *ops[2]
 			if mayShareObject(src) {
-				for _, c := range vis.contexts {
+				for _, c := range vis.getContexts(dst) {
 					vis.unifyByValue(MakeLocal(c, dst), MakeLocal(c, src))
 				}
 			}
@@ -445,21 +546,10 @@ func (vis *visitor) collectCalleeConstraints(common *ssa.CallCommon, fn *ssa.Fun
 	}
 
 	paramCstrs := make(map[ssa.Value][]ssa.Value)
-	// Add caller_reg -> {callee_reg} constraints for parameters.
+	// Add caller_arg -> {callee_parameter} constraints for parameters.
 	// For example, for g(a, b) and func g(x, y), add <a, g.x> and
 	// <b, g.y> into the constraints.
-	//
-	// Unifying the receivers may cause too many FPs.
-	// For example, func (t T) f() { ... }; o1.f() and o2.f(), the unification
-	// results in {o1, o2, t}, while o1 and o2 should not be unified.
-	// So the receiver unification is disabled here, which may lead to under-approximation
-	// and FNs of the checkers.
-	// TODO: support context-sensitivity to mitigate this issue.
-	start := 0
-	if fn.Signature.Recv() != nil {
-		start = 1
-	}
-	for i := start; i < len(common.Args); i++ {
+	for i := 0; i < len(common.Args); i++ {
 		arg, param := common.Args[i], fn.Params[i]
 		if mayShareObject(arg) && typeMayShareObject(param.Type()) {
 			paramCstrs[arg] = append(paramCstrs[arg], param)
@@ -516,7 +606,7 @@ func (vis *visitor) processHeapAccess(data ssa.Value, base ssa.Value, fd Field)
 		return
 	}
 	// unify instance field
-	for _, context := range vis.contexts {
+	for _, context := range vis.getContexts(base) {
 		vis.unifyField(context, base, fd, data)
 	}
 }
@@ -527,7 +617,7 @@ func (vis *visitor) processAddressToValue(addr ssa.Value, value ssa.Value) {
 		return
 	}
 	state := vis.state
-	for _, context := range vis.contexts {
+	for _, context := range vis.getContexts(addr) {
 		caddr := state.Insert(MakeReference(context, addr))
 		cvalue := state.Insert(MakeReference(context, value))
 		fmap := state.PartitionFieldMap(state.representative(caddr))
@@ -544,14 +634,14 @@ func (vis *visitor) processAddressToValue(addr ssa.Value, value ssa.Value) {
 	}
 }
 
-// Unify local regs "v1" and "v2" in all contexts.
-func (vis *visitor) unifyLocals(v1 ssa.Value, v2 ssa.Value) {
-	if !mayShareObject(v1) || !mayShareObject(v2) {
+// Unify local regs "u" and "v" in all contexts.
+func (vis *visitor) unifyLocals(u ssa.Value, v ssa.Value) {
+	if !mayShareObject(u) || !mayShareObject(v) {
 		return
 	}
 	state := vis.state
-	for _, context := range vis.contexts {
-		state.Unify(MakeReference(context, v1), MakeReference(context, v2))
+	for _, context := range vis.getContexts(u) {
+		state.Unify(MakeReference(context, u), MakeReference(context, v))
 	}
 }
 
@@ -747,6 +837,47 @@ func (vis *visitor) unifyByValue(ref1 Reference, ref2 Reference) {
 	}
 }
 
+// Return all the calling contexts of the function to which a value belongs.
+func (vis *visitor) getContexts(v ssa.Value) []*Context {
+	if fn := v.Parent(); fn != nil {
+		if ctxs, ok := vis.contexts[fn]; ok {
+			return ctxs
+		}
+	}
+	// Otherwise return an empty context, e.g. for a Global.
+	return []*Context{&emptyContext}
+}
+
+// Return whether the last K sites of two contexts are equal.
+func (vis *visitor) contextKEqual(c1 Context, c2 Context) bool {
+	l1, l2 := len(c1), len(c2)
+	for i := 0; i < vis.contextK; i++ {
+		if i >= l1 || i >= l2 { // no more to compare
+			return true
+		}
+		if c1[l1-i-1] != c2[l2-i-1] {
+			return false
+		}
+	}
+	// All the last K sites in the two contexts match.
+	return true
+}
+
+// Maps a callsite to the set of candidate callee functions.
+func mapCallees(cg *callgraph.Graph) map[*ssa.CallCommon][]*ssa.Function {
+	callees := make(map[*ssa.CallCommon][]*ssa.Function)
+	for f, node := range cg.Nodes {
+		if f == nil {
+			continue
+		}
+		for _, in := range node.In {
+			common := in.Site.Common()
+			callees[common] = append(callees[common], in.Callee.Func)
+		}
+	}
+	return callees
+}
+
 // Return whether a type can use the unify-by-reference semantics. If not then
 // it is assumed to be unifyByValue.
 func isUnifyByReference(tp types.Type) bool {


Test patch:
diff --git a/internal/pkg/earpointer/analysis_test.go b/internal/pkg/earpointer/analysis_test.go
index ade06c9..f102505 100644
--- a/internal/pkg/earpointer/analysis_test.go
+++ b/internal/pkg/earpointer/analysis_test.go
@@ -17,6 +17,7 @@ package earpointer_test
 import (
 	"fmt"
 	"sort"
+	"strconv"
 	"strings"
 	"testing"
 
@@ -27,8 +28,8 @@ import (
 	"golang.org/x/tools/go/ssa"
 )
 
-// Compiles the code and then runs the EAR pointer analysis.
-func runCode(code string) (*earpointer.Partitions, error) {
+// Compiles the code and then runs the EAR pointer analysis with a specific context K.
+func runCodeWithContext(code string, contextK int) (*earpointer.Partitions, error) {
 	pkg, err := buildSSA(code)
 	if err != nil {
 		return nil, fmt.Errorf("compilation failed: %s :\n %s", err, code)
@@ -42,6 +43,7 @@ func runCode(code string) (*earpointer.Partitions, error) {
 	}
 	ssainput := buildssa.SSA{Pkg: pkg, SrcFuncs: srcFuncs}
 	pass := analysis.Pass{ResultOf: map[*analysis.Analyzer]interface{}{buildssa.Analyzer: &ssainput}}
+	earpointer.Analyzer.Flags.Set("contextK", strconv.Itoa(contextK))
 	// Run the analysis.
 	partitions, err := earpointer.Analyzer.Run(&pass)
 	if err != nil {
@@ -50,6 +52,11 @@ func runCode(code string) (*earpointer.Partitions, error) {
 	return partitions.(*earpointer.Partitions), nil
 }
 
+// A shortcut with context K=0.
+func runCodeK0(code string) (*earpointer.Partitions, error) {
+	return runCodeWithContext(code, 0)
+}
+
 // Concatenates the "members -> fields" map items into a string.
 func concat(m map[string]string) string {
 	var pstrs []string
@@ -74,7 +81,7 @@ func TestFieldAddr(t *testing.T) {
 		*t0 = b
 		return
 	*/
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -106,7 +113,7 @@ func TestFieldAddr2(t *testing.T) {
 			*t1 = c
 			return
 	*/
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -138,7 +145,7 @@ func TestField(t *testing.T) {
 			t3 = t2.x [#0]                  *int
 			return
 	*/
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -173,7 +180,7 @@ func TestEmbeddedField(t *testing.T) {
 			t3 = t2.t [#0]                  T1
 			return
 	*/
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -221,7 +228,7 @@ func TestStructCopy(t *testing.T) { // mainly for coverage test
 			*t1 = t5
 			return
 	*/
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -282,7 +289,7 @@ func TestEmbeddedFieldClone(t *testing.T) {
 			*t1 = t5
 			return
 	*/
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -319,7 +326,7 @@ func TestIndexAddr(t *testing.T) {
 			*t0 = t2
 			return
 	*/
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -351,7 +358,7 @@ func TestIndex(t *testing.T) {
 			t3 = t2[0:int]                    *int
 			return
 	*/
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -396,7 +403,7 @@ func TestPhi(t *testing.T) {
 			t2 = print(t1)                   ()
 		...
 	*/
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -427,7 +434,7 @@ func TestStore(t *testing.T) {
 			*t2 = t4
 			return
 	*/
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -464,7 +471,7 @@ func TestDereference(t *testing.T) {
 			*x = t1
 			return
 	*/
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -499,7 +506,7 @@ func TestMapAccess(t *testing.T) {
 			t2 = m[i]             *T
 			return
 	*/
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -531,7 +538,7 @@ func TestLookUp(t *testing.T) {
 			t3 = *t1                            *int
 			return
 	*/
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -551,7 +558,7 @@ func TestConvert(t *testing.T) {
 		_ = (string)(a)
 	}
 	`
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -578,7 +585,7 @@ func TestTypeAssert(t *testing.T) {
 			t3 = typeassert a.(*int)          *int
 			return
 	*/
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -608,7 +615,7 @@ func TestChangeInterfaceOrType(t *testing.T) {
 			t3 = changetype interface{} <- I (a)      interface{}
 			return t3
 	*/
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -628,7 +635,7 @@ func TestBinOp(t *testing.T) {
 		_ = a + b
 	}
 	`
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -663,7 +670,7 @@ func TestRange(t *testing.T) {
 		3:                                   rangeiter.done P:1 S:0
 			return
 	*/
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -680,7 +687,7 @@ func TestChannel(t *testing.T) {
 		_ = <-ch  // Receive from ch.
 	}
 	`
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -706,7 +713,7 @@ func TestSlice(t *testing.T) {
 			t2 = slice t1[2:int:]          []int
 			return
 	*/
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -722,7 +729,7 @@ func TestMakeInterface(t *testing.T) {
 		return s
 	}
 	`
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -745,7 +752,7 @@ func TestUnimplemented(t *testing.T) {
 		panic("test")
 	}
 	`
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -774,7 +781,7 @@ func TestSimpleCall(t *testing.T) {
 		0:                                             entry P:0 S:0
 			return a
 	*/
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -811,7 +818,7 @@ func TestCallWithTupleReturn(t *testing.T) {
 			t0 = *x                                      T
 			return a, t0
 	*/
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -854,7 +861,7 @@ func TestClosureCall(t *testing.T) {
 			t6 = extract t3 #2                           int
 			return
 	*/
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -893,7 +900,7 @@ func TestRecursiveCall(t *testing.T) {
 			t3 = f(t2, t1)                       *int
 			return t3
 	*/
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -932,21 +939,19 @@ func TestMethod(t *testing.T) {
 				t1 = (A).g(t0, x)                       *int
 				return t1
 	*/
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
 	// f.t1, f.x and g.x are unified due to calling g().
 	// Note g.a is a struct receiver that won't be unified.
 	want := concat(map[string]string{
-		"{f.a}":                           "--> f.t0",
-		"{f.t0}":                          "[]",
+		"{f.a}":                           "--> A:g.a",
 		"{*A:g.t2,*A:g.x,A:g.x,f.t1,f.x}": "[]",
-		"{A:g.a}":                         "[]",
-		"{A:g.t0}":                        "--> A:g.a",
 		"{*A:g.a}":                        "[]",
-		"{*A:g.t0}":                       "--> *A:g.t1",
-		"{*A:g.t1}":                       "[]",
+		"{A:g.t0}":                        "--> A:g.a",
+		"{*A:g.t0}":                       "--> A:g.a",
+		"{*A:g.t1,A:g.a,f.t0}":            "[]",
 	})
 	if diff := cmp.Diff(want, state.String()); diff != "" {
 		t.Errorf("diff (-want +got):\n%s", diff)
@@ -978,7 +983,7 @@ func TestGoDefer(t *testing.T) {
 		1:                                       recover P:0 S:0
 			return
 	*/
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -1014,7 +1019,7 @@ func TestCallBuiltin(t *testing.T) {
 			t5 = print(d1, t3)                         ()
 			return
 	*/
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -1059,7 +1064,7 @@ func TestVariadicCall(t *testing.T) {
 			t1 = *t0                                           *int
 			return t1
 	*/
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -1104,7 +1109,7 @@ func TestMethodInvoke(t *testing.T) {
 			t1 = (*T1).f(t0, x)                     ()
 			return
 	*/
-	state, err := runCode(code)
+	state, err := runCodeK0(code)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -1113,14 +1118,182 @@ func TestMethodInvoke(t *testing.T) {
 	// Note that in "**T2", the first "*" is the synthesized ValueOf operator,
 	// and "*T2" is the receiver type.
 	want := concat(map[string]string{
-		"{**T2:f.arg0}":         "[T1->*T2:f.t0]",
-		"{*T1:f.arg0}":          "[]",
-		"{*T1:f.x,*T2:f.x,g.x}": "[]",
-		"{*T2:f.arg0}":          "--> **T2:f.arg0",
-		"{*T2:f.t0}":            "[]",
-		"{g.t0}":                "[]",
-		"{g.x2}":                "--> *g.x2",
-		"{*g.x2}":               "[T1->g.t0]",
+		"{**T2:f.arg0}":              "[T1->*T1:f.arg0]",
+		"{*T1:f.arg0,*T2:f.t0,g.t0}": "[]",
+		"{*T1:f.x,*T2:f.x,g.x}":      "[]",
+		"{*T2:f.arg0}":               "--> **T2:f.arg0",
+		"{*g.x2}":                    "[T1->*T1:f.arg0]",
+		"{g.x2}":                     "--> *g.x2",
+	})
+	if got := state.String(); got != want {
+		t.Errorf("\n  got: %s\n want: %s", got, want)
+	}
+}
+
+func TestMethodCallContextSensitive(t *testing.T) {
+	code := `package p
+	type A struct {}
+	func f(x *int) *int {
+		var a A
+		return a.g(x)
+	}
+	func (a *A) g(x *int) *int {
+		return x
+	}
+	`
+	/*
+		func f(x *int) *int:
+		0:                                               entry P:0 S:0
+			t0 = new A (a)                               *A
+			t1 = (*A).g(t0, x)                           *int
+			return t1
+	*/
+	state, err := runCodeWithContext(code, 1)
+	if err != nil {
+		t.Fatal(err)
+	}
+	want := concat(map[string]string{
+		"{[(*A).g(t0, x)]*A:g.a,f.t0}":     "[]",
+		"{[(*A).g(t0, x)]*A:g.x,f.t1,f.x}": "[]",
+	})
+	if diff := cmp.Diff(want, state.String()); diff != "" {
+		t.Errorf("diff (-want +got):\n%s", diff)
+	}
+}
+
+func TestCallContextSensitive(t *testing.T) {
+	code := `package p
+	func f(x, y *int, i int) *int {
+		if i > 0 {
+			return g(x)
+		}
+		return g(y)
+	}
+	func g(a *int) *int {
+		return a
+	}
+	`
+	/*
+		func f(x *int, y *int, i int) *int:
+		0:                                           entry P:0 S:2
+			t0 = i > 0:int                           bool
+			if t0 goto 1 else 2
+		1:                                           if.then P:1 S:0
+			t1 = g(x)                                *int
+			return t1
+		2:                                           if.done P:1 S:0
+			t2 = g(y)                                *int
+			return t2
+	*/
+	state, err := runCodeWithContext(code, 1)
+	if err != nil {
+		t.Fatal(err)
+	}
+	// g() is called at two call sites. f.x and f.y won't be unified.
+	want := concat(map[string]string{
+		"{[g(x)]g.a,f.t1,f.x}": "[]",
+		"{[g(y)]g.a,f.t2,f.y}": "[]",
+	})
+	if diff := cmp.Diff(want, state.String()); diff != "" {
+		t.Errorf("diff (-want +got):\n%s", diff)
+	}
+}
+
+func TestEmbeddedCallContextSensitive(t *testing.T) {
+	code := `package p
+	func f(x, y *int) {
+		g(x)
+		g(y)
+	}
+	func g(a *int) {
+		h(a)
+	}
+	func h(b *int) { }
+	`
+	/*
+		func f(x *int, y *int):
+		0:                                           entry P:0 S:0
+			t0 = g(x)                                ()
+			t1 = g(y)                                ()
+			return
+
+		func g(a *int):
+		0:                                           entry P:0 S:0
+			t0 = h(a)                                ()
+			return
+	*/
+	state, err := runCodeWithContext(code, 1)
+	if err != nil {
+		t.Fatal(err)
+	}
+	// When K=1, h.b will be unified with g.a regardless of the context of g.a.
+	want := concat(map[string]string{
+		"{[g(x)]g.a,[g(y)]g.a,[h(a)]h.b,f.x,f.y}": "[]",
+	})
+	if diff := cmp.Diff(want, state.String()); diff != "" {
+		t.Errorf("diff (-want +got):\n%s", diff)
+	}
+
+	state, _ = runCodeWithContext(code, 2)
+	// When K=2, g(x)'s calling h(a) is distinguished from g(y)'s calling h(a),
+	// i.e. h(a) is called in two different contexts: [g(x); h(a)] and [g(y); h(a)].
+	want = concat(map[string]string{
+		"{[g(x); h(a)]h.b,[g(x)]g.a,f.x}": "[]",
+		"{[g(y); h(a)]h.b,[g(y)]g.a,f.y}": "[]",
+	})
+	//	g() is called at two call sites.
+	if diff := cmp.Diff(want, state.String()); diff != "" {
+		t.Errorf("diff (-want +got):\n%s", diff)
+	}
+}
+
+func TestCallReturnContextSensitive(t *testing.T) {
+	code := `package p
+	func f(x, y, z *int) *int {
+		g(x)
+		h(y)
+		return g(z)
+	}
+	func g(a *int) *int {
+		return h(a)
+	}
+	func h(b *int) *int {
+		return b
+	}
+	`
+	/*
+		func f(x *int, y *int, z *int) *int:
+		0:                                                                entry P:0 S:0
+			t0 = g(x)                                                          *int
+			t1 = h(y)                                                          *int
+			t2 = g(z)                                                          *int
+			return t2
+
+		func g(a *int) *int:
+		0:                                                                entry P:0 S:0
+			t0 = h(a)                                                          *int
+			return t0
+	*/
+	state, err := runCodeWithContext(code, 1)
+	if err != nil {
+		t.Fatal(err)
+	}
+	// When K=1, h.b will be unified with g.a regardless of the context of g.a,
+	// which also resulting in the unfications of f.x, f.y and f.z.
+	want := concat(map[string]string{
+		"{[g(x)]g.a,[g(x)]g.t0,[g(z)]g.a,[g(z)]g.t0,[h(a)]h.b,f.t0,f.t2,f.x,f.z}": "[]",
+		"{[h(y)]h.b,f.t1,f.y}": "[]",
+	})
+	if diff := cmp.Diff(want, state.String()); diff != "" {
+		t.Errorf("diff (-want +got):\n%s", diff)
+	}
+
+	state, _ = runCodeWithContext(code, 2)
+	// When K=2, less unifications are performed. Now the two calls to h(b) are separate.
+	want = concat(map[string]string{
+		"{[g(x); h(a)]h.b,[g(x)]g.a,[g(x)]g.t0,f.t0,f.x}": "[]",
+		"{[g(z); h(a)]h.b,[g(z)]g.a,[g(z)]g.t0,f.t2,f.z}": "[]",
+		"{[h(y)]h.b,f.t1,f.y}":                            "[]",
 	})
 	if diff := cmp.Diff(want, state.String()); diff != "" {
 		t.Errorf("diff (-want +got):\n%s", diff)