Fix dynamic replacement runtime when generating replacements from multiple object files

The sections with the replacements get merged and we need to consider all replacements.

SR-10947
rdar://51913012

Author: aschwaighofer

lib/IRGen/GenDecl.cpp

@@ -1515,6 +1515,7 @@ void IRGenerator::emitDynamicReplacements() {
   autoReplacements.addInt32(1); // number of replacement entries.
   auto autoReplacementsArray = autoReplacements.beginArray();
   autoReplacementsArray.addRelativeAddress(var);
+  autoReplacementsArray.addInt32(0); // unused flags.
   autoReplacementsArray.finishAndAddTo(autoReplacements);
   auto autoReplVar = autoReplacements.finishAndCreateGlobal(
       "\x01l_auto_dynamic_replacements", IGM.getPointerAlignment(),

stdlib/public/runtime/MetadataLookup.cpp

@@ -2044,6 +2044,8 @@ class AutomaticDynamicReplacements
     for (auto &replacementEntry : getReplacementEntries())
       replacementEntry.enable();
   }
+
+  uint32_t getNumScopes() const { return numScopes; }
 };
 
 /// A map from original to replaced opaque type descriptor of a some type.
@@ -2083,26 +2085,63 @@ class AutomaticDynamicReplacementsSome
     for (auto &replacementEntry : getReplacementEntries())
       replacementEntry.enable(lock);
   }
+  uint32_t getNumEntries() const { return numEntries; }
 };
 
 } // anonymous namespace
 
 void swift::addImageDynamicReplacementBlockCallback(
     const void *replacements, uintptr_t replacementsSize,
     const void *replacementsSome, uintptr_t replacementsSomeSize) {
+
   auto *automaticReplacements =
       reinterpret_cast<const AutomaticDynamicReplacements *>(replacements);
+
   const AutomaticDynamicReplacementsSome *someReplacements = nullptr;
   if (replacementsSomeSize) {
     someReplacements =
         reinterpret_cast<const AutomaticDynamicReplacementsSome *>(
             replacementsSome);
   }
+
+  auto sizeOfCurrentEntry = sizeof(AutomaticDynamicReplacements) +
+                            (automaticReplacements->getNumScopes() *
+                             sizeof(AutomaticDynamicReplacementEntry));
+  auto sizeOfCurrentSomeEntry =
+      replacementsSomeSize == 0
+          ? 0
+          : sizeof(AutomaticDynamicReplacementsSome) +
+                (someReplacements->getNumEntries() *
+                 sizeof(DynamicReplacementSomeDescriptor));
+
   auto &lock = DynamicReplacementLock.get();
   lock.withLock([&] {
-    automaticReplacements->enableReplacements();
-    if (someReplacements)
+    auto endOfAutomaticReplacements =
+        ((const char *)automaticReplacements) + replacementsSize;
+    while (((const char *)automaticReplacements) < endOfAutomaticReplacements) {
+      automaticReplacements->enableReplacements();
+      automaticReplacements =
+          reinterpret_cast<const AutomaticDynamicReplacements *>(
+              ((const char *)automaticReplacements) + sizeOfCurrentEntry);
+      if ((const char*)automaticReplacements <  endOfAutomaticReplacements)
+        sizeOfCurrentEntry = sizeof(AutomaticDynamicReplacements) +
+                             (automaticReplacements->getNumScopes() *
+                              sizeof(AutomaticDynamicReplacementEntry));
+    }
+    if (!replacementsSomeSize)
+      return;
+    auto endOfSomeReplacements =
+        ((const char *)someReplacements) + replacementsSomeSize;
+    while (((const char *)someReplacements) < endOfSomeReplacements) {
       someReplacements->enableReplacements(lock);
+      someReplacements =
+          reinterpret_cast<const AutomaticDynamicReplacementsSome *>(
+              ((const char *)someReplacements) + sizeOfCurrentSomeEntry);
+      if  ((const char*) someReplacements < endOfSomeReplacements)
+        sizeOfCurrentSomeEntry = sizeof(AutomaticDynamicReplacementsSome) +
+                                 (someReplacements->getNumEntries() *
+                                  sizeof(DynamicReplacementSomeDescriptor));
+    }
   });
 }
 

test/IRGen/dynamic_replaceable.sil

@@ -17,9 +17,9 @@
 // CHECK:  %swift.dyn_repl_link_entry* @test_replacementTX
 // CHECK:  i32 0 }] }, section "__TEXT,__const"
 
-// CHECK: @"\01l_auto_dynamic_replacements" = private constant { i32, i32, [1 x i32] }
+// CHECK: @"\01l_auto_dynamic_replacements" = private constant { i32, i32, [2 x i32] }
 // CHECK: { i32 0, i32 1,
-// CHECK: [1 x i32] [{{.*}}@"\01l_unnamed_dynamic_replacements"{{.*}}]
+// CHECK: [2 x i32] [{{.*}}@"\01l_unnamed_dynamic_replacements"{{.*}}, i32 0]
 // CHECK: }, section "__TEXT, __swift5_replace, regular, no_dead_strip"
 
 // CHECK-LABEL: define swiftcc void @test_dynamically_replaceable()

test/Interpreter/dynamic_replacement_multifile/Inputs/dynamic_replacement_multi_file_A.swift

@@ -0,0 +1,35 @@
+dynamic func replaceable1() -> Int {
+   return 0
+}
+
+@_dynamicReplacement(for: replaceable1())
+func replaceable1_r() -> Int {
+  return 2
+}
+
+@available(macOS 9999, iOS 9999, tvOS 9999, watchOS 9999, *)
+dynamic func bar1(_ x: Int) -> some P {
+  return x
+}
+
+@available(macOS 9999, iOS 9999, tvOS 9999, watchOS 9999, *)
+@_dynamicReplacement(for: bar1(_:))
+func bar1_r(_ x: Int) -> some P {
+  return Pair()
+}
+
+@available(macOS 9999, iOS 9999, tvOS 9999, watchOS 9999, *)
+dynamic func bar2(_ x: Int) -> some P {
+  return x
+}
+
+@available(macOS 9999, iOS 9999, tvOS 9999, watchOS 9999, *)
+@_dynamicReplacement(for: bar2(_:))
+func bar2_r(_ x: Int) -> some P {
+  return Pair()
+}
+
+@_dynamicReplacement(for: replaceableInOtherFile())
+func replaceableInOtherFile_r() -> Int {
+  return 7
+}

test/Interpreter/dynamic_replacement_multifile/Inputs/dynamic_replacement_multi_file_B.swift

@@ -0,0 +1,19 @@
+dynamic func replaceable2() -> Int {
+   return 0
+}
+
+@_dynamicReplacement(for: replaceable2())
+func replaceable2_r() -> Int {
+  return 3
+}
+
+@available(macOS 9999, iOS 9999, tvOS 9999, watchOS 9999, *)
+dynamic func bar3(_ x: Int) -> some P {
+  return x
+}
+
+@available(macOS 9999, iOS 9999, tvOS 9999, watchOS 9999, *)
+@_dynamicReplacement(for: bar3(_:))
+func bar3_r(_ x: Int) -> some P {
+  return Pair()
+}

test/Interpreter/dynamic_replacement_multifile/main.swift

@@ -0,0 +1,61 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift -o %t/main %s %S/Inputs/dynamic_replacement_multi_file_A.swift %S/Inputs/dynamic_replacement_multi_file_B.swift -swift-version 5
+// RUN: %target-codesign %t/main
+// RUN: %target-run %t/main
+
+// REQUIRES: executable_test
+
+import StdlibUnittest
+
+dynamic func replaceable() -> Int {
+  return 0
+}
+
+dynamic func replaceableInOtherFile() -> Int {
+  return 0
+}
+
+@_dynamicReplacement(for: replaceable())
+func replaceable_r() -> Int {
+  return 1
+}
+
+protocol P {}
+
+extension Int : P {}
+
+struct Pair {
+  var x: Int64 = 0
+  var y: Int64 = 0
+}
+
+extension Pair : P {}
+
+@available(macOS 9999, iOS 9999, tvOS 9999, watchOS 9999, *)
+dynamic func bar(_ x: Int) -> some P {
+  return x
+}
+
+@available(macOS 9999, iOS 9999, tvOS 9999, watchOS 9999, *)
+@_dynamicReplacement(for: bar(_:))
+func bar_r(_ x: Int) -> some P {
+  return Pair()
+}
+
+var DynamicallyReplaceable = TestSuite("DynamicallyReplaceable")
+
+DynamicallyReplaceable.test("DynamicallyReplaceable") {
+  expectEqual(1, replaceable())
+  expectEqual(2, replaceable1())
+  expectEqual(3, replaceable2())
+  expectEqual(7, replaceableInOtherFile())
+  if #available(macOS 9999, iOS 9999, tvOS 9999, watchOS 9999, *) {
+    expectEqual(16, MemoryLayout.size(ofValue: bar(5)))
+    expectEqual(16, MemoryLayout.size(ofValue: bar1(5)))
+    expectEqual(16, MemoryLayout.size(ofValue: bar2(5)))
+    expectEqual(16, MemoryLayout.size(ofValue: bar3(5)))
+  }
+
+}
+
+runAllTests()