[noise] Use the shared PROGMEM table helper for the PSK (#20101)

esphome/codegen.py

@@ -40,6 +40,7 @@ from esphome.cpp_generator import (  # noqa: F401
     progmem_array,
     safe_exp,
     set_cpp_standard,
+    shared_progmem_array,
     statement,
     static_const_array,
     static_function,

esphome/components/api/__init__.py

@@ -601,7 +601,7 @@ async def to_code(config: ConfigType) -> None:
 
     if (encryption_config := config.get(CONF_ENCRYPTION, None)) is not None:
         if key := encryption_config.get(CONF_KEY):
-            cg.add(var.set_noise_psk(new_psk_progmem(config[CONF_ID], key)))
+            cg.add(var.set_noise_psk(new_psk_progmem(key)))
             cg.add_define("USE_API_NOISE_PSK_FROM_YAML")
         else:
             # No key provided, but encryption desired

esphome/components/esphome/ota/__init__.py

@@ -360,7 +360,7 @@ async def to_code(config: ConfigType) -> None:
         # Build time key: the ota keeps its own pointer so safe mode, which
         # has no api server, still has it
         cg.add_define("USE_OTA_ENCRYPTION")
-        cg.add(var.set_noise_psk(new_psk_progmem(config[CONF_ID], key)))
+        cg.add(var.set_noise_psk(new_psk_progmem(key)))
     elif CONF_ENCRYPTION in api_conf:
         # Runtime key: found in the api server, or in preferences in safe mode
         cg.add_define("USE_OTA_ENCRYPTION")

esphome/components/noise/__init__.py

@@ -5,12 +5,11 @@ from typing import Any
 import esphome.codegen as cg
 import esphome.config_validation as cv
 from esphome.const import CONF_ENCRYPTION, CONF_KEY
-from esphome.core import CORE, ID
+from esphome.core import CORE
 from esphome.cpp_generator import MockObj
 from esphome.types import ConfigType
 
 CODEOWNERS = ["@esphome/core"]
-DOMAIN = "noise"
 
 # Keep in sync with platformio.ini and esphome/idf_component.yml.
 # LIBSODIUM_VERSION must match the version noise-c pins in its manifests.
@@ -74,18 +73,13 @@ def static_encryption_key(conf: ConfigType) -> str | None:
     return (conf.get(CONF_ENCRYPTION) or {}).get(CONF_KEY) or None
 
 
-def new_psk_progmem(parent_id: ID, key: str) -> MockObj:
+def new_psk_progmem(key: str) -> MockObj:
     """Emit the decoded key as a PROGMEM array; the component keeps a pointer
     so the key never occupies RAM. Components sharing one key (api and ota)
     share the array."""
-    decoded = decode_encryption_key(key)
-    arrays: dict[bytes, MockObj] = CORE.data.setdefault(DOMAIN, {})
-    if (array := arrays.get(decoded)) is None:
-        array = arrays[decoded] = cg.progmem_array(
-            ID(f"{parent_id.id}_psk", is_declaration=True, type=cg.uint8),
-            list(decoded),
-        )
-    return array
+    return cg.shared_progmem_array(
+        "noise_psk", cg.uint8, list(decode_encryption_key(key))
+    )
 
 
 def encryption_schema(config: ConfigType | None) -> ConfigType:

esphome/cpp_generator.py

@@ -19,7 +19,12 @@ from esphome.core import (
     TimePeriodNanoseconds,
     TimePeriodSeconds,
 )
-from esphome.helpers import cpp_string_escape, indent, indent_all_but_first_and_last
+from esphome.helpers import (
+    cpp_string_escape,
+    ensure_unique_string,
+    indent,
+    indent_all_but_first_and_last,
+)
 from esphome.types import Expression, SafeExpType, TemplateArgsType
 from esphome.util import OrderedDict
 from esphome.yaml_util import ESPHomeDataBase
@@ -463,6 +468,36 @@ def progmem_array(id_, rhs) -> "MockObj":
     return obj
 
 
+def shared_progmem_array(
+    name: str, type_: "MockObjClass", rhs: SafeExpType, *, share: bool = True
+) -> "MockObj":
+    """Emit a global PROGMEM array once per distinct type and contents; later calls reuse it.
+
+    The array is ``static constexpr``, so elements must be constant expressions and lambdas
+    must be captureless. Its name is made unique against every config id and variable.
+    ``share=False`` always emits a new array, e.g. for lambdas that may keep static state.
+    """
+    from esphome.config import iter_ids
+    from esphome.config_validation import RESERVED_IDS
+
+    arrays: dict[str, MockObj] = CORE.data.setdefault("shared_progmem_array", {})
+    rhs = safe_exp(rhs)
+    key = f"{type_} {rhs}"
+    if share and (array := arrays.get(key)) is not None:
+        return array
+    used = {str(i) for i, _ in iter_ids(CORE.config)}
+    used |= {str(i) for i in CORE.variables}
+    used |= set(RESERVED_IDS) | CORE.loaded_integrations
+    id_ = ID(ensure_unique_string(name, used), is_declaration=True, type=type_)
+    # Global, so any scope can use it; anything a lambda references is already declared.
+    CORE.add_global(ProgmemAssignmentExpression(type_, id_, rhs))
+    array = MockObj(id_, ".")
+    CORE.register_variable(id_, array)
+    if share:
+        arrays[key] = array
+    return array
+
+
 def static_const_array(id_, rhs) -> "MockObj":
     rhs = safe_exp(rhs)
     obj = MockObj(id_, ".")

tests/component_tests/ota/test_esphome_ota.py

@@ -560,7 +560,7 @@ def test_encryption_offer_codegen(
     own_key = encrypted and "USE_OTA_ENCRYPTION_PROVISIONED" not in defines_present
     assert ("esphome_esphomeotacomponent_id->set_noise_psk(" in main_cpp) is own_key
     # The api shares the ota's array instead of emitting the same key twice
-    assert main_cpp.count("_psk[] PROGMEM") == (1 if own_key else 0)
+    assert main_cpp.count("noise_psk[] PROGMEM") == (1 if own_key else 0)
     assert ("set_auth_password(" in main_cpp) is ("USE_OTA_PASSWORD" in defines_present)
     # The noise transport source compiles only when the define is set
     assert FILTER_SOURCE_FILES() == ([] if encrypted else ["ota_esphome_noise.cpp"])

tests/unit_tests/test_cpp_generator.py

@@ -4,6 +4,7 @@ import math
 import pytest
 
 from esphome import cpp_generator as cg, cpp_types as ct
+from esphome.core import CORE, ID
 
 
 class TestExpressions:
@@ -808,3 +809,39 @@ async def test_templatable__lambda_with_std_string() -> None:
     result = await cg.templatable(lambda_obj, [], ct.std_string)
 
     assert isinstance(result, cg.LambdaExpression)
+
+
+class TestSharedProgmemArray:
+    def test_identical_contents_share_one_array(self) -> None:
+        CORE.config = {}
+        a = cg.shared_progmem_array("table", ct.uint8, [1, 2, 3])
+        b = cg.shared_progmem_array("table", ct.uint8, [1, 2, 3])
+        c = cg.shared_progmem_array("table", ct.uint8, [4])
+        assert a is b
+        assert str(a) != str(c)
+        assert sum("PROGMEM" in str(st) for st in CORE.global_statements) == 2
+        assert not any("PROGMEM" in str(st) for st in CORE.main_statements)
+
+    def test_share_false_always_emits_a_new_array(self) -> None:
+        CORE.config = {}
+        a = cg.shared_progmem_array("table", ct.uint8, [1], share=False)
+        b = cg.shared_progmem_array("table", ct.uint8, [1], share=False)
+        assert str(a) != str(b)
+        assert sum("PROGMEM" in str(st) for st in CORE.global_statements) == 2
+
+    def test_same_contents_different_type_are_separate(self) -> None:
+        CORE.config = {}
+        a = cg.shared_progmem_array("table", ct.uint8, [1])
+        b = cg.shared_progmem_array("table", ct.uint16, [1])
+        assert str(a) != str(b)
+
+    def test_name_avoids_config_ids(self) -> None:
+        CORE.config = {"sensor": [{"id": ID("table", is_declaration=True)}]}
+        array = cg.shared_progmem_array("table", ct.uint8, [1])
+        assert str(array) == "table_2"
+
+    def test_name_avoids_registered_variables(self) -> None:
+        CORE.config = {}
+        CORE.register_variable(ID("table", is_declaration=True), cg.MockObj("table"))
+        array = cg.shared_progmem_array("table", ct.uint8, [1])
+        assert str(array) == "table_2"