HasArgInCallExpr: support specify scope

ZhengHe-MD · Nov 5, 2019 · 3a27ff9 · 3a27ff9
1 parent 479635e
commit 3a27ff9
Show file tree

Hide file tree

Showing 4 changed files with 256 additions and 154 deletions.
diff --git a/func_body_test.go b/func_body_test.go
@@ -70,10 +70,10 @@ func TestHasStmtInsideFuncBody(t *testing.T) {
 			}, false},
 			{
 				&dst.AssignStmt{
-					Lhs:  []dst.Expr{dst.NewIdent("c")},
-					Tok:  token.DEFINE,
-					Rhs:  []dst.Expr{&dst.BasicLit{
-						Kind: token.INT,
+					Lhs: []dst.Expr{dst.NewIdent("c")},
+					Tok: token.DEFINE,
+					Rhs: []dst.Expr{&dst.BasicLit{
+						Kind:  token.INT,
 						Value: "1",
 					}},
 					Decs: dst.AssignStmtDecorations{},
@@ -106,10 +106,10 @@ func TestHasStmtInsideFuncBody(t *testing.T) {
 
 		df, _ := ParseSrcFileFromBytes([]byte(src))
 		stmt := &dst.AssignStmt{
-			Lhs:  []dst.Expr{dst.NewIdent("c")},
-			Tok:  token.DEFINE,
-			Rhs:  []dst.Expr{&dst.BasicLit{
-				Kind: token.INT,
+			Lhs: []dst.Expr{dst.NewIdent("c")},
+			Tok: token.DEFINE,
+			Rhs: []dst.Expr{&dst.BasicLit{
+				Kind:  token.INT,
 				Value: "3",
 			}},
 			Decs: dst.AssignStmtDecorations{},
@@ -224,10 +224,10 @@ func TestDeleteStmtFromFuncBody(t *testing.T) {
 
 		df, _ := ParseSrcFileFromBytes([]byte(src))
 		stmt := &dst.AssignStmt{
-			Lhs:  []dst.Expr{dst.NewIdent("c")},
-			Tok:  token.DEFINE,
-			Rhs:  []dst.Expr{&dst.BasicLit{
-				Kind: token.INT,
+			Lhs: []dst.Expr{dst.NewIdent("c")},
+			Tok: token.DEFINE,
+			Rhs: []dst.Expr{&dst.BasicLit{
+				Kind:  token.INT,
 				Value: "1",
 			}},
 			Decs: dst.AssignStmtDecorations{},
@@ -306,8 +306,8 @@ func TestDeleteSelectorExprFromFuncBody(t *testing.T) {
 		`
 
 		printSelector := &dst.SelectorExpr{
-			X:    dst.NewIdent("fmt"),
-			Sel:  dst.NewIdent("Printf"),
+			X:   dst.NewIdent("fmt"),
+			Sel: dst.NewIdent("Printf"),
 		}
 
 		df, _ := ParseSrcFileFromBytes([]byte(src))
@@ -352,8 +352,8 @@ func TestDeleteSelectorExprFromFuncBody(t *testing.T) {
 		`
 
 		printSelector := &dst.SelectorExpr{
-			X:    dst.NewIdent("log"),
-			Sel:  dst.NewIdent("Infof"),
+			X:   dst.NewIdent("log"),
+			Sel: dst.NewIdent("Infof"),
 		}
 
 		df, _ := ParseSrcFileFromBytes([]byte(src))
@@ -530,7 +530,7 @@ func TestAddStmtToFuncBodyRelativeTo(t *testing.T) {
 
 		cases := []struct {
 			direction    int
-			stmt 		 dst.Stmt
+			stmt         dst.Stmt
 			refStmt      dst.Stmt
 			expectedBody string
 		}{

diff --git a/func_call.go b/func_call.go
@@ -6,12 +6,16 @@ import (
 )
 
 // HasArgInCallExpr checks if the arguments of the function call has given arg
-func HasArgInCallExpr(df *dst.File, funcName string, arg dst.Expr) (ret bool) {
+func HasArgInCallExpr(df *dst.File, scope Scope, funcName string, arg dst.Expr) (ret bool) {
 	pre := func(c *dstutil.Cursor) bool {
 		node := c.Node()
 
 		switch node.(type) {
 		case *dst.CallExpr:
+			if !scope.IsInScope() {
+				return true
+			}
+
 			var found bool
 			nn := node.(*dst.CallExpr)
 			if ie, ok := nn.Fun.(*dst.Ident); ok && ie.Name == funcName {
@@ -30,11 +34,18 @@ func HasArgInCallExpr(df *dst.File, funcName string, arg dst.Expr) (ret bool) {
 				}
 				return false
 			}
+		default:
+			scope.TryEnterScope(node)
 		}
 		return true
 	}
 
-	dstutil.Apply(df, pre, nil)
+	post := func(c *dstutil.Cursor) bool {
+		scope.TryLeaveScope(c.Node())
+		return true
+	}
+
+	dstutil.Apply(df, pre, post)
 	return
 }
 
@@ -149,4 +160,4 @@ func SetMethodOnReceiver(df *dst.File, receiver, oldMethod, newMethod string) (m
 
 	dstutil.Apply(df, pre, nil)
 	return
-}
+}