// g++ -framework Foundation -W -Wall -Wno-unused-parameter llvmhom.mm `llvm-config --cxxflags --ldflags --libs core engine`
#include "llvm/Analysis/Verifier.h"
#include "llvm/Assembly/PrintModulePass.h"
#include "llvm/CallingConv.h"
#include "llvm/ExecutionEngine/ExecutionEngine.h"
#include "llvm/Function.h"
#include "llvm/Module.h"
#include "llvm/ModuleProvider.h"
#include "llvm/PassManager.h"
#include "llvm/Support/IRBuilder.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Target/TargetData.h"
#include "llvm/Transforms/Scalar.h"
#include "llvm/Type.h"


#import <Foundation/Foundation.h>
#import <objc/runtime.h>
#import <objc/message.h>


using namespace llvm;


@interface ArrayMapProxyLLVM : NSProxy
{
    NSArray *_array;
    int _count;
    int _index;
}
- (id)initWithArray:(NSArray *)array;
@end

@implementation ArrayMapProxyLLVM

static ExecutionEngine *ArrayMapProxyLLVMEngine;

// requires explicit namespace due to conflict with objc header Module type
static llvm::Module *ArrayMapProxyLLVMModule;

+ (void)initialize
{
    ArrayMapProxyLLVMModule = new llvm::Module("ArrayMapProxyLLVMDynamic");
    ArrayMapProxyLLVMEngine = ExecutionEngine::create(ArrayMapProxyLLVMModule);
}

+ (void)printModule
{
    PassManager PM;
    ModulePass *pmp = createPrintModulePass(&outs());
    PM.add(pmp);
    PM.run(*ArrayMapProxyLLVMModule);
}

static const IntegerType *intType(void)
{
    return IntegerType::get(sizeof(int) * CHAR_BIT);
}

static const IntegerType *charType(void)
{
    return IntegerType::get(CHAR_BIT);
}

static const IntegerType *intptrType(void)
{
    return IntegerType::get(sizeof(void *) * CHAR_BIT);
}

static const PointerType *idType(void)
{
    return PointerType::getUnqual(charType());
}

static const PointerType *selType(void)
{
    return PointerType::getUnqual(charType());
}

static const Type *LLVMTypeForObjCType(const char *type)
{
#define IF_ISTYPE(t) if(strcmp(@encode(t), type) == 0)
#define INT_TYPE(t) IF_ISTYPE(t) return IntegerType::get(sizeof(t) * CHAR_BIT)
#define PTR_TYPE(t) IF_ISTYPE(t) return PointerType::getUnqual(charType())
    INT_TYPE(char);
    INT_TYPE(short);
    INT_TYPE(int);
    INT_TYPE(long);
    INT_TYPE(long long);
    INT_TYPE(unsigned char);
    INT_TYPE(unsigned short);
    INT_TYPE(unsigned int);
    INT_TYPE(unsigned long);
    INT_TYPE(unsigned long long);
    IF_ISTYPE(float) return Type::FloatTy;
    IF_ISTYPE(double) return Type::DoubleTy;
    IF_ISTYPE(void) return Type::VoidTy;
    PTR_TYPE(char *);
    PTR_TYPE(id);
    PTR_TYPE(SEL);
    PTR_TYPE(Class);
    if(type[0] == '^') return PointerType::getUnqual(charType());
    
    return NULL;
}

static Value *PtrValue(void *ptr, IRBuilder<> &builder, const Type *type, const char *name)
{
    Value *intv = ConstantInt::get(intptrType(), (int64_t)ptr, 0);
    return builder.CreateIntToPtr(intv, type, name);
}

static Value *SELValue(SEL sel, IRBuilder<> &builder)
{
    return PtrValue(sel, builder, selType(), sel_getName(sel));
}

static Value *ClassValue(Class c, IRBuilder<> &builder)
{
    return PtrValue(c, builder, idType(), class_getName(c));
}

static Function *ObjcMsgSendFunction(void)
{
    static Function *f;
    if(!f)
    {
        std::vector<const Type *> msgSendArgTypes;
        msgSendArgTypes.push_back(idType());
        msgSendArgTypes.push_back(selType());
        FunctionType *msgSendType = FunctionType::get(idType(), msgSendArgTypes, true);
        f = Function::Create(msgSendType,
                             Function::ExternalLinkage,
                             "objc_msgSend",
                             ArrayMapProxyLLVMModule);
    }
    return f;
}

+ (Function *)_trampolineFunctionForSignature:(NSMethodSignature *)sig selector:(SEL)sel
{
    // convert the NSMethodSignature into an LLVM type array
    std::vector<const Type *> methodArgTypes;
    for(unsigned i = 0; i < [sig numberOfArguments]; i++)
        methodArgTypes.push_back(LLVMTypeForObjCType([sig getArgumentTypeAtIndex:i]));
    
    // create the trampoline function using the types from the method signature
    const Type *methodReturnType = LLVMTypeForObjCType([sig methodReturnType]);
    FunctionType *trampolineType = FunctionType::get(methodReturnType, methodArgTypes, false);
    Function *trampoline = (Function *)ArrayMapProxyLLVMModule->getOrInsertFunction(
        [NSStringFromSelector(sel) UTF8String],
        trampolineType);
    trampoline->setCallingConv(CallingConv::C);
    
    // get the 'self' and '_cmd' args as values, and name them
    // the rest we don't care about except to pass them along
    Function::arg_iterator args = trampoline->arg_begin();
    Value *selfarg = args++;
    selfarg->setName("self");
    Value *_cmdarg = args++;
    _cmdarg->setName("_cmd");
    
    // The function we want to reproduce is basically this:
    // - (id)trampoline {
    //     NSMutableArray *array = [NSMutableArray array];
    //     id obj;
    //     while((obj = [self _nextObject]))
    //         [array addObject:[obj trampoline]];
    //     return array;
    // }
    // but with the method argument types variable
    // in terms of LLVM code we want to do something like this:
    // entry:
    //  set up selectors
    //  array = [NSMutableArray array];
    //  go to loopstart
    // loopstart:
    //  obj = [self _nextObject]
    //  if obj == nil then go to return
    //  else go to loopbody
    // loopbody:
    //  result = [obj trampoline]
    //  [array addObject:result];
    //  goto loopstart
    // return:
    //  return array
    BasicBlock *entry = BasicBlock::Create("entry", trampoline);
    BasicBlock *loopstart = BasicBlock::Create("loopstart", trampoline);
    BasicBlock *loopbody = BasicBlock::Create("loopbody", trampoline);
    BasicBlock *ret = BasicBlock::Create("return", trampoline);
    
    // we'll be doing several message sends, so get a reference to objc_msgSend
    Function *msgsend = ObjcMsgSendFunction();
    
    IRBuilder<> builder(entry);
    // we need three selectors, array, addObject:, and _nextObject
    Value *arraySEL = SELValue(@selector(array), builder);
    Value *addObjectSEL = SELValue(@selector(addObject:), builder);
    Value *nextObjectSEL = SELValue(@selector(_nextObject), builder);
    
    // [NSMutableArray array]
    Value *nsmutablearray = ClassValue([NSMutableArray class], builder);
    Value *array = builder.CreateCall2(msgsend, nsmutablearray, arraySEL, "array");
    builder.CreateBr(loopstart);
    
    builder.SetInsertPoint(loopstart);
    // nextObject = [self _nextObject]
    Value *nextObject = builder.CreateCall2(msgsend, selfarg, nextObjectSEL, "nextObject");
    
    // convert the object to int and compare with 0 (there's probably a better way)
    Value *nextObjectInt = builder.CreatePtrToInt(nextObject, intptrType(), "nextObjectInt");
    Constant *zero = ConstantInt::get(intType(), 0, 1);
    Value *nextObjectIsNil = builder.CreateICmpEQ(nextObjectInt, zero, "nextObjectIsNil");
    // if it's 0, return, otherwise enter the loop body
    builder.CreateCondBr(nextObjectIsNil, ret, loopbody);
    
    builder.SetInsertPoint(loopbody);
    // send the object the same message with the same arguments as what we got
    // start by copying all of the arguments into a new vector, save the first
    // one (self) which we re-point to the new object
    Function::arg_iterator methodArgs = trampoline->arg_begin();
    std::vector<Value *> msgsendArgs;
    msgsendArgs.push_back(nextObject); methodArgs++;
    while(methodArgs != trampoline->arg_end())
        msgsendArgs.push_back(methodArgs++);
    // args copied, now make the call
    Value *result = builder.CreateCall(msgsend,
                                       msgsendArgs.begin(),
                                       msgsendArgs.end(),
                                       "result");
    // [array addObject:nextObject]
    builder.CreateCall3(msgsend, array, addObjectSEL, result);
    builder.CreateBr(loopstart);
    
    // return block, just return the array and we're done
    builder.SetInsertPoint(ret);
    builder.CreateRet(array);
    
    return trampoline;
}

+ (void)_optimizeFunction:(Function *)f
{
    static FunctionPassManager *fpm;
    if(!fpm)
    {
        ExistingModuleProvider *moduleProvider = new ExistingModuleProvider(ArrayMapProxyLLVMModule);
        fpm = new FunctionPassManager(moduleProvider);
        
        fpm->add(new TargetData(*ArrayMapProxyLLVMEngine->getTargetData()));
        fpm->add(createInstructionCombiningPass());
        fpm->add(createReassociatePass());
        fpm->add(createGVNPass());
        fpm->add(createCFGSimplificationPass());
    }
    
    fpm->run(*f);
}

+ (IMP)_trampolineMethodForSignature:(NSMethodSignature *)sig selector:(SEL)sel
{
    Function *f = [self _trampolineFunctionForSignature:sig selector:sel];
    [self _optimizeFunction:f];
    return (IMP)ArrayMapProxyLLVMEngine->getPointerToFunction(f);
}

- (NSMethodSignature *)methodSignatureForSelector:(SEL)sel
{
    return [[_array lastObject] methodSignatureForSelector:sel];
}

- (void)forwardInvocation:(NSInvocation *)inv
{
    SEL sel = [inv selector];
    id obj = [_array lastObject];
    Method method = class_getInstanceMethod(object_getClass(obj), sel);
    NSParameterAssert(method);
    
    const char *types = method_getTypeEncoding(method);
    NSMethodSignature *sig = [NSMethodSignature signatureWithObjCTypes:types];
    NSParameterAssert([sig methodReturnType][0] == '@');
    
    class_addMethod([self class],
                    sel,
                    [[self class] _trampolineMethodForSignature:sig selector:sel],
                    types);
    [inv invoke];
}

- (id)initWithArray:(NSArray *)array
{
    _array = array;
    _count = [_array count];
    return self;
}

- (id)_nextObject
{
    return (_index < _count
            ? (id)CFArrayGetValueAtIndex((CFArrayRef)_array, _index++)
            : nil);
}

@end

@interface ArrayMapProxyNormal : NSProxy
{
    NSArray *_array;
}
- (id)initWithArray:(NSArray *)array;
@end
@implementation ArrayMapProxyNormal
- (id)initWithArray:(NSArray *)array
{
    _array = array;
    return self;
}
- (NSMethodSignature *)methodSignatureForSelector:(SEL)sel
{
    return [[_array lastObject] methodSignatureForSelector:sel];
}
- (void)forwardInvocation:(NSInvocation *)inv
{
    NSMutableArray *newArray = [NSMutableArray array];
    for(id obj in _array)
    {
        id retval;
        [inv invokeWithTarget:obj];
        [inv getReturnValue:&retval];
        [newArray addObject:retval];
    }
    [inv setReturnValue:&newArray];
}
@end

@interface NSArray (DOExtensions)
- (id)mapLLVM;
- (id)mapNormal;
@end
@implementation NSArray (DOExtensions)
- (id)mapLLVM
{
    return [[[ArrayMapProxyLLVM alloc] initWithArray:self] autorelease];
}
- (id)mapNormal
{
    return [[[ArrayMapProxyNormal alloc] initWithArray:self] autorelease];
}
@end

@interface NSString (NOP)
- (id)nop;
- (id)nop:(int)x :(int)y :(int)z;
@end
@implementation NSString (Logging)
- (id)nop { return self; }
- (id)nop:(int)x :(int)y :(int)z
{
    NSParameterAssert(x == 1 && y == 2 && z == 3);
    return self;
}
@end


#define TIME(expr) do { \
    fprintf(stderr, "testing %s...", #expr); \
    /* let stuff happen a few times first for caching etc. */ \
    for(int i = 0; i < 10; i++) expr; \
    \
    NSTimeInterval totalTime = 0; \
    int iterations = 1; \
    while(totalTime < 5 && iterations < 2000000000) \
    { \
        iterations *= 5; \
        NSTimeInterval start = [NSDate timeIntervalSinceReferenceDate]; \
        NSAutoreleasePool *pool = [[NSAutoreleasePool alloc] init]; \
        for(int i = 0; i < iterations; i++) \
        { \
            expr; \
            if(!(i & 0xFF)) \
            { \
                [pool release]; \
                pool = [[NSAutoreleasePool alloc] init]; \
            } \
        } \
        [pool release]; \
        NSTimeInterval end = [NSDate timeIntervalSinceReferenceDate]; \
        totalTime = end - start; \
    } \
    fprintf(stderr, " %fus/call\n", totalTime * 1000000.0 / iterations); \
} while(0)

int main(int argc, char **argv)
{
    NSAutoreleasePool *pool = [NSAutoreleasePool new];
    
    NSArray *reallySmallTimeTestArray = [NSArray arrayWithObject: @"0"];
    NSMutableArray *smallTimeTestArray = [NSMutableArray array];
    for(int i = 0; i < 10; i++)
        [smallTimeTestArray addObject:[NSString stringWithFormat:@"%d", i]];
        
    NSMutableArray *largeTimeTestArray = [NSMutableArray array];
    for(int i = 0; i < 10000; i++)
        [largeTimeTestArray addObject:[NSString stringWithFormat:@"%d", i]];
    
    TIME([[reallySmallTimeTestArray mapLLVM] nop]);
    TIME([[reallySmallTimeTestArray mapNormal] nop]);
    TIME([[reallySmallTimeTestArray mapLLVM] nop:1 :2 :3]);
    TIME([[reallySmallTimeTestArray mapNormal] nop:1 :2 :3]);
    
    TIME([[smallTimeTestArray mapLLVM] nop]);
    TIME([[smallTimeTestArray mapNormal] nop]);
    TIME([[smallTimeTestArray mapLLVM] nop:1 :2 :3]);
    TIME([[smallTimeTestArray mapNormal] nop:1 :2 :3]);
    
    TIME([[largeTimeTestArray mapLLVM] nop]);
    TIME([[largeTimeTestArray mapNormal] nop]);
    TIME([[largeTimeTestArray mapLLVM] nop:1 :2 :3]);
    TIME([[largeTimeTestArray mapNormal] nop:1 :2 :3]);
    
    [pool release];
    return 0;
}