diff --git a/config/linux/aarch64/headers.txt b/config/linux/aarch64/headers.txt
index a40f4b8..6aa6681 100644
--- a/config/linux/aarch64/headers.txt
+++ b/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/utils/HdrGen/PrototypeTestGen/PrototypeTestGen.cpp b/utils/HdrGen/PrototypeTestGen/PrototypeTestGen.cpp
index 918b520..e6b4353 100644
--- a/utils/HdrGen/PrototypeTestGen/PrototypeTestGen.cpp
+++ b/utils/HdrGen/PrototypeTestGen/PrototypeTestGen.cpp
@@ -24,9 +24,21 @@
 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);