Block JIT

#import <Foundation/Foundation.h>

void eachrow(size_t total, size_t nrows, size_t ncols, char const*** rows, BOOL(^cb)(size_t i, size_t ncols, char const **row, size_t *lengths))
{
    size_t lengths[ncols];
    for (size_t k = 0; k < total; ++k)
    {
        size_t i = k % nrows;
        for (size_t j = 0; j < ncols; ++j)
        {
            lengths[j] = strlen(rows[i][j]);
        }
        if (!cb(k, ncols, rows[i], lengths)) break;
    }
}

#if 0
static void print_int(char const *s, size_t l)
{
    printf("%d, ", atoi(s));
}

static void print_string(char const *s, size_t l)
{
    printf("%.*s, ", (int)l, s);
}

static void print_float(char const *s, size_t l)
{
    printf("%f\n", atof(s));
}

int main(int argc, const char * argv[])
{
    __block struct
    {
        void (*fptrs[3])(char const *, size_t);
    }
    ctx;
    eachrow(1000000, 3, 3, (char const **[]){
        (char const *[]){ "1", "hello", "3.14" },
        (char const *[]){ "2", "pony",  "2.72" },
        (char const *[]){ "3", "world", "1.62" },
    }, ^(size_t i, size_t ncols, char const **row, size_t *lengths)
    {
        if (i == 0)
        {
            ctx.fptrs[0] = print_int;
            ctx.fptrs[1] = print_string;
            ctx.fptrs[2] = print_float;
        }
        
        for (size_t col = 0; col < ncols; ++col)
        {
            ctx.fptrs[col](row[col], lengths[col]);
        }
        
        return YES;
    });
    return 0;
}
#elif 0
int main(int argc, const char * argv[])
{
    __block struct
    {
        void *dispatch[3];
    }
    ctx;
    eachrow(1000000, 3, 3, (char const **[]){
        (char const *[]){ "1", "hello", "3.14" },
        (char const *[]){ "2", "pony",  "2.72" },
        (char const *[]){ "3", "world", "1.62" },
    }, ^(size_t i, size_t ncols, char const **row, size_t *lengths)
    {
        if (i == 0)
        {
            ctx.dispatch[0] = &&print_int;
            ctx.dispatch[1] = &&print_string;
            ctx.dispatch[2] = &&print_float;
        }
        
        for (size_t col = 0; col < ncols; ++col)
        {
            goto *ctx.dispatch[col];
        print_int:
            printf("%d, ", atoi(row[col]));
            continue;
        print_string:
            printf("%.*s, ", (int)lengths[col], row[col]);
            continue;
        print_float:
            printf("%f\n", atof(row[col]));
            continue;
        }
        
        return YES;
    });
    return 0;
}
#elif 0
// lots of gory details cribbed from
// http://llvm.org/svn/llvm-project/cfe/trunk/lib/CodeGen/CGBlocks.cpp
// http://llvm.org/svn/llvm-project/cfe/trunk/lib/CodeGen/CGBlocks.h
struct RowBlockDescriptor
{
    unsigned long reserved;
    unsigned long size;
    void *copy_func_helper_decl;
    void *destroy_func_decl;
///   void *block_method_encoding_address; // @encode for block literal signature.
///   void *block_layout_info;
};

struct RowBlock
{
    void  *isa;
    int    flags;
    int    reserved;
    BOOL (*invoke)(struct RowBlock *block, size_t i, size_t ncols, char const **row, size_t *lengths);
    struct RowBlockDescriptor *block_descriptor;
    
    int *kinds;
};

static BOOL process_row(struct RowBlock *block, size_t i, size_t ncols, char const **row, size_t *lengths)
{
    for (size_t col = 0; col < ncols; ++col)
    {
        switch (block->kinds[col])
        {
        case 0:
            printf("%d, ", atoi(row[col]));
            break;
        case 1:
            printf("%.*s, ", (int)lengths[col], row[col]);
            break;
        case 2:
            printf("%f\n", atof(row[col]));
            break;
        }
    }
    
    return YES;
}

static BOOL first_call(struct RowBlock *block, size_t i, size_t ncols, char const **row, size_t *lengths)
{
    block->kinds = malloc(3 * sizeof(int));
    block->kinds[0] = 0;
    block->kinds[1] = 1;
    block->kinds[2] = 2;

    block->invoke = process_row;
    return block->invoke(block, i, ncols, row, lengths);
    // should update destroy_func_decl?  need to free kinds somewhere
    // maybe just at the end of main
}

int main(int argc, const char * argv[])
{
    struct RowBlockDescriptor desc = {
        .reserved = 0,
        .size     = sizeof(struct RowBlock),
        .copy_func_helper_decl = NULL,
        .destroy_func_decl = NULL,
    };
    extern void *_NSConcreteStackBlock[];
    struct RowBlock block = {
        .isa              = _NSConcreteStackBlock,
        .flags            = 0,
        .reserved         = 0,
        .invoke           = first_call,
        .block_descriptor = &desc,
    };
    eachrow(1000000, 3, 3, (char const **[]){
        (char const *[]){ "1", "hello", "3.14" },
        (char const *[]){ "2", "pony",  "2.72" },
        (char const *[]){ "3", "world", "1.62" },
    }, (__bridge id)&block);
    return 0;
}
#elif 1
static void print_int(char const *s, size_t l)
{
    printf("%d, ", atoi(s));
}

static void print_string(char const *s, size_t l)
{
    printf("%.*s, ", (int)l, s);
}

static void print_float(char const *s, size_t l)
{
    printf("%f\n", atof(s));
}

// lots of gory details cribbed from
// http://llvm.org/svn/llvm-project/cfe/trunk/lib/CodeGen/CGBlocks.cpp
// http://llvm.org/svn/llvm-project/cfe/trunk/lib/CodeGen/CGBlocks.h
struct RowBlockDescriptor
{
    unsigned long reserved;
    unsigned long size;
    void *copy_func_helper_decl;
    void *destroy_func_decl;
///   void *block_method_encoding_address; // @encode for block literal signature.
///   void *block_layout_info;
};

struct RowBlock
{
    void  *isa;
    int    flags;
    int    reserved;
    BOOL (*invoke)(struct RowBlock *block, size_t i, size_t ncols, char const **row, size_t *lengths);
    struct RowBlockDescriptor *block_descriptor;
};

static BOOL jit_me_a_river(struct RowBlock *block, size_t i, size_t ncols, char const **row, size_t *lengths)
{
    size_t function_size = 4096;
    uint8_t *function = mmap(NULL, function_size, PROT_READ | PROT_WRITE, MAP_ANON | MAP_PRIVATE, 0, 0);
    assert(function);
    
    // on entry to this function, the registers are as follows
    // RDI    block
    // RSI    i
    // RDX    ncols
    // RCX    row
    // R8     lengths
    
    // to call the per-row functions, we must place
    // RDI    row[col]
    // RSI    lengths[col]
    
    // we don't care if the functions we call clobber
    // block, i or ncols.  We do care if they clobber
    // row or lengths, so we save those to r13 and r14,
    // which are callee-save.  However, we must also
    // save & restore those registers.
    
    size_t ins = 0;
#define PACK(b) function[ins++] = 0x##b
#define PACK2(a, b) PACK(a); PACK(b)
#define PACK3(a, b, c) PACK2(a, b); PACK(c)
#define PACK4(a, b, c, d) PACK3(a, b, c); PACK(d)
    
    // prologue (align stack, save callee-save r13 & r14, save caller-save rcx & r8)
    PACK(55); // pushq %rbp
    PACK2(41,55); // pushq %r13
    PACK2(41,56); // pushq %r14
    PACK3(49,89,cd); // movq %rcx, %r13
    PACK3(4d,89,c6); // movq %r8, %r14
    
    void *f[3] = { print_int, print_string, print_float };
    for (size_t col = 0; col < ncols; ++col)
    {
        // f[col](row[col], lengths[col]);
        assert(col <= 31); // else we'd need a different form of mov
        PACK3(49,8b,7d); function[ins++] = 8 * col; // movq $8*col(%r13), %rdi
        PACK3(49,8b,76); function[ins++] = 8 * col; // movq $8*col(%r14), %rsi
        ptrdiff_t delta = (uint8_t *)f[col] - (function + ins + 5);
        assert(delta <= INT_MAX && delta >= INT_MIN); // else we'd need an indirect call
        PACK(e8);
        function[ins++] = (delta >>  0);
        function[ins++] = (delta >>  8);
        function[ins++] = (delta >> 16);
        function[ins++] = (delta >> 24); // callq f[col]
    }
    
    // epilogue (reset stack, restore callee-save r13 & r14, return true)
    PACK2(b0,01); // movb $1, %al
    PACK2(41,5e); // popq %r14
    PACK2(41,5d); // popq %r13
    PACK(5d); // popq %rbp
    PACK(c3); // ret

#undef PACK4
#undef PACK3
#undef PACK
    
#if 0
    printf("JITted something at %p\n", function);
    for (size_t b = 0; b < ins; ++b)
    {
        printf("%02x ", function[b]);
    }
    printf("\n");
#endif
    
    mprotect(function, function_size, PROT_READ | PROT_EXEC);
    block->invoke = (void *)function;
    return block->invoke(block, i, ncols, row, lengths);
    // should update destroy_func_decl?  need to munmap function somewhere
    // maybe just at the end of main
}

int main(int argc, const char * argv[])
{
    struct RowBlockDescriptor desc = {
        .reserved = 0,
        .size     = sizeof(struct RowBlock),
        .copy_func_helper_decl = NULL,
        .destroy_func_decl = NULL,
    };
    extern void *_NSConcreteStackBlock[];
    struct RowBlock block = {
        .isa              = _NSConcreteStackBlock,
        .flags            = 0,
        .reserved         = 0,
        .invoke           = jit_me_a_river,
        .block_descriptor = &desc,
    };
    eachrow(1000000, 3, 3, (char const **[]){
        (char const *[]){ "1", "hello", "3.14" },
        (char const *[]){ "2", "pony",  "2.72" },
        (char const *[]){ "3", "world", "1.62" },
    }, (__bridge id)&block);
    return 0;
}
#endif