From adbbda407292589591a748b5c781e37b83398b17 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Mon, 7 Sep 2026 11:09:07 +0200 Subject: [PATCH] [api] Emit every encode call through one generator helper _encode_call() owns the cursor assignment and the _force suffix, so the convention lives in one place instead of at every emission site; the fixed32 fast path is an arm of the generic encode_content keyed by a per type value template. write_fixed32_le uses convert_little_endian instead of its own byte order switch. The integration test shares a StateWaiter from state_utils and leaves the disconnect to the fixture. --- esphome/components/api/proto.h | 10 +- script/api_protobuf/api_protobuf.py | 180 +++++++++++------- tests/integration/state_utils.py | 25 +++ .../integration/test_api_encode_boundaries.py | 52 +++-- .../api/test_api_protobuf_generator.py | 20 +- 5 files changed, 174 insertions(+), 113 deletions(-) diff --git a/esphome/components/api/proto.h b/esphome/components/api/proto.h index a76374aaef..f7dd065a68 100644 --- a/esphome/components/api/proto.h +++ b/esphome/components/api/proto.h @@ -413,14 +413,8 @@ class ProtoEncode { } /// Unaligned little-endian store; __builtin_memcpy stays inline even under -fno-builtin-memcpy. static inline void ESPHOME_ALWAYS_INLINE write_fixed32_le(uint8_t *__restrict__ pos, uint32_t value) { -#if __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__ - __builtin_memcpy(pos, &value, 4); -#else - pos[0] = static_cast(value); - pos[1] = static_cast(value >> 8); - pos[2] = static_cast(value >> 16); - pos[3] = static_cast(value >> 24); -#endif + const uint32_t le = convert_little_endian(value); + __builtin_memcpy(pos, &le, 4); } /// Write a precomputed tag byte + 32-bit value. Outlined on embedded: one copy beats inline stores per field. static PROTO_OUTLINE_FOR_SIZE uint8_t *write_tag_and_fixed32(uint8_t *__restrict__ pos PROTO_ENCODE_DEBUG_PARAM, diff --git a/script/api_protobuf/api_protobuf.py b/script/api_protobuf/api_protobuf.py index 2b8c688e84..3eb37f3b31 100755 --- a/script/api_protobuf/api_protobuf.py +++ b/script/api_protobuf/api_protobuf.py @@ -131,6 +131,12 @@ def force_str(force: bool) -> str: return str(force).lower() +def _encode_call(func: str, *args: str, force: bool = False) -> str: + """Emit one ProtoEncode call; every helper takes the cursor and returns it advanced.""" + suffix = "_force" if force else "" + return f"pos = ProtoEncode::{func}{suffix}({', '.join(('pos', *args))});" + + class TypeInfo(ABC): """Base class for all type information.""" @@ -264,14 +270,16 @@ class TypeInfo(ABC): # write_raw_byte(tag) + raw encode instead of the full encode_* method, # eliminating the zero-check branch and encode_field_raw indirection. # {value} is replaced with the actual field expression. - RAW_ENCODE_MAP: dict[str, str] = { - "encode_uint32": "pos = ProtoEncode::encode_varint_raw(pos, {value});", - "encode_uint64": "pos = ProtoEncode::encode_varint_raw_64(pos, {value});", - "encode_sint32": "pos = ProtoEncode::encode_varint_raw_short(pos, encode_zigzag32({value}));", - "encode_sint64": "pos = ProtoEncode::encode_varint_raw_64(pos, encode_zigzag64({value}));", - "encode_int64": "pos = ProtoEncode::encode_varint_raw_64(pos, static_cast({value}));", - "encode_bool": "pos = ProtoEncode::write_raw_byte(pos, {value} ? 0x01 : 0x00);", + RAW_ENCODE_MAP: dict[str, tuple[str, str]] = { + "encode_uint32": ("encode_varint_raw", "{value}"), + "encode_uint64": ("encode_varint_raw_64", "{value}"), + "encode_sint32": ("encode_varint_raw_short", "encode_zigzag32({value})"), + "encode_sint64": ("encode_varint_raw_64", "encode_zigzag64({value})"), + "encode_int64": ("encode_varint_raw_64", "static_cast({value})"), + "encode_bool": ("write_raw_byte", "{value} ? 0x01 : 0x00"), } + # Fixed32 value expression for the shared tag+fixed32 writer; None for other wire types + fixed32_value_template: str | None = None def _encode_with_precomputed_tag(self, value_expr: str) -> str | None: """Try to emit a precomputed-tag encode for a field. @@ -288,12 +296,17 @@ class TypeInfo(ABC): return None max_val = self.max_value # Only use RAW_ENCODE_MAP for forced fields or fields with max_value - raw_expr = None + raw = None if self.force or max_val is not None: - raw_expr = self.RAW_ENCODE_MAP.get(self.encode_func) - if raw_expr is None: + raw = self.RAW_ENCODE_MAP.get(self.encode_func) + if raw is None: return None - body = f"pos = ProtoEncode::write_raw_byte(pos, {tag});\n{raw_expr.format(value=value_expr)}" + func, arg = raw + body = ( + _encode_call("write_raw_byte", str(tag)) + + "\n" + + _encode_call(func, arg.format(value=value_expr)) + ) if self.force: return body # Non-forced with max_value: inline zero-check + raw encode @@ -314,14 +327,16 @@ class TypeInfo(ABC): return None # When max_len < 128, length varint is always 1 byte len_encode = ( - f"pos = ProtoEncode::write_raw_byte(pos, static_cast({len_expr}));" + _encode_call("write_raw_byte", f"static_cast({len_expr})") if max_len is not None and max_len < 128 - else f"pos = ProtoEncode::encode_varint_raw(pos, {len_expr});" + else _encode_call("encode_varint_raw", len_expr) ) - return ( - f"pos = ProtoEncode::write_raw_byte(pos, {tag});\n" - f"{len_encode}\n" - f"pos = ProtoEncode::encode_raw(pos, {data_expr}, {len_expr});" + return "\n".join( + ( + _encode_call("write_raw_byte", str(tag)), + len_encode, + _encode_call("encode_raw", data_expr, len_expr), + ) ) def _encode_fixed32_with_precomputed_tag(self, value_expr: str) -> str | None: @@ -330,22 +345,25 @@ class TypeInfo(ABC): if tag >= 128: return None if self.force: - return ( - f"pos = ProtoEncode::write_tag_and_fixed32(pos, {tag}, {value_expr});" - ) + return _encode_call("write_tag_and_fixed32", str(tag), value_expr) return ( f"if (uint32_t raw = {value_expr}; raw != 0) [[likely]] {{\n" - f" pos = ProtoEncode::write_tag_and_fixed32(pos, {tag}, raw);\n" + f" {_encode_call('write_tag_and_fixed32', str(tag), 'raw')}\n" "}" ) @property def encode_content(self) -> str: - if result := self._encode_with_precomputed_tag(f"this->{self.field_name}"): + value = f"this->{self.field_name}" + if result := self._encode_with_precomputed_tag(value): return result - if self.force: - return f"pos = ProtoEncode::{self.encode_func}_force(pos, {self.number}, this->{self.field_name});" - return f"pos = ProtoEncode::{self.encode_func}(pos, {self.number}, this->{self.field_name});" + if self.fixed32_value_template is not None and ( + result := self._encode_fixed32_with_precomputed_tag( + self.fixed32_value_template.format(value=value) + ) + ): + return result + return _encode_call(self.encode_func, str(self.number), value, force=self.force) encode_func = None @@ -650,13 +668,7 @@ class FloatType(FixedSizeTypeMixin, TypeInfo): encode_func = "encode_float" wire_type = WireType.FIXED32 # Uses wire type 5 - @property - def encode_content(self) -> str: - if result := self._encode_fixed32_with_precomputed_tag( - f"float_to_raw(this->{self.field_name})" - ): - return result - return super().encode_content + fixed32_value_template = "float_to_raw({value})" def dump(self, name: str) -> str: o = f'snprintf(buffer, sizeof(buffer), "%g", {name});\n' @@ -724,7 +736,7 @@ class UInt64Type(VarintTypeMixin, TypeInfo): if self.mac_address: return { **TypeInfo.RAW_ENCODE_MAP, - "encode_uint64": "pos = ProtoEncode::encode_varint_raw_48bit(pos, {value});", + "encode_uint64": ("encode_varint_raw_48bit", "{value}"), } return TypeInfo.RAW_ENCODE_MAP @@ -792,13 +804,7 @@ class Fixed32Type(FixedSizeTypeMixin, TypeInfo): o += "out.append(buffer);" return o - @property - def encode_content(self) -> str: - if result := self._encode_fixed32_with_precomputed_tag( - f"this->{self.field_name}" - ): - return result - return super().encode_content + fixed32_value_template = "{value}" def get_size_calculation(self, name: str, force: bool = False) -> str: field_id_size = self.calculate_field_id_size() @@ -872,9 +878,12 @@ class StringType(TypeInfo): f"this->{self.field_name}_ref_.size()", ): return result - if self.force: - return f"pos = ProtoEncode::encode_string_force(pos, {self.number}, this->{self.field_name}_ref_);" - return f"pos = ProtoEncode::encode_string(pos, {self.number}, this->{self.field_name}_ref_);" + return _encode_call( + "encode_string", + str(self.number), + f"this->{self.field_name}_ref_", + force=self.force, + ) def dump(self, name): # If name is 'it', this is a repeated field element - always use string @@ -972,7 +981,9 @@ class MessageType(TypeInfo): @property def encode_content(self) -> str: # Sub-message encoding needs buffer for backpatch/sync - return f"pos = ProtoEncode::{self.encode_func}(pos, buffer, {self.number}, this->{self.field_name});" + return _encode_call( + self.encode_func, "buffer", str(self.number), f"this->{self.field_name}" + ) @property def decode_length(self) -> str: @@ -1079,9 +1090,13 @@ class BytesType(TypeInfo): f"this->{self.field_name}_ptr_", f"this->{self.field_name}_len_" ): return result - if self.force: - return f"pos = ProtoEncode::encode_bytes_force(pos, {self.number}, this->{self.field_name}_ptr_, this->{self.field_name}_len_);" - return f"pos = ProtoEncode::encode_bytes(pos, {self.number}, this->{self.field_name}_ptr_, this->{self.field_name}_len_);" + return _encode_call( + "encode_bytes", + str(self.number), + f"this->{self.field_name}_ptr_", + f"this->{self.field_name}_len_", + force=self.force, + ) def dump(self, name: str) -> str: ptr_dump = f"format_hex_pretty(this->{self.field_name}_ptr_, this->{self.field_name}_len_)" @@ -1191,9 +1206,13 @@ class PointerToBytesBufferType(PointerToBufferTypeBase): f"this->{self.field_name}", f"this->{self.field_name}_len" ): return result - if self.force: - return f"pos = ProtoEncode::encode_bytes_force(pos, {self.number}, this->{self.field_name}, this->{self.field_name}_len);" - return f"pos = ProtoEncode::encode_bytes(pos, {self.number}, this->{self.field_name}, this->{self.field_name}_len);" + return _encode_call( + "encode_bytes", + str(self.number), + f"this->{self.field_name}", + f"this->{self.field_name}_len", + force=self.force, + ) @property def decode_length_content(self) -> str | None: @@ -1245,15 +1264,20 @@ class PointerToStringBufferType(PointerToBufferTypeBase): if max_len is not None and max_len < 128 and self.force: tag = self.calculate_tag() if tag < 128: - return f"pos = ProtoEncode::encode_short_string_force(pos, {tag}, this->{self.field_name});" + return _encode_call( + "encode_short_string_force", str(tag), f"this->{self.field_name}" + ) if result := self._encode_bytes_with_precomputed_tag( f"this->{self.field_name}.c_str()", f"this->{self.field_name}.size()", ): return result - if self.force: - return f"pos = ProtoEncode::encode_string_force(pos, {self.number}, this->{self.field_name});" - return f"pos = ProtoEncode::encode_string(pos, {self.number}, this->{self.field_name});" + return _encode_call( + "encode_string", + str(self.number), + f"this->{self.field_name}", + force=self.force, + ) @property def decode_length_content(self) -> str | None: @@ -1440,9 +1464,13 @@ class FixedArrayBytesType(TypeInfo): f"this->{self.field_name}", f"this->{self.field_name}_len", max_len=max_len ): return result - if self.force: - return f"pos = ProtoEncode::encode_bytes_force(pos, {self.number}, this->{self.field_name}, this->{self.field_name}_len);" - return f"pos = ProtoEncode::encode_bytes(pos, {self.number}, this->{self.field_name}, this->{self.field_name}_len);" + return _encode_call( + "encode_bytes", + str(self.number), + f"this->{self.field_name}", + f"this->{self.field_name}_len", + force=self.force, + ) def dump(self, name: str) -> str: return f"out.append(format_hex_pretty({name}, {name}_len));" @@ -1539,10 +1567,8 @@ class EnumType(VarintTypeMixin, TypeInfo): @property def encode_content(self) -> str: value_expr = f"static_cast(this->{self.field_name})" - if self.force: - return f"pos = ProtoEncode::{self.encode_func}_force(pos, {self.number}, {value_expr});" - return ( - f"pos = ProtoEncode::{self.encode_func}(pos, {self.number}, {value_expr});" + return _encode_call( + self.encode_func, str(self.number), value_expr, force=self.force ) def dump(self, name: str) -> str: @@ -1722,9 +1748,9 @@ def _generate_inline_encode_block( lines = [] lines.append(f"auto &sub_msg = {element};") - lines.append(f"pos = ProtoEncode::write_raw_byte(pos, {tag});") + lines.append(_encode_call("write_raw_byte", str(tag))) lines.append("uint8_t *len_pos = pos;") - lines.append("pos = ProtoEncode::reserve_byte(pos);") + lines.append(_encode_call("reserve_byte")) # Generate inline field encoding for each sub-message field for field in sub_desc.field: @@ -1796,15 +1822,22 @@ class FixedArrayRepeatedType(TypeInfo): def _encode_element(self, element: str) -> str: """Helper to generate encode statement for a single element.""" if isinstance(self._ti, EnumType): - return f"pos = ProtoEncode::{self._ti.encode_func}_force(pos, {self.number}, static_cast({element}));" + return _encode_call( + self._ti.encode_func, + str(self.number), + f"static_cast({element})", + force=True, + ) # Repeated message elements use encode_sub_message (force=true is default) if isinstance(self._ti, MessageType): if _is_inline_encode(self._ti.cpp_type): return _generate_inline_encode_block( self.number, self._ti.cpp_type, element ) - return f"pos = ProtoEncode::encode_sub_message(pos, buffer, {self.number}, {element});" - return f"pos = ProtoEncode::{self._ti.encode_func}_force(pos, {self.number}, {element});" + return _encode_call( + "encode_sub_message", "buffer", str(self.number), element + ) + return _encode_call(self._ti.encode_func, str(self.number), element, force=True) @property def cpp_type(self) -> str: @@ -2156,11 +2189,18 @@ class RepeatedTypeInfo(TypeInfo): def _encode_element_call(self, element: str) -> str: """Helper to generate encode call for a single element.""" if isinstance(self._ti, EnumType): - return f"pos = ProtoEncode::{self._ti.encode_func}_force(pos, {self.number}, static_cast({element}));" + return _encode_call( + self._ti.encode_func, + str(self.number), + f"static_cast({element})", + force=True, + ) # Repeated message elements use encode_sub_message (force=true is default) if isinstance(self._ti, MessageType): - return f"pos = ProtoEncode::encode_sub_message(pos, buffer, {self.number}, {element});" - return f"pos = ProtoEncode::{self._ti.encode_func}_force(pos, {self.number}, {element});" + return _encode_call( + "encode_sub_message", "buffer", str(self.number), element + ) + return _encode_call(self._ti.encode_func, str(self.number), element, force=True) @property def encode_content(self) -> str: @@ -2169,7 +2209,7 @@ class RepeatedTypeInfo(TypeInfo): # Special handling for const char* elements (when container_no_template contains "const char") if "const char" in self._container_no_template: o = f"for (const char *it : *this->{self.field_name}) {{\n" - o += f" pos = ProtoEncode::{self._ti.encode_func}_force(pos, {self.number}, it, strlen(it));\n" + o += f" {_encode_call(self._ti.encode_func, str(self.number), 'it', 'strlen(it)', force=True)}\n" else: o = f"for (const auto &it : *this->{self.field_name}) {{\n" o += f" {self._encode_element_call('it')}\n" diff --git a/tests/integration/state_utils.py b/tests/integration/state_utils.py index 9c0debbc5c..68403ea12c 100644 --- a/tests/integration/state_utils.py +++ b/tests/integration/state_utils.py @@ -57,6 +57,31 @@ async def wait_for_state( return await asyncio.wait_for(future, timeout=timeout) +class StateWaiter: + """Route one state subscription to any number of predicate waits.""" + + def __init__(self) -> None: + self._waiters: list[ + tuple[Callable[[EntityState], bool], asyncio.Future[EntityState]] + ] = [] + + def on_state(self, state: EntityState) -> None: + for predicate, future in list(self._waiters): + if not future.done() and predicate(state): + future.set_result(state) + + async def expect( + self, predicate: Callable[[EntityState], bool], timeout: float = 5.0 + ) -> EntityState: + """Wait for the next state matching ``predicate``.""" + entry = (predicate, asyncio.get_running_loop().create_future()) + self._waiters.append(entry) + try: + return await asyncio.wait_for(entry[1], timeout) + finally: + self._waiters.remove(entry) + + def find_entity[T: EntityInfo]( entities: list[EntityInfo], object_id_substring: str, diff --git a/tests/integration/test_api_encode_boundaries.py b/tests/integration/test_api_encode_boundaries.py index 7441437390..6ab7bc7476 100644 --- a/tests/integration/test_api_encode_boundaries.py +++ b/tests/integration/test_api_encode_boundaries.py @@ -1,12 +1,11 @@ """Encode paths at their branch boundaries: zero skipped float, fixed32 state, negative int32, -length prefixes of two varint bytes, two byte field tags and a clean disconnect.""" +length prefixes of two varint bytes and two byte field tags.""" from __future__ import annotations import asyncio from aioesphomeapi import ( - EntityState, NumberState, SelectInfo, SensorInfo, @@ -15,7 +14,7 @@ from aioesphomeapi import ( ) import pytest -from .state_utils import InitialStateHelper, require_entity +from .state_utils import InitialStateHelper, StateWaiter, require_entity from .types import APIClientConnectedFactory, RunCompiledFunction LONG_OPTION = ( @@ -30,12 +29,12 @@ async def test_api_encode_boundaries( run_compiled: RunCompiledFunction, api_client_connected: APIClientConnectedFactory, ) -> None: - loop = asyncio.get_running_loop() async with run_compiled(yaml_config), api_client_connected() as client: - device_info = await client.device_info() + device_info, (entities, _) = await asyncio.gather( + client.device_info(), client.list_entities_services() + ) assert device_info.suggested_area == "Kitchen" - entities, _ = await client.list_entities_services() sensor = require_entity(entities, "zero_then_value", SensorInfo) assert sensor.accuracy_decimals == -2 select = require_entity(entities, "long_option_select", SelectInfo) @@ -45,25 +44,9 @@ async def test_api_encode_boundaries( number = require_entity(entities, "negative_number") button = require_entity(entities, "publish_values") - sensor_value: asyncio.Future[float] = loop.create_future() - text_value: asyncio.Future[str] = loop.create_future() initial = InitialStateHelper(entities) - - def on_state(state: EntityState) -> None: - if ( - isinstance(state, SensorState) - and state.key == sensor.key - and not sensor_value.done() - ): - sensor_value.set_result(state.state) - elif ( - isinstance(state, TextSensorState) - and state.key == text.key - and not text_value.done() - ): - text_value.set_result(state.state) - - client.subscribe_states(initial.on_state_wrapper(on_state)) + waiter = StateWaiter() + client.subscribe_states(initial.on_state_wrapper(waiter.on_state)) await initial.wait_for_initial_states() # A float of exactly zero is skipped on the wire and must still read as 0.0, not missing @@ -75,8 +58,19 @@ async def test_api_encode_boundaries( assert first_number.state == -123.5 client.button_command(button.key) - assert await asyncio.wait_for(sensor_value, 5.0) == 12.5 - assert await asyncio.wait_for(text_value, 5.0) == "y" * 200 - - # DisconnectRequest and DisconnectResponse carry no fields - await client.disconnect() + await asyncio.gather( + waiter.expect( + lambda s: ( + isinstance(s, SensorState) + and s.key == sensor.key + and s.state == 12.5 + ) + ), + waiter.expect( + lambda s: ( + isinstance(s, TextSensorState) + and s.key == text.key + and s.state == "y" * 200 + ) + ), + ) diff --git a/tests/unit_tests/components/api/test_api_protobuf_generator.py b/tests/unit_tests/components/api/test_api_protobuf_generator.py index 763a764894..b872e300ba 100644 --- a/tests/unit_tests/components/api/test_api_protobuf_generator.py +++ b/tests/unit_tests/components/api/test_api_protobuf_generator.py @@ -118,10 +118,9 @@ def test_message_id_above_maximum_is_rejected() -> None: validate_message_id(MAX_MESSAGE_ID + 1, "TooBigMessage") -def _encode_field( - field_type: int, number: int = 1, force: bool = False, repeated: bool = False -) -> str: - """Return the encode statement the generator emits for one encode-only field.""" +def _field( + field_type: int, number: int = 1, *, force: bool = False, repeated: bool = False +) -> descriptor_pb2.FieldDescriptorProto: field = descriptor_pb2.FieldDescriptorProto( name="value", number=number, type=field_type ) @@ -129,8 +128,17 @@ def _encode_field( field.label = descriptor_pb2.FieldDescriptorProto.LABEL_REPEATED if force: field.options.Extensions[pb.force] = True - ti = create_field_type_info(field, needs_decode=False, needs_encode=True) - return ti.encode_content + return field + + +def _encode_field( + field_type: int, number: int = 1, force: bool = False, repeated: bool = False +) -> str: + """Return the encode statement the generator emits for one encode-only field.""" + field = _field(field_type, number, force=force, repeated=repeated) + return create_field_type_info( + field, needs_decode=False, needs_encode=True + ).encode_content SCALAR_TYPES = [