Simplify the incremental reading of remote descriptors

Author: tbkka

include/swift/Remote/MetadataReader.h

@@ -1017,152 +1017,81 @@ class MetadataReader {
           address, reinterpret_cast<const TargetContextDescriptor<Runtime> *>(
                        cached->second.get()));
 
-    // Read the flags to figure out what kind of descriptor this is
-    ContextDescriptorFlags flags;
-    if (!Reader->readBytes(RemoteAddress(address), (uint8_t*)&flags,
-                           sizeof(flags))) {
-      return nullptr;
-    }
-    TypeContextDescriptorFlags typeFlags(flags.getKindSpecificFlags());
-    uint64_t baseSize = 0;
-    uint64_t genericHeaderSize = sizeof(GenericContextDescriptorHeader);
-    uint64_t metadataInitSize = 0;
-    bool hasVTable = false;
-
-    auto readMetadataInitSize = [&]() -> unsigned {
-      switch (typeFlags.getMetadataInitialization()) {
-      case TypeContextDescriptorFlags::NoMetadataInitialization:
-        return 0;
-      case TypeContextDescriptorFlags::SingletonMetadataInitialization:
-        // FIXME: classes
-        return sizeof(TargetSingletonMetadataInitialization<Runtime>);
-      case TypeContextDescriptorFlags::ForeignMetadataInitialization:
-        return sizeof(TargetForeignMetadataInitialization<Runtime>);
+    auto buffer = (uint8_t *)nullptr;
+    unsigned available = 0; // Size of data in buffer
+
+    // In the first read, just grab the initial flag value...
+    unsigned sizeEstimate = sizeof(ContextDescriptorFlags);
+    while (available < sizeEstimate) {
+      // The first `available` bytes have already been read, so grow
+      // the buffer and read more bytes to fulfill `sizeEstimate`
+      auto newbuffer = (uint8_t *)realloc(buffer, sizeEstimate);
+      if (!newbuffer) {
+        free(buffer);
+        return nullptr;
       }
-      return 0;
-    };
-
-    // Compute the size of the base object based on the kind of descriptor
-    unsigned baseSize = 0;
-    switch (auto kind = flags.getKind()) {
-    case ContextDescriptorKind::Module:
-      baseSize = sizeof(TargetModuleContextDescriptor<Runtime>);
-      break;
-    case ContextDescriptorKind::Extension:
-      baseSize = sizeof(TargetExtensionContextDescriptor<Runtime>);
-      break;
-    case ContextDescriptorKind::Anonymous:
-      baseSize = sizeof(TargetAnonymousContextDescriptor<Runtime>);
-      break;
-    case ContextDescriptorKind::Class:
-      baseSize = sizeof(TargetClassDescriptor<Runtime>);
-      break;
-    case ContextDescriptorKind::Enum:
-      baseSize = sizeof(TargetEnumDescriptor<Runtime>);
-      break;
-    case ContextDescriptorKind::Struct:
-      baseSize = sizeof(TargetStructDescriptor<Runtime>);
-      break;
-    case ContextDescriptorKind::Protocol:
-      baseSize = sizeof(TargetProtocolDescriptor<Runtime>);
-      break;
-    case ContextDescriptorKind::OpaqueType:
-      baseSize = sizeof(TargetOpaqueTypeDescriptor<Runtime>);
-      break;
-    default:
-      // We don't know about this kind of context.
-      return nullptr;
-    }
-
-    // We don't know what trailing objects might be present, so
-    // round up our initial estimate a bit.
-    unsigned sizeEstimate = std::max(baseSize, 128U);
-    auto buffer = (uint8_t *)malloc(sizeEstimate);
-    unsigned bufferSize = 0; // Size of data in buffer
-
-    while (true) {
-      // Read more of the object
-      // The first `bufferSize` bytes have already been read, so just
-      // grab the next section to fulfill `sizeEstimate`
-      if (!Reader->readBytes(RemoteAddress(address + bufferSize),
-                             buffer + bufferSize,
-                             sizeEstimate - bufferSize)) {
+      buffer = newbuffer;
+      if (!Reader->readBytes(RemoteAddress(address + available),
+                             buffer + available,
+                             sizeEstimate - available)) {
         free(buffer);
         return nullptr;
       }
-      bufferSize = sizeEstimate;
-      // Note that `buffer` includes at least the base object, so
-      // it's safe to query the object as long as we don't
-      // try to actually access trailing objects that might not exist.
-
-      // Based on what we have so far, figure out whether we have
-      // all the trailing data yet.
-      unsigned revisedEstimate = 0;
+      available = sizeEstimate;
+
+      // Based on the data so far, update our guess of the full descriptor size.
+
+      // All objects start with a flags value, and our first read always gets
+      // that much data, so this is safe.
+      ContextDescriptorFlags flags;
+      memcpy(&flags, buffer, sizeof(flags));
       switch (auto kind = flags.getKind()) {
-      // For fixed-size descriptors without trailing data, the baseSize is always correct
+
+      // Fixed-size descriptors without trailing data have fixed size.
       case ContextDescriptorKind::Module:
-      case ContextDescriptorKind::Extension: {
-        revisedEstimate = baseSize;
+        sizeEstimate = sizeof(TargetModuleContextDescriptor<Runtime>);
         break;
-      }
+      case ContextDescriptorKind::Extension:
+        sizeEstimate = sizeof(TargetExtensionContextDescriptor<Runtime>);
+        break;
+
       // For types that use trailing objects, ask the trailing object logic to
-      // look at what we have so far and tell us whether we're done or not
+      // look at what we have so far and tell us whether we're done or not.
       case ContextDescriptorKind::Anonymous: {
-        revisedEstimate = TargetAnonymousContextDescriptor<Runtime>::totalSizeOfPartialObject(buffer, bufferSize);
+        sizeEstimate = TargetAnonymousContextDescriptor<Runtime>::totalSizeOfPartialObject(buffer, available);
         break;
       }
       case ContextDescriptorKind::Class: {
-        revisedEstimate = TargetClassDescriptor<Runtime>::totalSizeOfPartialObject(buffer, bufferSize);
+        sizeEstimate = TargetClassDescriptor<Runtime>::totalSizeOfPartialObject(buffer, available);
         break;
       }
       case ContextDescriptorKind::Enum: {
-        revisedEstimate = TargetEnumDescriptor<Runtime>::totalSizeOfPartialObject(buffer, bufferSize);
+        sizeEstimate = TargetEnumDescriptor<Runtime>::totalSizeOfPartialObject(buffer, available);
         break;
       }
       case ContextDescriptorKind::Struct: {
-        revisedEstimate = TargetStructDescriptor<Runtime>::totalSizeOfPartialObject(buffer, bufferSize);
+        sizeEstimate = TargetStructDescriptor<Runtime>::totalSizeOfPartialObject(buffer, available);
         break;
       }
       case ContextDescriptorKind::Protocol: {
-        revisedEstimate = TargetProtocolDescriptor<Runtime>::totalSizeOfPartialObject(buffer, bufferSize);
+        sizeEstimate = TargetProtocolDescriptor<Runtime>::totalSizeOfPartialObject(buffer, available);
         break;
       }
       case ContextDescriptorKind::OpaqueType: {
-        revisedEstimate = TargetOpaqueTypeDescriptor<Runtime>::totalSizeOfPartialObject(buffer, bufferSize);
+        sizeEstimate = TargetOpaqueTypeDescriptor<Runtime>::totalSizeOfPartialObject(buffer, available);
         break;
       }
+
       // We don't know about this kind of context.
       default:
         return nullptr;
       }
-
-      // If we have everything, then clean up the buffer and exit the loop
-      if (revisedEstimate <= bufferSize) {
-        auto trueSize = revisedEstimate;
-        // If the final true object size is a lot smaller than the allocated
-        // buffer, then it's worth a little more work to save memory.
-        if (trueSize < bufferSize / 2 && trueSize + 64 < bufferSize) {
-          auto newbuffer = (uint8_t *)realloc(buffer, trueSize);
-          if (!newbuffer) {
-            free(buffer);
-            return nullptr;
-          }
-          buffer = newbuffer;
-        }
-        bufferSize = trueSize;
-        break;
-      }
-
-      // We didn't get everything, increase the working size estimate and try again
-      sizeEstimate = std::max(sizeEstimate + 128, revisedEstimate);
-      auto newbuffer = (uint8_t *)realloc(buffer, sizeEstimate);
-      if (!newbuffer) {
-        free(buffer);
-        return nullptr;
-      }
-      buffer = newbuffer;
     }
 
+    // After exiting loop above, `sizeEstimate` is the true size of the
+    // descriptor with all trailing objects and buffer holds at least that much
+    // data.
+
     // Insert the final object into the descriptor cache and return it
     auto descriptor
       = reinterpret_cast<TargetContextDescriptor<Runtime> *>(buffer);