1 files changed, 43 insertions, 17 deletions
diff --git a/gocomplete/pkgs.go b/gocomplete/pkgs.go
index b223ea9..356a0fa 100644
--- a/gocomplete/pkgs.go
+++ b/gocomplete/pkgs.go
@@ -4,37 +4,59 @@ import (
 	"bytes"
 	"encoding/json"
 	"os/exec"
+	"path/filepath"
+	"regexp"
 	"strings"
 
 	"github.com/posener/complete"
 )
 
-const goListFormat = `'{"name": "{{.Name}}", "dir": "{{.Dir}}"}'`
+const goListFormat = `{"Name": "{{.Name}}", "Path": "{{.Dir}}", "FilesString": "{{.GoFiles}}"}`
 
-func predictPackages(packageName string) complete.Predictor {
-	return complete.PredictFunc(func(a complete.Args) (prediction []string) {
-		dir := a.Directory()
-		dir = strings.TrimRight(dir, "/.") + "/..."
+// regexp matches a main function
+var reMainFunc = regexp.MustCompile("^main$")
 
-		pkgs := listPackages(dir)
+func predictPackages(a complete.Args) (prediction []string) {
+	dir := a.Directory()
+	pkgs := listPackages(dir)
 
-		files := make([]string, 0, len(pkgs))
-		for _, p := range pkgs {
-			if packageName != "" && p.Name != packageName {
-				continue
+	files := make([]string, 0, len(pkgs))
+	for _, p := range pkgs {
+		files = append(files, p.Path)
+	}
+	return complete.PredictFilesSet(files).Predict(a)
+}
+
+func predictRunnableFiles(a complete.Args) (prediction []string) {
+	dir := a.Directory()
+	pkgs := listPackages(dir)
+
+	files := []string{}
+	for _, p := range pkgs {
+		// filter non main pacakges
+		if p.Name != "main" {
+			continue
+		}
+		for _, f := range p.Files {
+			path := filepath.Join(p.Path, f)
+			if len(functionsInFile(path, reMainFunc)) > 0 {
+				files = append(files, path)
 			}
-			files = append(files, p.Path)
 		}
-		return complete.PredictFilesSet(files).Predict(a)
-	})
+	}
+	complete.Log("FILES: %s", files)
+	return complete.PredictFilesSet(files).Predict(a)
 }
 
 type pack struct {
-	Name string
-	Path string
+	Name        string
+	Path        string
+	FilesString string
+	Files       []string
 }
 
 func listPackages(dir string) (pkgs []pack) {
+	dir = strings.TrimRight(dir, "/") + "/..."
 	out, err := exec.Command("go", "list", "-f", goListFormat, dir).Output()
 	if err != nil {
 		return
@@ -42,9 +64,13 @@ func listPackages(dir string) (pkgs []pack) {
 	lines := bytes.Split(out, []byte("\n"))
 	for _, line := range lines {
 		var p pack
-		if err := json.Unmarshal(line, &p); err == nil {
-			pkgs = append(pkgs, p)
+		err := json.Unmarshal(line, &p)
+		if err != nil {
+			continue
 		}
+		// parse the FileString from a string "[file1 file2 file3]" to a list of files
+		p.Files = strings.Split(strings.Trim(p.FilesString, "[]"), " ")
+		pkgs = append(pkgs, p)
 	}
 	return
 }