Merge pull request #3 from gcpug/fix-far-block

Fix issue#1

README.md

@@ -29,7 +29,6 @@ $ zagane github.com/gcpug/spshovel/...
 ### unstopiter
 
 `unstopiter` finds spanner.RowIterator which is not calling Stop method such as below code.
-But it currently has some bugs ([#1](https://github.com/gcpug/zagane/issues/1)).
 
 ```go
 _, _ = client.Single().Query(ctx, stmt).Next()

passes/unstopiter/testdata/src/a/a.go

@@ -12,3 +12,25 @@ func f(ctx context.Context, client *spanner.Client) {
 	client.Single().Query(ctx, stmt).Stop()        // OK
 	defer client.Single().Query(ctx, stmt).Stop()  // OK
 }
+
+func g(ctx context.Context, client *spanner.Client) {
+	stmt := spanner.Statement{SQL: `SELECT 1`}
+	iter1 := client.Single().Query(ctx, stmt) // want "iterator must be stop"
+	if iter1 == nil {
+		iter1.Stop()
+	}
+
+	iter2 := client.Single().Query(ctx, stmt) // OK
+	if iter2 == nil {
+		iter2.Stop()
+	}
+	iter2.Stop()
+}
+
+func h(ctx context.Context, client *spanner.Client) {
+	stmt := spanner.Statement{SQL: `SELECT 1`}
+	iter := client.Single().Query(ctx, stmt) // want "iterator must be stop"
+	if iter == nil {
+		defer iter.Stop()
+	}
+}

passes/unstopiter/unstopiter.go

@@ -91,10 +91,8 @@ func (r *runner) unstop(b *ssa.BasicBlock, i int) bool {
 		return false
 	}
 
-	for _, nb := range b.Succs {
-		if r.callStopIn(nb.Instrs, call) {
-			return false
-		}
+	if r.callStopInSuccs(b, call, map[*ssa.BasicBlock]bool{}) {
+		return false
 	}
 
 	return true
@@ -115,3 +113,22 @@ func (r *runner) callStopIn(instrs []ssa.Instruction, call *ssa.Call) bool {
 	}
 	return false
 }
+
+func (r *runner) callStopInSuccs(b *ssa.BasicBlock, call *ssa.Call, done map[*ssa.BasicBlock]bool) bool {
+	if done[b] {
+		return false
+	}
+	done[b] = true
+
+	if len(b.Succs) == 0 {
+		return false
+	}
+
+	for _, s := range b.Succs {
+		if !r.callStopIn(s.Instrs, call) && !r.callStopInSuccs(s, call, done) {
+			return false
+		}
+	}
+
+	return true
+}