Morton code encoder generator. Generates C code that encodes the X dimension of a 3D morton code. Other dimensions can be obtained by shifting by 1 or 2.

morton_encoder_generator.c

#include <stdlib.h>
#include <stdio.h>
#include <stdint.h>
#include <stdbool.h>
#include <inttypes.h>
#include <limits.h>
#include <assert.h>

typedef struct op_s op_t;

struct op_s {
    enum {
        AND,
        OR,
        LSHIFT,
        RSHIFT,
        VAL,
        VAR
    } tag : 8;
    bool generated : 1;
    const char* name;
    union {
        char var;
        uint64_t val;
        struct {
            op_t* left;
            op_t* right;
        } bin_op;
    };
};

static inline uint64_t encode_slow(uint64_t x) {
    uint64_t res = 0;
    for (int i = 0; x; ++i, x >>= 1)
        res |= (x & 1) << (3 * i);
    return res;
}

static inline uint64_t encode_fast(uint64_t x) {
    const size_t bits = 32;
    uint64_t mask = (UINT64_C(1) << bits) - 1;
    for (size_t n = bits; n > 1; n /= 2) {
        mask = (mask | (mask << n)) & ~(mask << (n / 2));
        x = (x | (x << n)) & mask;
    }
    return x;
}

static inline op_t* bin_op(op_t* left, op_t* right, unsigned tag) {
    op_t* op = calloc(1, sizeof(op_t));
    op->bin_op.left = left;
    op->bin_op.right = right;
    op->tag = tag;
    return op;
}

static inline op_t* val_op(uint64_t val) {
    op_t* op = calloc(1, sizeof(op_t));
    op->val = val;
    op->tag = VAL;
    return op;
}

static inline op_t* var_op(char var) {
    op_t* op = calloc(1, sizeof(op_t));
    op->var = var;
    op->tag = VAR;
    return op;
}

static inline op_t* rshift_op(op_t* left, op_t* right) {
    return bin_op(left, right, RSHIFT);
}

static inline op_t* lshift_op(op_t* left, op_t* right) {
    return bin_op(left, right, LSHIFT);
}

static inline void reorder_ops(op_t** left, op_t** right) {
    if ((*left)->tag == VAL) {
        op_t* tmp = *left;
        *left = *right;
        *right = tmp;
    }
}

static inline op_t* and_op(op_t* left, op_t* right) {
    reorder_ops(&left, &right);
    return bin_op(left, right, AND);
}

static inline op_t* or_op(op_t* left, op_t* right) {
    reorder_ops(&left, &right);
    return bin_op(left, right, OR);
}

static const char* generate_name(void) {
    static size_t count = 0;
    char* name = malloc(32);
    sprintf(name, "inst%zu", count++);
    return name;
}

static void print_op(op_t*);

static void generate_op(op_t* x) {
    if (x->generated)
        return;
    static const char* ops[] = {
        "&",
        "|",
        "<<",
        ">>"
    };
    switch (x->tag) {
        case AND:
        case OR:
        case LSHIFT:
        case RSHIFT:
            generate_op(x->bin_op.left);
            generate_op(x->bin_op.right);
            x->name = generate_name();    
            printf("    uint64_t %s = ", x->name);
            print_op(x->bin_op.left);
            printf(" %s ", ops[x->tag]);
            print_op(x->bin_op.right);
            printf(";\n");
            break;
    }
    x->generated = true;
}

static void print_op(op_t* x) {
    switch (x->tag) {
        case VAR: printf("%c", x->var); break;
        case VAL: printf("0x%"PRIx64, x->val); break;
        case AND:
        case OR:
        case LSHIFT:
        case RSHIFT:
            generate_op(x);
            printf("%s", x->name);
            break;
        default:
            printf("<ERROR>");
            break;
    }
}

static inline op_t* generate_encoder(op_t* x, size_t bits) {
    size_t q = 1;
    while (q < bits) q <<= 1;
    bits = q;

    uint64_t mask = (UINT64_C(1) << bits) - 1;
    for (size_t n = bits; n > 1; n /= 2) {
        mask = (mask | (mask << n)) & ~(mask << (n / 2));
        x = and_op(or_op(x, lshift_op(x, val_op(n))), val_op(mask));
    }
    return x;
}

int main(int argc, char** argv) {
    printf("#include <stdint.h>\n"
           "#include <inttypes.h>\n"
           "#include <stdio.h>\n"
           "#include <stdlib.h>\n"
           "uint64_t morton(uint64_t x) {\n");
    op_t* encoder = generate_encoder(var_op('x'), 8);
    generate_op(encoder);
    printf("    return ");
    print_op(encoder);
    printf(";\n}\n");
    printf("int main(int argc, char** argv) {\n"
           "    if (argc != 2)\n"
           "        return 1;\n"
           "    printf(\"%\"PRIu64\"\\n\", morton(strtoull(argv[1], NULL, 10)));\n"
           "    return 0;\n"
           "}\n");
    return 0;
}