Cod4 Shader Compiler

/*
Copyright (c) 2015 Otso O.

This software is provided 'as-is', without any express or implied
warranty. In no event will the authors be held liable for any damages
arising from the use of this software.

Permission is granted to anyone to use this software for any purpose,
including commercial applications, and to alter it and redistribute it
freely, subject to the following restrictions:

1. The origin of this software must not be misrepresented; you must not
   claim that you wrote the original software. If you use this software
   in a product, an acknowledgement in the product documentation would be
   appreciated but is not required.
2. Altered source versions must be plainly marked as such, and must not be
   misrepresented as being the original software.
3. This notice may not be removed or altered from any source distribution.
*/

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <time.h>

#ifdef _DEBUG
#   define _CRTDBG_MAP_ALLOC
#   include <stdlib.h>
#   include <crtdbg.h>
#endif // _DEBUG

int isDefined(void const *ptr)
{
    return (*(uint32_t*)ptr != 0);
}

#include <windows.h>
int colorPrint(uint8_t color, char const *fmt, ...)
{
    int result;
    va_list arg;

    HANDLE handle = 0;
    handle = GetStdHandle(STD_OUTPUT_HANDLE);

    va_start(arg, fmt);
    if(isDefined(&handle))
    {
        SetConsoleTextAttribute(handle, color);
        result = vprintf(fmt, arg);
        SetConsoleTextAttribute(handle, 0x07);
    } else
        result = vprintf(fmt, arg);

    va_end(arg);
    return result;
}

uint32_t swapByteOrder32(uint32_t uint)
{
    return ((uint & 0x000000ff) << 24) |
           ((uint & 0x0000ff00) << 8)  |
           ((uint & 0x00ff0000) >> 8)  |
           ((uint & 0xff000000) >> 24);
}

uint32_t generateHashValue(char const *str)
{
    uint32_t esi, eax, ebx, ecx;
    ebx = ebx ^ ebx;
    for(esi = 0; esi < strlen(str); ++esi)
    {
        ecx = str[esi];
        eax = ebx;
        eax = eax << 5;
        ebx = ebx + eax;
        ebx = ebx ^ ecx;
    }
    return ebx;
}

enum COLORS {BLACK, DARK_BLUE, DARK_GREEN, DARK_CYAN, DARK_RED, DARK_MAGENTA, DARK_YELLOW,
             GRAY, DARK_GRAY, BLUE, GREEN, CYAN, RED, MAGENTA, YELLOW, WHITE};

enum SHADER_TYPES {VERTEX_SHADER = 0x01, PIXEL_SHADER = 0x02};

typedef struct Shader
{
    uint8_t  type;
    uint32_t hlsl_file;
    uint32_t bin_file;
} Shader;

int updateShaders(Shader **shaders, uint32_t num_new);
int getNumShaders(FILE *file, uint32_t *num_vs, uint32_t *num_ps);
int backupShaders(Shader const *shaders, uint32_t ofs, uint32_t num_vs, uint32_t num_ps);
int getUniqueShaders(Shader **shaders, uint32_t num_new, uint32_t num_vs, uint32_t num_ps);
int updateNumShaders(Shader const *shaders, uint32_t num_new, uint32_t *num_vs, uint32_t *num_ps);
int sortShaders(void const *pa, void const *pb);
int compileShaders(int argc, char const **argv);
int updateShaderSize(char const *filename);

int main(int argc, char **argv)
{
    int i;

    Shader *shaders = 0;

    if(argc < 2)
    {
        printf("USAGE: shader_tool.exe <shader_name> <shader_name2> <shader_name3> ...\n");
        colorPrint(YELLOW, "WARNING: The use of special characters in shader names will result in undefined behavior!\n");
        return 0;
    }

    shaders = malloc((argc - 1) * sizeof(Shader));
    if(!isDefined(&shaders))
    {
        colorPrint(RED, "ERROR: Failed to allocate %u bytes of memory!\n", ((argc - 1) * sizeof(Shader)));
        return 1;
    }

    printf("#----------------------------------------------------------------------------#\n");
    printf("|                  -- Call of Duty 4 - Shader Tool v0.1 --                   |\n");
    printf("|                         Copyright (c) 2015 Otso O.                         |\n");
    printf("#----------#----------------------------#------------------------------------#\n");
    printf("| Hash     | Shader name for techniques | File name in raw\\shader_bin        |\n");
    printf("#----------#----------------------------#------------------------------------#\n");
    for(i = 0; i < (argc - 1); ++i)
    {
        int x;

        char const *arg = argv[i + 1];

        for(x = 1; x < argc; ++x)
        {
            if(strcmp(argv[x], arg) == 0 && x != (i + 1))
            {
                free(shaders);
                printf("#----------#----------------------------#------------------------------------#\n");
                colorPrint(RED, "ERROR: Encountered a duplicate shader name '%s'!\n", arg);
                colorPrint(RED, "Remove the duplicate shader name from command line arguments and try again.\n");
                return 1;
            }
        }

        shaders[i].type      = 0;
        shaders[i].hlsl_file = generateHashValue(arg);
        shaders[i].bin_file  = 0xFFFFFFFF - swapByteOrder32(shaders[i].hlsl_file);
        printf("| %.8X | %-24.24s%s | vs_3_0_%.8x or ps_3_0_%.8x |\n",
                   swapByteOrder32(shaders[i].hlsl_file), arg, strlen(arg) > 24 ? ".." : "  ",
                   shaders[i].bin_file, shaders[i].bin_file);
    }
    printf("#----------#----------------------------#------------------------------------#\n");

    qsort(shaders, (argc - 1), sizeof(Shader), sortShaders);
    updateShaders(&shaders, (argc - 1));
    free(shaders);

    compileShaders(argc, argv);

#ifdef _DEBUG
    system("pause");
    _CrtDumpMemoryLeaks();
#endif // _DEBUG

    return 0;
}

int updateShaders(Shader **shaders, uint32_t num_new)
{
    int i, j;

    FILE *file      = 0;
    Shader *temp    = 0;
    void *ptr       = 0;
    uint32_t ofs    = num_new;
    uint32_t num_vs = 0;
    uint32_t num_ps = 0;
    uint32_t total  = 0;
    int count       = 0;

    printf("\nUpdating 'shader_names' file ...\n");
    file = fopen("shader_names", "rb");
    if(!isDefined(&file))
    {
        colorPrint(RED, "ERROR: Failed to open file 'shader_names' for reading!\n");
        colorPrint(RED, "Make sure you are running shader_tool.exe in the raw\\shader_bin\\ directory.\n");
        return 0;
    }
    if(!getNumShaders(file, &num_vs, &num_ps))
    {
        fclose(file);
        return 0;
    }

    total = (num_vs + num_ps + num_new);
    temp = realloc(*shaders, total * sizeof(Shader));
    if(!isDefined(&temp))
    {
        fclose(file);
        colorPrint(RED, "ERROR: Failed to allocate %u bytes of memory!\n", total * sizeof(Shader));
        return 0;
    }

    *shaders = temp;
    ptr = ((uint32_t)*shaders + (ofs * sizeof(Shader)));
    memset(ptr, 0, ((num_vs + num_ps) * sizeof(Shader)));

    fseek(file, sizeof(uint32_t), SEEK_SET);
    for(i = 0; i < (num_vs + num_ps); ++i)
    {
        if(i == num_vs)
            fseek(file, sizeof(uint32_t), SEEK_CUR);

        if(i < num_vs)
            (*shaders)[i + ofs].type = VERTEX_SHADER;
        else
            (*shaders)[i + ofs].type = PIXEL_SHADER;

        if(!fread(&(*shaders)[i + ofs].hlsl_file, 1, sizeof(uint32_t), file) || !fread(&(*shaders)[i + ofs].bin_file, 1, sizeof(uint32_t), file))
        {
            fclose(file);
            colorPrint(RED, "ERROR: Failed to read shader #%d!\n", i);
            return 0;
        }
    }

    fclose(file);
    if(!backupShaders(*shaders, ofs, num_vs, num_ps))
    {
        colorPrint(RED, "ERROR: Failed to backup the current 'shader_names' file!\n");
        return 0;
    }

    file = fopen("shader_names", "wb");
    if(!isDefined(&file))
    {
        colorPrint(RED, "ERROR: Failed to open file 'shader_names' for writing!\n");
        colorPrint(RED, "Make sure you are running shader_tool.exe as administrator!\n");
        return 0;
    }

    count = getUniqueShaders(shaders, num_new, num_vs, num_ps);
    printf("%d new shaders\n", count);
    updateNumShaders(*shaders, num_new, &num_vs, &num_ps);

    fwrite(&num_vs, sizeof(uint32_t), 1, file);
    for(i = 0; i < total; ++i)
    {
        if((*shaders)[i].type == VERTEX_SHADER)
        {
            for(j = 0; j < num_new; ++j)
            {
                if((*shaders)[i].hlsl_file > (*shaders)[j].hlsl_file && (*shaders)[j].type == 0x00)
                {
                    (*shaders)[j].type = 0xAA;
                    fwrite(&(*shaders)[j].hlsl_file, sizeof(uint32_t), 1, file);
                    fwrite(&(*shaders)[j].bin_file, sizeof(uint32_t), 1, file);
                }
            }
            fwrite(&(*shaders)[i].hlsl_file, sizeof(uint32_t), 1, file);
            fwrite(&(*shaders)[i].bin_file, sizeof(uint32_t), 1, file);
        } else
        if((*shaders)[i].type == 0x00 && (*shaders)[i].hlsl_file < (*shaders)[num_new].hlsl_file)
        {
            (*shaders)[i].type = 0xAA;
            fwrite(&(*shaders)[i].hlsl_file, sizeof(uint32_t), 1, file);
            fwrite(&(*shaders)[i].bin_file, sizeof(uint32_t), 1, file);
        }
    }

    fwrite(&num_ps, sizeof(uint32_t), 1, file);
    for(i = 0; i < total; ++i)
    {
        if((*shaders)[i].type == PIXEL_SHADER)
        {
            for(j = 0; j < num_new; ++j)
            {
                if((*shaders)[i].hlsl_file > (*shaders)[j].hlsl_file && (*shaders)[j].type == 0xAA)
                {
                    (*shaders)[j].type = 0xBB;
                    fwrite(&(*shaders)[j].hlsl_file, sizeof(uint32_t), 1, file);
                    fwrite(&(*shaders)[j].bin_file, sizeof(uint32_t), 1, file);
                }
            }
            fwrite(&(*shaders)[i].hlsl_file, sizeof(uint32_t), 1, file);
            fwrite(&(*shaders)[i].bin_file, sizeof(uint32_t), 1, file);
        } else
        if((*shaders)[i].type == 0xAA && (*shaders)[i].hlsl_file < (*shaders)[num_vs].hlsl_file)
        {
            (*shaders)[i].type = 0xBB;
            fwrite(&(*shaders)[i].hlsl_file, sizeof(uint32_t), 1, file);
            fwrite(&(*shaders)[i].bin_file, sizeof(uint32_t), 1, file);
        }
    }

    fclose(file);
    return count;
}

int getNumShaders(FILE *file, uint32_t *num_vs, uint32_t *num_ps)
{
    uint32_t fsize = 0;

    fseek(file, 0, SEEK_END);
    fsize = ftell(file);
    fseek(file, 0, SEEK_SET);

    if(fsize == 0)
    {
        colorPrint(RED, "ERROR: File 'shader_names' is empty!\n");
        colorPrint(RED, "Please restore the original 'shader_names' file and try again.\n");
        return 0;
    }
    if(fread(num_vs, sizeof(uint32_t), 1, file))
    {
        if(fseek(file, 8 * (*num_vs) + sizeof(uint32_t), SEEK_SET) == 0 && fread(num_ps, sizeof(uint32_t), 1, file))
        {
            if(fsize == ((2 * sizeof(uint32_t)) + (((*num_vs) + (*num_ps)) * 8)))
            {
                printf("%u vertex shaders\n%u pixel shaders\n", (*num_vs), (*num_ps));
                return 1;
            }

            colorPrint(RED, "ERROR: File 'shader_names' is corrupt! Expected size of %u bytes, got %u bytes instead.\n",
                      ((2 * sizeof(uint32_t)) + (((*num_vs) + (*num_ps)) * 8)), fsize);
        } else
            colorPrint(RED, "ERROR: Could not read the number of pixel shaders!\n");
    } else
        colorPrint(RED, "ERROR: Could not read the number of vertex shaders!\n");

    colorPrint(RED, "Please restore the original 'shader_names' file and try again.\n");
    return 0;
}

int backupShaders(Shader const *shaders, uint32_t ofs, uint32_t num_vs, uint32_t num_ps)
{
    int i;
    time_t now;
    struct tm *t;
    char filename[64];

    FILE *file = 0;

    now = time(0);
    t = localtime(&now);
    strftime(filename, sizeof(filename), "backups\\shader_names_%d%m%y%H%M%S.bak", t);

    file = fopen(filename, "wb");
    if(!isDefined(&file))
    {
        colorPrint(RED, "ERROR: Failed to create file '%s'!\n", filename);
        colorPrint(RED, "Make sure you are running shader_tool.exe as administrator!\n");
        return 0;
    }

    fwrite(&num_vs, sizeof(uint32_t), 1, file);
    for(i = 0; i < (num_vs + num_ps); ++i)
    {
        if(i == num_vs)
            fwrite(&num_ps, sizeof(uint32_t), 1, file);

        if(!fwrite(&shaders[i + ofs].hlsl_file, sizeof(uint32_t), 1, file) || !fwrite(&shaders[i + ofs].bin_file, sizeof(uint32_t), 1, file))
        {
            colorPrint(RED, "ERROR: Failed to write shader #%d!\n", i);
            fclose(file);
            return 0;
        }
    }
    fclose(file);
    return 1;
}

int getUniqueShaders(Shader **shaders, uint32_t num_new, uint32_t num_vs, uint32_t num_ps)
{
    int i, j;

    int uniques    = 0;
    int duplicates = 0;

    for(i = 0; i < num_new; ++i)
    {
        for(j = 0; j < (num_vs + num_ps); ++j)
        {
            if((*shaders)[i].hlsl_file == (*shaders)[j + num_new].hlsl_file || (*shaders)[i].bin_file == (*shaders)[j + num_new].bin_file)
            {
                (*shaders)[i].type = 0xFF;
                (*shaders)[i].hlsl_file = 0xDEADBEEF;
                (*shaders)[i].bin_file = 0xDEADBEEF;
                ++duplicates;
            }
        }
    }
    if(duplicates > 0)
        colorPrint(YELLOW, "WARNING: %d of the new shaders already exist in 'shader_names' file!\n", duplicates);

    uniques = num_new - duplicates;
    return uniques;
}

int updateNumShaders(Shader const *shaders, uint32_t num_new, uint32_t *num_vs, uint32_t *num_ps)
{
    int i;

    for(i = 0; i < num_new; ++i)
    {
        if(shaders[i].type != 0xFF)
        {
            ++(*num_vs);
            ++(*num_ps);
        }
    }
    return 1;
}

int sortShaders(void const *pa, void const *pb)
{
    Shader *a = (Shader*)pa;
    Shader *b = (Shader*)pb;
    return a->hlsl_file >= b->hlsl_file ? 1 : 0;
}

int compileShaders(int argc, char const **argv)
{
    int i;

    printf("\nCompiling shaders with fxc.exe ...\n");
    for(i = 0; i < (argc - 1); ++i)
    {
        char shader[32];
        char cmd[256];

        char const *arg = argv[i + 1];
        uint32_t bin_file = 0xFFFFFFFF - swapByteOrder32(generateHashValue(arg));

        sprintf(shader, "vs_3_0_%.8x", bin_file);
        sprintf(cmd, "shader_src\\fxc.exe shader_src\\vs_3_0_%s.hlsl /T vs_3_0 /E vs_main /Fo vs_3_0_%.8x > nul", arg, bin_file);
        system(cmd);
        updateShaderSize(shader);

        sprintf(shader, "ps_3_0_%.8x", bin_file);
        sprintf(cmd, "shader_src\\fxc.exe shader_src\\ps_3_0_%s.hlsl /T ps_3_0 /E ps_main /Fo ps_3_0_%.8x > nul", arg, bin_file);
        system(cmd);
        updateShaderSize(shader);
    }

    return 1;
}

int updateShaderSize(char const *filename)
{
    FILE *file = 0;

    file = fopen(filename, "rb");
    if(isDefined(&file))
    {
        uint8_t *buffer;

        uint32_t fsize = 0;

        fseek(file, 0, SEEK_END);
        fsize = ftell(file);
        fseek(file, 0, SEEK_SET);

        buffer = malloc(fsize * sizeof(uint8_t));
        if(isDefined(&buffer))
        {
            fread(buffer, sizeof(uint8_t), fsize, file);
            fclose(file);

            fopen(filename, "wb");
            if(isDefined(&file))
            {
                fwrite(&fsize, sizeof(uint32_t), 1, file);
                fwrite(buffer, sizeof(uint8_t), fsize, file);
                printf("%s, %u bytes\n", filename, fsize);
                fclose(file);
            } else {
                colorPrint(RED, "ERROR: Failed to open file '%s' for writing!\n", filename);
                colorPrint(RED, "Make sure you are running shader_tool.exe as administrator!\n");
            }
            free(buffer);
        } else {
            fclose(file);
            colorPrint(RED, "ERROR: Failed to allocate %u bytes of memory!\n", fsize);
        }
    }
    return 1;
}

/*
00464EC0  /$ 55             PUSH EBP
00464EC1  |. 8BEC           MOV EBP,ESP
00464EC3  |. 53             PUSH EBX
00464EC4  |. 56             PUSH ESI
00464EC5  |. 57             PUSH EDI
00464EC6  |. 8B7D 08        MOV EDI,DWORD PTR SS:[EBP+8]
00464EC9  |. 33DB           XOR EBX,EBX
00464ECB  |. 381F           CMP BYTE PTR DS:[EDI],BL
00464ECD  |. 8BF7           MOV ESI,EDI
00464ECF  |. 74 5B          JE SHORT linker_p.00464F2C
00464ED1  |> 8A06           /MOV AL,BYTE PTR DS:[ESI]
00464ED3  |. 3C 41          |CMP AL,41
00464ED5  |. 7C 20          |JL SHORT linker_p.00464EF7
00464ED7  |. 3C 5A          |CMP AL,5A
00464ED9  |. 7F 1C          |JG SHORT linker_p.00464EF7
00464EDB  |. 57             |PUSH EDI                                ; /Arg6
00464EDC  |. 68 9C044F00    |PUSH linker_p.004F049C                  ; |Arg5 = 004F049C ASCII "(*pos < 'A' || *pos > 'Z')"
00464EE1  |. 68 EC554D00    |PUSH linker_p.004D55EC                  ; |Arg4 = 004D55EC ASCII "%s
        (name) = %s"
00464EE6  |. 6A 00          |PUSH 0                                  ; |Arg3 = 00000000
00464EE8  |. 6A 2A          |PUSH 2A                                 ; |Arg2 = 0000002A
00464EEA  |. 68 5C044F00    |PUSH linker_p.004F045C                  ; |Arg1 = 004F045C ASCII "C:\trees\cod3-pc\cod3-modtools\cod3src\src\gfx_d3d\r_utils.cpp"
00464EEF  |. E8 AC970200    |CALL linker_p.0048E6A0                  ; \linker_p.0048E6A0
00464EF4  |. 83C4 18        |ADD ESP,18
00464EF7  |> 803E 5C        |CMP BYTE PTR DS:[ESI],5C
00464EFA  |. 75 1C          |JNZ SHORT linker_p.00464F18
00464EFC  |. 57             |PUSH EDI                                ; /Arg6
00464EFD  |. 68 3C044F00    |PUSH linker_p.004F043C                  ; |Arg5 = 004F043C ASCII "(*pos != '\\' || *pos == '/')"
00464F02  |. 68 EC554D00    |PUSH linker_p.004D55EC                  ; |Arg4 = 004D55EC ASCII "%s
        (name) = %s"
00464F07  |. 6A 00          |PUSH 0                                  ; |Arg3 = 00000000
00464F09  |. 6A 2B          |PUSH 2B                                 ; |Arg2 = 0000002B
00464F0B  |. 68 5C044F00    |PUSH linker_p.004F045C                  ; |Arg1 = 004F045C ASCII "C:\trees\cod3-pc\cod3-modtools\cod3src\src\gfx_d3d\r_utils.cpp"
00464F10  |. E8 8B970200    |CALL linker_p.0048E6A0                  ; \linker_p.0048E6A0
00464F15  |. 83C4 18        |ADD ESP,18
00464F18  |> 0FBE0E         |MOVSX ECX,BYTE PTR DS:[ESI]
00464F1B  |. 8BC3           |MOV EAX,EBX
00464F1D  |. C1E0 05        |SHL EAX,5
00464F20  |. 03D8           |ADD EBX,EAX
00464F22  |. 83C6 01        |ADD ESI,1
00464F25  |. 33D9           |XOR EBX,ECX
00464F27  |. 803E 00        |CMP BYTE PTR DS:[ESI],0
00464F2A  |.^75 A5          \JNZ SHORT linker_p.00464ED1
00464F2C  |> 5F             POP EDI
00464F2D  |. 5E             POP ESI
00464F2E  |. 8BC3           MOV EAX,EBX
00464F30  |. 5B             POP EBX
00464F31  |. 5D             POP EBP
00464F32  \. C3             RETN
*/