构建带记忆化的 Block 代理

#Original (English)

Source: https://www.mikeash.com/pyblog/friday-qa-2011-11-11-building-a-memoizing-block-proxy.html

Last time, I talked about my crazy hack that misuses the Objective-C message forwarding machinery to do block proxying. This allows writing code that interposes in front of an arbitrary block to intercept its arguments, manipulate its return value, etc. Today, I want to present an exanmple of using this hack which almost verges on the practical. Specifically, I’m going to discuss how to use it to build a generalized block memoization facility.

What the Heck is Memoization?In short, memoization is just a cache in front of a function. Imagine a pure function, which is to say, no side effects, and you always get the same result for the same arguments. Now imagine that this function is expensive to compute, and called multiple times with the same arguments. If the result is always the same with the same arguments, it’s a waste to constantly recompute it.

The standard answer to this is a cache. Whenever the function computes a result, it can put that result into a dictionary that’s keyed off the arguments. Whenever the function is called, it can check the dictionary first to see if the result is already in there. If it is, it can return the previously computed value.

Memoization is the process of performing this caching in a wrapper function. Instead of putting cache code into every function that fits this profile, you write the cache code once, as a wrapper around arbitrary functions. The wrapper has no idea how the function works, but it knows enough to cache return values and bypass the function if a particular set of arguments has already been computed.

CodeAs before, the code is available on GitHub here:

https://github.com/mikeash/MABlockForwarding

ApproachThe previous adventures eventually resulted in a MAForwardingBlock which takes a block and an interposer, and returns a new block. When the new block is called, the interposer is first called with an NSInvocation representing the call. The interposer can then manipulate the NSInvocation at will, optionally call through to the original block, and then give the NSInvocation’s return value back to the caller.

The goal here is to write an MAMemoize function. This function takes a block, and returns a new block which wraps the original. The new block memoizes the original, caching values based on the arguments passed in. MAMemoize will use MAForwardingBlock to interpose in front of the original block, and use an interposer block to perform the memoization.

In order to accomplish this, the interposer needs to unpack all of the arguments of the NSInvocation into an object which can be used as a dictionary key. It also needs to be able to translate the return value to and from an object that can be used as a dictionary value.

Once it’s able to do this, the rest is simple. Keep a dictionary which contains the memory. Unpack the arguments, and check to see if the dictionary has an entry for them. If it does, grab the value for those arguments and return it. Otherwise, call through to the original block, then fetch the return value and stash it in the memory dictionary before returning.

ImplementationTo start out with, we’ll define the MAMemoize function. It takes and return id because it deals with arbitrary blocks, not just a particular type:

1
    id MAMemoize(id block) {

Next up, create the memory dictionary. This is just a local variable, but it will be captured by the interposer block, so it ends up persisting. Blocks can be interesting like that. You can think of a block as being an object with just one method, where the captured variables are instance variables.

1
        NSMutableDictionary *memory = [NSMutableDictionary dictionary];

Now we call MAForwardingBlock and give it our interposer. The interposer block takes an NSInvocation and a block that calls through to the original block. Here’s the beginning of the call:

1
        return MAForwardingBlock(^(NSInvocation *inv, void (^call)(void)) {

All of the code starting from here runs when the returned block is called. The inv argument contains the NSInvocation of that call, and the call block calls the original function.

The first thing to do is fetch the arguments into an object that can be used as a dictionary key. In this case, I chose to use an NSArray. Although this is probably not a very efficient dictionary key, this code isn’t very efficient anyway and it does the job fine. The array will be filled out with objects representing each argument.

Fetching the arguments also requires knowing the invocation’s method signature, so we grab that at the same time:

1
            NSMethodSignature *sig = [inv methodSignature];
2
            NSMutableArray *args = [NSMutableArray array];

Now, loop through the arguments and add them to the array. Note that the loop starts at argument 1, because argument 0 is the implicit block pointer argument and doesn’t need to be included in the cache:

1
            for(unsigned i = 1; i < [sig numberOfArguments]; i++)
2
            {

For the case where the argument is an object, the task is easy: fetch it using getArgument:atIndex: and then stash it in the array. For non-object arguments, life gets more complicated. In fact, it’s not possible to handle every possible type, since things like pointer types simply can’t be introspected well enough to understand how they’re being used.

My approach is to special-case C strings, since they seem like a common parameter type. For all other parameter types, fetch the raw argument into an NSData and treat that as the argument, not even making an attempt to interpret the meaning. In practice, this means that all primitive values will work, as will some structs and the rare pointers, like SEL which can be compared using == and persist for the lifetime of the app.

To start with, I create an arg local variable which will hold the extracted object, whatever it is. I also grab the string describing the type of this argument:

1
                id arg = nil;
2
                const char *type = [sig getArgumentTypeAtIndex: i];

The first thing is to check to see if this argument is an object. We can do this by comparing the first character of ‘type’ with the first character of ‘@encode(id)’:

1
                if(type[0] == @encode(id)[0])
2
                {

For this case, just grab the argument into arg. I add a small additional flourish where, if the argument conforms to NSCopying, it’s copied before being placed into the array:

1
                    [inv getArgument: &arg atIndex: i];
2
                    if([arg conformsToProtocol: @protocol(NSCopying)])
3
                        arg = [arg copy];
4
                }

Next, check for a C string (any char * parameter is assumed to be a C string) using the same technique. If it is one, pull its contents into an NSData:

1
                else if(type[0] == @encode(char *)[0])
2
                {
3
                    char *str;
4
                    [inv getArgument: &str atIndex: i];
5
                    arg = [NSData dataWithBytes: str length: strlen(str)];
6
                }

Finally, for the fallback case, pull the argument directly into an NSData. The NSGetSizeAndAlignment function can tell us how large the argument is and thus how large the NSData needs to be:

1
                else
2
                {
3
                    NSUInteger size;
4
                    NSGetSizeAndAlignment(type, &size, NULL);
5
                    arg = [NSMutableData dataWithLength: size];
6
                    [inv getArgument: [arg mutableBytes] atIndex: i];
7
                }

Now we have this argument in arg. Before adding it to the array, since NSArray doesn’t like nil, check for that and use NSNull instead:

1
                if(!arg)
2
                    arg = [NSNull null];

Now add the argument to the array and continue looping until all arguments are handled:

1
                [args addObject: arg];
2
            }

Next, we check memory to see if the arguments are in it. If they are, we extract the return value and put it into the NSInvocation. If they aren’t, we call the original block, then extract the return value from the `NSInvocation. In both cases, we need the method’s return type, and we need to know if it’s an object, a C string, or something else, so we precompute those:

1
            const char *type = [sig methodReturnType];
2
            BOOL isObj = type[0] == @encode(id)[0];
3
            BOOL isCStr = type[0] == @encode(char *)[0];

Next up, check memory. This is done in a @synchronized block so that this code is thread safe:

1
            id result;
2
            @synchronized(memory)
3
            {
4
                result = [[[memory objectForKey: args] retain] autorelease];
5
            }

If the arguments exist in the dictionary, then result will contain the object representation of the return value. Otherwise, result will be nil. First, let’s look at the case where there is no entry in memory:

1
            if(!result)
2
            {

Since there’s no saved value, the first thing to do is call the original block to compute the return value:

1
                call();

Once this completes, the return value will be contained within the NSInvocation. We need to fetch it into an object so it can be stored in memory. If the return value is an object, this is really easy. Just use getReturnValue: and done:

1
                if(isObj)
2
                {
3
                    [inv getReturnValue: &result];
4
                }

If it’s a C string, then we need to fetch the string pointer, then convert that to an NSData. Note that unlike the argument saving code, I add 1 to the length of the string in order to copy the string’s terminating NUL byte as well. Since this pointer needs to be returned to the caller, and the caller will expect the terminating NUL, this saved additional redundant conversion:

1
                else if(isCStr)
2
                {
3
                    char *str;
4
                    [inv getReturnValue: &str];
5
                    result = str ? [NSData dataWithBytes: str length: strlen(str) + 1] : NULL;
6
                }

Finally, the fallback case. Like before, we fetch it into an NSData, using NSGetSizeAndAlignment to figure out the necessary length:

1
                else
2
                {
3
                    NSUInteger size;
4
                    NSGetSizeAndAlignment(type, &size, NULL);
5
                    result = [NSMutableData dataWithLength: size];
6
                    [inv getReturnValue: [result mutableBytes]];
7
                }

And again, if the object is nil, replace it with NSNull:

1
                if(!result)
2
                    result = [NSNull null];

Now that we have the result, store it in memory. At this point, we’re done with this particular code path. The proper return value is already in the NSInvocation and will be returned to the caller:

1
                @synchronized(memory)
2
                {
3
                    [memory setObject: result forKey: args];
4
                }
5
            }

Now let’s look at the case where a result was found in memory. In that case, all we have to do is stuff the contents of result into the NSInvocation using setReturnValue:. Here, we do the reverse transformation for nil, checking result for NSNull and changing it to nil:

1
            else
2
            {
3
                if(result == [NSNull null])
4
                    result = nil;

If the return type is an object, a direct call to setReturnValue: gets the job done:

1
                if(isObj)
2
                {
3
                    [inv setReturnValue: &result];
4
                }

If it’s a C string, we have to fetch the string pointer into a local variable, then set that:

1
                else if(isCStr)
2
                {
3
                    const char *str = [result bytes];
4
                    [inv setReturnValue: &str];
5
                }

Otherwise, result is an NSData of the appropriate length containing the return value. We can just pass the pointer straight into the invocation, and that’s all there is to it!

1
                else
2
                {
3
                    [inv setReturnValue: [result mutableBytes]];
4
                }
5
            }
6
        }, block);
7
    }

For those of you who have lost track, the block at the end of this is the original block which was passed into MAMemoize, and tells MABlockForwarding which block to wrap.

ExampleAs an example of memoization, let’s look at a simple, computationally-expensive block:

1
    __block uint64_t (^fib)(int) = ^uint64_t (int n) {
2
        if(n <= 1)
3
            return 1;
4
        else
5
            return fib(n - 1) + fib(n - 2);
6
    }

This fib block computes the well known Fibonacci function using recursion. This particular approach is extremely slow. The two recursive calls in the else clause result in an exponential growth of computation. Each of the two recursive calls spawns two more recursive calls, which in turn each spawn two more recursive calls, etc.

However, almost all of the calls are redundant. If you analyze it, you’ll see that fib(5) calls fib(4) and fib(3), but then fib(4) just calls fib(3) again. Memoizing fib will make it much faster, since subsequent calls to fib(3) will be cached rather than recomputed. There are, of course, much better ways compute the Fibonacci function altogether, but this still serves as an interesting example.

To memoize fib, we just call MAMemoize and assign the result back to fib:

1
    fib = MAMemoize(fib);

Since fib is declared as __block, the recursive calls pick up the memoized version, and the result is massively faster on large values than the original.

ConclusionThis memoization wrapper is not entirely useful. For one thing, it’s based on my colossal block proxying hack, which really can’t be used in code you intend to ship. Even ignoring that, all of the games played with NSInvocation are slow and this is a big problem for code that is, essentially, an optimization.

Despite this, it’s a great learning experience. Even if it’s not practical, this memoization wrapper illustrates how to deal with NSInvocation in complicated situations and deal with some of the variety of argument and return types that you’re likely to encounter in the wild. Beyond that, this kind of hacking is just good fun.

That wraps things up for today. Keep those suggestions coming. Friday Q&A is (usually) driven by reader-submitted topics, so if you have a subject you’d like to see covered, send it in!