BCn Texture Compression

#ifndef __TEXTURE_COMPRESSOR__H__
#define __TEXTURE_COMPRESSOR__H__

#include <math.h>

#define BC_ROUNDING_BIAS

namespace Engine
{
    class TextureCompressor
    {
    public:
        enum class Format
        {
            BC1 = 0,
            BC2,
            BC3,
            BC4,
            BC5
        };

        enum class Quality
        {
            NORMAL = 0,
            HIGH
        };

    private:
        static unsigned char mLookUp8To5[256][2];
        static unsigned char mLookUp8To6[256][2];
        static float mMidPoints5[32];
        static float mMidPoints6[64];
        static unsigned char mExpand5[32];
        static unsigned char mExpand6[64];
        static unsigned char mQuantGTab[256 + 16];
        static unsigned char mQuantRBTab[256 + 16];
        static bool mQuantInitialized;

        static int Mul8Bit(int a, int b)
        {
            int t = a * b + 128;
            return (t + (t >> 8)) >> 8;
        }

        static void From16Bit(unsigned char* out, unsigned short v)
        {
            int r = (v & 0xF800) >> 11;
            int g = (v & 0x07E0) >> 5;
            int b = (v & 0x001F) >> 0;

            // Expand to 8-bit via bit replication
            out[0] = (r * 33) >> 2;
            out[1] = (g * 65) >> 4;
            out[2] = (b * 33) >> 2;
            out[3] = 0;
        }

        static unsigned short As16Bit(int r, int g, int b)
        {
            return (unsigned short)((Mul8Bit(r, 31) << 11) + (Mul8Bit(g, 63) << 5) + Mul8Bit(b, 31));
        }

        static int Lerp13(int a, int b)
        {
#ifdef BC_ROUNDING_BIAS
            return a + Mul8Bit(b - a, 0x55);
#else
            return (2 * a + b) / 3;
#endif
        }

        static void Lerp13RGB(unsigned char* out, unsigned char* p1, unsigned char* p2)
        {
            out[0] = (unsigned char)Lerp13(p1[0], p2[0]);
            out[1] = (unsigned char)Lerp13(p1[1], p2[1]);
            out[2] = (unsigned char)Lerp13(p1[2], p2[2]);
        }

        static void EvalColors(unsigned char* color, unsigned short c0, unsigned short c1)
        {
            From16Bit(color + 0, c0);
            From16Bit(color + 4, c1);
            Lerp13RGB(color + 8, color + 0, color + 4);
            Lerp13RGB(color + 12, color + 4, color + 0);
        }

        static void PrecomputeTables()
        {
            for (int i = 0; i < 32; i++)
            {
                mExpand5[i] = (i << 3) | (i >> 2);
            }

            for (int i = 0; i < 64; i++)
            {
                mExpand6[i] = (i << 2) | (i >> 4);
            }

            for (int i = 0; i < 256 + 16; i++)
            {
                int v = i - 8 < 0 ? 0 : i - 8 > 255 ? 255 : i - 8;
                mQuantRBTab[i] = mExpand5[Mul8Bit(v, 31)];
                mQuantGTab[i] = mExpand6[Mul8Bit(v, 63)];
            }

            mQuantInitialized = true;
        }

        static void DitherBlock(unsigned char* dest, unsigned char* block)
        {
            if (mQuantInitialized == false)
            {
                PrecomputeTables();
            }

            int err[8];
            int* ep1 = err;
            int* ep2 = err + 4;
            int* et;

            // Process channels separately and dither
            for (int channel = 0; channel < 3; channel++)
            {
                unsigned char* bp = block + channel;
                unsigned char* dp = dest + channel;

                unsigned char* quant = (channel == 1) ? mQuantGTab + 8 : mQuantRBTab + 8;

                err[0] = err[1] = err[2] = err[3] = err[4] = err[5] = err[6] = err[7] = 0;

                for (int i = 0; i < 4; i++)
                {
                    dp[0] = quant[bp[0] + ((3 * ep2[1] + 5 * ep2[0]) >> 4)];
                    ep1[0] = bp[0] - dp[0];
                    dp[4] = quant[bp[4] + ((7 * ep1[0] + 3 * ep2[2] + 5 * ep2[1] + ep2[0]) >> 4)];
                    ep1[1] = bp[4] - dp[4];
                    dp[8] = quant[bp[8] + ((7 * ep1[1] + 3 * ep2[3] + 5 * ep2[2] + ep2[1]) >> 4)];
                    ep1[2] = bp[8] - dp[8];
                    dp[12] = quant[bp[12] + ((7 * ep1[2] + 5 * ep2[3] + ep2[2]) >> 4)];
                    ep1[3] = bp[12] - dp[12];
                    bp += 16;
                    dp += 16;

                    // Swap
                    et = ep1;
                    ep1 = ep2;
                    ep2 = et;
                }
            }
        }

        static unsigned int MatchColorsBlock(unsigned char* block, unsigned char* color, bool dither)
        {
            // Calculate vector for color line
            int dir[3] = { color[0 * 4 + 0] - color[1 * 4 + 0], color[0 * 4 + 1] - color[1 * 4 + 1], color[0 * 4 + 2] - color[1 * 4 + 2] };

            // Calculate colors projected on this line
            int stops[4];
            for (int i = 0; i < 4; i++)
            {
                stops[i] = color[i * 4 + 0] * dir[0] + color[i * 4 + 1] * dir[1] + color[i * 4 + 2] * dir[2];
            }

            // Calculate ideal projection of each pixel color on the line
            int dots[16];
            for (int i = 0; i < 16; i++)
            {
                dots[i] = block[i * 4 + 0] * dir[0] + block[i * 4 + 1] * dir[1] + block[i * 4 + 2] * dir[2];
            }

            // At this point all colors are projected on the line, we have just 4 colors linearly on the given line. Determine which is the closest color for given projection
            int c0 = (stops[1] + stops[3]);
            int half = (stops[3] + stops[2]);
            int c3 = (stops[2] + stops[0]);

            // Resulting mask
            int mask = 0;

            if (!dither)
            {
                for (int i = 15; i >= 0; i--)
                {
                    int dot = dots[i] * 2;
                    mask <<= 2;

                    if (dot < half)
                    {
                        mask |= (dot < c0) ? 1 : 3;
                    }
                    else
                    {
                        mask |= (dot < c3) ? 2 : 0;
                    }
                }
            }
            else
            {
                // Floyd-Steinberg dithering
                int err[8];
                int* ep1 = err;
                int* ep2 = err + 4;

                int* dp = dots;

                c0 <<= 3;
                half <<= 3;
                c3 <<= 3;

                for (int i = 0; i < 8; i++)
                {
                    err[i] = 0;
                }

                for (int y = 0; y < 4; y++)
                {
                    int dot, lmask, step;

                    dot = (dp[0] << 4) + (3 * ep2[1] + 5 * ep2[0]);
                    if (dot < half)
                        step = (dot < c0) ? 1 : 3;
                    else
                        step = (dot < c3) ? 2 : 0;
                    ep1[0] = dp[0] - stops[step];
                    lmask = step;

                    dot = (dp[1] << 4) + (7 * ep1[0] + 3 * ep2[2] + 5 * ep2[1] + ep2[0]);
                    if (dot < half)
                        step = (dot < c0) ? 1 : 3;
                    else
                        step = (dot < c3) ? 2 : 0;
                    ep1[1] = dp[1] - stops[step];
                    lmask |= step << 2;

                    dot = (dp[2] << 4) + (7 * ep1[1] + 3 * ep2[3] + 5 * ep2[2] + ep2[1]);
                    if (dot < half)
                        step = (dot < c0) ? 1 : 3;
                    else
                        step = (dot < c3) ? 2 : 0;
                    ep1[2] = dp[2] - stops[step];
                    lmask |= step << 4;

                    dot = (dp[3] << 4) + (7 * ep1[2] + 5 * ep2[3] + ep2[2]);
                    if (dot < half)
                        step = (dot < c0) ? 1 : 3;
                    else
                        step = (dot < c3) ? 2 : 0;
                    ep1[3] = dp[3] - stops[step];
                    lmask |= step << 6;

                    dp += 4;
                    mask |= lmask << (y * 8);
                    { int* et = ep1; ep1 = ep2; ep2 = et; } // swap
                }

            }

            return mask;
        }

        static void OptimizeColorBlock(unsigned char* block, unsigned short* max16, unsigned short* min16)
        {
            // Find minimum, maximum and mean color per channel
            int mean_value[3];
            int min_value[3];
            int max_value[3];

            for (int channel = 0; channel < 3; channel++)
            {
                int minv, maxv, muv;

                muv = maxv = minv = block[channel];
                for (int i = 4; i < 64; i += 4)
                {
                    muv += block[i + channel];

                    if (block[i + channel] < minv)
                    {
                        minv = block[i + channel];
                    }
                    else if (block[i + channel] > maxv)
                    {
                        maxv = block[i + channel];
                    }
                }

                mean_value[channel] = (muv + 8) >> 4;
                min_value[channel] = minv;
                max_value[channel] = maxv;
            }

            // Calculate elements of covariance matrix (symmetrical matrix - therefore calculating only bottom triangle of it):
            // R*R G*R B*R
            // R*G G*G B*G
            // R*B G*B B*B
            int covariance_int[6] = { 0 };
            for (int i = 0; i < 16; i++)
            {
                int r = block[i * 4 + 0] - mean_value[0];
                int g = block[i * 4 + 1] - mean_value[1];
                int b = block[i * 4 + 2] - mean_value[2];

                covariance_int[0] += r * r;
                covariance_int[1] += r * g;
                covariance_int[2] += r * b;
                covariance_int[3] += g * g;
                covariance_int[4] += g * b;
                covariance_int[5] += b * b;
            }

            float covariance_float[6] = { 0.0f };
            for (int i = 0; i < 6; i++)
            {
                covariance_float[i] = (float)covariance_int[i] / 255.0f;
            }

            // Find principal axis based on power interation
            float axis[3] = { (float)(max_value[0] - min_value[0]), (float)(max_value[1] - min_value[1]), (float)(max_value[2] - min_value[2]) };

            int ITERATIONS = 1;
            for (int iter = 0; iter < ITERATIONS; iter++)
            {
                float r = axis[0] * covariance_float[0] + axis[1] * covariance_float[1] + axis[2] * covariance_float[2];
                float g = axis[0] * covariance_float[1] + axis[1] * covariance_float[3] + axis[2] * covariance_float[4];
                float b = axis[0] * covariance_float[2] + axis[1] * covariance_float[4] + axis[2] * covariance_float[5];

                axis[0] = r;
                axis[1] = g;
                axis[2] = b;
            }

            // Calculate magnitude of strongest principal axis component
            float magnitude = fabsf(axis[0]);
            if (fabsf(axis[1]) > magnitude)
            {
                magnitude = fabsf(axis[1]);
            }
            if (fabsf(axis[2]) > magnitude)
            {
                magnitude = fabsf(axis[2]);
            }

            // Calculate coefficients for principal axis (in integer - range 0 - 1000), if strongest component of principal axis is too weak, default to luminance
            int vector[3] = { 299, 587, 114 };
            if (magnitude >= 4.0f)
            {
                magnitude = 512.0f / magnitude;
                vector[0] = (int)(axis[0] * magnitude);
                vector[1] = (int)(axis[1] * magnitude);
                vector[2] = (int)(axis[2] * magnitude);
            }

            // Pick colors at extreme points
            unsigned char* minp = block;
            unsigned char* maxp = block;

            int mind = block[0] * vector[0] + block[1] * vector[1] + block[2] * vector[2];
            int maxd = mind;
            for (int i = 1; i < 16; i++)
            {
                int dot = block[i * 4 + 0] * vector[0] + block[i * 4 + 1] * vector[1] + block[i * 4 + 2] * vector[2];

                if (dot < mind)
                {
                    mind = dot;
                    minp = block + i * 4;
                }

                if (dot > maxd)
                {
                    maxd = dot;
                    maxp = block + i * 4;
                }
            }

            *min16 = As16Bit(minp[0], minp[1], minp[2]);
            *max16 = As16Bit(maxp[0], maxp[1], maxp[2]);
        }

        static unsigned short Quantize5(float x)
        {
            unsigned short q;
            x = x < 0.0f ? 0.0f : x > 1.0f ? 1.0f : x;
            q = (unsigned short)(x * 31.0f);
            q += (x > mMidPoints5[q]);
            return q;
        }

        static unsigned short Quantize6(float x)
        {
            unsigned short q;
            x = x < 0.0f ? 0.0f : x > 1.0f ? 1.0f : x;
            q = (unsigned short)(x * 63.0f);
            q += (x > mMidPoints6[q]);
            return q;
        }

        static int RefineBlock(unsigned char* block, unsigned short* max16, unsigned short* min16, unsigned int mask)
        {
            int oldmin = *min16;
            int oldmax = *max16;

            // Check whether all pixels have same idx
            if ((mask ^ (mask << 2)) < 4)
            {
                int r = 8;
                int g = 8;
                int b = 8;

                for (int i = 0; i < 16; i++)
                {
                    r += block[i * 4 + 0];
                    g += block[i * 4 + 1];
                    b += block[i * 4 + 2];
                }

                r >>= 4;
                g >>= 4;
                b >>= 4;

                *max16 = (mLookUp8To5[r][0] << 11) | (mLookUp8To6[g][0] << 5) | mLookUp8To5[b][0];
                *min16 = (mLookUp8To5[r][1] << 11) | (mLookUp8To6[g][1] << 5) | mLookUp8To5[b][1];
            }
            else
            {
                int at1[3] = { 0, 0, 0 };
                int at2[3] = { 0, 0, 0 };
                unsigned int cm = mask;

                int w1table[4] = { 3, 0, 2 ,1 };
                int prods[4] = { 0x090000, 0x000900, 0x040102, 0x010402 };

                int accumulate = 0;

                for (int i = 0; i < 16; i++, cm >>= 2)
                {
                    int step = cm & 3;
                    int w1 = w1table[step];

                    int r = block[i * 4 + 0];
                    int g = block[i * 4 + 1];
                    int b = block[i * 4 + 2];

                    accumulate += prods[step];

                    at1[0] += w1 * r;
                    at1[1] += w1 * g;
                    at1[2] += w1 * b;
                    at2[0] += r;
                    at2[1] += g;
                    at2[2] += b;
                }

                at2[0] = 3 * at2[0] - at1[0];
                at2[1] = 3 * at2[1] - at1[1];
                at2[2] = 3 * at2[2] - at1[2];

                int xx = accumulate >> 16;
                int yy = (accumulate >> 8) & 0xff;
                int xy = (accumulate >> 0) & 0xff;

                float f = 3.0f / 255.0f / (xx * yy - xy * xy);

                *max16 = Quantize5((at1[0] * yy - at2[0] * xy) * f) << 11;
                *max16 |= Quantize6((at1[1] * yy - at2[1] * xy) * f) << 5;
                *max16 |= Quantize5((at1[2] * yy - at2[2] * xy) * f) << 0;

                *min16 = Quantize5((at2[0] * xx - at1[0] * xy) * f) << 11;
                *min16 |= Quantize6((at2[1] * xx - at1[1] * xy) * f) << 5;
                *min16 |= Quantize5((at2[2] * xx - at1[2] * xy) * f) << 0;
            }

            return (oldmin != *min16) || (oldmax != *max16);
        }

        static void CompressColorBlock(unsigned char* dest, unsigned char* src, bool dithering)
        {
            unsigned int mask = 0;
            unsigned short max16 = 0;
            unsigned short min16 = 0;
;
            unsigned char dblock[16 * 4] = { 0 };

            // Check if block is single constant color
            int i = 0;
            for (i = 1; i < 16; i++)
            {
                //if (((unsigned int*)src)[i] != ((unsigned int*)src)[0])
                if ((src[i * 4 + 0] != src[0]) ||
                    (src[i * 4 + 1] != src[1]) ||
                    (src[i * 4 + 2] != src[2]))
                {
                    break;
                }
            }

            if (i == 16)
            {
                // Constant color, use just that one
                int r = src[0];
                int g = src[1];
                int b = src[2];

                max16 = (mLookUp8To5[r][0] << 11) | (mLookUp8To6[g][0] << 5) | mLookUp8To5[b][0];
                min16 = (mLookUp8To5[r][1] << 11) | (mLookUp8To6[g][1] << 5) | mLookUp8To5[b][1];

                mask = 0xAAAAAAAA;
            }
            else
            {
                if (dithering)
                {
                    DitherBlock(dblock, src);
                }

                OptimizeColorBlock(dithering ? dblock : src, &max16, &min16);

                if (max16 == min16)
                {
                    mask = 0x0;

                    min16 = 0;
                    max16 = 0;
                }
                else
                {
                    unsigned char color[4 * 4];
                    EvalColors(color, max16, min16);
                    mask = MatchColorsBlock(src, color, dithering);
                }

                for (i = 0; i < 2; i++)
                {
                    unsigned int lastmask = mask;

                    if (RefineBlock(dithering ? dblock : src, &max16, &min16, mask))
                    {
                        if (max16 == min16)
                        {
                            mask = 0;
                            break;
                        }
                        else
                        {
                            unsigned char color[4 * 4];
                            EvalColors(color, max16, min16);
                            mask = MatchColorsBlock(src, color, dithering);
                        }
                    }

                    if (mask == lastmask)
                    {
                        break;
                    }
                }
            }

            if (max16 < min16)
            {
                unsigned short t = min16;
                min16 = max16;
                max16 = t;

                mask ^= 0x55555555;
            }

            dest[0] = (unsigned char)(max16);
            dest[1] = (unsigned char)(max16 >> 8);
            dest[2] = (unsigned char)(min16);
            dest[3] = (unsigned char)(min16 >> 8);
            dest[4] = (unsigned char)(mask);
            dest[5] = (unsigned char)(mask >> 8);
            dest[6] = (unsigned char)(mask >> 16);
            dest[7] = (unsigned char)(mask >> 24);
        }

        static void SampleAlphaBlock(unsigned char* dest, unsigned char* src, int stride)
        {
            // Resample alpha to be 4 bit per pixel
            for (int i = 0; i < 16; i += 2)
            {
                int a1 = (src[i * stride]) >> 4;
                int a2 = (src[(i + 1) * stride]) >> 4;

                unsigned char alpha = (unsigned char)((a1 << 4) + a2);

                dest[i / 2] = alpha;
            }
        }

        static void CompressAlphaBlock(unsigned char* dest, unsigned char* src, int stride)
        {
            // Find min and max value
            int minval = src[0];
            int maxval = src[0];

            for (int i = 1; i < 16; i++)
            {
                if (src[i * stride] < minval)
                {
                    minval = src[i * stride];
                }
                else if (src[i * stride] > maxval)
                {
                    maxval = src[i * stride];
                }
            }

            // Encode min and max value in result
            dest[0] = (unsigned char)maxval;
            dest[1] = (unsigned char)minval;
            dest += 2;

            // Determine color indices (mask)
            int dist = maxval - minval;
            int dist2 = dist * 2;
            int dist4 = dist * 4;

            int bias = (dist < 8) ? (dist - 1) : (dist / 2 + 2);
            bias -= minval * 7;

            int bits = 0;
            int mask = 0;

            for (int i = 0; i < 16; i++)
            {
                int a = src[i * stride] * 7 + bias;

                int ind, t;

                // Select indext (linear scale lerp factor between 0 (val = min) and 7 (val = max)
                t = (a >= dist4) ? -1 : 0; ind = t & 4; a -= dist4 & t;
                t = (a >= dist2) ? -1 : 0; ind += t & 2; a -= dist2 & t;
                ind += (a >= dist);

                // Linear scale into block compression index
                ind = -ind & 7;
                ind ^= (2 > ind);

                // Write index
                mask |= ind << bits;
                if ((bits += 3) >= 8)
                {
                    *dest++ = (unsigned char)mask;
                    mask >>= 8;
                    bits -= 8;
                }
            }
        }

    public:
        static void CompressBlock(unsigned char* dest, unsigned char* src, Format format, bool dithering)
        {
            if (format == Format::BC1)
            {
                CompressColorBlock(dest, src, dithering);
            }
            else if (format == Format::BC2)
            {
                SampleAlphaBlock(dest, src + 3, 4);
                CompressColorBlock(dest + 8, src, dithering);
            }
            else if (format == Format::BC3)
            {
                CompressAlphaBlock(dest, src + 3, 4);
                CompressColorBlock(dest + 8, src, dithering);
            }
            else if (format == Format::BC4)
            {
                CompressAlphaBlock(dest, src, 4);
            }
            else if (format == Format::BC5)
            {
                CompressAlphaBlock(dest, src, 4);
                CompressAlphaBlock(dest + 8, src + 1, 4);
            }
        }

        static void Compress(size_t width, size_t height, unsigned char* dest, unsigned char* src, Format format, size_t stride, bool alpha, bool dithering)
        {
            int channels = alpha ? 4 : 3;

            unsigned char* ptr = dest;

            for (int y = 0; y < height; y += 4)
            {
                for (int x = 0; x < width; x += 4)
                {
                    unsigned char compressed[16 * 4] = { 0 };
                    unsigned char original[16 * 4] = { 0 };

                    for (int i = 0; i < 4; i++)
                    {
                        for (int j = 0; j < 4; j++)
                        {
                            original[(i * 4 + j) * 4 + 0] = src[((y + i) * width + (x + j)) * channels + 0];
                            original[(i * 4 + j) * 4 + 1] = src[((y + i) * width + (x + j)) * channels + 1];
                            original[(i * 4 + j) * 4 + 2] = src[((y + i) * width + (x + j)) * channels + 2];
                            original[(i * 4 + j) * 4 + 3] = (alpha ? src[((y + i) * width + (x + j)) * channels + 3] : 255);
                        }
                    }

                    CompressBlock(compressed, original, format, dithering);

                    if (format == Format::BC1 || format == Format::BC4)
                    {
                        for (int i = 0; i < 8; i++)
                        {
                            (*ptr) = compressed[i];
                            ptr++;
                        }
                    }
                    else
                    {
                        for (int i = 0; i < 16; i++)
                        {
                            (*ptr) = compressed[i];
                            ptr++;
                        }
                    }
                }
            }
        }
    };
}

#endif