[Reflection] Cache the various DataQuery values.

We can end up querying the pointer size and similar things very frequently, which can be slow. These values don't change for any given reader, so call them once and cache the values locally instead.

rdar://150221008

Author: mikeash

include/swift/Remote/CMemoryReader.h

@@ -42,38 +42,16 @@ namespace remote {
 class CMemoryReader final : public MemoryReader {
   MemoryReaderImpl Impl;
 
-  uint64_t ptrauthMask;
-
-  uint64_t getPtrauthMask() {
-    if (ptrauthMask == 0) {
-      int success;
-      if (Impl.PointerSize == 4) {
-        uint32_t ptrauthMask32 = 0;
-        success = queryDataLayout(DataLayoutQueryType::DLQ_GetPtrAuthMask,
-                                  nullptr, &ptrauthMask32);
-        ptrauthMask = ptrauthMask32;
-      } else if (Impl.PointerSize == 8) {
-        success = queryDataLayout(DataLayoutQueryType::DLQ_GetPtrAuthMask,
-                                  nullptr, &ptrauthMask);
-      } else {
-        success = 0;
-      }
-
-      if (!success)
-        ptrauthMask = ~0ull;
-    }
-    return ptrauthMask;
-  }
-
   // Check to see if an address has bits outside the ptrauth mask. This suggests
   // that we're likely failing to strip a signed pointer when reading from it.
   bool hasSignatureBits(RemoteAddress address) {
     uint64_t addressData = address.getAddressData();
-    return addressData != (addressData & getPtrauthMask());
+    uint64_t mask = getPtrAuthMask().value_or(~uint64_t(0));
+    return addressData != (addressData & mask);
   }
 
 public:
-  CMemoryReader(MemoryReaderImpl Impl) : Impl(Impl), ptrauthMask(0) {
+  CMemoryReader(MemoryReaderImpl Impl) : Impl(Impl) {
     assert(this->Impl.queryDataLayout && "No queryDataLayout implementation");
     assert(this->Impl.getStringLength && "No stringLength implementation");
     assert(this->Impl.readBytes && "No readBytes implementation");

include/swift/Remote/MemoryReader.h

@@ -22,6 +22,7 @@
 #include "swift/SwiftRemoteMirror/MemoryReaderInterface.h"
 #include <optional>
 
+#include <cstdint>
 #include <cstdlib>
 #include <cstring>
 #include <functional>
@@ -37,7 +38,80 @@ namespace remote {
 /// This abstraction presents memory as if it were a read-only
 /// representation of the address space of a remote process.
 class MemoryReader {
+  uint8_t cachedPointerSize = 0;
+  uint8_t cachedSizeSize = 0;
+  uint64_t cachedPtrAuthMask = 0;
+  uint8_t cachedObjCReservedLowBits = 0xff;
+  uint64_t cachedLeastValidPointerValue = 0;
+  uint8_t cachedObjCInteropIsEnabled = 0xff;
+
+protected:
+  virtual bool queryDataLayout(DataLayoutQueryType type, void *inBuffer,
+                               void *outBuffer) = 0;
+
 public:
+  std::optional<uint8_t> getPointerSize() {
+    if (cachedPointerSize == 0) {
+      if (!queryDataLayout(DLQ_GetPointerSize, nullptr, &cachedPointerSize))
+        return std::nullopt;
+    }
+    return cachedPointerSize;
+  }
+
+  std::optional<uint8_t> getSizeSize() {
+    if (cachedSizeSize == 0) {
+      if (!queryDataLayout(DLQ_GetSizeSize, nullptr, &cachedSizeSize))
+        return std::nullopt;
+    }
+    return cachedSizeSize;
+  }
+
+  std::optional<uint64_t> getPtrAuthMask() {
+    if (cachedPtrAuthMask == 0) {
+      auto ptrSize = getPointerSize();
+      if (!ptrSize)
+        return std::nullopt;
+
+      if (ptrSize.value() == sizeof(uint64_t)) {
+        if (!queryDataLayout(DLQ_GetPtrAuthMask, nullptr, &cachedPtrAuthMask))
+          return std::nullopt;
+      } else if (ptrSize.value() == sizeof(uint32_t)) {
+        uint32_t mask;
+        if (!queryDataLayout(DLQ_GetPtrAuthMask, nullptr, &mask))
+          return std::nullopt;
+        cachedPtrAuthMask = mask;
+      }
+    }
+    return cachedPtrAuthMask;
+  }
+
+  std::optional<uint8_t> getObjCReservedLowBits() {
+    if (cachedObjCReservedLowBits == 0xff) {
+      if (!queryDataLayout(DLQ_GetObjCReservedLowBits, nullptr,
+                           &cachedObjCReservedLowBits))
+        return std::nullopt;
+    }
+    return cachedObjCReservedLowBits;
+  }
+
+  std::optional<uint64_t> getLeastValidPointerValue() {
+    if (cachedLeastValidPointerValue == 0) {
+      if (!queryDataLayout(DLQ_GetLeastValidPointerValue, nullptr,
+                           &cachedLeastValidPointerValue))
+        return std::nullopt;
+    }
+    return cachedLeastValidPointerValue;
+  }
+
+  std::optional<bool> getObjCInteropIsEnabled() {
+    if (cachedObjCInteropIsEnabled == 0xff) {
+      if (!queryDataLayout(DLQ_GetObjCInteropIsEnabled, nullptr,
+                           &cachedObjCInteropIsEnabled))
+        return std::nullopt;
+    }
+    return cachedObjCInteropIsEnabled;
+  }
+
   /// A convenient name for the return type from readBytes.
   using ReadBytesResult =
       std::unique_ptr<const void, std::function<void(const void *)>>;
@@ -46,9 +120,6 @@ class MemoryReader {
   using ReadObjResult =
       std::unique_ptr<const T, std::function<void(const void *)>>;
 
-  virtual bool queryDataLayout(DataLayoutQueryType type, void *inBuffer,
-                               void *outBuffer) = 0;
-
   /// Look up the given public symbol name in the remote process.
   virtual RemoteAddress getSymbolAddress(const std::string &name) = 0;
 
@@ -209,50 +280,37 @@ class MemoryReader {
   // index (counting from 0).
   bool readHeapObjectExtraInhabitantIndex(RemoteAddress address,
                                           int *extraInhabitantIndex) {
-    uint8_t PointerSize;
-    if (!queryDataLayout(DataLayoutQueryType::DLQ_GetPointerSize,
-                         nullptr, &PointerSize)) {
-      return false;
-    }
-    uint64_t LeastValidPointerValue;
-    if (!queryDataLayout(DataLayoutQueryType::DLQ_GetLeastValidPointerValue,
-                         nullptr, &LeastValidPointerValue)) {
-      return false;
-    }
-    uint8_t ObjCReservedLowBits;
-    if (!queryDataLayout(DataLayoutQueryType::DLQ_GetObjCReservedLowBits,
-                         nullptr, &ObjCReservedLowBits)) {
+    auto PointerSize = getPointerSize();
+    auto LeastValidPointerValue = getLeastValidPointerValue();
+    auto ObjCReservedLowBits = getObjCReservedLowBits();
+
+    if (!PointerSize || !LeastValidPointerValue || !ObjCReservedLowBits)
       return false;
-    }
+
     uint64_t RawPointerValue;
-    if (!readInteger(address, PointerSize, &RawPointerValue)) {
+    if (!readInteger(address, PointerSize.value(), &RawPointerValue)) {
       return false;
     }
-    if (RawPointerValue >= LeastValidPointerValue) {
+    if (RawPointerValue >= LeastValidPointerValue.value()) {
       *extraInhabitantIndex = -1; // Valid value, not an XI
     } else {
-      *extraInhabitantIndex = (RawPointerValue >> ObjCReservedLowBits);
+      *extraInhabitantIndex = (RawPointerValue >> ObjCReservedLowBits.value());
     }
     return true;
   }
 
   bool readFunctionPointerExtraInhabitantIndex(RemoteAddress address,
                                                int *extraInhabitantIndex) {
-    uint8_t PointerSize;
-    if (!queryDataLayout(DataLayoutQueryType::DLQ_GetPointerSize,
-                         nullptr, &PointerSize)) {
+    auto PointerSize = getPointerSize();
+    auto LeastValidPointerValue = getLeastValidPointerValue();
+    if (!PointerSize || !LeastValidPointerValue)
       return false;
-    }
-    uint64_t LeastValidPointerValue;
-    if (!queryDataLayout(DataLayoutQueryType::DLQ_GetLeastValidPointerValue,
-                         nullptr, &LeastValidPointerValue)) {
-      return false;
-    }
+
     uint64_t RawPointerValue;
-    if (!readInteger(address, PointerSize, &RawPointerValue)) {
+    if (!readInteger(address, PointerSize.value(), &RawPointerValue)) {
       return false;
     }
-    if (RawPointerValue >= LeastValidPointerValue) {
+    if (RawPointerValue >= LeastValidPointerValue.value()) {
       *extraInhabitantIndex = -1; // Valid value, not an XI
     } else {
       *extraInhabitantIndex = RawPointerValue;

include/swift/Remote/MetadataReader.h

@@ -422,12 +422,8 @@ class MetadataReader {
   }
 
   StoredPointer queryPtrAuthMask() {
-    StoredPointer QueryResult;
-    if (Reader->queryDataLayout(DataLayoutQueryType::DLQ_GetPtrAuthMask,
-                                nullptr, &QueryResult)) {
-      return QueryResult;
-    }
-    return ~StoredPointer(0);
+    auto QueryResult = Reader->getPtrAuthMask();
+    return QueryResult.value_or(~StoredPointer(0));
   }
 
   template <class... T>

lib/StaticMirror/ObjectFileContext.cpp

@@ -565,16 +565,18 @@ std::unique_ptr<ReflectionContextHolder> makeReflectionContextForObjectFiles(
   const std::vector<const ObjectFile *> &objectFiles, bool ObjCInterop) {
   auto Reader = std::make_shared<ObjectMemoryReader>(objectFiles);
 
-  uint8_t pointerSize;
-  Reader->queryDataLayout(DataLayoutQueryType::DLQ_GetPointerSize, nullptr,
-                          &pointerSize);
+  auto pointerSize = Reader->getPointerSize();
+  if (!pointerSize) {
+    fputs("unable to get target pointer size\n", stderr);
+    abort();
+  }
 
-  switch (pointerSize) {
+  switch (pointerSize.value()) {
   case 4:
 #define MAKE_CONTEXT(INTEROP, PTRSIZE)                                         \
   makeReflectionContextForMetadataReader<                                      \
       External<INTEROP<RuntimeTarget<PTRSIZE>>>>(std::move(Reader),            \
-                                                 pointerSize)
+                                                 pointerSize.value())
 #if SWIFT_OBJC_INTEROP
     if (ObjCInterop)
       return MAKE_CONTEXT(WithObjCInterop, 4);