Search code examples
goabstract-syntax-tree

Use go ast to get surrounding function name from position


What's the best way to get the surrounding function name from a file position via the source code alone (no runtime checks)?

For example say I have some code:

func MyFunc() {
    doSomething() // package/file.go:215:15
}

And I have the position of doSomething, at package/file.go:215:15, is there a way to easily get MyFunc?


Solution

  • Questions containing the phrase "the best" are always hard to answer since this very much depends on your requirements, which you have not named.

    I thought of 2 ways to solve it this, both with their own pro and con list:


    Quick and dirty

    The quick and dirty approach would be to just loop over every line of code and look for the most recent occurrence of a function deceleration before we hit the line we are interested in. Which should be the function containing our line.

    package main
    
    import (
        "bufio"
        "fmt"
        "os"
        "regexp"
        "strconv"
        "strings"
    )
    
    func main() {
        if len(os.Args) < 2 {
            usage()
        }
    
        loc := strings.Split(os.Args[1], ":")
        if len(loc) != 2 {
            usage()
        }
    
        filePath := loc[0]
        lineStr := loc[1]
        targetLine, err := strconv.Atoi(lineStr)
        if err != nil {
            fmt.Println(err.Error())
            usage()
        }
    
        f, err := os.Open(filePath)
        if err != nil {
            fmt.Println(err.Error())
            usage()
        }
        defer f.Close()
    
        lineScanner := bufio.NewScanner(f)
        line := 0
        var lastFunc string
        for lineScanner.Scan() {
            m := funcName.FindStringSubmatch(lineScanner.Text())
            if len(m) > 0 {
                lastFunc = m[1]
            }
    
            if line == targetLine {
                fmt.Println(lastFunc)
                return
            }
    
            line++
        }
    }
    
    func usage() {
        fmt.Fprintf(os.Stderr, "Usage: %s {file:line}\n", os.Args[0])
        os.Exit(1)
    }
    
    // Look for a func followed by anything and ending at a `(` or ` `(space).
    var funcName = regexp.MustCompile(`func ([^ (]+)`)
    

    Pros:

    • Quick (relative to building a AST)
    • Easy to understand (no need to know how ASTs and Depth-First-Search works)
    • Works when go code contains errors on which the parser fails (which can be what you want for some types of linters or tools).

    Cons:

    • Assumptions about the code, we assume that ( and are the only chars that end a func declaration, and we assume the code is formatted so that there are never 2 func declerations on 1 line.

    Parsing, AST, Depth First Search

    The more "spec correct" approach is to parse the go code using the "official" go parser which the go compiler also uses. This results in an AST(Abstract Syntax Tree) which we can traverse using the DFS(Depth First Search) algorithm to find the most specific AST node which contains our location, along the way also finding the latest function decleration.

    package main
    
    import (
        "fmt"
        "go/ast"
        "go/parser"
        "go/token"
        "os"
        "regexp"
        "strconv"
        "strings"
    )
    
    func main() {
        if len(os.Args) < 2 {
            usage()
        }
    
        var pos token.Position
    
        loc := strings.Split(os.Args[1], ":")
        if len(loc) >= 2 {
            pos.Filename = loc[0]
            line, err := strconv.Atoi(loc[1])
            if err != nil {
                fmt.Println(err.Error())
                usage()
            }
            pos.Line = line
        } else {
            usage()
        }
    
        if len(loc) >= 3 {
            col, err := strconv.Atoi(loc[2])
            if err != nil {
                fmt.Println(err.Error())
                usage()
            }
            pos.Column = col
        }
    
        file, err := os.Open(pos.Filename)
        if err != nil {
            fmt.Println(err.Error())
            usage()
        }
    
        fset := token.NewFileSet()
        f, err := parser.ParseFile(fset, "", file, 0)
        if err != nil {
            fmt.Println(err.Error())
            usage()
        }
    
        var lastFunc *ast.FuncDecl
        ast.Inspect(f, func(n ast.Node) bool {
            if n == nil {
                return false
            }
    
            // Store the most specific function declaration
            if funcDecl, ok := n.(*ast.FuncDecl); ok {
                lastFunc = funcDecl
            }
    
            start := fset.Position(n.Pos())
            end := fset.Position(n.End())
    
            // Don't traverse nodes which don't contain the target line
            if start.Line > pos.Line || end.Line < pos.Line {
                return false
            }
    
            // If node starts and stops on the same line
            if start.Line == pos.Line && end.Line == pos.Line {
                // Don't traverse nodes which don't contain the target column
                if start.Column > pos.Column || end.Column < pos.Column {
                    return false
                }
            }
    
            // Note, the very last node to be traversed is our target node
    
            return true
        })
    
        if lastFunc != nil {
            fmt.Println(lastFunc.Name.String())
        }
    }
    
    func usage() {
        fmt.Fprintf(os.Stderr, "Usage: %s {file:line:column}\n", os.Args[0])
        os.Exit(1)
    }
    

    Pros:

    • Spec correct, the parser handles all edge cases, unicode chars, ect.

    Cons:

    • Slower since the parser has do to a lot of work, building the AST tree, ect.
    • Requires more knowledge to understand / work with
    • If the go file contains errors, the parser will error out on those, thus not giving a result