如何在 Objective-C 中實現協議擴展

02-01

Swift 中的協議擴展為 iOS 開發帶來了非常多的可能性，它為我們提供了一種類似多重繼承的功能，幫助我們減少一切可能導致重複代碼的地方。

關於 Protocol Extension

在 Swift 中比較出名的 Then 就是使用了協議擴展為所有的 AnyObject 添加方法，而且不需要調用 runtime 相關的 API，其實現簡直是我見過最簡單的開源框架之一：

public protocol Then {}nnextension Then where Self: AnyObject {n public func then(@noescape block: Self -> Void) -> Self {n block(self)n return selfn }n}nnextension NSObject: Then {}n

只有這麼幾行代碼，就能為所有的 NSObject 添加下面的功能：

let titleLabel = UILabel().then {n $0.textColor = .blackColor()n $0.textAlignment = .Centern}n

這裡沒有調用任何的 runtime 相關 API，也沒有在 NSObject 中進行任何的方法聲明，甚至 protocol Then {} 協議本身都只有一個大括弧，整個 Then 框架就是基於協議擴展來實現的。

在 Objective-C 中同樣有協議，但是這些協議只是相當於介面，遵循某個協議的類只表明實現了這些介面，每個類都需要對這些介面有單獨的實現，這就很可能會導致重複代碼的產生。

而協議擴展可以調用協議中聲明的方法，以及 where Self: AnyObject 中的 AnyObject 的類/實例方法，這就大大提高了可操作性，便於開發者寫出一些意想不到的擴展。

如果讀者對 Protocol Extension 興趣或者不了解協議擴展，可以閱讀最後的 Reference 了解相關內容。

ProtocolKit

其實協議擴展的強大之處就在於它能為遵循協議的類添加一些方法的實現，而不只是一些介面，而今天為各位讀者介紹的 ProtocolKit 就實現了這一功能，為遵循協議的類添加方法。

ProtocolKit 的使用

我們先來看一下如何使用 ProtocolKit，首先定義一個協議：

@protocol TestProtocolnn@requirednn- (void)fizz;nn@optionalnn- (void)buzz;nn@endn

在協議中定義了兩個方法，必須實現的方法 fizz 以及可選實現 buzz，然後使用 ProtocolKit 提供的介面 defs 來定義協議中方法的實現了：

@defs(TestProtocol)nn- (void)buzz {n NSLog(@"Buzz");n}nn@endn

這樣所有遵循 TestProtocol 協議的對象都可以調用 buzz 方法，哪怕它們沒有實現：

上面的 XXObject 雖然沒有實現 buzz 方法，但是該方法仍然成功執行了。

ProtocolKit 的實現

ProtocolKit 的主要原理仍然是 runtime 以及宏的；通過宏的使用來隱藏類的聲明以及實現的代碼，然後在 main 函數運行之前，將類中的方法實現載入到內存，使用 runtime 將實現注入到目標類中。

如果你對上面的原理有所疑惑也不是太大的問題，這裡只是給你一個 ProtocolKit 原理的簡單描述，讓你了解它是如何工作的。

ProtocolKit 中有兩條重要的執行路線：

_pk_extension_load 將協議擴展中的方法實現載入到了內存
_pk_extension_inject_entry 負責將擴展協議注入到實現協議的類

載入實現

首先要解決的問題是如何將方法實現載入到內存中，這裡可以先了解一下上面使用到的 defs 介面，它其實只是一個調用了其它宏的超級宏這名字是我編的：

#define defs _pk_extensionnn#define _pk_extension($protocol) _pk_extension_imp($protocol, _pk_get_container_class($protocol))nn#define _pk_extension_imp($protocol, $container_class) n protocol $protocol; n @interface $container_class : NSObject <$protocol> @end n @implementation $container_class n + (void)load { n _pk_extension_load(@protocol($protocol), $container_class.class); n } nn#define _pk_get_container_class($protocol) _pk_get_container_class_imp($protocol, __COUNTER__)n#define _pk_get_container_class_imp($protocol, $counter) _pk_get_container_class_imp_concat(__PKContainer_, $protocol, $counter)n#define _pk_get_container_class_imp_concat($a, $b, $c) $a ## $b ## _ ## $cn

使用 defs 作為介面的是因為它是一個保留的 keyword，Xcode 會將它渲染成與 @property 等其他關鍵字相同的顏色。

上面的這一坨宏並不需要一個一個來分析，只需要看一下最後展開會變成什麼：

@protocol TestProtocol; nn@interface __PKContainer_TestProtocol_0 : NSObject <TestProtocol>nn@endnn@implementation __PKContainer_TestProtocol_0nn+ (void)load {n _pk_extension_load(@protocol(TestProtocol), __PKContainer_TestProtocol_0.class); n}n

根據上面宏的展開結果，這裡可以介紹上面的一坨宏的作用：

defs 這貨沒什麼好說的，只是 _pk_extension 的別名，為了提供一個更加合適的名字作為介面
_pk_extension 向 _pk_extension_imp 中傳入 $protocol 和 _pk_get_container_class($protocol) 參數
- _pk_get_container_class 的執行生成一個類名，上面生成的類名就是 __PKContainer_TestProtocol_0，這個類名是 __PKContainer_、 $protocol 和 __COUNTER__ 拼接而成的（__COUNTER__ 只是一個計數器，可以理解為每次調用時加一）
_pk_extension_imp 會以傳入的類名生成一個遵循當前 $protocol 協議的類，然後在 + load 方法中執行 _pk_extension_load 載入擴展協議

通過宏的運用成功隱藏了 __PKContainer_TestProtocol_0 類的聲明以及實現，還有 _pk_extension_load 函數的調用：

void _pk_extension_load(Protocol *protocol, Class containerClass) {nn pthread_mutex_lock(&protocolsLoadingLock);nn if (extendedProtcolCount >= extendedProtcolCapacity) {n size_t newCapacity = 0;n if (extendedProtcolCapacity == 0) {n newCapacity = 1;n } else {n newCapacity = extendedProtcolCapacity << 1;n }n allExtendedProtocols = realloc(allExtendedProtocols, sizeof(*allExtendedProtocols) * newCapacity);n extendedProtcolCapacity = newCapacity;n }nn ...nn pthread_mutex_unlock(&protocolsLoadingLock);n}n

ProtocolKit 使用了 protocolsLoadingLock 來保證靜態變數 allExtendedProtocols 以及 extendedProtcolCount extendedProtcolCapacity 不會因為線程競爭導致問題：

allExtendedProtocols 用於保存所有的 PKExtendedProtocol 結構體
後面的兩個變數確保數組不會越界，並在數組滿的時候，將內存佔用地址翻倍

方法的後半部分會在靜態變數中尋找或創建傳入的 protocol 對應的 PKExtendedProtocol 結構體：

size_t resultIndex = SIZE_T_MAX;nfor (size_t index = 0; index < extendedProtcolCount; ++index) {n if (allExtendedProtocols[index].protocol == protocol) {n resultIndex = index;n break;n }n}nnif (resultIndex == SIZE_T_MAX) {n allExtendedProtocols[extendedProtcolCount] = (PKExtendedProtocol){n .protocol = protocol,n .instanceMethods = NULL,n .instanceMethodCount = 0,n .classMethods = NULL,n .classMethodCount = 0,n };n resultIndex = extendedProtcolCount;n extendedProtcolCount++;n}nn_pk_extension_merge(&(allExtendedProtocols[resultIndex]), containerClass);n

這裡調用的 _pk_extension_merge 方法非常重要，不過在介紹 _pk_extension_merge 之前，首先要了解一個用於保存協議擴展信息的私有結構體 PKExtendedProtocol：

typedef struct {n Protocol *__unsafe_unretained protocol;n Method *instanceMethods;n unsigned instanceMethodCount;n Method *classMethods;n unsigned classMethodCount;n} PKExtendedProtocol;n

PKExtendedProtocol 結構體中保存了協議的指針、實例方法、類方法、實例方法數以及類方法數用於框架記錄協議擴展的狀態。

回到 _pk_extension_merge 方法，它會將新的擴展方法追加到 PKExtendedProtocol 結構體的數組 instanceMethods 以及 classMethods 中：

void _pk_extension_merge(PKExtendedProtocol *extendedProtocol, Class containerClass) {n // Instance methodsn unsigned appendingInstanceMethodCount = 0;n Method *appendingInstanceMethods = class_copyMethodList(containerClass, &appendingInstanceMethodCount);n Method *mergedInstanceMethods = _pk_extension_create_merged(extendedProtocol->instanceMethods,n extendedProtocol->instanceMethodCount,n appendingInstanceMethods,n appendingInstanceMethodCount);n free(extendedProtocol->instanceMethods);n extendedProtocol->instanceMethods = mergedInstanceMethods;n extendedProtocol->instanceMethodCount += appendingInstanceMethodCount;nn // Class methodsn ...n}n

因為類方法的追加與實例方法幾乎完全相同，所以上述代碼省略了向結構體中的類方法追加方法的實現代碼。

實現中使用 class_copyMethodList 從 containerClass 拉出方法列表以及方法數量；通過 _pk_extension_create_merged 返回一個合併之後的方法列表，最後在更新結構體中的 instanceMethods 以及 instanceMethodCount 成員變數。

_pk_extension_create_merged 只是重新 malloc 一塊內存地址，然後使用 memcpy 將所有的方法都複製到了這塊內存地址中，最後返回首地址：

Method *_pk_extension_create_merged(Method *existMethods, unsigned existMethodCount, Method *appendingMethods, unsigned appendingMethodCount) {nn if (existMethodCount == 0) {n return appendingMethods;n }n unsigned mergedMethodCount = existMethodCount + appendingMethodCount;n Method *mergedMethods = malloc(mergedMethodCount * sizeof(Method));n memcpy(mergedMethods, existMethods, existMethodCount * sizeof(Method));n memcpy(mergedMethods + existMethodCount, appendingMethods, appendingMethodCount * sizeof(Method));n return mergedMethods;n}n

這一節的代碼從使用宏生成的類中抽取方法實現，然後以結構體的形式載入到內存中，等待之後的方法注入。

注入方法實現

注入方法的時間點在 main 函數執行之前議實現的注入並不是在 + load 方法 + initialize 方法調用時進行的，而是使用的編譯器指令(compiler directive) __attribute__((constructor)) 實現的：

__attribute__((constructor)) static void _pk_extension_inject_entry(void);n

使用上述編譯器指令的函數會在 shared library 載入的時候執行，也就是 main 函數之前，可以看 StackOverflow 上的這個問題 How exactly does attribute((constructor)) work?。

__attribute__((constructor)) static void _pk_extension_inject_entry(void) {n #1：加鎖n unsigned classCount = 0;n Class *allClasses = objc_copyClassList(&classCount);nn @autoreleasepool {n for (unsigned protocolIndex = 0; protocolIndex < extendedProtcolCount; ++protocolIndex) {n PKExtendedProtocol extendedProtcol = allExtendedProtocols[protocolIndex];n for (unsigned classIndex = 0; classIndex < classCount; ++classIndex) {n Class class = allClasses[classIndex];n if (!class_conformsToProtocol(class, extendedProtcol.protocol)) {n continue;n }n _pk_extension_inject_class(class, extendedProtcol);n }n }n }n #2：解鎖並釋放 allClasses、allExtendedProtocolsn}n

_pk_extension_inject_entry 會在 main 執行之前遍歷內存中的所有 Class（整個遍歷過程都是在一個自動釋放池中進行的），如果某個類遵循了allExtendedProtocols 中的協議，調用 _pk_extension_inject_class 向類中注射（inject）方法實現：

static void _pk_extension_inject_class(Class targetClass, PKExtendedProtocol extendedProtocol) {nn for (unsigned methodIndex = 0; methodIndex < extendedProtocol.instanceMethodCount; ++methodIndex) {n Method method = extendedProtocol.instanceMethods[methodIndex];n SEL selector = method_getName(method);nn if (class_getInstanceMethod(targetClass, selector)) {n continue;n }nn IMP imp = method_getImplementation(method);n const char *types = method_getTypeEncoding(method);n class_addMethod(targetClass, selector, imp, types);n }nn #1: 注射類方法n}n

如果類中沒有實現該實例方法就會通過 runtime 中的 class_addMethod 注射該實例方法；而類方法的注射有些不同，因為類方法都是保存在元類中的，而一些類方法由於其特殊地位最好不要改變其原有實現，比如 + load 和 + initialize 這兩個類方法就比較特殊，如果想要了解這兩個方法的相關信息，可以在 Reference 中查看相關的信息。

Class targetMetaClass = object_getClass(targetClass);nfor (unsigned methodIndex = 0; methodIndex < extendedProtocol.classMethodCount; ++methodIndex) {n Method method = extendedProtocol.classMethods[methodIndex];n SEL selector = method_getName(method);nn if (selector == @selector(load) || selector == @selector(initialize)) {n continue;n }n if (class_getInstanceMethod(targetMetaClass, selector)) {n continue;n }nn IMP imp = method_getImplementation(method);n const char *types = method_getTypeEncoding(method);n class_addMethod(targetMetaClass, selector, imp, types);n}n

實現上的不同僅僅在獲取元類、以及跳過 + load 和 + initialize 方法上。

總結

ProtocolKit 通過宏和 runtime 實現了類似協議擴展的功能，其實現代碼總共也只有 200 多行，還是非常簡潔的；在另一個叫做 libextobjc 的框架中也實現了類似的功能，有興趣的讀者可以查看 EXTConcreteProtocol.h · libextobjc 這個文件。

Reference

Protocols · Apple Doc
EXTConcreteProtocol.h · libextobjc
__attribute__ · NSHipster
你真的了解 load 方法么？
懶惰的 initialize 方法

Github Repo：iOS-Source-Code-Analyze
Follow: Draveness · Github
原文鏈接：如何在 Objective-C 中實現協議擴展