Add support to modify chaining behavior of dynamic replacements

include/swift/ABI/Metadata.h

@@ -4127,6 +4127,8 @@ class DynamicReplacementDescriptor {
   RelativeDirectPointer<DynamicReplacementChainEntry, false> chainEntry;
   uint32_t flags;
 
+  enum : uint32_t { EnableChainingMask = 0x1 };
+
 public:
   /// Enable this replacement by changing the function's replacement chain's
   /// root entry.
@@ -4143,6 +4145,8 @@ class DynamicReplacementDescriptor {
   void disableReplacement() const;
 
   uint32_t getFlags() const { return flags; }
+
+  bool shouldChain() const { return (flags & EnableChainingMask); }
 };
 
 /// A collection of dynamic replacement records.

include/swift/AST/IRGenOptions.h

@@ -187,6 +187,9 @@ class IRGenOptions {
   /// Instrument code to generate profiling information.
   unsigned GenerateProfile : 1;
 
+  /// Enable chaining of dynamic replacements.
+  unsigned EnableDynamicReplacementChaining : 1;
+
   /// Path to the profdata file to be used for PGO, or the empty string.
   std::string UseProfile = "";
 
@@ -219,7 +222,8 @@ class IRGenOptions {
         ValueNames(false), EnableReflectionMetadata(true),
         EnableReflectionNames(true), EnableClassResilience(false),
         EnableResilienceBypass(false), UseIncrementalLLVMCodeGen(true),
-        UseSwiftCall(false), GenerateProfile(false), CmdArgs(),
+        UseSwiftCall(false), GenerateProfile(false),
+        EnableDynamicReplacementChaining(false), CmdArgs(),
         SanitizeCoverage(llvm::SanitizerCoverageOptions()),
         TypeInfoFilter(TypeInfoDumpFilter::All) {}
 

include/swift/Option/FrontendOptions.td

@@ -362,6 +362,11 @@ def enable_implicit_dynamic : Flag<["-"], "enable-implicit-dynamic">,
   Flags<[FrontendOption, NoInteractiveOption, HelpHidden]>,
   HelpText<"Add 'dynamic' to all declarations">;
 
+def enable_dynamic_replacement_chaining :
+  Flag<["-"], "enable-dynamic-replacement-chaining">,
+  Flags<[FrontendOption, NoInteractiveOption, HelpHidden]>,
+  HelpText<"Enable chaining of dynamic replacements">;
+
 def enable_nskeyedarchiver_diagnostics :
   Flag<["-"], "enable-nskeyedarchiver-diagnostics">,
   HelpText<"Diagnose classes with unstable mangled names adopting NSCoding">;

lib/Frontend/CompilerInvocation.cpp

@@ -989,6 +989,9 @@ static bool ParseIRGenArgs(IRGenOptions &Opts, ArgList &Args,
 
   Opts.PrintInlineTree |= Args.hasArg(OPT_print_llvm_inline_tree);
 
+  Opts.EnableDynamicReplacementChaining |=
+      Args.hasArg(OPT_enable_dynamic_replacement_chaining);
+
   Opts.UseSwiftCall = Args.hasArg(OPT_enable_swiftcall);
 
   // This is set to true by default.

lib/IRGen/GenDecl.cpp

@@ -1288,7 +1288,7 @@ void IRGenerator::emitDynamicReplacements() {
   //     RelativeIndirectablePointer<KeyEntry, false> replacedFunctionKey;
   //     RelativeDirectPointer<void> newFunction;
   //     RelativeDirectPointer<LinkEntry> replacement;
-  //     uint32_t flags; // unused.
+  //     uint32_t flags; // shouldChain.
   //   }[0]
   // };
   ConstantInitBuilder builder(IGM);
@@ -1315,7 +1315,8 @@ void IRGenerator::emitDynamicReplacements() {
     replacement.addRelativeAddress(newFnPtr); // direct relative reference.
     replacement.addRelativeAddress(
         replacementLinkEntry); // direct relative reference.
-    replacement.addInt32(0); // unused flags.
+    replacement.addInt32(
+        Opts.EnableDynamicReplacementChaining ? 1 : 0);
     replacement.finishAndAddTo(replacementsArray);
   }
   replacementsArray.finishAndAddTo(replacementScope);

stdlib/public/runtime/MetadataLookup.cpp

@@ -1565,6 +1565,13 @@ void DynamicReplacementDescriptor::enableReplacement() const {
     }
   }
 
+  // Unlink the previous entry if we are not chaining.
+  if (!shouldChain() && chainRoot->next) {
+    auto *previous = chainRoot->next;
+    chainRoot->next = previous->next;
+    chainRoot->implementationFunction = previous->implementationFunction;
+  }
+
   // First populate the current replacement's chain entry.
   auto *currentEntry =
       const_cast<DynamicReplacementChainEntry *>(chainEntry.get());

test/Interpreter/Inputs/dynamic_replacement_chaining_A.swift

@@ -0,0 +1,6 @@
+public struct Impl {
+  public init() {}
+  dynamic public func foo() -> Int {
+     return 1
+  }
+}

test/Interpreter/Inputs/dynamic_replacement_chaining_B.swift

@@ -0,0 +1,8 @@
+import A
+
+extension Impl {
+  @_dynamicReplacement(for: foo())
+  func repl() -> Int {
+    return foo() + 1
+  }
+}

test/Interpreter/dynamic_replacement_chaining.swift

@@ -0,0 +1,54 @@
+// First build without chaining.
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift-dylib(%t/libA.%target-dylib-extension) -module-name A -emit-module -emit-module-path %t -swift-version 5 %S/Inputs/dynamic_replacement_chaining_A.swift
+// RUN: %target-build-swift-dylib(%t/libB.%target-dylib-extension) -I%t -L%t -lA -Xlinker -rpath -Xlinker %t -module-name B -emit-module -emit-module-path %t -swift-version 5 %S/Inputs/dynamic_replacement_chaining_B.swift
+// RUN: %target-build-swift-dylib(%t/libC.%target-dylib-extension) -I%t -L%t -lA -Xlinker -rpath -Xlinker %t -module-name C -emit-module -emit-module-path %t -swift-version 5 %S/Inputs/dynamic_replacement_chaining_B.swift
+// RUN: %target-build-swift -I%t -L%t -lA -o %t/main -Xlinker -rpath -Xlinker %t %s -swift-version 5
+// RUN: %target-codesign %t/main %t/libA.%target-dylib-extension %t/libB.%target-dylib-extension %t/libC.%target-dylib-extension
+// RUN: %target-run %t/main %t/libA.%target-dylib-extension %t/libB.%target-dylib-extension %t/libC.%target-dylib-extension
+
+// Now build with chaining enabled.
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift-dylib(%t/libA.%target-dylib-extension) -module-name A -emit-module -emit-module-path %t -swift-version 5 %S/Inputs/dynamic_replacement_chaining_A.swift
+// RUN: %target-build-swift-dylib(%t/libB.%target-dylib-extension) -I%t -L%t -lA -Xlinker -rpath -Xlinker %t -module-name B -emit-module -emit-module-path %t -swift-version 5 %S/Inputs/dynamic_replacement_chaining_B.swift -Xfrontend -enable-dynamic-replacement-chaining
+// RUN: %target-build-swift-dylib(%t/libC.%target-dylib-extension) -I%t -L%t -lA -Xlinker -rpath -Xlinker %t -module-name C -emit-module -emit-module-path %t -swift-version 5 %S/Inputs/dynamic_replacement_chaining_B.swift -Xfrontend -enable-dynamic-replacement-chaining
+// RUN: %target-build-swift -I%t -L%t -lA -DCHAINING -o %t/main -Xlinker -rpath -Xlinker %t %s -swift-version 5
+// RUN: %target-codesign %t/main %t/libA.%target-dylib-extension %t/libB.%target-dylib-extension %t/libC.%target-dylib-extension
+// RUN: %target-run %t/main %t/libA.%target-dylib-extension %t/libB.%target-dylib-extension %t/libC.%target-dylib-extension
+
+
+import A
+
+import StdlibUnittest
+
+#if os(Linux)
+  import Glibc
+  let dylibSuffix = "so"
+#else
+  import Darwin
+  let dylibSuffix = "dylib"
+#endif
+
+var DynamicallyReplaceable = TestSuite("DynamicallyReplaceableChaining")
+
+
+DynamicallyReplaceable.test("DynamicallyReplaceable") {
+  var executablePath = CommandLine.arguments[0]
+  executablePath.removeLast(4)
+
+#if os(Linux)
+	_ = dlopen("libB."+dylibSuffix, RTLD_NOW)
+	_ = dlopen("libC."+dylibSuffix, RTLD_NOW)
+#else
+	_ = dlopen(executablePath+"libB."+dylibSuffix, RTLD_NOW)
+	_ = dlopen(executablePath+"libC."+dylibSuffix, RTLD_NOW)
+#endif
+
+#if CHAINING
+  expectEqual(3, Impl().foo())
+#else
+  expectEqual(2, Impl().foo())
+#endif
+}
+
+runAllTests()