Swift 中的協(xié)議擴(kuò)展為 iOS 開發(fā)帶來(lái)了非常多的可能性，它為我們提供了一種類似多重繼承的功能，幫助我們減少一切可能導(dǎo)致重復(fù)代碼的地方。

關(guān)于 Protocol Extension

在 Swift 中比較出名的 Then 就是使用了協(xié)議擴(kuò)展為所有的 AnyObject 添加方法，而且不需要調(diào)用 runtime 相關(guān)的 API，其實(shí)現(xiàn)簡(jiǎn)直是我見過(guò)最簡(jiǎn)單的開源框架之一：

public protocol Then {}

extension Then where Self: AnyObject {
    public func then(@noescape block: Self -> Void) -> Self {
        block(self)
        return self
    }
}

extension NSObject: Then {}

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

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

這里沒(méi)有調(diào)用任何的 runtime 相關(guān) API，也沒(méi)有在 NSObject 中進(jìn)行任何的方法聲明，甚至 protocol Then {} 協(xié)議本身都只有一個(gè)大括號(hào)，整個(gè) Then 框架就是基于協(xié)議擴(kuò)展來(lái)實(shí)現(xiàn)的。

在 Objective-C 中同樣有協(xié)議，但是這些協(xié)議只是相當(dāng)于接口，遵循某個(gè)協(xié)議的類只表明實(shí)現(xiàn)了這些接口，每個(gè)類都需要對(duì)這些接口有單獨(dú)的實(shí)現(xiàn)，這就很可能會(huì)導(dǎo)致重復(fù)代碼的產(chǎn)生。

而協(xié)議擴(kuò)展可以調(diào)用協(xié)議中聲明的方法，以及 where Self: AnyObject 中的 AnyObject 的類/實(shí)例方法，這就大大提高了可操作性，便于開發(fā)者寫出一些意想不到的擴(kuò)展。

如果讀者對(duì) Protocol Extension 興趣或者不了解協(xié)議擴(kuò)展，可以閱讀最后的 Reference 了解相關(guān)內(nèi)容。

ProtocolKit

其實(shí)協(xié)議擴(kuò)展的強(qiáng)大之處就在于它能為遵循協(xié)議的類添加一些方法的實(shí)現(xiàn)，而不只是一些接口，而今天為各位讀者介紹的 ProtocolKit 就實(shí)現(xiàn)了這一功能，為遵循協(xié)議的類添加方法。

ProtocolKit 的使用

我們先來(lái)看一下如何使用 ProtocolKit，首先定義一個(gè)協(xié)議：

@protocol TestProtocol

@required

- (void)fizz;

@optional

- (void)buzz;

@end

在協(xié)議中定義了兩個(gè)方法，必須實(shí)現(xiàn)的方法 fizz 以及可選實(shí)現(xiàn) buzz，然后使用 ProtocolKit 提供的接口 defs 來(lái)定義協(xié)議中方法的實(shí)現(xiàn)了：

@defs(TestProtocol)

- (void)buzz {
    NSLog(@"Buzz");
}

@end

這樣所有遵循 TestProtocol 協(xié)議的對(duì)象都可以調(diào)用 buzz 方法，哪怕它們沒(méi)有實(shí)現(xiàn)：

protocol-demo

上面的 XXObject 雖然沒(méi)有實(shí)現(xiàn) buzz 方法，但是該方法仍然成功執(zhí)行了。

ProtocolKit 的實(shí)現(xiàn)

ProtocolKit 的主要原理仍然是 runtime 以及宏的；通過(guò)宏的使用來(lái)隱藏類的聲明以及實(shí)現(xiàn)的代碼，然后在 main 函數(shù)運(yùn)行之前，將類中的方法實(shí)現(xiàn)加載到內(nèi)存，使用 runtime 將實(shí)現(xiàn)注入到目標(biāo)類中。

如果你對(duì)上面的原理有所疑惑也不是太大的問(wèn)題，這里只是給你一個(gè) ProtocolKit 原理的簡(jiǎn)單描述，讓你了解它是如何工作的。

ProtocolKit 中有兩條重要的執(zhí)行路線：

• _pk_extension_load 將協(xié)議擴(kuò)展中的方法實(shí)現(xiàn)加載到了內(nèi)存
• _pk_extension_inject_entry 負(fù)責(zé)將擴(kuò)展協(xié)議注入到實(shí)現(xiàn)協(xié)議的類

加載實(shí)現(xiàn)

首先要解決的問(wèn)題是如何將方法實(shí)現(xiàn)加載到內(nèi)存中，這里可以先了解一下上面使用到的 defs 接口，它其實(shí)只是一個(gè)調(diào)用了其它宏的超級(jí)宏這名字是我編的：

#define defs _pk_extension

#define _pk_extension($protocol) _pk_extension_imp($protocol, _pk_get_container_class($protocol))

#define _pk_extension_imp($protocol, $container_class) \
    protocol $protocol; \
    @interface $container_class : NSObject <$protocol> @end \
    @implementation $container_class \
    + (void)load { \
        _pk_extension_load(@protocol($protocol), $container_class.class); \
    } \

#define _pk_get_container_class($protocol) _pk_get_container_class_imp($protocol, __COUNTER__)
#define _pk_get_container_class_imp($protocol, $counter) _pk_get_container_class_imp_concat(__PKContainer_, $protocol, $counter)
#define _pk_get_container_class_imp_concat($a, $b, $c) $a ## $b ## _ ## $c

使用 defs 作為接口的是因?yàn)樗且粋€(gè)保留的 keyword，Xcode 會(huì)將它渲染成與 @property 等其他關(guān)鍵字相同的顏色。

上面的這一坨宏并不需要一個(gè)一個(gè)來(lái)分析，只需要看一下最后展開會(huì)變成什么：

@protocol TestProtocol; 

@interface __PKContainer_TestProtocol_0 : NSObject <TestProtocol>

@end

@implementation __PKContainer_TestProtocol_0

+ (void)load {
    _pk_extension_load(@protocol(TestProtocol), __PKContainer_TestProtocol_0.class); 
}

根據(jù)上面宏的展開結(jié)果，這里可以介紹上面的一坨宏的作用：

• defs 這貨沒(méi)什么好說(shuō)的，只是 _pk_extension 的別名，為了提供一個(gè)更加合適的名字作為接口
• _pk_extension 向 _pk_extension_imp 中傳入 $protocol 和 _pk_get_container_class($protocol) 參數(shù)
  • _pk_get_container_class 的執(zhí)行生成一個(gè)類名，上面生成的類名就是 __PKContainer_TestProtocol_0，這個(gè)類名是 __PKContainer_、 $protocol 和 __COUNTER__ 拼接而成的（__COUNTER__ 只是一個(gè)計(jì)數(shù)器，可以理解為每次調(diào)用時(shí)加一）
• _pk_extension_imp 會(huì)以傳入的類名生成一個(gè)遵循當(dāng)前 $protocol 協(xié)議的類，然后在 + load 方法中執(zhí)行 _pk_extension_load 加載擴(kuò)展協(xié)議

通過(guò)宏的運(yùn)用成功隱藏了 __PKContainer_TestProtocol_0 類的聲明以及實(shí)現(xiàn)，還有 _pk_extension_load 函數(shù)的調(diào)用：

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

    pthread_mutex_unlock(&protocolsLoadingLock);
}

ProtocolKit 使用了 protocolsLoadingLock 來(lái)保證靜態(tài)變量 allExtendedProtocols 以及 extendedProtcolCount extendedProtcolCapacity 不會(huì)因?yàn)榫€程競(jìng)爭(zhēng)導(dǎo)致問(wèn)題：

• allExtendedProtocols 用于保存所有的 PKExtendedProtocol 結(jié)構(gòu)體
• 后面的兩個(gè)變量確保數(shù)組不會(huì)越界，并在數(shù)組滿的時(shí)候，將內(nèi)存占用地址翻倍

方法的后半部分會(huì)在靜態(tài)變量中尋找或創(chuàng)建傳入的 protocol 對(duì)應(yīng)的 PKExtendedProtocol 結(jié)構(gòu)體：

size_t resultIndex = SIZE_T_MAX;
for (size_t index = 0; index < extendedProtcolCount; ++index) {
    if (allExtendedProtocols[index].protocol == protocol) {
        resultIndex = index;
        break;
    }
}

if (resultIndex == SIZE_T_MAX) {
    allExtendedProtocols[extendedProtcolCount] = (PKExtendedProtocol){
        .protocol = protocol,
        .instanceMethods = NULL,
        .instanceMethodCount = 0,
        .classMethods = NULL,
        .classMethodCount = 0,
    };
    resultIndex = extendedProtcolCount;
    extendedProtcolCount++;
}

_pk_extension_merge(&(allExtendedProtocols[resultIndex]), containerClass);

這里調(diào)用的 _pk_extension_merge 方法非常重要，不過(guò)在介紹 _pk_extension_merge 之前，首先要了解一個(gè)用于保存協(xié)議擴(kuò)展信息的私有結(jié)構(gòu)體 PKExtendedProtocol：

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

PKExtendedProtocol 結(jié)構(gòu)體中保存了協(xié)議的指針、實(shí)例方法、類方法、實(shí)例方法數(shù)以及類方法數(shù)用于框架記錄協(xié)議擴(kuò)展的狀態(tài)。

回到 _pk_extension_merge 方法，它會(huì)將新的擴(kuò)展方法追加到 PKExtendedProtocol 結(jié)構(gòu)體的數(shù)組 instanceMethods 以及 classMethods 中：

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

因?yàn)轭惙椒ǖ淖芳优c實(shí)例方法幾乎完全相同，所以上述代碼省略了向結(jié)構(gòu)體中的類方法追加方法的實(shí)現(xiàn)代碼。

實(shí)現(xiàn)中使用 class_copyMethodList 從 containerClass 拉出方法列表以及方法數(shù)量；通過(guò) _pk_extension_create_merged 返回一個(gè)合并之后的方法列表，最后在更新結(jié)構(gòu)體中的 instanceMethods 以及 instanceMethodCount 成員變量。

_pk_extension_create_merged 只是重新 malloc 一塊內(nèi)存地址，然后使用 memcpy 將所有的方法都復(fù)制到了這塊內(nèi)存地址中，最后返回首地址：

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

這一節(jié)的代碼從使用宏生成的類中抽取方法實(shí)現(xiàn)，然后以結(jié)構(gòu)體的形式加載到內(nèi)存中，等待之后的方法注入。

注入方法實(shí)現(xiàn)

注入方法的時(shí)間點(diǎn)在 main 函數(shù)執(zhí)行之前議實(shí)現(xiàn)的注入并不是在 + load 方法 + initialize 方法調(diào)用時(shí)進(jìn)行的，而是使用的編譯器指令(compiler directive) __attribute__((constructor)) 實(shí)現(xiàn)的：

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

使用上述編譯器指令的函數(shù)會(huì)在 shared library 加載的時(shí)候執(zhí)行，也就是 main 函數(shù)之前，可以看 StackOverflow 上的這個(gè)問(wèn)題 How exactly does attribute((constructor)) work?。

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

_pk_extension_inject_entry 會(huì)在 main 執(zhí)行之前遍歷內(nèi)存中的所有 Class（整個(gè)遍歷過(guò)程都是在一個(gè)自動(dòng)釋放池中進(jìn)行的），如果某個(gè)類遵循了allExtendedProtocols 中的協(xié)議，調(diào)用 _pk_extension_inject_class 向類中注射（inject）方法實(shí)現(xiàn)：

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

如果類中沒(méi)有實(shí)現(xiàn)該實(shí)例方法就會(huì)通過(guò) runtime 中的 class_addMethod 注射該實(shí)例方法；而類方法的注射有些不同，因?yàn)轭惙椒ǘ际潜４嬖谠愔械模恍╊惙椒ㄓ捎谄涮厥獾匚蛔詈貌灰淖兤湓袑?shí)現(xiàn)，比如 + load 和 + initialize 這兩個(gè)類方法就比較特殊，如果想要了解這兩個(gè)方法的相關(guān)信息，可以在 Reference 中查看相關(guān)的信息。

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

實(shí)現(xiàn)上的不同僅僅在獲取元類、以及跳過(guò) + load 和 + initialize 方法上。

總結(jié)

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

Reference

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

Github Repo：iOS-Source-Code-Analyze

Follow: Draveness · Github

原文鏈接：http://draveness.me/protocol-extension/