Merge branch 'main' into method-receiver-closure

compiler/cmd/coroc/main.go

@@ -3,8 +3,11 @@ package main
 import (
 	"flag"
 	"fmt"
+	"io"
+	"log"
 	"os"
 	"runtime/debug"
+	"runtime/pprof"
 
 	"github.com/stealthrocket/coroutine/compiler"
 )
@@ -17,9 +20,26 @@ USAGE:
 
 OPTIONS:
   -h, --help      Show this help information
+  -l, --list      List all files that would be compiled
   -v, --version   Show the compiler version
+
+ADVANCED OPTIONS:
+  -cpuprofile     Write CPU profile to file
+  -memprofile     Write memory profile to file
 `
 
+var (
+	showVersion   bool
+	onlyListFiles bool
+	cpuProfile    string
+	memProfile    string
+)
+
+func boolFlag(ptr *bool, short, long string) {
+	flag.BoolVar(ptr, short, false, "")
+	flag.BoolVar(ptr, long, false, "")
+}
+
 func main() {
 	if err := run(); err != nil {
 		fmt.Fprintf(os.Stderr, "error: %v\n", err)
@@ -30,17 +50,38 @@ func main() {
 func run() error {
 	flag.Usage = func() { println(usage[1:]) }
 
-	var showVersion bool
-	flag.BoolVar(&showVersion, "v", false, "")
-	flag.BoolVar(&showVersion, "version", false, "")
-
+	boolFlag(&showVersion, "v", "version")
+	boolFlag(&onlyListFiles, "l", "list")
+	flag.StringVar(&cpuProfile, "cpuprofile", "", "")
+	flag.StringVar(&memProfile, "memprofile", "", "")
 	flag.Parse()
 
 	if showVersion {
 		fmt.Println(version())
 		return nil
 	}
 
+	if memProfile != "" {
+		f, err := os.Create(memProfile)
+		if err != nil {
+			return err
+		}
+		defer f.Close()
+		defer pprof.WriteHeapProfile(f)
+	}
+
+	if cpuProfile != "" {
+		f, err := os.Create(cpuProfile)
+		if err != nil {
+			return err
+		}
+		defer f.Close()
+		if err := pprof.StartCPUProfile(f); err != nil {
+			return err
+		}
+		defer pprof.StopCPUProfile()
+	}
+
 	path := flag.Arg(0)
 	if path == "" {
 		// If the compiler was invoked via go generate, the GOFILE
@@ -55,7 +96,13 @@ func run() error {
 		}
 	}
 
-	return compiler.Compile(path)
+	if onlyListFiles {
+		log.SetOutput(io.Discard)
+	}
+
+	return compiler.Compile(path,
+		compiler.OnlyListFiles(onlyListFiles),
+	)
 }
 
 func version() (version string) {

compiler/compile.go

@@ -43,10 +43,9 @@ func Compile(path string, options ...Option) error {
 	return c.compile(path)
 }
 
-// Option configures the compiler.
-type Option func(*compiler)
-
 type compiler struct {
+	onlyListFiles bool
+
 	prog         *ssa.Program
 	generics     map[*ssa.Function][]*ssa.Function
 	coroutinePkg *packages.Package
@@ -183,6 +182,17 @@ func (c *compiler) compile(path string) error {
 		pkgColors[fn] = color
 	}
 
+	if c.onlyListFiles {
+		cwd, _ := os.Getwd()
+		for pkg := range colorsByPkg {
+			for _, filePath := range pkg.GoFiles {
+				relPath, _ := filepath.Rel(cwd, filePath)
+				fmt.Println(relPath)
+			}
+		}
+		return nil
+	}
+
 	// Before mutating packages, we need to ensure that packages exist in a
 	// location where mutations can be made safely (without affecting other
 	// builds).
@@ -219,6 +229,7 @@ func (c *compiler) compile(path string) error {
 		// Reject packages outside ./vendor.
 		return fmt.Errorf("cannot mutate package %s (%s) safely. Please vendor dependencies: go mod vendor", p.PkgPath, dir)
 	}
+
 	if len(needVendoring) > 0 {
 		log.Printf("vendoring GOROOT packages")
 		newRoot := filepath.Join(moduleDir, "goroot")

compiler/options.go

@@ -0,0 +1,10 @@
+package compiler
+
+// Option configures the compiler.
+type Option func(*compiler)
+
+func OnlyListFiles(enabled bool) Option {
+	return func(c *compiler) {
+		c.onlyListFiles = enabled
+	}
+}