esphome/esphome/platformio/pch.py.script

import hashlib
import os
from pathlib import Path
import posixpath
import re
import shlex
import subprocess
import traceback

# pylint: disable=E0602
Import("env")  # noqa: F821
_projenv_error = None
try:
    Import("projenv")  # noqa: F821
except Exception as err:  # noqa: BLE001 -- not exported under -t nobuild
    projenv = None
    _projenv_error = err

# Precompile the src force-includes plus defines.h (which pulls in
# Arduino.h on Arduino platforms) and force-include the result into C++ src
# compiles only; those TUs already include this content first, so their
# preprocessed output is unchanged. Post script: the final src flags exist,
# nothing compiled yet. Registration is gated host-side (the platform's
# __init__.py, pch_enabled()). Keep the header tail, ccache values,
# include-closure recipe, and the checksum/failed-marker stamp flow in sync
# with build_helpers/pch.py.

# Compiler failures that clear on their own must not latch the .failed marker
_TRANSIENT_ERRORS = ("No space left", "Cannot allocate", "Resource temporarily")

_INCLUDE_RE = re.compile(rb'^\s*#\s*include\s+["<]([^">]+)[">]', re.MULTILINE)
_CORE_HEADER = "esphome/core/defines.h"


def _include_closure(src_dir: Path, roots: list) -> dict:
    """Quoted-include closure: src-relative name -> contents (mirror of
    build_helpers/pch.py)."""
    seen = {}
    stack = [(name, "") for name in roots]
    while stack:
        name, from_dir = stack.pop()
        for candidate in (f"{from_dir}/{name}" if from_dir else name, name):
            rel = posixpath.normpath(candidate)
            if not rel.startswith("..") and (src_dir / rel).is_file():
                break
        else:
            continue
        if rel in seen:
            continue
        try:
            data = (src_dir / rel).read_bytes()
        except OSError as err:
            # If stat also fails the identity is unknown: the OSError
            # propagates to the outer handler, which skips the pch
            print(f"ESPHome: could not read {rel} for the pch checksum: {err}")
            st = (src_dir / rel).stat()
            data = f"<unreadable:{st.st_mtime_ns}:{st.st_size}>".encode()
        seen[rel] = data
        parent = posixpath.dirname(rel)
        stack.extend(
            (inc.decode(errors="surrogateescape"), parent)
            for inc in _INCLUDE_RE.findall(data)
        )
    return seen


def _shell_arg(element) -> str:
    """One compiler argv from one SCons element, matching the real spawn:
    SCons whole-quotes spaced elements, the shell unquotes the rest. On
    Windows there is no POSIX shell pass and shlex would eat path
    backslashes."""
    arg = str(element)
    if " " in arg or os.name == "nt":
        return arg.replace('\\"', '"')
    if not arg.strip():
        return arg
    try:
        tokens = shlex.split(arg)
    except ValueError as err:
        print(f"ESPHome: could not lex flag {arg!r} for the pch: {err}")
        return arg
    if len(tokens) != 1:
        # The shell-quoting model is wrong for this element; say so rather
        # than surfacing only as downstream pch warnings
        print(f"ESPHome: passing flag {arg!r} through unlexed for the pch")
        return arg
    return tokens[0]


def _compile_gch(cxx, flags, header: Path, gch: Path, proj_dir: Path):
    """Compile the .gch, then probe that the toolchain can load it back
    (GCC 10 on macOS arm64 builds one it then rejects per-process: "had
    text segment at different address"). Returns a deterministic error
    string or None; OSError propagates for transient handling."""
    result = subprocess.run(  # noqa: PLW1510
        [cxx, "-x", "c++-header", *flags, "-c", str(header), "-o", str(gch)],
        cwd=proj_dir,
        capture_output=True,
        text=True,
    )
    if result.returncode < 0:
        # Signal-killed (OOM, ^C) is environmental; raising OSError routes
        # it to the transient no-marker path
        raise OSError(f"compiler killed by signal {-result.returncode}")
    if result.returncode != 0:
        return result.stderr
    probe = subprocess.run(  # noqa: PLW1510
        [
            cxx,
            *flags,
            "-MF",
            os.devnull,
            "-Winvalid-pch",
            "-include",
            str(header),
            "-fsyntax-only",
            "-x",
            "c++",
            "-",
        ],
        cwd=proj_dir,
        input="",
        capture_output=True,
        text=True,
    )
    if probe.returncode != 0 or ".gch" in probe.stderr:
        return f"toolchain cannot load the pch: {probe.stderr.strip()}"
    return None


def _read_stamp(path: Path) -> str:
    """A corrupt sidecar must read as stale, not kill the pch forever."""
    try:
        return path.read_text(encoding="utf-8").strip()
    except (OSError, UnicodeDecodeError):
        return ""


def _setup_pch() -> None:
    if projenv is None:
        # Expected under -t nobuild; anything else must leave a trail
        print(f"ESPHome: projenv unavailable ({_projenv_error}); skipping pch")
        return
    # Project root, not $BUILD_DIR: SCons compiles run with the project dir
    # as cwd, so "-include esphome_pch.h" resolves here as a relative path.
    # An absolute path would put the per-device build path on every compile
    # command and defeat cross-device ccache sharing.
    proj_dir = Path(env.subst("$PROJECT_DIR"))  # noqa: F821
    src_dir = Path(env.subst("$PROJECT_SRC_DIR"))  # noqa: F821
    header = proj_dir / "esphome_pch.h"
    gch = Path(f"{header}.gch")
    sum_path = Path(f"{gch}.sum")

    cxx = projenv.subst("$CXX")  # noqa: F821
    # The header holds the -include entries itself, so the .gch compile must
    # not see them; consumers keep theirs, which the .gch then satisfies.
    flags = []
    include_headers = []
    flag_it = iter(
        _shell_arg(element)
        for element in projenv.subst_list("$CXXFLAGS $CCFLAGS $_CCCOMCOM")[0]  # noqa: F821
    )
    for tok in flag_it:
        if tok == "-include":
            include_headers.append(next(flag_it, ""))
        else:
            flags.append(tok)
    if any(not name for name in include_headers):
        print("ESPHome: build_src_flags has a trailing -include; skipping pch")
        return
    # Fold only names that resolve under src/: consumers keep their own
    # -include entries, so an unguarded user header folded here would be
    # included twice. An unfolded header simply stays consumer-only and
    # ccache hashes it directly off the command line.
    folded = [
        name
        for name in include_headers
        # An absolute name would sneak past src_dir /: keep it consumer-only
        if not Path(name).is_absolute() and (src_dir / name).is_file()
    ]
    if unfolded := [n for n in include_headers if n not in folded]:
        print(f"ESPHome: not precompiling non-src force-includes: {unfolded}")
    content = "".join(f'#include "{name}"\n' for name in (*folded, _CORE_HEADER))

    digest = hashlib.sha256()
    digest.update(content.encode(errors="surrogateescape"))
    digest.update(cxx.encode(errors="surrogateescape"))
    # Mirror CCACHE_BASEDIR: strip the per-device build path so identical
    # configs produce identical .sum files and share cache entries
    flags_id = " ".join(flags)
    if basedir := os.environ.get("CCACHE_BASEDIR"):
        flags_id = flags_id.replace(basedir, "")
    digest.update(flags_id.encode(errors="surrogateescape"))
    # GCC never validates a .gch against its source headers, and PlatformIO
    # package paths carry no version, so a package bump must invalidate here
    platform = env.PioPlatform()  # noqa: F821
    for package in sorted(platform.packages):
        try:
            version = platform.get_package_version(package)
        except KeyError:
            # Only trust KeyError as "absent" when the package really is not
            # installed; an unresolved manifest must not hash as a constant
            if platform.get_package(package) is not None:
                print(f"ESPHome: skipping precompiled header: no version for {package}")
                return
            version = None  # absent optional package
        except Exception as err:  # noqa: BLE001
            # Without trustworthy package identity a stale .gch could be
            # reused across upgrades; skip the pch instead
            print(f"ESPHome: skipping precompiled header: {err}")
            return
        digest.update(f"{package}={version}".encode())
        digest.update(b"\0")
    closure = _include_closure(src_dir, [*folded, _CORE_HEADER])
    for rel in sorted(closure):
        digest.update(rel.encode(errors="surrogateescape"))
        digest.update(closure[rel])
        digest.update(b"\0")
    # Project-local include dirs outside src (e.g. rp2's lwip_override)
    # hold generated headers the src closure cannot see; hash them so an
    # ESPHome-side change invalidates an existing build dir
    prev = ""
    for tok in flags:
        inc = tok[2:] if tok.startswith("-I") and len(tok) > 2 else ""
        if prev == "-I":
            inc = tok
        prev = tok
        if not inc:
            continue
        inc_dir = Path(inc)
        if not (
            inc_dir.is_dir()
            and inc_dir.is_relative_to(proj_dir)
            and not inc_dir.is_relative_to(src_dir)
            # lib_deps trees are not part of the prefix closure today (the
            # roots resolve under src/ only); walking them would read every
            # library file on every build for nothing
            and not inc_dir.is_relative_to(proj_dir / ".piolibdeps")
            and not inc_dir.is_relative_to(proj_dir / ".pioenvs")
        ):
            continue
        headers = (
            p
            for p in inc_dir.rglob("*")
            if p.is_file() and p.suffix in (".h", ".hpp", ".hh", ".inc")
        )
        for local in sorted(headers):
            try:
                data = local.read_bytes()
            except OSError as err:
                # mtime/size keep a changed-but-unreadable header shifting
                # the digest; a stat failure propagates and skips the pch
                print(f"ESPHome: could not read {local} for the pch checksum: {err}")
                st = local.stat()
                data = f"<unreadable:{st.st_mtime_ns}:{st.st_size}>".encode()
            digest.update(str(local.relative_to(proj_dir)).encode())
            digest.update(data)
            digest.update(b"\0")
    checksum = digest.hexdigest()

    # The ccache .sum sidecar doubles as the freshness stamp
    if (
        not header.is_file()
        or not gch.is_file()
        or not sum_path.is_file()
        or (_read_stamp(sum_path) != checksum)
    ):
        failed_marker = Path(f"{gch}.failed")
        if _read_stamp(failed_marker) == checksum:
            print(
                "ESPHome: skipping precompiled header (previous attempt "
                f"failed); delete {failed_marker.name} to retry"
            )
            return
        header.write_text(content, encoding="utf-8")
        try:
            error = _compile_gch(cxx, flags, header, gch, proj_dir)
        except OSError as err:
            # Transient spawn/IO failure: no marker, retry next build
            print(f"ESPHome: precompiled header compile did not run: {err}")
            gch.unlink(missing_ok=True)
            sum_path.unlink(missing_ok=True)
            return
        if error is not None:
            print("ESPHome: precompiled header failed; compiling without it")
            print(error)
            gch.unlink(missing_ok=True)
            sum_path.unlink(missing_ok=True)
            if any(m in error for m in _TRANSIENT_ERRORS):
                # Resource exhaustion clears on its own; retry next build
                return
            # Skip retries until a flag/header/platform change alters the checksum
            failed_marker.write_text(checksum + "\n", encoding="utf-8")
            return
        failed_marker.unlink(missing_ok=True)
        sum_path.write_text(checksum + "\n", encoding="utf-8")

    # projenv["ENV"] aliases os.environ under PlatformIO, so these reach
    # framework/library TUs too; only time_macros affects non-pch TUs (the
    # trade-off ccache_pch_env documents). User-set values win.
    for key, value in (
        ("CCACHE_SLOPPINESS", "pch_defines,time_macros"),
        ("CCACHE_PCH_EXTSUM", "true"),
    ):
        if key not in os.environ:
            projenv["ENV"][key] = value  # noqa: F821
    sloppiness = os.environ.get("CCACHE_SLOPPINESS")
    if sloppiness is not None:
        tokens = {tok.strip() for tok in sloppiness.split(",")}
        missing = [t for t in ("pch_defines", "time_macros") if t not in tokens]
        if missing:
            # Without these ccache declines every pch-consuming compile;
            # union rather than override so the user's own tokens survive
            projenv["ENV"]["CCACHE_SLOPPINESS"] = ",".join((sloppiness, *missing))  # noqa: F821
            print(f"ESPHome: adding {','.join(missing)} to CCACHE_SLOPPINESS for the pch")

    # Prepended so it is processed before the build_src_flags -include
    # entries: GCC only uses a .gch while no other tokens have been seen.
    # The relative name also reaches "pio run -t idedata" output; external
    # consumers replaying cxx_flags must run from the project dir.
    projenv.Prepend(CXXFLAGS=["-Winvalid-pch", "-include", header.name])  # noqa: F821
    print("ESPHome: Compiling with precompiled header")


try:
    _setup_pch()
except Exception:  # noqa: BLE001 -- a speedup must never break the build
    # Stable marker: an unexpected error, unlike the expected skip prints
    print("ESPHome: pch internal error; compiling without it")
    traceback.print_exc()