main: add -print-allocs flag that lets you print all heap allocations

This flag, if set, is a regexp for function names. If there are heap
allocations in the matching function names, these heap allocations will
be printed with an explanation why the heap allocation exists (and why
the object can't be stack allocated).

6 files changed, 225 insertions, 16 deletions

diff --git a/compileopts/options.go b/compileopts/options.go
index 7570fefe5..8a1dea997 100644
--- a/compileopts/options.go
+++ b/compileopts/options.go
@@ -2,6 +2,7 @@ package compileopts
 
 import (
 	"fmt"
+	"regexp"
 	"strings"
 )
 
@@ -27,6 +28,7 @@ type Options struct {
 	PrintCommands bool
 	Debug         bool
 	PrintSizes    string
+	PrintAllocs   *regexp.Regexp // regexp string
 	PrintStacks   bool
 	Tags          string
 	WasmAbi       string
diff --git a/main.go b/main.go
index 3d9bd6e75..50b12abd3 100644
--- a/main.go
+++ b/main.go
@@ -13,6 +13,7 @@ import (
 	"os/exec"
 	"os/signal"
 	"path/filepath"
+	"regexp"
 	"runtime"
 	"strings"
 	"sync/atomic"
@@ -901,6 +902,7 @@ func main() {
 	target := flag.String("target", "", "LLVM target | .json file with TargetSpec")
 	printSize := flag.String("size", "", "print sizes (none, short, full)")
 	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")
 	nodebug := flag.Bool("no-debug", false, "disable DWARF debug symbol generation")
 	ocdOutput := flag.Bool("ocd-output", false, "print OCD daemon output during debug")
@@ -941,6 +943,14 @@ func main() {
 		fmt.Fprintln(os.Stderr, err)
 		os.Exit(1)
 	}
+	var printAllocs *regexp.Regexp
+	if *printAllocsString != "" {
+		printAllocs, err = regexp.Compile(*printAllocsString)
+		if err != nil {
+			fmt.Fprintln(os.Stderr, err)
+			os.Exit(1)
+		}
+	}
 	options := &compileopts.Options{
 		Target:        *target,
 		Opt:           *opt,
@@ -953,6 +963,7 @@ func main() {
 		Debug:         !*nodebug,
 		PrintSizes:    *printSize,
 		PrintStacks:   *printStacks,
+		PrintAllocs:   printAllocs,
 		PrintCommands: *printCommands,
 		Tags:          *tags,
 		GlobalValues:  globalVarValues,
diff --git a/transform/allocs.go b/transform/allocs.go
index 9cf15a1dc..00d318310 100644
--- a/transform/allocs.go
+++ b/transform/allocs.go
@@ -6,6 +6,10 @@ package transform
 // interprocedural escape analysis.
 
 import (
+	"fmt"
+	"go/token"
+	"regexp"
+
 	"tinygo.org/x/go-llvm"
 )
 
@@ -20,7 +24,10 @@ const maxStackAlloc = 256
 // whenever possible. It relies on the LLVM 'nocapture' flag for interprocedural
 // escape analysis, and within a function looks whether an allocation can escape
 // to the heap.
-func OptimizeAllocs(mod llvm.Module) {
+// If printAllocs is non-nil, it indicates the regexp of functions for which a
+// heap allocation explanation should be printed (why the object can't be stack
+// allocated).
+func OptimizeAllocs(mod llvm.Module, printAllocs *regexp.Regexp, logger func(token.Position, string)) {
 	allocator := mod.NamedFunction("runtime.alloc")
 	if allocator.IsNil() {
 		// nothing to optimize
@@ -32,14 +39,21 @@ func OptimizeAllocs(mod llvm.Module) {
 	builder := mod.Context().NewBuilder()
 
 	for _, heapalloc := range getUses(allocator) {
+		logAllocs := printAllocs != nil && printAllocs.MatchString(heapalloc.InstructionParent().Parent().Name())
 		if heapalloc.Operand(0).IsAConstant().IsNil() {
 			// Do not allocate variable length arrays on the stack.
+			if logAllocs {
+				logAlloc(logger, heapalloc, "size is not constant")
+			}
 			continue
 		}
 
 		size := heapalloc.Operand(0).ZExtValue()
 		if size > maxStackAlloc {
 			// The maximum size for a stack allocation.
+			if logAllocs {
+				logAlloc(logger, heapalloc, fmt.Sprintf("object size %d exceeds maximum stack allocation size %d", size, maxStackAlloc))
+			}
 			continue
 		}
 
@@ -68,7 +82,15 @@ func OptimizeAllocs(mod llvm.Module) {
 			bitcast = uses[0]
 		}
 
-		if mayEscape(bitcast) {
+		if at := valueEscapesAt(bitcast); !at.IsNil() {
+			if logAllocs {
+				atPos := getPosition(at)
+				msg := "escapes at unknown line"
+				if atPos.Line != 0 {
+					msg = fmt.Sprintf("escapes at line %d", atPos.Line)
+				}
+				logAlloc(logger, heapalloc, msg)
+			}
 			continue
 		}
 		// The pointer value does not escape.
@@ -97,9 +119,9 @@ func OptimizeAllocs(mod llvm.Module) {
 	}
 }
 
-// mayEscape returns whether the value might escape. It returns true if it might
-// escape, and false if it definitely doesn't. The value must be an instruction.
-func mayEscape(value llvm.Value) bool {
+// valueEscapesAt returns the instruction where the given value may escape and a
+// nil llvm.Value if it definitely doesn't. The value must be an instruction.
+func valueEscapesAt(value llvm.Value) llvm.Value {
 	uses := getUses(value)
 	for _, use := range uses {
 		if use.IsAInstruction().IsNil() {
@@ -107,13 +129,13 @@ func mayEscape(value llvm.Value) bool {
 		}
 		switch use.InstructionOpcode() {
 		case llvm.GetElementPtr:
-			if mayEscape(use) {
-				return true
+			if at := valueEscapesAt(use); !at.IsNil() {
+				return at
 			}
 		case llvm.BitCast:
 			// A bitcast escapes if the casted-to value escapes.
-			if mayEscape(use) {
-				return true
+			if at := valueEscapesAt(use); !at.IsNil() {
+				return at
 			}
 		case llvm.Load:
 			// Load does not escape.
@@ -121,21 +143,27 @@ func mayEscape(value llvm.Value) bool {
 			// Store only escapes when the value is stored to, not when the
 			// value is stored into another value.
 			if use.Operand(0) == value {
-				return true
+				return use
 			}
 		case llvm.Call:
 			if !hasFlag(use, value, "nocapture") {
-				return true
+				return use
 			}
 		case llvm.ICmp:
 			// Comparing pointers don't let the pointer escape.
 			// This is often a compiler-inserted nil check.
 		default:
 			// Unknown instruction, might escape.
-			return true
+			return use
 		}
 	}
 
 	// Checked all uses, and none let the pointer value escape.
-	return false
+	return llvm.Value{}
+}
+
+// logAlloc prints a message to stderr explaining why the given object had to be
+// allocated on the heap.
+func logAlloc(logger func(token.Position, string), allocCall llvm.Value, reason string) {
+	logger(getPosition(allocCall), "object allocated on the heap: "+reason)
 }
diff --git a/transform/allocs_test.go b/transform/allocs_test.go
index 9176b2ff9..e6327827a 100644
--- a/transform/allocs_test.go
+++ b/transform/allocs_test.go
@@ -1,12 +1,129 @@
 package transform_test
 
 import (
+	"go/token"
+	"go/types"
+	"io/ioutil"
+	"path/filepath"
+	"regexp"
+	"sort"
+	"strconv"
+	"strings"
 	"testing"
 
+	"github.com/tinygo-org/tinygo/compileopts"
+	"github.com/tinygo-org/tinygo/compiler"
+	"github.com/tinygo-org/tinygo/loader"
 	"github.com/tinygo-org/tinygo/transform"
+	"tinygo.org/x/go-llvm"
 )
 
 func TestAllocs(t *testing.T) {
 	t.Parallel()
-	testTransform(t, "testdata/allocs", transform.OptimizeAllocs)
+	testTransform(t, "testdata/allocs", func(mod llvm.Module) {
+		transform.OptimizeAllocs(mod, nil, nil)
+	})
+}
+
+type allocsTestOutput struct {
+	filename string
+	line     int
+	msg      string
+}
+
+func (out allocsTestOutput) String() string {
+	return out.filename + ":" + strconv.Itoa(out.line) + ": " + out.msg
+}
+
+// Test with a Go file as input (for more accurate tests).
+func TestAllocs2(t *testing.T) {
+	t.Parallel()
+
+	target, err := compileopts.LoadTarget("i686--linux")
+	if err != nil {
+		t.Fatal("failed to load target:", err)
+	}
+	config := &compileopts.Config{
+		Options: &compileopts.Options{},
+		Target:  target,
+	}
+	compilerConfig := &compiler.Config{
+		Triple:             config.Triple(),
+		GOOS:               config.GOOS(),
+		GOARCH:             config.GOARCH(),
+		CodeModel:          config.CodeModel(),
+		RelocationModel:    config.RelocationModel(),
+		Scheduler:          config.Scheduler(),
+		FuncImplementation: config.FuncImplementation(),
+		AutomaticStackSize: config.AutomaticStackSize(),
+		Debug:              true,
+	}
+	machine, err := compiler.NewTargetMachine(compilerConfig)
+	if err != nil {
+		t.Fatal("failed to create target machine:", err)
+	}
+
+	// Load entire program AST into memory.
+	lprogram, err := loader.Load(config, []string{"./testdata/allocs2.go"}, config.ClangHeaders, types.Config{
+		Sizes: compiler.Sizes(machine),
+	})
+	if err != nil {
+		t.Fatal("failed to create target machine:", err)
+	}
+	err = lprogram.Parse()
+	if err != nil {
+		t.Fatal("could not parse", err)
+	}
+
+	// Compile AST to IR.
+	program := lprogram.LoadSSA()
+	pkg := lprogram.MainPkg()
+	mod, errs := compiler.CompilePackage("allocs2.go", pkg, program.Package(pkg.Pkg), machine, compilerConfig, false)
+	if errs != nil {
+		for _, err := range errs {
+			t.Error(err)
+		}
+		return
+	}
+
+	// Run functionattrs pass, which is necessary for escape analysis.
+	pm := llvm.NewPassManager()
+	defer pm.Dispose()
+	pm.AddInstructionCombiningPass()
+	pm.AddFunctionAttrsPass()
+	pm.Run(mod)
+
+	// Run heap to stack transform.
+	var testOutputs []allocsTestOutput
+	transform.OptimizeAllocs(mod, regexp.MustCompile("."), func(pos token.Position, msg string) {
+		testOutputs = append(testOutputs, allocsTestOutput{
+			filename: filepath.Base(pos.Filename),
+			line:     pos.Line,
+			msg:      msg,
+		})
+	})
+	sort.Slice(testOutputs, func(i, j int) bool {
+		return testOutputs[i].line < testOutputs[j].line
+	})
+	testOutput := ""
+	for _, out := range testOutputs {
+		testOutput += out.String() + "\n"
+	}
+
+	// Load expected test output (the OUT: lines).
+	testInput, err := ioutil.ReadFile("./testdata/allocs2.go")
+	if err != nil {
+		t.Fatal("could not read test input:", err)
+	}
+	var expectedTestOutput string
+	for i, line := range strings.Split(strings.ReplaceAll(string(testInput), "\r\n", "\n"), "\n") {
+		if idx := strings.Index(line, " // OUT: "); idx > 0 {
+			msg := line[idx+len(" // OUT: "):]
+			expectedTestOutput += "allocs2.go:" + strconv.Itoa(i+1) + ": " + msg + "\n"
+		}
+	}
+
+	if testOutput != expectedTestOutput {
+		t.Errorf("output does not match expected output:\n%s", testOutput)
+	}
 }
diff --git a/transform/optimizer.go b/transform/optimizer.go
index cf570a2d4..dca4a2bda 100644
--- a/transform/optimizer.go
+++ b/transform/optimizer.go
@@ -3,6 +3,8 @@ package transform
 import (
 	"errors"
 	"fmt"
+	"go/token"
+	"os"
 
 	"github.com/tinygo-org/tinygo/compileopts"
 	"github.com/tinygo-org/tinygo/compiler/ircheck"
@@ -65,7 +67,7 @@ func Optimize(mod llvm.Module, config *compileopts.Config, optLevel, sizeLevel i
 		OptimizeMaps(mod)
 		OptimizeStringToBytes(mod)
 		OptimizeReflectImplements(mod)
-		OptimizeAllocs(mod)
+		OptimizeAllocs(mod, nil, nil)
 		err := LowerInterfaces(mod, sizeLevel)
 		if err != nil {
 			return []error{err}
@@ -86,7 +88,9 @@ func Optimize(mod llvm.Module, config *compileopts.Config, optLevel, sizeLevel i
 		goPasses.Run(mod)
 
 		// Run TinyGo-specific interprocedural optimizations.
-		OptimizeAllocs(mod)
+		OptimizeAllocs(mod, config.Options.PrintAllocs, func(pos token.Position, msg string) {
+			fmt.Fprintln(os.Stderr, pos.String()+": "+msg)
+		})
 		OptimizeStringToBytes(mod)
 		OptimizeStringEqual(mod)
 
diff --git a/transform/testdata/allocs2.go b/transform/testdata/allocs2.go
new file mode 100644
index 000000000..f01ad167c
--- /dev/null
+++ b/transform/testdata/allocs2.go
@@ -0,0 +1,47 @@
+package main
+
+func main() {
+	n1 := 5
+	derefInt(&n1)
+
+	// This should eventually be modified to not escape.
+	n2 := 6 // OUT: object allocated on the heap: escapes at line 9
+	returnIntPtr(&n2)
+
+	s1 := make([]int, 3)
+	readIntSlice(s1)
+
+	s2 := [3]int{}
+	readIntSlice(s2[:])
+
+	// This should also be modified to not escape.
+	s3 := make([]int, 3) // OUT: object allocated on the heap: escapes at line 19
+	returnIntSlice(s3)
+
+	_ = make([]int, getUnknownNumber()) // OUT: object allocated on the heap: size is not constant
+
+	s4 := make([]byte, 300) // OUT: object allocated on the heap: object size 300 exceeds maximum stack allocation size 256
+	readByteSlice(s4)
+}
+
+func derefInt(x *int) int {
+	return *x
+}
+
+func returnIntPtr(x *int) *int {
+	return x
+}
+
+func readIntSlice(s []int) int {
+	return s[1]
+}
+
+func readByteSlice(s []byte) byte {
+	return s[1]
+}
+
+func returnIntSlice(s []int) []int {
+	return s
+}
+
+func getUnknownNumber() int