Swift5 Method Swizzling 方法交换

最近看了一下关于Swift5的特性，发现了一个叫 @_dynamicReplacement(for:) 的新特性比较有趣，所以研究了一下。
新建一个Swift 5.0的工程，不能用playground！
然后输入以下代码，这里有一个key word叫dynamic，这个关键字在Swift3中就出现了，代表动态派发，只有加了dynamic修饰的方法才能被@_dynamicReplacement(for:) ：

import Foundation

class Test {
    dynamic func testA() {
        print("testA")
    }
}

extension Test {
    @_dynamicReplacement(for: testA())
    func testB() {
        print("testB")
    }
}

Test().testA()

运行一下发现控制台输出了：

testB
Program ended with exit code: 0

( ⊙ o ⊙ )啊？我们期待已久的 Method Swizzling 又回来了？

看到这里，我们不禁要与OC的Method Swizzling做一个比较：
先写一个OC的类：

//Test.h
#import <Foundation/Foundation.h>

NS_ASSUME_NONNULL_BEGIN

@interface Test : NSObject
- (void)testA;
- (void)testB;
- (void)testC;
@end

NS_ASSUME_NONNULL_END

//Test.m
#import "Test.h"
#import <objc/runtime.h>

@implementation Test

+ (void)load {
    static dispatch_once_t onceToken;
    dispatch_once(&onceToken, ^{
        Class class = [self class];
        SEL testASelector = @selector(testA);
        SEL testBSelector = @selector(testB);
        Method testAMethod = class_getInstanceMethod(class, testASelector);
        Method testBMethod = class_getInstanceMethod(class, testBSelector);
        method_exchangeImplementations(testAMethod, testBMethod);

        SEL testCSelector = @selector(testC);
        Method testCMethod = class_getInstanceMethod(class, testCSelector);
        method_exchangeImplementations(testAMethod, testCMethod);
    });
}

- (void)testA {
    NSLog(@"A");
}

- (void)testB {
    NSLog(@"B");
    [self testB];
}

- (void)testC {
    NSLog(@"C");
    [self testC];
}
@end

//main.m
#import <Foundation/Foundation.h>
#import "Test.h"

int main(int argc, const char * argv[]) {
    @autoreleasepool {
        Test * t = [[Test alloc] init];
        [t testA];
    }
    return 0;
}

上面的程序很明显，运行结果是：

2019-09-26 23:12:39.513796+0800 OC[87404:218117] C
2019-09-26 23:12:39.514014+0800 OC[87404:218117] B
2019-09-26 23:12:39.514036+0800 OC[87404:218117] A
Program ended with exit code: 0

OC中方法交换是一个一个来的，所以结果会输出 CBA，因为 testA 方法首先被替换成了 testB ，其替换后的结果又被替换成了 testC 。进而由于 Swizze 的方法都调用了原方法，所以会输出 CBA。

Method Swizzling 在 Runtime 中的原理就是方法指针的交换。
由于 OC 对于实例方法的存储方式是以方法实例表，那么我们只要能够访问到其指定的方法实例，修改 imp 指针对应的指向，再对引用计数和内存开辟等于 Class 相关的信息做一次更新就实现了 Method Swizzling。

对于这个例子，我们用Swift的@_dynamicReplacement(for:) 写一遍：

import Foundation

class Test {
    dynamic func testA() {
        print("A")
    }
}

extension Test {
    @_dynamicReplacement(for: testA())
    func testB() {
        print("B")
        testA()
    }
}

extension Test {
    @_dynamicReplacement(for: testA())
    func testC() {
        print("C")
        testA()
    }
}

Test().testA()

从视觉角度上来看，通过对 Swift Functions 的显式声明，我们完成了对于 Method Swizzling 的实现。
但是跑一下代码，发现运行结果却是：

C
A
Program ended with exit code: 0

为什么结果只显示两个？我交换了一下两个 extension 的顺序继续尝试，其打印结果又变成了 BA 。所以可以断定在执行顺序上，后声明的将会生效。那么应该如何实现这种连环 Hook 的场景呢？经过多次的尝试，从代码层面应该是不可能办到了，总会以最后一次的方法为准。

按照正常程序员的逻辑，如果我们在重构一个模块的代码，新的模块代码无论从功能还是效率上，都应该优于之前的方式、覆盖之前所有的逻辑场景。如果 Swift 支持这种连环修改的场景，那这个新的 Feature 放出其实是功能不完备的！于是我们开始翻看 Swift 这个 Feature 的 PR 代码，来一探 Dynamic Method Replacement 的原理。

首先来看这个 Dynamic Method Replacement 特性的 Issue-20333[2]，作者上来就贴了两段很有意思的代码：

/// 片段一
// Module A
struct Foo {
 dynamic func bar() {}
}
// Module B
extension Foo {
  @_dynamicReplacement(for: bar())
  func barReplacement() {
    ...
    // Calls previously active implementation of bar()
    bar()
  }
}

/// 片段二
dynamic_replacement_scope AGroupOfReplacements {
   extension Foo {
     func replacedFunc() {}
   }
   extension AnotherType {
     func replacedFunc() {}
   }
}

AGroupOfReplacements.enable()
...
AGroupOfReplacements.disable()

大概意思就是，他希望这种动态替换的特性，通过一些关键字标记和内部标记，兼容动态替换和启用开关。既然他们有规划 enable 和 disable 这两个方法来控制启用，那么就来检索它们的实现。

通过关键字搜索，我在 MetadataLookup.cpp 这个文件的 L1523-L1536 中找到了 enable 和 disable 两个方法的实现源码：

// Metadata.h#L4390-L4394:
// https://github.com/apple/swift/blob/659c49766be5e5cfa850713f43acc4a86f347fd8/include/swift/ABI/Metadata.h#L4390-L4394

/// dynamic replacement functions 的链表实现
/// 只有一个 next 指针和 imp 指针
struct DynamicReplacementChainEntry {
  void *implementationFunction;
  DynamicReplacementChainEntry *next;
};

// MetadataLookup.cpp#L1523-L1563
// https://github.com/aschwaighofer/swift/blob/fff13330d545b914d069aad0ef9fab2b4456cbdd/stdlib/public/runtime/MetadataLookup.cpp#L1523-L1563
void DynamicReplacementDescriptor::enableReplacement() const {
    // 拿到根节点
  auto *chainRoot = const_cast<DynamicReplacementChainEntry *>(
      replacedFunctionKey->root.get());

  // 通过遍历链表来保证这个方法是 enabled 的
  for (auto *curr = chainRoot; curr != nullptr; curr = curr->next) {
    if (curr == chainEntry.get()) {
            // 如果在 Replacement 链中发现了这个方法，说明已经 enable，中断操作
      swift::swift_abortDynamicReplacementEnabling();
    }
  }

  // 将 Root 节点的 imp 保存到 current，并将 current 头插
  auto *currentEntry =
      const_cast<DynamicReplacementChainEntry *>(chainEntry.get());
  currentEntry->implementationFunction = chainRoot->implementationFunction;
  currentEntry->next = chainRoot->next;

  // Root 继续进行头插操作
  chainRoot->next = chainEntry.get();
    // Root 的 imp 换成了 replacement 实现
  chainRoot->implementationFunction = replacementFunction.get();
}

// 同理 disable 做逆操作
void DynamicReplacementDescriptor::disableReplacement() const {
  const auto *chainRoot = replacedFunctionKey->root.get();
  auto *thisEntry =
      const_cast<DynamicReplacementChainEntry *>(chainEntry.get());

  // Find the entry previous to this one.
  auto *prev = chainRoot;
  while (prev && prev->next != thisEntry)
    prev = prev->next;
  if (!prev) {
    swift::swift_abortDynamicReplacementDisabling();
    return;
  }

  // Unlink this entry.
  auto *previous = const_cast<DynamicReplacementChainEntry *>(prev);
  previous->next = thisEntry->next;
  previous->implementationFunction = thisEntry->implementationFunction;
}

我们发现 Swift 中处理每一个 dynamic 方法，会为其建立一个 dynamicReplacement 链表来记录实现记录。
那么也就是说不管我们对原来的 dynamic 做了多少次 @_dynamicReplacement ，其实现原则上都会被记录下来，这一点也说明Swift并未抛弃OC中连环hook的场景。
但是调用方法后的执行代码我始终没有找到对应的逻辑，所以无法判断 Swift 在调用时机做了哪些事情。

但是总会有大神想出办法，比如 Whirlwind ,既然我们无法找到调用的实现，那么另辟蹊径：既然 Swift 已经通过链式记录了所有的实现，那么在单元测试的时候应该会有这种逻辑测试。

在根据关键字和文件后缀搜索了大量的单元测试文件后，我们发现了这个文件dynamic_replacement_chaining.swift 。我们注意到 L13 的执行命令：

// RUN: %target-build-swift-dylib(%t/%target-library-name(B)) -I%t -L%t -lA %target-rpath(%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

在命令参数中增加了 -Xfrontend -enable-dynamic-replacement-chaining ，第一反应：这个东西像 Build Settings 中的 Flags。
翻看了 Build Settings 中所有的 Compile Flags，将其尝试写入 Other Swift Flags 中：

重新运行一遍，奇迹发生了：

说明我们的实验和猜想是正确的，Swift 在处理 dynamic 是将所有实现的 imp 保存，并且也有办法根据记录的链表来触发实现。

@_dynamicReplacement 虽然在 Swift 5 中就已经带入了 Swift 中，但是在官方论坛和官方仓库中并未找到 Release 日志痕迹。
这个 PR 虽然已经开了一年之久，苹果在今年的 WWDC 之前就偷偷的 Merge 到了 master 分支上，不得不猜，Apple 是为了实现 SwiftUI，而专门定制的 Feature。