解析片段并找到所有顶级定义

Question

The service I'm writing receive code snippets and process them, the snippets could either be complement program or a fragment, if it is a fragment, I need to add the enclosing main function. For example, the snippet:

var v int
v = 3
fmt.Println(v)

should be classified as a fragment, and add main to it:

func main() {
    var v int
    v = 3
    fmt.Println(v)
}

If the snippet is:

package main
import "fmt"
func main() {
    fmt.Println("hello")
}

then no modification should be done.

The way I'm doing now is run the go parser against the snippet:

var fset *token.FileSet
file, err := parser.ParseFile(fset, "stdin", code, 0)
if err != nil {
        // add function
        code = fmt.Sprintf("func main() {
%s
}", code)
        // ...
}

This works for the 1st snippet above, however it fails if the fragment has a main function after some other declarations, e.g.

type S struct {
    a int
}
func main() {
    fmt.Println("foo")
}

I also try to look into the file returned by ParseFile, check the Decls, but it looks it will stop parsing after the 1st error, so Decls is nil in this case. So my question is is there a robust way to handle this?

PS. The inclusion of package clause and the required imports are not relevant because I'm feeding the processed code into golang.org/x/tools/imports anyway.

Answer 1

The dumbest thing that might work is after reading in the file (to a buffer most likely), is to do a string search for func main(){.

If it isn't formatted already, you might need to change that to a regex with whitespaces, but it should be pretty straight-forward.