Extending Methods in a Category by Method Swizzling

Probably you have been in the situation to being writing a category and at some point exclaim “F**K! It’s a category!”. Yes, and probably was because you accidentally tried to override a method without remembering that effectively, you were writing in a category.

In the largest percentage of situations, when you need to override a method, you want to write a subclass; when instead you want to write a category for a class, it is because you probably need to have some new methods or properties that you’d like to be used from every subclass.

There is only another situation which could confuse you, and this is when you are in both the previous cases, together.

The problem

Let’s suppose you want to add some generic code in -layoutSubviews. The first thing which comes in your mind is:

1
2
3
4
5
6

- (void)layoutSubviews
{
    [super layoutSubviews];

    //Your extra code
}

Great.

Wait..let’s suppose also you would like every subclass of UIView to implement an extra code, like for instance logging something or checking a property before do something else.
Make this by subclassing would be time-consuming and would implies code repetition as you should create a subclass and override the -layoutSubviews method for UIView, UIImageView, UIControl, UILabel and so on, as well as set these subclasses for any new and old class…

Exactly, it is unthinkable.

When we need to add methods or properties to a class, we normally write a single Category for the higher superclass we need to extend. In this case that class would be UIView. But, we still have a problem: we don’t want to add a new method, we want to extend an existing method.
We know that this is not possible, as a category is a class extension, and implementing an existing method like -layoutSubviews, would results in a method repetition. As figurative and easy example think just to something like this:

1
2
3
4
5
6
7
8
9
10
11
12
13
14

//UIView standard code...

- (void)layoutSubviews
{
    //..default UIView code
}

//...UIView standard code

//My layoutSubviews method
- (void)layoutSubviews
{
    //..my code
}

So, what can we do?

Playing around the problem with Method Swizzling

Method Swizzling is a technique with which, roughly speaking, you can swipe two methods making the selector of the first method calls the second method and vice-versa.

The idea is to create a new method and swap his implementation with the -layoutSubviews one. Therefore, calling -layoutSubviews, will be effectively executed the new method with our new code, and in the new method, we will also call his selector, which instead will executes the -layoutSubviews original method.

Calm down, I know, it’s not easy to explain and to understand with these few lines above. So I created a couple of diagrams to have in mind a visual image of the situation:

The normal condition we have is the following:

The condition we want to pull out, is instead this:

Being not possible insert extra code in the -layoutSubviews method, in a category, we have to redirect the flow through our custom method, in order to manipulate it.

Now, it’s the time to have a look at the explained code step by step.

Let’s start by saying that Apple offers a set of C functions, useful to manipulate properties, ivars, classes, methods and so on, at runtime. We are gonna see some of these in order to sort our situation.

STEP 1

Add the new method

The firs thing to do, is to add a new (empty for now) method to our category. This method, will be the one containing our code extension or our part of code we want to be executed when -layoutSubviews is called:

1
2
3
4
5
6
7
8

@implementation UIView (MG_layoutSubviews)

- (void)mg_layoutSubviews
{

}

@end

STEP 2

Write the swizzling method - First part

The method containing the swizzling code, needs to be a class method. We will see the reason later. We must to remeber that we are doing this to have an extended -layoutSubviews default implementation:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

//..category code

+ (void)mg_extendsLayoutSubviews
{
    //PART 1
    Class thisClass = self;

    //layoutSubviews selector, method, implementation
    SEL layoutSubviewsSEL = @selector(layoutSubviews);
    Method layoutSubviewsMethod = class_getInstanceMethod(thisClass, layoutSubviewsSEL);
    IMP layoutSubviewsIMP = method_getImplementation(layoutSubviewsMethod);

    //mg_layoutSubviews selector, method, implementation
    SEL mg_layoutSubviewsSEL = @selector(mg_layoutSubviews);
    Method mg_layoutSubviewsMethod = class_getInstanceMethod(thisClass, mg_layoutSubviewsSEL);
    IMP mg_layoutSubviewsIMP = method_getImplementation(mg_layoutSubviewsMethod);

    //....
}

//..category code

class_getInstanceMethod return the method which corresponds to the implementation of the selector and the class specified, or NULL if the specified class or its superclasses do not contain an instance method with the specified selector.

method_getImplementation return a function pointer of type IMP.

It is important to underline that each class has a dispatch table, and as Apple says:

This table has entries that associate method selectors with the class-specific addresses of the methods they identify.

So basically, Method is the glue between SEL and IMP, and at this point we have completely access to these information.

STEP 3

Write the swizzling method - Second part

Now, we need to exchange the two implementations:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

//..category code

+ (void)mg_extendsLayoutSubviews
{
    //PART 1
    Class thisClass = self;

    //layoutSubviews selector, method, implementation
    SEL layoutSubviewsSEL = @selector(layoutSubviews);
    Method layoutSubviewsMethod = class_getInstanceMethod(thisClass, layoutSubviewsSEL);
    IMP layoutSubviewsIMP = method_getImplementation(layoutSubviewsMethod);

    //mg_layoutSubviews selector, method, implementation
    SEL mg_layoutSubviewsSEL = @selector(mg_layoutSubviews);
    Method mg_layoutSubviewsMethod = class_getInstanceMethod(thisClass, mg_layoutSubviewsSEL);
    IMP mg_layoutSubviewsIMP = method_getImplementation(mg_layoutSubviewsMethod);

    //PART 2
    //Try to add the method layoutSubviews with the new implementation (if already exists it'll return NO)
    BOOL wasMethodAdded = class_addMethod(thisClass, layoutSubviewsSEL, mg_layoutSubviewsIMP, method_getTypeEncoding(mg_layoutSubviewsMethod));

    if (wasMethodAdded) {
        //Just set the new selector points to the original layoutSubviews method
        class_replaceMethod(thisClass, mg_layoutSubviewsSEL, layoutSubviewsIMP, method_getTypeEncoding(layoutSubviewsMethod));
    } else {
        method_exchangeImplementations(layoutSubviewsMethod, mg_layoutSubviewsMethod);
    }
}

//..category code

What’s happening in this code it is not really hard as could to seem.
The first thing to do is to try to add the method -layoutSubviews. In fact, also if this method is in the default implementation of UIView, it could not be implemented in the subclass of UIView which the developer is gonna write. It means, that could not exists an associated implementation for that subclass, but just for the superclass.

Using the function BOOL class_addMethod ( Class cls, SEL name, IMP imp, const char *types ), we can add in a Class a method associated to an implementation and called by a selector. Therefore, if the method will results to be not implemented, the returned var will be YES; if instead there will be an error due for example the fact just explained above, the resulted var will be equals to NO.

In this code, we are adding a method called by the -layoutSubviews selector, but with the implementation of our new -mg_layoutSubviews method. The last parameter is an array of characters that describe the types of the arguments to the method.
Please note that thisClass is equals to self, but remember that we are in a class method (+), and it means that self is exactly the class UIView and not an her instance.

Now, coming back to the variable result:

case YES, with the same approach, we have just to replace the other method associating the -mg_layoutSubviews selector to the implementation of the -layoutSubviews method.

case NO, means that the method and his implementation already exist, and we have just to exchange the two implementations between them.

STEP 4

Call back -layoutSubviews (by the swizzled selector)

Even if we exchanged the two implementations, we are still missing a part of our final diagram above: the arrow which calls the -layoutSubviews implementation.
In fact, at the moment, UIView will calls -layoutSubviews, which due the swizzling, it is associated with our new method. But we still have that method empty.

What we need, is close the circle executing also the original -layoutSubviews implementation. Just remember, that now that implementation can be executed calling the new selector -mg_layoutSubviews. So let’s write it by code:

1
2
3
4
5
6
7
8
9
10

//..category code

- (void)mg_layoutSubviews
{
    [self mg_layoutSuviews];

    //..our extra code
}

//..category code

Now the circle is closed and we have our method, effectively, “extended”.

Finalize

As saw above, we wanted to use the class method +mg_extendsLayoutSubviews to write our swizzling code. So now, we have to call it in order to launch all this process.

The best method in which do it, is +load. This because that method is executed when the class is initially loaded.
Basically we are saying to the class to exchange his SEL -> Method (and therefore IMP) association in the dispatch table, and we are saying this directly when the class is loaded:

1
2
3
4
5
6
7

+ (void)load
{
    static dispatch_once_t onceToken;
    dispatch_once(&onceToken, ^{
        [self mg_createNewLayoutSubviews];
    });
}

Because we are changing this global state of the class, we need to do it extremely careful. For this reason, we use the dispatch_once function, which offers us the atomicity, and the guarantee that the code inside will be executed only once.

Being called by the class method +load, this explains why +mg_createNewLayoutSubviews is a class method too.

Conclusion

In the 90% of situations, you will not need to use this approach, but sometimes could be really useful. Just be sure you have not other solutions before to proceed, because since this is an extremely powerful way, it could be also extremely dangerous if you don’t pay the maximum attention.

Follow me on Twitter!

Enjoy ;)