Initial refactor of source package

Major source cleanup, still not finished. Removes gross control flow.

proc/stack.go

@@ -1,8 +1,6 @@
 package proc
 
-import (
+import "encoding/binary"
-	"encoding/binary"
-)
 
 // Takes an offset from RSP and returns the address of the
 // instruction the currect function is going to return to.
@@ -78,7 +76,6 @@ func (dbp *Process) stacktrace(pc, sp uint64, depth int) ([]Location, error) {
 		if fn != nil && fn.Name == "runtime.goexit" {
 			break
 		}
-
 	}
 	return locations, nil
 }

source/source.go

@@ -26,11 +26,11 @@ func (n NoNodeError) Error() string {
 }
 
 // Returns the first node at the given file:line.
-func (s *Searcher) FirstNodeAt(fname string, line int) (ast.Node, error) {
+func (s *Searcher) FirstNodeAt(fname string, line int) (*ast.File, ast.Node, error) {
 	var node ast.Node
 	f, err := s.parse(fname)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	ast.Inspect(f, func(n ast.Node) bool {
 		if n == nil {
@@ -44,9 +44,9 @@ func (s *Searcher) FirstNodeAt(fname string, line int) (ast.Node, error) {
 		return true
 	})
 	if node == nil {
-		return nil, NoNodeError{f: fname, l: line}
+		return nil, nil, NoNodeError{f: fname, l: line}
 	}
-	return node, nil
+	return f, node, nil
 }
 
 type Done string
@@ -58,245 +58,43 @@ func (d Done) Error() string {
 // Returns all possible lines that could be executed after the given file:line,
 // within the same source file.
 func (s *Searcher) NextLines(fname string, line int) (lines []int, err error) {
-	var found bool
+	parsedFile, n, err := s.FirstNodeAt(fname, line)
-	n, err := s.FirstNodeAt(fname, line)
 	if err != nil {
 		return lines, nil
 	}
-	defer func() {
-		if e := recover(); e != nil {
-			e = e.(Done)
-			nl := make([]int, 0, len(lines))
-			fnd := make(map[int]bool)
-			for _, l := range lines {
-				if _, ok := fnd[l]; !ok {
-					fnd[l] = true
-					nl = append(nl, l)
-				}
-			}
-			lines = nl
-		}
-	}()
 
 	switch x := n.(type) {
-	// Check if we are at an 'if' statement.
+	// Follow if statements
-	//
-	// If we are at an 'if' statement, employ the following algorithm:
-	//    * Follow all 'else if' statements, appending their line number
-	//    * Follow any 'else' statement if it exists, appending the line
-	//      number of the statement following the 'else'.
-	//    * If there is no 'else' statement, append line of first statement
-	//      following the entire 'if' block.
 	case *ast.IfStmt:
-		var rbrace int
+		lines = removeDuplicateLines(s.parseIfStmtBlock(x, parsedFile))
-		p := x.Body.List[0].Pos()
+		return
-		pos := s.fileset.Position(p)
-		lines = append(lines, pos.Line)
-
-		if x.Else == nil {
-			// Grab first line after entire 'if' block
-			rbrace = s.fileset.Position(x.Body.Rbrace).Line
-			f, err := s.parse(fname)
-			if err != nil {
-				return nil, err
-			}
-			ast.Inspect(f, func(n ast.Node) bool {
-				if n == nil {
-					return true
-				}
-				pos := s.fileset.Position(n.Pos())
-				if rbrace < pos.Line {
-					lines = append(lines, pos.Line)
-					panic(Done("done"))
-				}
-				return true
-			})
-		} else {
-			// Follow any 'else' statements
-			for {
-				if stmt, ok := x.Else.(*ast.IfStmt); ok {
-					pos := s.fileset.Position(stmt.Pos())
-					lines = append(lines, pos.Line)
-					x = stmt
-					continue
-				}
-				if x.Else != nil {
-					pos := s.fileset.Position(x.Else.Pos())
-					ast.Inspect(x, func(n ast.Node) bool {
-						if found {
-							panic(Done("done"))
-						}
-						if n == nil {
-							return false
-						}
-						p := s.fileset.Position(n.Pos())
-						if pos.Line < p.Line {
-							lines = append(lines, p.Line)
-							found = true
-							return false
-						}
-						return true
-					})
-				}
-			}
-		}
 
 	// Follow case statements.
-	//
 	// Append line for first statement following each 'case' condition.
 	case *ast.SwitchStmt:
+		var switchEnd int
 		ast.Inspect(x, func(n ast.Node) bool {
 			if stmt, ok := n.(*ast.SwitchStmt); ok {
-				ast.Inspect(stmt, func(n ast.Node) bool {
+				switchEnd = s.fileset.Position(stmt.End()).Line
-					if stmt, ok := n.(*ast.CaseClause); ok {
+				return true
-						p := stmt.Body[0].Pos()
+			}
-						pos := s.fileset.Position(p)
+			if switchEnd < s.fileset.Position(x.Pos()).Line {
-						lines = append(lines, pos.Line)
+				return false
-						return false
+			}
-					}
+			if stmt, ok := n.(*ast.CaseClause); ok {
-					return true
+				p := stmt.Body[0].Pos()
-				})
+				pos := s.fileset.Position(p)
-				panic(Done("done"))
+				lines = append(lines, pos.Line)
+				return false
 			}
 			return true
 		})
+		lines = removeDuplicateLines(lines)
+		return
 	// Default case - find next source line.
-	//
-	// We are not at a branch, employ the following algorithm:
-	//    * Traverse tree, storing any loop as a parent
-	//    * Find next source line after the given line
-	//    * Check and see if we've passed the scope of any parent we've
-	//      stored. If so, pop them off the stack. The last parent that
-	//      is left get's appending to our list of lines since we could
-	//      end up at the top of the loop again.
 	default:
-		var (
+		lines = removeDuplicateLines(s.parseDefaultBlock(x, parsedFile, line))
-			parents              []*ast.BlockStmt
+		return
-			parentBlockBeginLine int
-			deferEndLine         int
-		)
-		f, err := s.parse(fname)
-		if err != nil {
-			return nil, err
-		}
-		ast.Inspect(f, func(n ast.Node) bool {
-			if found {
-				panic(Done("done"))
-			}
-			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())
-				parentBlockBeginLine = pos.Line
-			}
-
-			if _, ok := n.(*ast.GenDecl); ok {
-				return true
-			}
-
-			if dn, ok := n.(*ast.DeferStmt); ok && line < pos.Line {
-				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())
-				if beginpos.Line < endpos.Line {
-					return true
-				}
-			}
-
-			// Check to see if we've found the "next" line.
-			if line < pos.Line {
-				if _, ok := n.(*ast.BlockStmt); ok {
-					return true
-				}
-				var (
-					parent        *ast.BlockStmt
-					parentEndLine int
-				)
-				for len(parents) > 0 {
-					parent = parents[len(parents)-1]
-
-					// Grab the line number of the right brace of the parent block.
-					parentEndLine = s.fileset.Position(parent.Rbrace).Line
-
-					// Check to see if we're still within the parents block.
-					// If we are, we're done and that is our parent.
-					if parentEndLine > line {
-						parentBlockBeginLine = s.fileset.Position(parent.Pos()).Line
-						break
-					}
-					// If we weren't, and there is only 1 parent, we no longer have one.
-					if len(parents) == 1 {
-						parent = nil
-						break
-					}
-					// Remove that parent from the stack.
-					parents = parents[0 : len(parents)-1]
-				}
-				if parent != nil {
-					var (
-						endfound   bool
-						beginFound bool
-						beginLine  int
-					)
-
-					ast.Inspect(f, func(n ast.Node) bool {
-						if n == nil || endfound {
-							return false
-						}
-						if _, ok := n.(*ast.BlockStmt); ok {
-							return true
-						}
-						pos := s.fileset.Position(n.Pos())
-						if parentBlockBeginLine < pos.Line && !beginFound {
-							beginFound = true
-							beginLine = pos.Line
-							return true
-						}
-						if parentEndLine < pos.Line {
-							if _, ok := n.(*ast.FuncDecl); !ok {
-								lines = append(lines, beginLine, pos.Line)
-							}
-							endfound = true
-							return false
-						}
-						return true
-					})
-					lines = append(lines, parentBlockBeginLine)
-				}
-				switch n.(type) {
-				case *ast.BranchStmt, *ast.FuncDecl:
-				default:
-					lines = append(lines, pos.Line)
-				}
-				found = true
-				return false
-			}
-			return true
-		})
-		if len(lines) == 0 && 0 < len(parents) {
-			parent := parents[len(parents)-1]
-			lbrace := s.fileset.Position(parent.Lbrace).Line
-			pos := s.fileset.Position(parent.List[0].Pos())
-			lines = append(lines, lbrace, pos.Line)
-		}
 	}
 	return lines, nil
 }
@@ -313,3 +111,196 @@ func (s *Searcher) parse(fname string) (*ast.File, error) {
 	s.visited[fname] = f
 	return f, nil
 }
+
+func (s *Searcher) nextLineAfter(parsedFile *ast.File, line int) (nextLine int) {
+	var done bool
+	ast.Inspect(parsedFile, func(n ast.Node) bool {
+		if done || n == nil {
+			return false
+		}
+		pos := s.fileset.Position(n.Pos())
+		if line < pos.Line {
+			nextLine = pos.Line
+			done = true
+			return false
+		}
+		return true
+	})
+	return
+}
+
+// If we are at an 'if' statement, employ the following algorithm:
+//    * Follow all 'else if' statements, appending their line number
+//    * Follow any 'else' statement if it exists, appending the line
+//      number of the statement following the 'else'.
+//    * If there is no 'else' statement, append line of first statement
+//      following the entire 'if' block.
+func (s *Searcher) parseIfStmtBlock(ifRoot *ast.IfStmt, parsedFile *ast.File) []int {
+	var (
+		rbrace     int
+		ifStmtLine = s.fileset.Position(ifRoot.Body.List[0].Pos()).Line
+		lines      = []int{ifStmtLine}
+	)
+
+	for {
+		if ifRoot.Else == nil {
+			// Grab first line after entire 'if' block
+			rbrace = s.fileset.Position(ifRoot.Body.Rbrace).Line
+			return append(lines, s.nextLineAfter(parsedFile, rbrace))
+		}
+		// Continue following 'else if' branches.
+		if elseStmt, ok := ifRoot.Else.(*ast.IfStmt); ok {
+			lines = append(lines, s.fileset.Position(elseStmt.Pos()).Line)
+			ifRoot = elseStmt
+			continue
+		}
+		// Grab next line after final 'else'.
+		pos := s.fileset.Position(ifRoot.Else.Pos())
+		return append(lines, s.nextLineAfter(parsedFile, pos.Line))
+	}
+}
+
+// We are not at a branch, employ the following algorithm:
+//    * Traverse tree, storing any loop as a parent
+//    * Find next source line after the given line
+//    * Check and see if we've passed the scope of any parent we've
+//      stored. If so, pop them off the stack. The last parent that
+//      is left get's appending to our list of lines since we could
+//      end up at the top of the loop again.
+func (s *Searcher) parseDefaultBlock(ifRoot ast.Node, parsedFile *ast.File, line int) []int {
+	var (
+		found                bool
+		lines                []int
+		parents              []*ast.BlockStmt
+		parentBlockBeginLine int
+		deferEndLine         int
+	)
+	ast.Inspect(parsedFile, func(n ast.Node) bool {
+		if found || n == nil {
+			return false
+		}
+
+		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())
+			parentBlockBeginLine = pos.Line
+		}
+
+		if _, ok := n.(*ast.GenDecl); ok {
+			return true
+		}
+
+		if dn, ok := n.(*ast.DeferStmt); ok && line < pos.Line {
+			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())
+			if beginpos.Line < endpos.Line {
+				return true
+			}
+		}
+
+		// Check to see if we've found the "next" line.
+		if line < pos.Line {
+			if _, ok := n.(*ast.BlockStmt); ok {
+				return true
+			}
+			var (
+				parent        *ast.BlockStmt
+				parentEndLine int
+			)
+			for len(parents) > 0 {
+				parent = parents[len(parents)-1]
+
+				// Grab the line number of the right brace of the parent block.
+				parentEndLine = s.fileset.Position(parent.Rbrace).Line
+
+				// Check to see if we're still within the parents block.
+				// If we are, we're done and that is our parent.
+				if parentEndLine > line {
+					parentBlockBeginLine = s.fileset.Position(parent.Pos()).Line
+					break
+				}
+				// If we weren't, and there is only 1 parent, we no longer have one.
+				if len(parents) == 1 {
+					parent = nil
+					break
+				}
+				// Remove that parent from the stack.
+				parents = parents[0 : len(parents)-1]
+			}
+			if parent != nil {
+				var (
+					endfound   bool
+					beginFound bool
+					beginLine  int
+				)
+
+				ast.Inspect(parsedFile, func(n ast.Node) bool {
+					if n == nil || endfound {
+						return false
+					}
+					if _, ok := n.(*ast.BlockStmt); ok {
+						return true
+					}
+					pos := s.fileset.Position(n.Pos())
+					if parentBlockBeginLine < pos.Line && !beginFound {
+						beginFound = true
+						beginLine = pos.Line
+						return true
+					}
+					if parentEndLine < pos.Line {
+						if _, ok := n.(*ast.FuncDecl); !ok {
+							lines = append(lines, beginLine, pos.Line)
+						}
+						endfound = true
+						return false
+					}
+					return true
+				})
+				lines = append(lines, parentBlockBeginLine)
+			}
+			switch n.(type) {
+			case *ast.BranchStmt, *ast.FuncDecl:
+			default:
+				lines = append(lines, pos.Line)
+			}
+			found = true
+			return false
+		}
+		return true
+	})
+	if len(lines) == 0 && 0 < len(parents) {
+		parent := parents[len(parents)-1]
+		lbrace := s.fileset.Position(parent.Lbrace).Line
+		pos := s.fileset.Position(parent.List[0].Pos())
+		lines = append(lines, lbrace, pos.Line)
+	}
+	return lines
+}
+
+func removeDuplicateLines(lines []int) []int {
+	nl := make([]int, 0, len(lines))
+	fnd := make(map[int]bool)
+	for _, l := range lines {
+		if _, ok := fnd[l]; !ok {
+			fnd[l] = true
+			nl = append(nl, l)
+		}
+	}
+	return nl
+}

source/source_test.go

@@ -12,7 +12,7 @@ func TestTokenAtLine(t *testing.T) {
 		tf, _ = filepath.Abs("../_fixtures/testvisitorprog.go")
 		v     = New()
 	)
-	n, err := v.FirstNodeAt(tf, 8)
+	_, n, err := v.FirstNodeAt(tf, 8)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -52,9 +52,9 @@ func TestNextLines(t *testing.T) {
 			fmt.Println(lines)
 			t.Fatalf("did not get correct number of lines back expected %d got %d for test case %d", len(c.nextlines), len(lines), i+1)
 		}
-		for i, l := range lines {
+		for j, l := range lines {
-			if l != c.nextlines[i] {
+			if l != c.nextlines[j] {
-				t.Fatalf("expected index %d to be %d got %d", i, c.nextlines[i], l)
+				t.Fatalf("expected index %d to be %d got %d for case %d", j, c.nextlines[j], l, i+1)
 			}
 		}
 	}