From d90bb66dd931a45cae3399dc76942bf8e412fa7b Mon Sep 17 00:00:00 2001
From: Siva Chandra <sivachandra@google.com>
Date: Wed, 27 Jan 2021 11:08:31 -0800
Subject: [PATCH] [libc] Include only the relevant header files in the
 integration test.

---
 libc/config/linux/aarch64/headers.txt              |  1 +
 .../HdrGen/PrototypeTestGen/PrototypeTestGen.cpp   | 14 +++++++++++++-
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/libc/config/linux/aarch64/headers.txt b/libc/config/linux/aarch64/headers.txt
index a40f4b827607..6aa66813e940 100644
--- a/libc/config/linux/aarch64/headers.txt
+++ b/libc/config/linux/aarch64/headers.txt
@@ -1,5 +1,6 @@
 set(PUBLIC_HEADERS
     libc.include.errno
     libc.include.math
+    libc.include.stdlib
     libc.include.sys_syscall
 )
diff --git a/libc/utils/HdrGen/PrototypeTestGen/PrototypeTestGen.cpp b/libc/utils/HdrGen/PrototypeTestGen/PrototypeTestGen.cpp
index 918b520d3910..e6b4353ca8d2 100644
--- a/libc/utils/HdrGen/PrototypeTestGen/PrototypeTestGen.cpp
+++ b/libc/utils/HdrGen/PrototypeTestGen/PrototypeTestGen.cpp
@@ -24,9 +24,21 @@ llvm::cl::list<std::string>
 bool TestGeneratorMain(llvm::raw_ostream &OS, llvm::RecordKeeper &records) {
   OS << "#include \"TypeTraits.h\"\n";
   llvm_libc::APIIndexer G(records);
-  for (const auto &header : G.PublicHeaders)
+  std::unordered_set<std::string> headerFileSet;
+  for (const auto &entrypoint : EntrypointNamesOption) {
+    auto match = G.FunctionToHeaderMap.find(entrypoint);
+    if (match == G.FunctionToHeaderMap.end()) {
+      llvm::errs() << "ERROR: entrypoint '" << entrypoint
+                   << "' could not be found in spec in any public header\n";
+      return true;
+    }
+    headerFileSet.insert(match->second);
+  }
+  for (const auto &header : headerFileSet)
     OS << "#include <" << header << ">\n";
+
   OS << '\n';
+
   OS << "int main() {\n";
   for (const auto &entrypoint : EntrypointNamesOption) {
     auto match = G.FunctionSpecMap.find(entrypoint);