From c638f03b3c4264689a487043f1f99b8309d0aaac Mon Sep 17 00:00:00 2001
From: Ayke van Laethem <aykevanlaethem@gmail.com>
Date: Wed, 22 Sep 2021 14:00:51 +0200
Subject: [PATCH] main: add -p flag to set parallelism

This is very useful for debugging.
---
 builder/build.go       |  4 ++--
 builder/jobs.go        | 15 ++++++++++++---
 builder/library.go     |  2 +-
 compileopts/options.go |  1 +
 main.go                |  3 +++
 5 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/builder/build.go b/builder/build.go
index 18c9d8a3..73a0f257 100644
--- a/builder/build.go
+++ b/builder/build.go
@@ -462,7 +462,7 @@ func Build(pkgName, outpath string, config *compileopts.Config, action func(Buil
 	outext := filepath.Ext(outpath)
 	if outext == ".o" || outext == ".bc" || outext == ".ll" {
 		// Run jobs to produce the LLVM module.
-		err := runJobs(programJob)
+		err := runJobs(programJob, config.Options.Parallelism)
 		if err != nil {
 			return err
 		}
@@ -686,7 +686,7 @@ func Build(pkgName, outpath string, config *compileopts.Config, action func(Buil
 	// Run all jobs to compile and link the program.
 	// Do this now (instead of after elf-to-hex and similar conversions) as it
 	// is simpler and cannot be parallelized.
-	err = runJobs(linkJob)
+	err = runJobs(linkJob, config.Options.Parallelism)
 	if err != nil {
 		return err
 	}
diff --git a/builder/jobs.go b/builder/jobs.go
index 57326bfa..16bd214e 100644
--- a/builder/jobs.go
+++ b/builder/jobs.go
@@ -71,7 +71,16 @@ func (job *compileJob) readyToRun() bool {
 // It runs all jobs in the order of the dependencies slice, depth-first.
 // Therefore, if some jobs are preferred to run before others, they should be
 // ordered as such in the job dependencies.
-func runJobs(job *compileJob) error {
+func runJobs(job *compileJob, parallelism int) error {
+	if parallelism == 0 {
+		// Have a default, if the parallelism isn't set. This is useful for
+		// tests.
+		parallelism = runtime.NumCPU()
+	}
+	if parallelism < 1 {
+		return fmt.Errorf("-p flag must be at least 1, provided -p=%d", parallelism)
+	}
+
 	// Create a slice of jobs to run, where all dependencies are run in order.
 	jobs := []*compileJob{}
 	addedJobs := map[*compileJob]struct{}{}
@@ -94,7 +103,7 @@ func runJobs(job *compileJob) error {
 	defer close(workerChan)
 
 	// Start a number of workers.
-	for i := 0; i < runtime.NumCPU(); i++ {
+	for i := 0; i < parallelism; i++ {
 		if jobRunnerDebug {
 			fmt.Println("## starting worker", i)
 		}
@@ -109,7 +118,7 @@ func runJobs(job *compileJob) error {
 	for {
 		// If there are free workers, try starting a new job (if one is
 		// available). If it succeeds, try again to fill the entire worker pool.
-		if numRunningJobs < runtime.NumCPU() {
+		if numRunningJobs < parallelism {
 			jobToRun := nextJob(jobs)
 			if jobToRun != nil {
 				// Start job.
diff --git a/builder/library.go b/builder/library.go
index 0e09fa44..511348c4 100644
--- a/builder/library.go
+++ b/builder/library.go
@@ -52,7 +52,7 @@ func (l *Library) Load(config *compileopts.Config, tmpdir string) (dir string, e
 	if err != nil {
 		return "", err
 	}
-	err = runJobs(job)
+	err = runJobs(job, config.Options.Parallelism)
 	return filepath.Dir(job.result), err
 }
 
diff --git a/compileopts/options.go b/compileopts/options.go
index ec47a84f..fd041645 100644
--- a/compileopts/options.go
+++ b/compileopts/options.go
@@ -31,6 +31,7 @@ type Options struct {
 	DumpSSA         bool
 	VerifyIR        bool
 	PrintCommands   func(cmd string, args ...string)
+	Parallelism     int // -p flag
 	Debug           bool
 	PrintSizes      string
 	PrintAllocs     *regexp.Regexp // regexp string
diff --git a/main.go b/main.go
index 00b5b9bb..9e53d2f9 100644
--- a/main.go
+++ b/main.go
@@ -1118,6 +1118,7 @@ func main() {
 	printStacks := flag.Bool("print-stacks", false, "print stack sizes of goroutines")
 	printAllocsString := flag.String("print-allocs", "", "regular expression of functions for which heap allocations should be printed")
 	printCommands := flag.Bool("x", false, "Print commands")
+	parallelism := flag.Int("p", runtime.GOMAXPROCS(0), "the number of build jobs that can run in parallel")
 	nodebug := flag.Bool("no-debug", false, "strip debug information")
 	ocdCommandsString := flag.String("ocd-commands", "", "OpenOCD commands, overriding target spec (can specify multiple separated by commas)")
 	ocdOutput := flag.Bool("ocd-output", false, "print OCD daemon output during debug")
@@ -1190,6 +1191,7 @@ func main() {
 		PrintIR:         *printIR,
 		DumpSSA:         *dumpSSA,
 		VerifyIR:        *verifyIR,
+		Parallelism:     *parallelism,
 		Debug:           !*nodebug,
 		PrintSizes:      *printSize,
 		PrintStacks:     *printStacks,
@@ -1278,6 +1280,7 @@ func main() {
 		}
 		defer os.RemoveAll(tmpdir)
 		config := &compileopts.Config{
+			Options: options,
 			Target: &compileopts.TargetSpec{
 				Triple: *target,
 			},