From e67b5a78d0b4fd4d06ddc347141eb4e19c7d4f10 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Sun, 22 Mar 2026 20:51:40 -1000
Subject: [PATCH] [esp32] Patch DRAM segment for testing mode to fix grouped
 component test overflow (#15102)

---
 esphome/components/esp32/iram_fix.py.script | 70 ++++++++++++---------
 1 file changed, 42 insertions(+), 28 deletions(-)

diff --git a/esphome/components/esp32/iram_fix.py.script b/esphome/components/esp32/iram_fix.py.script
index 0d23f9a81b..b656d7d1a6 100644
--- a/esphome/components/esp32/iram_fix.py.script
+++ b/esphome/components/esp32/iram_fix.py.script
@@ -4,12 +4,40 @@ import re
 # pylint: disable=E0602
 Import("env")  # noqa
 
-# IRAM size for testing mode (2MB - large enough to accommodate grouped tests)
-TESTING_IRAM_SIZE = 0x200000
+# Memory sizes for testing mode (large enough to accommodate grouped tests)
+TESTING_IRAM_SIZE = 0x200000  # 2MB
+TESTING_DRAM_SIZE = 0x200000  # 2MB
+
+
+def patch_segment(content, segment_name, new_size):
+    """Patch a memory segment's length in linker script content.
+
+    Handles both single-line and multi-line segment definitions, e.g.:
+      iram0_0_seg (RX) : org = 0x40080000, len = 0x20000 + 0x0
+    or split across lines:
+      dram0_0_seg (RW) : org = 0x3FFB0000 + 0xdb5c,
+                                       len = 0x2c200 - 0xdb5c
+
+    Args:
+        content: Full linker script content as string
+        segment_name: Name of the segment (e.g., 'iram0_0_seg')
+        new_size: New size as integer
+
+    Returns:
+        Tuple of (new_content, was_patched)
+    """
+    # Match segment name through to "len = <value>" allowing newlines between org and len
+    pattern = rf'({re.escape(segment_name)}\s*\([^)]*\)\s*:\s*org\s*=\s*.+?,\s*len\s*=\s*)(\S+[^\n]*)'
+    if match := re.search(pattern, content, re.DOTALL):
+        replacement = f"{match.group(1)}{new_size:#x}"
+        new_content = content[:match.start()] + replacement + content[match.end():]
+        if new_content != content:
+            return new_content, True
+    return content, False
 
 
 def patch_idf_linker_script(source, target, env):
-    """Patch ESP-IDF linker script to increase IRAM size for testing mode."""
+    """Patch ESP-IDF linker script to increase IRAM and DRAM size for testing mode."""
     # Check if we're in testing mode by looking for the define
     build_flags = env.get("BUILD_FLAGS", [])
     testing_mode = any("-DESPHOME_TESTING_MODE" in flag for flag in build_flags)
@@ -34,36 +62,22 @@ def patch_idf_linker_script(source, target, env):
         print(f"ESPHome: Error reading linker script: {e}")
         return
 
-    # Check if this file contains iram0_0_seg
-    if 'iram0_0_seg' not in content:
-        print(f"ESPHome: Warning - iram0_0_seg not found in {memory_ld}")
-        return
+    patches = []
 
-    # Look for iram0_0_seg definition and increase its length
-    # ESP-IDF format can be:
-    #   iram0_0_seg (RX) : org = 0x40080000, len = 0x20000 + 0x0
-    # or more complex with nested parentheses:
-    #   iram0_0_seg (RX) : org = (0x40370000 + 0x4000), len = (((0x403CB700 - (0x40378000 - 0x3FC88000)) - 0x3FC88000) + 0x8000 - 0x4000)
-    # We want to change len to TESTING_IRAM_SIZE for testing
+    content, patched = patch_segment(content, 'iram0_0_seg', TESTING_IRAM_SIZE)
+    if patched:
+        patches.append(f"IRAM={TESTING_IRAM_SIZE:#x}")
 
-    # Use a more robust approach: find the line and manually parse it
-    lines = content.split('\n')
-    for i, line in enumerate(lines):
-        if 'iram0_0_seg' in line and 'len' in line:
-            # Find the position of "len = " and replace everything after it until the end of the statement
-            match = re.search(r'(iram0_0_seg\s*\([^)]*\)\s*:\s*org\s*=\s*(?:\([^)]+\)|0x[0-9a-fA-F]+)\s*,\s*len\s*=\s*)(.+?)(\s*)$', line)
-            if match:
-                lines[i] = f"{match.group(1)}{TESTING_IRAM_SIZE:#x}{match.group(3)}"
-                break
+    content, patched = patch_segment(content, 'dram0_0_seg', TESTING_DRAM_SIZE)
+    if patched:
+        patches.append(f"DRAM={TESTING_DRAM_SIZE:#x}")
 
-    updated = '\n'.join(lines)
-
-    if updated != content:
+    if patches:
         with open(memory_ld, "w") as f:
-            f.write(updated)
-        print(f"ESPHome: Patched IRAM size to {TESTING_IRAM_SIZE:#x} in {memory_ld} for testing mode")
+            f.write(content)
+        print(f"ESPHome: Patched {', '.join(patches)} in {memory_ld} for testing mode")
     else:
-        print(f"ESPHome: Warning - could not patch iram0_0_seg in {memory_ld}")
+        print(f"ESPHome: Warning - could not patch memory segments in {memory_ld}")
 
 
 # Hook into the build process before linking