Handle defer blocks when next'ing

_fixtures/testvisitorprog.go

@@ -63,3 +63,10 @@ func decltest() {
 	var baz = 9
 	fmt.Println(foo, baz)
 }
+
+func defertest() {
+	defer func() {
+		fmt.Println("this is a useless defer")
+	}()
+	fmt.Println("I should get here")
+}

source/source.go

@@ -162,6 +162,7 @@ func (s *Searcher) NextLines(fname string, line int) (lines []int, err error) {
 		var (
 			parents              []*ast.BlockStmt
 			parentBlockBeginLine int
+			deferEndLine         int
 		)
 		f, err := s.parse(fname)
 		if err != nil {
@@ -174,6 +175,17 @@ func (s *Searcher) NextLines(fname string, line int) (lines []int, err error) {
 			if n == nil {
 				return true
 			}
+
+			pos := s.fileset.Position(n.Pos())
+			if line < pos.Line && deferEndLine != 0 {
+				p := s.fileset.Position(n.Pos())
+				if deferEndLine < p.Line {
+					found = true
+					lines = append(lines, p.Line)
+					return false
+				}
+			}
+
 			if stmt, ok := n.(*ast.ForStmt); ok {
 				parents = append(parents, stmt.Body)
 				pos := s.fileset.Position(stmt.Pos())
@@ -184,6 +196,13 @@ func (s *Searcher) NextLines(fname string, line int) (lines []int, err error) {
 				return true
 			}
 
+			if dn, ok := n.(*ast.DeferStmt); ok {
+				fmt.Println("defer")
+				endpos := s.fileset.Position(dn.End())
+				deferEndLine = endpos.Line
+				return false
+			}
+
 			if st, ok := n.(*ast.DeclStmt); ok {
 				beginpos := s.fileset.Position(st.Pos())
 				endpos := s.fileset.Position(st.End())
@@ -193,7 +212,6 @@ func (s *Searcher) NextLines(fname string, line int) (lines []int, err error) {
 			}
 
 			// Check to see if we've found the "next" line.
-			pos := s.fileset.Position(n.Pos())
 			if line < pos.Line {
 				if _, ok := n.(*ast.BlockStmt); ok {
 					return true

source/source_test.go

@@ -40,6 +40,7 @@ func TestNextLines(t *testing.T) {
 		{57, []int{55}},
 		{30, []int{32}},
 		{62, []int{63}},
+		{67, []int{71}},
 	}
 	for i, c := range cases {
 		lines, err := v.NextLines(tf, c.line)