From 6375c754a7417ac7a2c27d3a41de22b314d64fcc Mon Sep 17 00:00:00 2001
From: sudowanderer <21277644+sudowanderer@users.noreply.github.com>
Date: Thu, 2 Jul 2026 17:16:37 +0800
Subject: [PATCH] builder: support ESP flash size target option

---
 builder/build.go      |  2 +-
 builder/esp.go        | 31 ++++++++++++++++++-
 builder/esp_test.go   | 70 +++++++++++++++++++++++++++++++++++++++++++
 compileopts/target.go |  1 +
 4 files changed, 102 insertions(+), 2 deletions(-)
 create mode 100644 builder/esp_test.go

diff --git a/builder/build.go b/builder/build.go
index dabbe7544e..ce98095199 100644
--- a/builder/build.go
+++ b/builder/build.go
@@ -1080,7 +1080,7 @@ func Build(pkgName, outpath, tmpdir string, config *compileopts.Config) (BuildRe
 		// Special format for the ESP family of chips (parsed by the ROM
 		// bootloader).
 		result.Binary = filepath.Join(tmpdir, "main"+outext)
-		err := makeESPFirmwareImage(result.Executable, result.Binary, outputBinaryFormat)
+		err := makeESPFirmwareImage(result.Executable, result.Binary, outputBinaryFormat, config.Target)
 		if err != nil {
 			return result, err
 		}
diff --git a/builder/esp.go b/builder/esp.go
index a480b4d1ea..7b84319bde 100644
--- a/builder/esp.go
+++ b/builder/esp.go
@@ -16,6 +16,8 @@ import (
 	"os"
 	"sort"
 	"strings"
+
+	"github.com/tinygo-org/tinygo/compileopts"
 )
 
 type espImageSegment struct {
@@ -23,6 +25,29 @@ type espImageSegment struct {
 	data []byte
 }
 
+var espFlashSizes = map[string]uint8{
+	"1MB":   0x00,
+	"2MB":   0x10,
+	"4MB":   0x20,
+	"8MB":   0x30,
+	"16MB":  0x40,
+	"32MB":  0x50,
+	"64MB":  0x60,
+	"128MB": 0x70,
+}
+
+func setESPFlashSize(spiSpeedSize uint8, target *compileopts.TargetSpec) (uint8, error) {
+	if target == nil || target.ESPFlashSize == "" {
+		return spiSpeedSize, nil
+	}
+	size := strings.ToUpper(target.ESPFlashSize)
+	spiSize, ok := espFlashSizes[size]
+	if !ok {
+		return 0, fmt.Errorf("invalid esp-flash-size %q", size)
+	}
+	return spiSize | (spiSpeedSize & 0x0f), nil
+}
+
 // makeESPFirmwareImage converts an input ELF file to an image file for an ESP32 or
 // ESP8266 chip. This is a special purpose image format just for the ESP chip
 // family, and is parsed by the on-chip mask ROM bootloader.
@@ -31,7 +56,7 @@ type espImageSegment struct {
 // https://github.com/espressif/esptool/wiki/Firmware-Image-Format
 // https://github.com/espressif/esp-idf/blob/8fbb63c2a701c22ccf4ce249f43aded73e134a34/components/bootloader_support/include/esp_image_format.h#L58
 // https://github.com/espressif/esptool/blob/master/esptool.py
-func makeESPFirmwareImage(infile, outfile, format string) error {
+func makeESPFirmwareImage(infile, outfile, format string, target *compileopts.TargetSpec) error {
 	inf, err := elf.Open(infile)
 	if err != nil {
 		return err
@@ -118,6 +143,10 @@ func makeESPFirmwareImage(infile, outfile, format string) error {
 	// Image header.
 	switch chip {
 	case "esp32", "esp32c3", "esp32s3", "esp32c6":
+		spiSpeedSize, err = setESPFlashSize(spiSpeedSize, target)
+		if err != nil {
+			return err
+		}
 		// Header format:
 		// https://github.com/espressif/esp-idf/blob/v4.3/components/bootloader_support/include/esp_app_format.h#L71
 		// Note: not adding a SHA256 hash as the binary is modified by
diff --git a/builder/esp_test.go b/builder/esp_test.go
new file mode 100644
index 0000000000..42ec9ef1f4
--- /dev/null
+++ b/builder/esp_test.go
@@ -0,0 +1,70 @@
+package builder
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/tinygo-org/tinygo/compileopts"
+)
+
+func TestESPFlashSize(t *testing.T) {
+	tests := []struct {
+		name             string
+		defaultSpeedSize uint8
+		target           *compileopts.TargetSpec
+		wantSpeedSize    uint8
+	}{
+		{
+			name:             "keeps esp32c3 default",
+			defaultSpeedSize: 0x1f,
+			target:           &compileopts.TargetSpec{},
+			wantSpeedSize:    0x1f,
+		},
+		{
+			name:             "sets esp32c3 4MB size nibble",
+			defaultSpeedSize: 0x1f,
+			target:           &compileopts.TargetSpec{ESPFlashSize: "4MB"},
+			wantSpeedSize:    0x2f,
+		},
+		{
+			name:             "keeps esp32c6 default frequency nibble",
+			defaultSpeedSize: 0x10,
+			target:           &compileopts.TargetSpec{ESPFlashSize: "4MB"},
+			wantSpeedSize:    0x20,
+		},
+		{
+			name:             "normalizes lowercase size",
+			defaultSpeedSize: 0x1f,
+			target:           &compileopts.TargetSpec{ESPFlashSize: "8mb"},
+			wantSpeedSize:    0x3f,
+		},
+		{
+			name:             "allows nil target",
+			defaultSpeedSize: 0x1f,
+			target:           nil,
+			wantSpeedSize:    0x1f,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			spiSpeedSize, err := setESPFlashSize(test.defaultSpeedSize, test.target)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if spiSpeedSize != test.wantSpeedSize {
+				t.Fatalf("unexpected spi speed/size: got %#x, want %#x", spiSpeedSize, test.wantSpeedSize)
+			}
+		})
+	}
+}
+
+func TestESPFlashSizeInvalid(t *testing.T) {
+	_, err := setESPFlashSize(0x1f, &compileopts.TargetSpec{ESPFlashSize: "6MB"})
+	if err == nil {
+		t.Fatal("expected error, got nil")
+	}
+	if !strings.Contains(err.Error(), "esp-flash-size") {
+		t.Fatalf("unexpected error: got %q, want it to contain %q", err, "esp-flash-size")
+	}
+}
diff --git a/compileopts/target.go b/compileopts/target.go
index 70c7047462..0f62953846 100644
--- a/compileopts/target.go
+++ b/compileopts/target.go
@@ -58,6 +58,7 @@ type TargetSpec struct {
 	FlashFilename    string   `json:"msd-firmware-name,omitempty"`
 	UF2FamilyID      string   `json:"uf2-family-id,omitempty"`
 	BinaryFormat     string   `json:"binary-format,omitempty"`
+	ESPFlashSize     string   `json:"esp-flash-size,omitempty"`
 	OpenOCDInterface string   `json:"openocd-interface,omitempty"`
 	OpenOCDTarget    string   `json:"openocd-target,omitempty"`
 	OpenOCDTransport string   `json:"openocd-transport,omitempty"`