Allow larger systems to have a larger max stack alloc

Fixes #3331

6 files changed, 19 insertions, 15 deletions

diff --git a/builder/build.go b/builder/build.go
index 5fc6f44ee..b9965c1e8 100644
--- a/builder/build.go
+++ b/builder/build.go
@@ -197,6 +197,7 @@ func Build(pkgName, outpath, tmpdir string, config *compileopts.Config) (BuildRe
 		Scheduler:          config.Scheduler(),
 		AutomaticStackSize: config.AutomaticStackSize(),
 		DefaultStackSize:   config.StackSize(),
+		MaxStackAlloc:      config.MaxStackAlloc(),
 		NeedsStackObjects:  config.NeedsStackObjects(),
 		Debug:              !config.Options.SkipDWARF, // emit DWARF except when -internal-nodwarf is passed
 	}
diff --git a/compileopts/config.go b/compileopts/config.go
index 2260de2d2..e1689d3ba 100644
--- a/compileopts/config.go
+++ b/compileopts/config.go
@@ -189,6 +189,15 @@ func (c *Config) StackSize() uint64 {
 	return c.Target.DefaultStackSize
 }
 
+// MaxStackAlloc returns the size of the maximum allocation to put on the stack vs heap.
+func (c *Config) MaxStackAlloc() uint64 {
+	if c.StackSize() > 32*1024 {
+		return 1024
+	}
+
+	return 256
+}
+
 // RP2040BootPatch returns whether the RP2040 boot patch should be applied that
 // calculates and patches in the checksum for the 2nd stage bootloader.
 func (c *Config) RP2040BootPatch() bool {
diff --git a/compiler/compiler.go b/compiler/compiler.go
index 432412461..2ebd5a9f2 100644
--- a/compiler/compiler.go
+++ b/compiler/compiler.go
@@ -53,6 +53,7 @@ type Config struct {
 	Scheduler          string
 	AutomaticStackSize bool
 	DefaultStackSize   uint64
+	MaxStackAlloc      uint64
 	NeedsStackObjects  bool
 	Debug              bool // Whether to emit debug information in the LLVM module.
 }
@@ -1997,8 +1998,8 @@ func (b *builder) createExpr(expr ssa.Value) (llvm.Value, error) {
 	case *ssa.Alloc:
 		typ := b.getLLVMType(expr.Type().Underlying().(*types.Pointer).Elem())
 		size := b.targetData.TypeAllocSize(typ)
-		// Move all "large" allocations to the heap.  This value is also transform.maxStackAlloc.
-		if expr.Heap || size > 256 {
+		// Move all "large" allocations to the heap.
+		if expr.Heap || size > b.MaxStackAlloc {
 			// Calculate ^uintptr(0)
 			maxSize := llvm.ConstNot(llvm.ConstInt(b.uintptrType, 0, false)).ZExtValue()
 			if size > maxSize {
diff --git a/transform/allocs.go b/transform/allocs.go
index 908d72e7b..67ca1ea43 100644
--- a/transform/allocs.go
+++ b/transform/allocs.go
@@ -13,14 +13,6 @@ import (
 	"tinygo.org/x/go-llvm"
 )
 
-// maxStackAlloc is the maximum size of an object that will be allocated on the
-// stack. Bigger objects have increased risk of stack overflows and thus will
-// always be heap allocated.
-//
-// TODO: tune this, this is just a random value.
-// This value is also used in the compiler when translating ssa.Alloc nodes.
-const maxStackAlloc = 256
-
 // OptimizeAllocs tries to replace heap allocations with stack allocations
 // whenever possible. It relies on the LLVM 'nocapture' flag for interprocedural
 // escape analysis, and within a function looks whether an allocation can escape
@@ -28,7 +20,7 @@ const maxStackAlloc = 256
 // 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)) {
+func OptimizeAllocs(mod llvm.Module, printAllocs *regexp.Regexp, maxStackAlloc uint64, logger func(token.Position, string)) {
 	allocator := mod.NamedFunction("runtime.alloc")
 	if allocator.IsNil() {
 		// nothing to optimize
diff --git a/transform/allocs_test.go b/transform/allocs_test.go
index 59a5b14e2..7f7ff5b75 100644
--- a/transform/allocs_test.go
+++ b/transform/allocs_test.go
@@ -17,7 +17,7 @@ import (
 func TestAllocs(t *testing.T) {
 	t.Parallel()
 	testTransform(t, "testdata/allocs", func(mod llvm.Module) {
-		transform.OptimizeAllocs(mod, nil, nil)
+		transform.OptimizeAllocs(mod, nil, 256, nil)
 	})
 }
 
@@ -47,7 +47,7 @@ func TestAllocs2(t *testing.T) {
 
 	// Run heap to stack transform.
 	var testOutputs []allocsTestOutput
-	transform.OptimizeAllocs(mod, regexp.MustCompile("."), func(pos token.Position, msg string) {
+	transform.OptimizeAllocs(mod, regexp.MustCompile("."), 256, func(pos token.Position, msg string) {
 		testOutputs = append(testOutputs, allocsTestOutput{
 			filename: filepath.Base(pos.Filename),
 			line:     pos.Line,
diff --git a/transform/optimizer.go b/transform/optimizer.go
index 42acc2ddc..e73e97600 100644
--- a/transform/optimizer.go
+++ b/transform/optimizer.go
@@ -64,7 +64,8 @@ func Optimize(mod llvm.Module, config *compileopts.Config) []error {
 		// Run TinyGo-specific optimization passes.
 		OptimizeStringToBytes(mod)
 		OptimizeReflectImplements(mod)
-		OptimizeAllocs(mod, nil, nil)
+		maxStackSize := config.MaxStackAlloc()
+		OptimizeAllocs(mod, nil, maxStackSize, nil)
 		err = LowerInterfaces(mod, config)
 		if err != nil {
 			return []error{err}
@@ -84,7 +85,7 @@ func Optimize(mod llvm.Module, config *compileopts.Config) []error {
 		}
 
 		// Run TinyGo-specific interprocedural optimizations.
-		OptimizeAllocs(mod, config.Options.PrintAllocs, func(pos token.Position, msg string) {
+		OptimizeAllocs(mod, config.Options.PrintAllocs, maxStackSize, func(pos token.Position, msg string) {
 			fmt.Fprintln(os.Stderr, pos.String()+": "+msg)
 		})
 		OptimizeStringToBytes(mod)