Centralize stamp injection between the cc toolchain and PostMark.

PiperOrigin-RevId: 903199755
Change-Id: I56023825cfdb0ece1f99ad0c7e83ff6789fe9c16

Author: bigelephant29

cc/common/cc_helper_internal.bzl

@@ -44,6 +44,87 @@ CPP_SOURCE_TYPE_HEADER = "HEADER"
 CPP_SOURCE_TYPE_SOURCE = "SOURCE"
 CPP_SOURCE_TYPE_CLIF_INPUT_PROTO = "CLIF_INPUT_PROTO"
 
+def get_fdo_build_stamp(cpp_configuration, fdo_context, feature_configuration):
+    """Returns the FDO build stamp.
+
+    Args:
+      cpp_configuration: The C++ configuration.
+      fdo_context: The FDO context.
+      feature_configuration: The feature configuration.
+
+    Returns:
+      The FDO build stamp string, or None if FDO is not enabled.
+    """
+    branch_fdo_profile = getattr(fdo_context, "branch_fdo_profile", None)
+    if branch_fdo_profile:
+        branch_fdo_mode = branch_fdo_profile.branch_fdo_mode
+        if branch_fdo_mode == "auto_fdo":
+            return "AFDO" if feature_configuration.is_enabled("autofdo") else None
+        if branch_fdo_mode == "xbinary_fdo":
+            return "XFDO" if feature_configuration.is_enabled("xbinaryfdo") else None
+        if branch_fdo_mode == "llvm_cs_fdo" or cpp_configuration.cs_fdo_instrument():
+            return "CSFDO"
+    if branch_fdo_profile or cpp_configuration.fdo_instrument():
+        return "FDO"
+    return None
+
+def get_linkstamp_stamps(
+        cc_toolchain,
+        feature_configuration,
+        label_replacement,
+        output_replacement,
+        additional_linkstamp_defines):
+    """Returns a dict of stamps for linkstamp compilation/PostMark.
+
+    Args:
+      cc_toolchain: The C++ toolchain provider.
+      feature_configuration: The feature configuration.
+      label_replacement: String to replace ${LABEL} in linkstamp defines.
+      output_replacement: String to replace ${OUTPUT_PATH} in linkstamp defines.
+      additional_linkstamp_defines: A list of additional defines for linkstamp compilation.
+
+    Returns:
+      A dictionary of linkstamp defines.
+    """
+    fdo_build_stamp = get_fdo_build_stamp(
+        cc_toolchain._cpp_configuration,
+        cc_toolchain._fdo_context,
+        feature_configuration,
+    )
+    stamps = {
+        "GPLATFORM": cc_toolchain.toolchain_id,
+        "BUILD_COVERAGE_ENABLED": "1" if feature_configuration.is_enabled("coverage") else "0",
+        # G3_TARGET_NAME is a C string literal that normally contain the label of the target
+        # being linked.  However, they are set differently when using shared native deps. In
+        # that case, a single .so file is shared by multiple targets, and its contents cannot
+        # depend on which target(s) were specified on the command line.  So in that case we
+        # have to use the (obscure) name of the .so file instead, or more precisely the path of
+        # the .so file relative to the workspace root.
+        "G3_TARGET_NAME": label_replacement,
+        # G3_BUILD_TARGET is a C string literal containing the output of this
+        # link.  (An undocumented and untested invariant is that G3_BUILD_TARGET is the
+        # location of the executable, either absolutely, or relative to the directory part of
+        # BUILD_INFO.)
+        "G3_BUILD_TARGET": output_replacement,
+    }
+    if fdo_build_stamp:
+        stamps["BUILD_FDO_TYPE"] = fdo_build_stamp
+
+    if feature_configuration.is_enabled("thin_lto"):
+        stamps["BUILD_LTO_TYPE"] = "thin"
+
+    if additional_linkstamp_defines:
+        for d in additional_linkstamp_defines:
+            # TODO: b/503100490 - This replacement looks redundant, but we should remove this in a separate CL.
+            d = d.replace("${LABEL}", label_replacement).replace("${OUTPUT_PATH}", output_replacement)
+            if "=" in d:
+                k, v = d.split("=", 1)
+                stamps[k] = v
+            else:
+                stamps[d] = "1"
+
+    return stamps
+
 # LINT.IfChange(forked_exports)
 
 CREATE_COMPILE_ACTION_API_ALLOWLISTED_PACKAGES = [("", "devtools/rust/cc_interop"), ("", "third_party/crubit"), ("", "tools/build_defs/clif")]

cc/private/compile/compile_build_variables.bzl

@@ -15,7 +15,7 @@
 All build variables we create for various `CppCompileAction`s
 """
 
-load("//cc/common:cc_helper_internal.bzl", "extensions", _PRIVATE_STARLARKIFICATION_ALLOWLIST = "PRIVATE_STARLARKIFICATION_ALLOWLIST")
+load("//cc/common:cc_helper_internal.bzl", "extensions", "get_fdo_build_stamp", "get_linkstamp_stamps", _PRIVATE_STARLARKIFICATION_ALLOWLIST = "PRIVATE_STARLARKIFICATION_ALLOWLIST")
 load("//cc/private:cc_internal.bzl", _cc_internal = "cc_internal")
 load("//cc/private/rules_impl:native_cc_common.bzl", _cc_common_internal = "native_cc_common")
 
@@ -154,7 +154,7 @@ def create_compile_variables(
     common_vars = _setup_common_compile_build_variables_internal(
         feature_configuration = feature_configuration,
         is_using_memprof = getattr(fdo_context, "memprof_profile_artifact", None) != None,
-        fdo_build_stamp = _get_fdo_build_stamp(cpp_configuration, fdo_context, feature_configuration),
+        fdo_build_stamp = get_fdo_build_stamp(cpp_configuration, fdo_context, feature_configuration),
         variables_extension = variables_extension,
         includes = includes or [],
         include_dirs = include_directories or depset(),
@@ -194,7 +194,7 @@ def setup_common_compile_build_variables(
     common_vars = _setup_common_compile_build_variables_internal(
         feature_configuration = feature_configuration,
         is_using_memprof = getattr(fdo_context, "memprof_profile_artifact", None) != None,
-        fdo_build_stamp = _get_fdo_build_stamp(cpp_configuration, fdo_context, feature_configuration),
+        fdo_build_stamp = get_fdo_build_stamp(cpp_configuration, fdo_context, feature_configuration),
         variables_extension = variables_extension,
         include_dirs = cc_compilation_context.includes,
         quote_include_dirs = cc_compilation_context.quote_includes,
@@ -259,20 +259,6 @@ def _setup_common_compile_build_variables_internal(
         result[_VARS.EXTERNAL_INCLUDE_PATHS] = external_include_dirs
     return _cc_internal.cc_toolchain_variables(vars = result)
 
-def _get_fdo_build_stamp(cpp_configuration, fdo_context, feature_configuration):
-    branch_fdo_profile = getattr(fdo_context, "branch_fdo_profile", None)
-    if branch_fdo_profile:
-        branch_fdo_mode = branch_fdo_profile.branch_fdo_mode
-        if branch_fdo_mode == "auto_fdo":
-            return "AFDO" if feature_configuration.is_enabled("autofdo") else None
-        if branch_fdo_mode == "xbinary_fdo":
-            return "XFDO" if feature_configuration.is_enabled("xbinaryfdo") else None
-        if branch_fdo_mode == "llvm_cs_fdo" or cpp_configuration.cs_fdo_instrument():
-            return "CSFDO"
-    if branch_fdo_profile or cpp_configuration.fdo_instrument():
-        return "FDO"
-    return None
-
 # Note: this method is side-effect free, callers should add fdo inputs to
 # cc_compile_action_builder themselves
 def get_specific_compile_build_variables(
@@ -573,12 +559,6 @@ def get_linkstamp_compile_variables(
         action_name = "linkstamp-compile",
     ):
         fail("Action 'linkstamp-compile' is not configured.")
-    fdo_build_stamp = _get_fdo_build_stamp(
-        cc_toolchain._cpp_configuration,
-        cc_toolchain._fdo_context,
-        feature_configuration,
-    )
-    code_coverage_enabled = feature_configuration.is_enabled("coverage")
     copts = get_copts(
         language = "c++",  # The only language that receives special treatment is "objc".
         cpp_configuration = cc_toolchain._cpp_configuration,
@@ -593,8 +573,7 @@ def get_linkstamp_compile_variables(
         output_replacement = output_replacement,
         additional_linkstamp_defines = additional_linkstamp_defines,
         cc_toolchain = cc_toolchain,
-        fdo_build_stamp = fdo_build_stamp,
-        code_coverage_enabled = code_coverage_enabled,
+        feature_configuration = feature_configuration,
     )
     return create_compile_variables(
         feature_configuration = feature_configuration,
@@ -608,37 +587,31 @@ def get_linkstamp_compile_variables(
         user_compile_flags = copts,
     )
 
+def _stamps_to_defines(stamps):
+    """Converts a dictionary of stamps to a list of C preprocessor defines."""
+    defines = []
+    for k, v in stamps.items():
+        if v.isdigit() or k == "BUILD_LTO_TYPE":
+            # BUILD_LTO_TYPE is a special case because of the AS_STRING(x) call in "builddata_globals.cc".
+            # Normally, a string set as -DKEY_NAME="VALUE" is interpreted as KEY_NAME="VALUE",
+            # but for this specific case, it's KEY_NAME="\"VALUE\"" because of AS_STRING(x).
+            defines.append("{}={}".format(k, v))
+        else:
+            defines.append('{}="{}"'.format(k, v))
+    return defines
+
 def _compute_all_linkstamp_defines(
         label_replacement,
         output_replacement,
         additional_linkstamp_defines,
         cc_toolchain,
-        fdo_build_stamp,
-        code_coverage_enabled):
+        feature_configuration):
     """Computes defines for linkstamp compilation."""
-    defines = [
-        'GPLATFORM="' + cc_toolchain.toolchain_id + '"',
-        "BUILD_COVERAGE_ENABLED=" + ("1" if code_coverage_enabled else "0"),
-        # G3_TARGET_NAME is a C string literal that normally contain the label of the target
-        # being linked.  However, they are set differently when using shared native deps. In
-        # that case, a single .so file is shared by multiple targets, and its contents cannot
-        # depend on which target(s) were specified on the command line.  So in that case we
-        # have to use the (obscure) name of the .so file instead, or more precisely the path of
-        # the .so file relative to the workspace root.
-        'G3_TARGET_NAME="${LABEL}"',
-        # G3_BUILD_TARGET is a C string literal containing the output of this
-        # link.  (An undocumented and untested invariant is that G3_BUILD_TARGET is the
-        # location of the executable, either absolutely, or relative to the directory part of
-        # BUILD_INFO.)
-        'G3_BUILD_TARGET="${OUTPUT_PATH}"',
-    ]
-    if additional_linkstamp_defines:
-        defines.extend(additional_linkstamp_defines)
-
-    if fdo_build_stamp:
-        defines.append('BUILD_FDO_TYPE="' + fdo_build_stamp + '"')
-
-    return [
-        define.replace("${LABEL}", label_replacement).replace("${OUTPUT_PATH}", output_replacement)
-        for define in defines
-    ]
+    stamps = get_linkstamp_stamps(
+        cc_toolchain = cc_toolchain,
+        feature_configuration = feature_configuration,
+        label_replacement = label_replacement,
+        output_replacement = output_replacement,
+        additional_linkstamp_defines = additional_linkstamp_defines,
+    )
+    return _stamps_to_defines(stamps)

cc/private/rules_impl/cc_binary_impl.bzl

@@ -679,7 +679,9 @@ def cc_binary_impl(ctx, additional_linkopts, force_linkstatic = False):
             cc_toolchain,
             binary,
             output_binary_for_linking,
+            feature_configuration = feature_configuration,
             additional_stamp_infos = additional_stamp_infos,
+            additional_linkstamp_defines = [],
         )
 
     cc_linking_outputs_binary_library = cc_linking_outputs_binary.library_to_link

cc/private/toolchain/unix_cc_toolchain_config.bzl

@@ -1641,6 +1641,7 @@ def _impl(ctx):
                     ),
                 ],
             ),
+            # TODO: b/503100490 - We should centralize the stamp injection and remove this.
             flag_set(
                 actions = [ACTION_NAMES.linkstamp_compile],
                 flag_groups = [flag_group(flags = ["-DBUILD_LTO_TYPE=thin"])],