[Reflection] Fix the ownsAddress function to claim heap-allocated Metadatas.

ownsAddress was a simple range check on images, but that won't find Metadatas that get allocated on the heap. If an address isn't found, try reading it as a Metadata and doing a range check on the type context descriptor too.

rdar://problem/60981575

Author: mikeash

include/swift/Reflection/ReflectionContext.h

@@ -541,15 +541,35 @@ class ReflectionContext
   }
 
   /// Returns true if the address falls within a registered image.
-  bool ownsAddress(RemoteAddress Address) {
+  bool ownsAddressRaw(RemoteAddress Address) {
     for (auto Range : imageRanges) {
       auto Start = std::get<0>(Range);
       auto End = std::get<1>(Range);
       if (Start.getAddressData() <= Address.getAddressData()
           && Address.getAddressData() < End.getAddressData())
         return true;
     }
-  
+
+    return false;
+  }
+
+  /// Returns true if the address is known to the reflection context.
+  /// Currently, that means that either the address falls within a registered
+  /// image, or the address points to a Metadata whose type context descriptor
+  /// is within a registered image.
+  bool ownsAddress(RemoteAddress Address) {
+    if (ownsAddressRaw(Address))
+      return true;
+
+    // This is usually called on a Metadata address which might have been
+    // on the heap. Try reading it and looking up its type context descriptor
+    // instead.
+    if (auto Metadata = readMetadata(Address.getAddressData()))
+      if (auto DescriptorAddress =
+          super::readAddressOfNominalTypeDescriptor(Metadata, true))
+        if (ownsAddressRaw(RemoteAddress(DescriptorAddress)))
+          return true;
+
     return false;
   }
 

include/swift/Remote/MetadataReader.h

@@ -1682,32 +1682,16 @@ class MetadataReader {
     return nullptr;
   }
 
-private:
-  template <template <class R> class M>
-  MetadataRef _readMetadata(StoredPointer address) {
-    return _readMetadata(address, sizeof(M<Runtime>));
-  }
-
-  MetadataRef _readMetadata(StoredPointer address, size_t sizeAfter) {
-    auto size = sizeAfter;
-    uint8_t *buffer = (uint8_t *) malloc(size);
-    if (!Reader->readBytes(RemoteAddress(address), buffer, size)) {
-      free(buffer);
-      return nullptr;
-    }
-
-    auto metadata = reinterpret_cast<TargetMetadata<Runtime>*>(buffer);
-    MetadataCache.insert(std::make_pair(address, OwnedMetadataRef(metadata)));
-    return MetadataRef(address, metadata);
-  }
-
   StoredPointer
   readAddressOfNominalTypeDescriptor(MetadataRef &metadata,
                                      bool skipArtificialSubclasses = false) {
     switch (metadata->getKind()) {
     case MetadataKind::Class: {
       auto classMeta = cast<TargetClassMetadata<Runtime>>(metadata);
       while (true) {
+        if (!classMeta->isTypeMetadata())
+          return 0;
+
         StoredSignedPointer descriptorAddressSigned = classMeta->getDescriptionAsSignedPointer();
         StoredPointer descriptorAddress = stripSignedPointer(descriptorAddressSigned);
 
@@ -1753,7 +1737,26 @@ class MetadataReader {
       return 0;
     }
   }
-  
+
+private:
+  template <template <class R> class M>
+  MetadataRef _readMetadata(StoredPointer address) {
+    return _readMetadata(address, sizeof(M<Runtime>));
+  }
+
+  MetadataRef _readMetadata(StoredPointer address, size_t sizeAfter) {
+    auto size = sizeAfter;
+    uint8_t *buffer = (uint8_t *) malloc(size);
+    if (!Reader->readBytes(RemoteAddress(address), buffer, size)) {
+      free(buffer);
+      return nullptr;
+    }
+
+    auto metadata = reinterpret_cast<TargetMetadata<Runtime>*>(buffer);
+    MetadataCache.insert(std::make_pair(address, OwnedMetadataRef(metadata)));
+    return MetadataRef(address, metadata);
+  }
+
   /// Returns Optional(ParentContextDescriptorRef()) if there's no parent descriptor.
   /// Returns None if there was an error reading the parent descriptor.
   Optional<ParentContextDescriptorRef>

include/swift/SwiftRemoteMirror/SwiftRemoteMirror.h

@@ -116,7 +116,7 @@ SWIFT_REMOTE_MIRROR_LINKAGE
 int
 swift_reflection_ownsObject(SwiftReflectionContextRef ContextRef, uintptr_t Object);
 
-/// Returns whether the given address is within an image added to this
+/// Returns whether the given address is associated with an image added to this
 /// library. Images must be added using swift_reflection_addImage, not
 /// swift_reflection_addReflectionInfo, for this function to work
 /// properly. If addReflectionInfo is used, the return value will always