From dff00eaa7886fdfaef3a6e02abd8df6aaf16a218 Mon Sep 17 00:00:00 2001
From: Matt Hurd <matthurd@google.com>
Date: Fri, 3 Apr 2026 15:57:36 -0700
Subject: [PATCH] Validate selective binary loading in builds

PiperOrigin-RevId: 894290467
---
 BUILD                                   | 13 ++++++
 xprof/pywrap/BUILD                      | 11 ++++-
 xprof/pywrap/_pywrap_profiler_plugin.py | 57 +++++++++++++++++++++++++
 xprof/pywrap/profiler_wrapper_test.py   |  6 +++
 4 files changed, 86 insertions(+), 1 deletion(-)

diff --git a/BUILD b/BUILD
index e2b5c046c..6cfa5fe91 100644
--- a/BUILD
+++ b/BUILD
@@ -1,3 +1,4 @@
+load("@bazel_skylib//rules:common_settings.bzl", "bool_flag")
 load("@repository_configuration//:repository_config.bzl", "PROFILER_REQUIREMENTS_FILE")
 load("@rules_python//python:pip.bzl", "compile_pip_requirements")
 
@@ -15,6 +16,18 @@ exports_files([
     "rollup.config.js",
 ])
 
+bool_flag(
+    name = "enable_module_summary",
+    build_setting_default = False,
+    visibility = ["//visibility:public"],
+)
+
+config_setting(
+    name = "module_summary_enabled",
+    flag_values = {":enable_module_summary": "True"},
+    visibility = ["//visibility:public"],
+)
+
 py_library(
     name = "expect_tensorflow_installed",
     # This is a dummy rule used as a tensorflow dependency in open-source.
diff --git a/xprof/pywrap/BUILD b/xprof/pywrap/BUILD
index 9e0452721..78c87f26e 100644
--- a/xprof/pywrap/BUILD
+++ b/xprof/pywrap/BUILD
@@ -22,6 +22,9 @@ cc_binary(
     ] + select({
         ":dbg_build": ["-g"],
         "//conditions:default": [],
+    }) + select({
+        "//xprof:module_summary_enabled": ["-DENABLE_MODULE_SUMMARY"],
+        "//conditions:default": [],
     }),
     linkopts = select({
         "@xla//xla/tsl:windows": [
@@ -43,7 +46,13 @@ cc_binary(
         "@org_xprof//xprof/convert:tool_options",
         "@rules_python//python/cc:current_py_cc_headers",
         "@xla//xla/tsl/profiler/rpc/client:capture_profile",
-    ],
+    ] + select({
+        "//xprof:module_summary_enabled": [
+            "@org_xprof//internal/module_summary:module_summary_c_api",
+            "@org_xprof//internal/module_summary:refactored_tool_lib",
+        ],
+        "//conditions:default": [],
+    }),
 )
 
 py_library(
diff --git a/xprof/pywrap/_pywrap_profiler_plugin.py b/xprof/pywrap/_pywrap_profiler_plugin.py
index 2257b4c45..b6cbc74f9 100644
--- a/xprof/pywrap/_pywrap_profiler_plugin.py
+++ b/xprof/pywrap/_pywrap_profiler_plugin.py
@@ -505,3 +505,60 @@ def initialize_stubs(worker_service_addresses: str) -> None:
   _lib.InitializeStubs(
       worker_service_addresses.encode() if worker_service_addresses else None
   )
+
+
+if hasattr(_lib, "CreateModuleSummary"):
+  _lib.CreateModuleSummary.argtypes = [ctypes.c_char_p]
+  _lib.CreateModuleSummary.restype = ctypes.c_void_p
+
+  _lib.GetTotalCount.argtypes = [ctypes.c_void_p]
+  _lib.GetTotalCount.restype = ctypes.c_int
+
+  _lib.GetUniqueResourceCount.argtypes = [ctypes.c_void_p]
+  _lib.GetUniqueResourceCount.restype = ctypes.c_int
+
+  _lib.GetHistogramBuckets.argtypes = [ctypes.c_void_p]
+  _lib.GetHistogramBuckets.restype = ctypes.c_int
+
+  _lib.GetBucketKey.argtypes = [ctypes.c_void_p, ctypes.c_int]
+  _lib.GetBucketKey.restype = ctypes.c_int
+
+  _lib.GetBucketValue.argtypes = [ctypes.c_void_p, ctypes.c_int]
+  _lib.GetBucketValue.restype = ctypes.c_int
+
+  _lib.FreeModuleSummary.argtypes = [ctypes.c_void_p]
+  _lib.FreeModuleSummary.restype = None
+
+  def get_module_digest(filename: str) -> dict[str, Any]:
+    """Gets a module digest."""
+    _ensure_initialized()
+    handle = _lib.CreateModuleSummary(filename.encode("utf-8"))
+    if not handle:
+      return {"success": False}
+
+    try:
+      result = {
+          "success": True,
+          "total_count": _lib.GetTotalCount(handle),
+          "unique_resource_count": _lib.GetUniqueResourceCount(handle),
+      }
+
+      resources = {}
+      num_resources = _lib.GetHistogramBuckets(handle)
+      for i in range(num_resources):
+        resource = _lib.GetBucketKey(handle, i)
+        count = _lib.GetBucketValue(handle, i)
+        resources[resource] = count
+
+      result["resource_histogram"] = resources
+      return result
+    finally:
+      _lib.FreeModuleSummary(handle)
+
+else:
+
+  def get_module_digest(filename: str) -> dict[str, Any]:
+    del filename
+    raise NotImplementedError(
+        "get_module_digest is not supported in this build"
+    )
diff --git a/xprof/pywrap/profiler_wrapper_test.py b/xprof/pywrap/profiler_wrapper_test.py
index 9d3813d3d..858a346c7 100644
--- a/xprof/pywrap/profiler_wrapper_test.py
+++ b/xprof/pywrap/profiler_wrapper_test.py
@@ -59,6 +59,12 @@ def test_check_error_raises_runtime_error(self):
     with self.assertRaisesRegex(RuntimeError, r"^INVALID_ARGUMENT"):
       profiler_wrapper_plugin._check_error(err)
 
+  def test_get_module_digest(self):
+    if hasattr(profiler_wrapper_plugin, "get_module_digest"):
+      test_file = self.create_tempfile().full_path
+      result = profiler_wrapper_plugin.get_module_digest(test_file)
+      self.assertFalse(result["success"])
+
 
 if __name__ == "__main__":
   absltest.main()