[RemoteMirrors] Change MemoryReader::readBytes to return a unique_ptr instead of a tuple so that the destructor handles freeing the memory.

Author: mikeash

include/swift/Reflection/ReflectionContext.h

@@ -34,6 +34,7 @@
 #include <set>
 #include <vector>
 #include <unordered_map>
+#include <utility>
 
 #if defined(__APPLE__) && defined(__MACH__)
 #ifndef __LP64__
@@ -75,7 +76,9 @@ class ReflectionContext
 
   std::unordered_map<typename super::StoredPointer, const TypeInfo *> Cache;
 
-  std::vector<std::function<void()>> freeFuncs;
+  /// All buffers we need to keep around long term. This will automatically free them
+  /// when this object is destroyed.
+  std::vector<MemoryReader::ReadBytesResult> savedBuffers;
   std::vector<std::tuple<RemoteAddress, RemoteAddress>> dataSegments;
 
 public:
@@ -95,11 +98,6 @@ class ReflectionContext
   ReflectionContext(const ReflectionContext &other) = delete;
   ReflectionContext &operator=(const ReflectionContext &other) = delete;
 
-  ~ReflectionContext() {
-    for (auto f : freeFuncs)
-      f();
-  }
-
   MemoryReader &getReader() {
     return *this->Reader;
   }
@@ -115,50 +113,38 @@ class ReflectionContext
 
 #if defined(__APPLE__) && defined(__MACH__)
   bool addImage(RemoteAddress ImageStart) {
-    const void *Buf;
-    std::function<void()> FreeFunc;
-
-    std::tie(Buf, FreeFunc) =
-        this->getReader().readBytes(ImageStart, sizeof(MachHeader));
-    if (Buf == nullptr)
+    auto Buf = this->getReader().readBytes(ImageStart, sizeof(MachHeader));
+    if (!Buf)
       return false;
 
-    auto Header = reinterpret_cast<MachHeader *>(Buf);
-
+    auto Header = reinterpret_cast<MachHeader *>(Buf.get());
     if (Header->magic != MH_MAGIC && Header->magic != MH_MAGIC_64) {
-      FreeFunc();
       return false;
     }
-
     auto Length = Header->sizeofcmds;
-    FreeFunc();
 
     // Read the commands.
-    std::tie(Buf, FreeFunc) = this->getReader().readBytes(ImageStart, Length);
-
-    if (Buf == nullptr)
+    Buf = this->getReader().readBytes(ImageStart, Length);
+    if (!Buf)
       return false;
 
     // Find the TEXT segment and figure out where the end is.
-    Header = reinterpret_cast<MachHeader *>(Buf);
+    Header = reinterpret_cast<MachHeader *>(Buf.get());
     unsigned long TextSize;
     auto *TextSegment = getsegmentdata(Header, "__TEXT", &TextSize);
-    if (TextSegment == nullptr) {
-      FreeFunc();
+    if (TextSegment == nullptr)
       return false;
-    }
+    
     auto TextEnd =
-        TextSegment - reinterpret_cast<const uint8_t *>(Buf) + TextSize;
-    FreeFunc();
+        TextSegment - reinterpret_cast<const uint8_t *>(Buf.get()) + TextSize;
 
     // Read everything including the TEXT segment.
-    std::tie(Buf, FreeFunc) = this->getReader().readBytes(ImageStart, TextEnd);
-
-    if (Buf == nullptr)
+    Buf = this->getReader().readBytes(ImageStart, TextEnd);
+    if (!Buf)
       return false;
 
     // Read all the metadata parts.
-    Header = reinterpret_cast<MachHeader *>(Buf);
+    Header = reinterpret_cast<MachHeader *>(Buf.get());
 
     // The docs say "not all sections may be present." We'll succeed if ANY of
     // them are present. Not sure if that's the right thing to do.
@@ -173,12 +159,10 @@ class ReflectionContext
 
     bool success = FieldMd.second || AssocTyMd.second || BuiltinTyMd.second ||
                    CaptureMd.second || TyperefMd.second || ReflStrMd.second;
-    if (!success) {
-      FreeFunc();
-      return 0;
-    }
+    if (!success)
+      return false;
 
-    auto LocalStartAddress = reinterpret_cast<uintptr_t>(Buf);
+    auto LocalStartAddress = reinterpret_cast<uintptr_t>(Buf.get());
     auto RemoteStartAddress =
         reinterpret_cast<uint64_t>(ImageStart.getAddressData());
 
@@ -194,17 +178,18 @@ class ReflectionContext
 
     this->addReflectionInfo(info);
 
-    freeFuncs.push_back(FreeFunc);
-
     unsigned long DataSize;
     auto *DataSegment = getsegmentdata(Header, "__DATA", &DataSize);
     if (DataSegment != nullptr) {
-      auto DataSegmentStart = DataSegment - reinterpret_cast<const uint8_t *>(Buf)
+      auto DataSegmentStart = DataSegment - reinterpret_cast<const uint8_t *>(Buf.get())
                             + ImageStart.getAddressData();
       auto DataSegmentEnd = DataSegmentStart + DataSize;
       dataSegments.push_back(std::make_tuple(RemoteAddress(DataSegmentStart),
                                              RemoteAddress(DataSegmentEnd)));
     }
+    
+    savedBuffers.push_back(std::move(Buf));
+
     return true;
   }
 #endif // defined(__APPLE__) && defined(__MACH__)

include/swift/Remote/CMemoryReader.h

@@ -62,31 +62,26 @@ class CMemoryReader final : public MemoryReader {
     if (!length)
       return false;
 
-    const void *ptr;
-    std::function<void()> freeFunc;
-    std::tie(ptr, freeFunc) = readBytes(address, length);
-    if (ptr != nullptr) {
-      dest = std::string(reinterpret_cast<const char *>(ptr), length);
-      freeFunc();
-      return true;
-    } else {
+    auto Buf = readBytes(address, length);
+    if (!Buf)
       return false;
-    }
+    
+    dest = std::string(reinterpret_cast<const char *>(Buf.get()), length);
+    return true;
   }
 
-  std::tuple<const void *, std::function<void()>>
-    readBytes(RemoteAddress address, uint64_t size) override {
+  ReadBytesResult readBytes(RemoteAddress address, uint64_t size) override {
       void *FreeContext;
       auto Ptr = Impl.readBytes(Impl.reader_context, address.getAddressData(), size,
                                 &FreeContext);
 
       auto Free = Impl.free;
       if (Free == nullptr)
-        return std::make_tuple(Ptr, []{});
+        return ReadBytesResult(Ptr, [](const void *) {});
 
       auto ReaderContext = Impl.reader_context;
-      auto freeLambda = [=]{ Free(ReaderContext, Ptr, FreeContext); };
-      return std::make_tuple(Ptr, freeLambda);
+      auto freeLambda = [=](const void *Ptr) { Free(ReaderContext, Ptr, FreeContext); };
+      return ReadBytesResult(Ptr, freeLambda);
   }
 };
 

include/swift/Remote/InProcessMemoryReader.h

@@ -42,9 +42,8 @@ class InProcessMemoryReader final : public MemoryReader {
     return true;
   }
 
-  std::tuple<const void *, std::function<void()>>
-    readBytes(RemoteAddress address, uint64_t size) override {
-    return std::make_tuple(address.getLocalPointer<void>(), []{});
+  ReadBytesResult readBytes(RemoteAddress address, uint64_t size) override {
+    return ReadBytesResult(address.getLocalPointer<void>(), [](const void *) {});
   }
 };
 

include/swift/Remote/MemoryReader.h

@@ -21,6 +21,7 @@
 #include "swift/Remote/RemoteAddress.h"
 
 #include <functional>
+#include <memory>
 #include <string>
 #include <tuple>
 
@@ -33,6 +34,9 @@ namespace remote {
 /// representation of the address space of a remote process.
 class MemoryReader {
 public:
+  /// A convenient name for the return type from readBytes.
+  using ReadBytesResult = std::unique_ptr<const void, std::function<void(const void *)>>;
+
   /// Return the size of an ordinary pointer in the remote process, in bytes.
   virtual uint8_t getPointerSize() = 0;
 
@@ -47,7 +51,7 @@ class MemoryReader {
   ///
   /// Returns false if the operation failed.
   virtual bool readString(RemoteAddress address, std::string &dest) = 0;
-
+  
   /// Attempts to read an integer from the given address in the remote
   /// process.
   ///
@@ -65,15 +69,17 @@ class MemoryReader {
   ///
   /// NOTE: subclasses MUST override at least one of the readBytes functions. The default
   /// implementation calls through to the other one.
-  virtual std::tuple<const void *, std::function<void()>>
+  virtual ReadBytesResult
     readBytes(RemoteAddress address, uint64_t size) {
-    void *buffer = malloc(size);
-    bool success = readBytes(address, reinterpret_cast<uint8_t *>(buffer), size);
+    auto *Buf = malloc(size);
+    ReadBytesResult Result(Buf, [](const void *ptr) {
+      free(const_cast<void *>(ptr));
+    });
+    bool success = readBytes(address, reinterpret_cast<uint8_t *>(Buf), size);
     if (!success) {
-      free(buffer);
-      return std::make_tuple(nullptr, []{});
+      Result.reset();
     }
-    return std::make_tuple(buffer, [=]{ free(buffer); });
+    return Result;
   }
 
   /// Attempts to read 'size' bytes from the given address in the
@@ -84,16 +90,12 @@ class MemoryReader {
   /// NOTE: subclasses MUST override at least one of the readBytes functions. The default
   /// implementation calls through to the other one.
   virtual bool readBytes(RemoteAddress address, uint8_t *dest, uint64_t size) {
-    const void *ptr;
-    std::function<void()> freeFunc;
-    std::tie(ptr, freeFunc) = readBytes(address, size);
-    if (ptr != nullptr) {
-      memcpy(dest, ptr, size);
-      freeFunc();
-      return true;
-    } else {
+    auto Ptr = readBytes(address, size);
+    if (!Ptr)
       return false;
-    }
+    
+    memcpy(dest, Ptr.get(), size);
+    return true;
   }
 
   virtual ~MemoryReader() = default;

tools/swift-reflection-dump/swift-reflection-dump.cpp

@@ -196,14 +196,12 @@ class ObjectMemoryReader : public MemoryReader {
     return false;
   }
 
-  std::tuple<const void *, std::function<void()>>
-    readBytes(RemoteAddress address, uint64_t size) override {
+  ReadBytesResult readBytes(RemoteAddress address, uint64_t size) override {
     if (!isAddressValid(address, size))
-      return std::make_tuple(nullptr, []{});
+      return ReadBytesResult(nullptr, [](const void *){});
 
     // TODO: Account for offset in ELF binaries
-    auto ptr = (const void *)address.getAddressData();
-    return std::make_tuple(ptr, []{});
+    return ReadBytesResult((const void *)address.getAddressData(), [](const void *) {});
   }
 
   bool readString(RemoteAddress address, std::string &dest) override {