maxwell-dump.c

#define _CRT_SECURE_NO_WARNINGS
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdbool.h>
#include <string.h>

#ifdef _WIN32
#include <Windows.h>
#endif

#define PROGRAM_TYPE_VERTEX 0
#define PROGRAM_TYPE_TESS_CONTROL 1
#define PROGRAM_TYPE_TESS_EVAL 2
#define PROGRAM_TYPE_GEOMETRY 3
#define PROGRAM_TYPE_FRAGMENT 4
#define PROGRAM_TYPE_COMPUTE 5

#define SIZEOF_METADATA 621
#define SIZEOF_CONST_BUFFER_KEY 12
#define SIZEOF_BOUND_SAMPLER_KEY 12
#define SIZEOF_BINDLESS_SAMPLER_KEY 16

typedef uint8_t u8;
typedef uint32_t u32;
typedef uint64_t u64;

u32 read_u8(FILE *file) {
    u8 value;
    if (fread(&value, sizeof(u8), 1, file) != 1) {
        fprintf(stderr, "Failed to read file\n");
        exit(EXIT_FAILURE);
    }
    return value;
}

u32 read_u32(FILE *file) {
    u32 value;
    if (fread(&value, sizeof(u32), 1, file) != 1) {
        fprintf(stderr, "Failed to read file\n");
        exit(EXIT_FAILURE);
    }
    return value;
}

u64 read_u64(FILE *file) {
    u64 value;
    if (fread(&value, sizeof(u64), 1, file) != 1) {
        fprintf(stderr, "Failed to read file\n");
        exit(EXIT_FAILURE);
    }
    return value;
}

const char* program_type_name(u32 program_type) {
    switch (program_type) {
    case PROGRAM_TYPE_VERTEX:
        return "VS";
    case PROGRAM_TYPE_TESS_CONTROL:
        return "TCS";
    case PROGRAM_TYPE_TESS_EVAL:
        return "TES";
    case PROGRAM_TYPE_GEOMETRY:
        return "GS";
    case PROGRAM_TYPE_FRAGMENT:
        return "FS";
    case PROGRAM_TYPE_COMPUTE:
        return "CS";
    default:
        fprintf(stderr, "Invalid program type %d\n", program_type);
        exit(EXIT_FAILURE);
    }
}

u32 program_offset(u32 program_type) {
    if (program_type != PROGRAM_TYPE_COMPUTE) {
        return 10;
    }
    return 0;
}

void load_and_save_dump(FILE* input, const char* output_dir, bool code) {
    u32 program_type = read_u32(input);
    u32 program_code_size = read_u32(input);
    u32 program_code_size_b = read_u32(input);

    u64 *program_code = malloc(program_code_size * sizeof(u64));
    if (!program_code)
        goto fail_malloc_1;

    u64 *program_code_b = malloc(program_code_size_b * sizeof(u64));
    if (!program_code_b && program_code_size_b)
        goto fail_malloc_2;

    if (fread(program_code, sizeof(u64), program_code_size, input) !=
            program_code_size) {
        goto fail_fread;
    }
    if (fread(program_code_b, sizeof(u64), program_code_size_b, input) !=
            program_code_size_b) {
        goto fail_fread;
    }

    u64 unique_identifier = read_u64(input);

    if (fseek(input, SIZEOF_METADATA, SEEK_CUR) != 0) {
        fprintf(stderr, "Failed to fseek()\n");
        exit(EXIT_FAILURE);
    }

    size_t offset = 0;
    offset += SIZEOF_CONST_BUFFER_KEY * read_u32(input);
    offset += SIZEOF_BOUND_SAMPLER_KEY * read_u32(input);
    offset += SIZEOF_BINDLESS_SAMPLER_KEY * read_u32(input);
    if (fseek(input, offset, SEEK_CUR) != 0) {
        fprintf(stderr, "Failed to fseek()\n");
        exit(EXIT_FAILURE);
    }

    const char* prefix = program_type_name(program_type);
    char path[256];
    if (snprintf(path, sizeof(path), "%s/%s%016llX.bin", output_dir, prefix,
            unique_identifier) >= sizeof(path)) {
        fprintf(stderr, "File path got truncated %s\n", path);
        goto finish;
    }

    FILE *output = fopen(path, "wb");
    if (!output) {
        fprintf(stderr, "Failed to write file %s\n", path);
        goto finish;
    }
    if (code) {
        const u32 offset = program_offset(program_type);
        (void)fwrite(program_code + offset, sizeof(u64),
                        program_code_size - offset, output);

        for (u64 i = program_code_size - offset; i % 4 != 0; ++i) {
            u64 zero = 0;
            (void)fwrite(&zero, sizeof(u64), 1, output);
        }
    } else {
        (void)fwrite(program_code, sizeof(u64), program_code_size, output);
    }
    fclose(output);

    if (program_type == 0 && program_code_size_b > 0) {
        fprintf(stderr, "Program dump B not implemented\n");
        goto finish;
    }

    fprintf(stdout, "Dumped program %s%016llX\n", prefix, unique_identifier);
finish:
    free(program_code);
    free(program_code_b);
    return;

fail_malloc_2:
    free(program_code);
fail_malloc_1:
    fprintf(stderr, "Failed to malloc()\n");
    exit(EXIT_FAILURE);

fail_fread:
    free(program_code);
    free(program_code_b);
    fprintf(stderr, "Failed to fread()\n");
    exit(EXIT_FAILURE);
}

int main(int argc, char **argv) {
    if (argc < 4) {
        fprintf(stderr, "Usage: %s <input file> <output dir> <code/full>\n",
                        argv[0]);
        return EXIT_FAILURE;
    }

    const char* output_dir = argv[2];
#ifdef _WIN32
    if (CreateDirectoryA(output_dir, NULL) == 0) {
        if (GetLastError() != ERROR_ALREADY_EXISTS) {
            fprintf(stderr, "Failed to create directory %s\n", output_dir);
            return EXIT_FAILURE;
        }
    }
#endif

    const char *code_or_full = argv[3];
    bool code;
    if (strcmp(code_or_full, "code") == 0) {
        code = true;
    } else if (strcmp(code_or_full, "full") == 0) {
        code = false;
    } else {
        fprintf(stderr, "Unknown dump mode %s\n", code_or_full);
        return EXIT_FAILURE;
    }

    const char* input_file = argv[1];
    FILE *input = fopen(input_file, "rb");
    if (!input) {
        fprintf(stderr, "Failed to open file %s\n", input_file);
        return EXIT_FAILURE;
    }

    fseek(input, 0, SEEK_END);
    long input_size = ftell(input);
    rewind(input);

    u32 version = read_u32(input);
    if (version != 20) {
        fprintf(stderr, "Unimplemented transferable version %d\n", version);
        return EXIT_FAILURE;
    }

    while (ftell(input) < input_size) {
        load_and_save_dump(input, output_dir, code);
    }

    return EXIT_SUCCESS;
}