path: root/library/cpp/codecs/float_huffman.cpp

#include "float_huffman.h"

#include <util/generic/array_ref.h> 
#include <util/generic/bitops.h> 
#include <util/generic/cast.h> 
#include <util/generic/yexception.h> 
#include <util/system/unaligned_mem.h> 
#include <util/system/yassert.h> 
#include <util/stream/format.h> 
 
namespace NCodecs::NFloatHuff { 
    namespace { 
        struct THuffEntry { 
            ui32 CodeBase; 
            ui16 Prefix; 
            int PrefLength; 
            int CodeLength; 
            int TotalLength; 
            ui32 Mask; 
            ui64 Offset; 
 
            THuffEntry() = default; 
 
            constexpr THuffEntry(ui32 codeBase, ui16 prefix, int prefLength, int codeLength) 
                : CodeBase(codeBase) 
                , Prefix(prefix) 
                , PrefLength(prefLength) 
                , CodeLength(codeLength) 
                , TotalLength(prefLength + codeLength) 
                , Mask(Mask64(codeLength)) 
                , Offset(-(ui64(codeBase) << prefLength) + prefix) 
            {} 
 
            bool Fit(ui32 code) const { 
                return code >= CodeBase && code < CodeBase + (1ULL << CodeLength); 
            } 
        }; 
 
        // NB. There is a typo in the penultimate line (34 instead of 24). It was there from the very 
        // first commit and we cannot fix it without breaking all the clients. 
        constexpr THuffEntry entries[16] = { 
            {0x00000000, 0x01, 1, 0}, // Only +0.0f, 1 bit, prefix                     [1] 
            {0x3f800000, 0x0e, 4, 0}, // Only +1.0f, 4 bits, prefix                    [0111] 
            {0x3f700000, 0x08, 5, 20}, // [0.9375, 1.0), 25 bits, prefix               [00010] 
            {0x3f000000, 0x00, 5, 20}, // [0.5, 0.5625), 25 bits, prefx                [00000] 
            {0x3f400000, 0x06, 6, 20}, // [0.75, 0.8125), 26 bits, prefix              [011000] 
            {0x3f500000, 0x22, 6, 20}, // [0.8125, 0.875), 26 bits, prefix             [010001] 
            {0x3f200000, 0x02, 6, 20}, // [0.625, 0.6875), 26 bits, prefix             [010000] 
            {0x3f100000, 0x38, 6, 20}, // [0.5625, 0.625), 26 bits, prefix             [000111] 
            {0x3f600000, 0x18, 6, 20}, // [0.875, 0.9375), 26 bits, prefix             [000110] 
            {0x3f300000, 0x30, 6, 20}, // [0.6875, 0.75), 26 bits, prefix              [000011] 
            {0x3e800000, 0x10, 6, 20}, // [0.25, 0.28125), 26 bits, prefix             [000010] 
            {0x3e000000, 0x04, 3, 24}, // [0.125, 0.5), 27 bits, prefix                [001] 
            {0x3d000000, 0x0a, 4, 24}, // [0.03125, 0.125), 28 bits, prefix            [0101] 
            {0x3c000000, 0x12, 5, 24}, // [0.0078125, 0.03125), 29 bits, prefix        [01001] 
            {0x3b000000, 0x26, 6, 34}, // [0.001953125, end of range), 40 bits, prefix [011001] 
            {0x00000000, 0x16, 5, 32}, // whole range, 37 bits, prefix                 [01101] 
        };

        [[noreturn]] Y_NO_INLINE void ThrowInvalidOffset(size_t size, size_t byteOffset) {
            ythrow yexception() << 
                "Decompression error: requested decoding 8 bytes past end of input buffer of " << size << " bytes size at position " << byteOffset << ". "; 
        } 

        struct THuffInfo { 
            constexpr THuffInfo() { 
                for (size_t i = 0; i < 64; ++i) { 
                    bool init = false; 
                    for (size_t j = 0; j != 16; ++j) { 
                        ui16 prefix = i & Mask64(entries[j].PrefLength); 
                        if (entries[j].Prefix == prefix) { 
                            init = true; 
                            DecodeLookup[i] = entries[j]; 
                            break; 
                        } 
                    } 
                    Y_ASSERT(init); 
                } 
 
                for (ui32 i = 0; i < (1 << 12); ++i) { 
                    // First two entries (+0.0f and +1.0f) are not present in the lookup, they are handled separately 
                    for (int value = 2; value < 16; ++value) { 
                        if (entries[value].Fit(i << 20)) { 
                            EncodeLookup[i] = value; 
                            break; 
                        } 
                    } 
                } 
            } 
 
            std::pair<ui64, int> GetCode(ui32 value) const { 
                // Zeros are handled separately in the main loop 
                Y_ASSERT(value != 0); 
 
                if (value == 0x3f800000) { 
                    return {0x0e, 4}; 
                } 
 
                const auto& entry = entries[EncodeLookup[value >> 20]]; 
 
                return { 
                    (ui64(value) << entry.PrefLength) + entry.Offset, 
                    entry.TotalLength 
                }; 
            } 
 
            THuffEntry DecodeLookup[64]; 
            ui8 EncodeLookup[1 << 12]; 
        }; 
 
        const THuffInfo huffInfo; 
        /// End Of Stream 
        const ui32 EOS = ui32(-1); 
    } 
 
    TString Encode(TArrayRef<const float> factors) { 
        TString result; 
        result.resize((factors.size() + 1) * 40 / 8 + 8, 0); // Max code length is 40 bits 
        int usedBits = 0; 
        ui64 buffer = 0; 
        char* writePtr = result.begin(); 
 
        auto writeBits = [&](ui64 code, int size) { 
            const auto bitsToTransfer = Min(size, 64 - usedBits); 
            buffer |= (code << usedBits); 
            usedBits += bitsToTransfer; 
            if (usedBits == 64) { 
                memcpy(writePtr, &buffer, 8); 
                usedBits = size - bitsToTransfer; 
                if (bitsToTransfer != 64) { 
                    buffer = code >> bitsToTransfer; 
                } else { 
                    buffer = 0; 
                } 
                writePtr += 8; 
            } 
        }; 
 
        for (size_t i = 0; i != factors.size();) { 
            if (BitCast<ui32>(factors[i]) == 0) { 
                int zeroCount = 1; 
                for (;;) { 
                    ++i; 
                    if (i == factors.size() || BitCast<ui32>(factors[i]) != 0) { 
                        break; 
                    } 
                    ++zeroCount; 
                } 
                for (; zeroCount >= 64; zeroCount -= 64) { 
                    writeBits(ui64(-1), 64); 
                } 
                writeBits(Mask64(zeroCount), zeroCount); 
            } else { 
                const auto [code, codeSize] = huffInfo.GetCode(BitCast<ui32>(factors[i])); 
                writeBits(code, codeSize); 
                ++i; 
            } 
        } 
        // Write EOS. 
        // We use precomputed constants instead of the following: 
        //      auto [code, codeSize] = huffInfo.GetCode(EOS); 
        //      writeBits(code, codeSize); 
        writeBits(211527139302, 40); 
        memcpy(writePtr, &buffer, 8); 
        result.resize(writePtr - result.begin() + usedBits / 8 + 8); 
 
        return result; 
    } 
 
    TDecoder::TDecoder(TStringBuf data) 
        : State{ 
            .Workspace = data.size() < 8 ? ThrowInvalidOffset(data.size(), 0), 0 : ReadUnaligned<ui64>(data.data()), 
            .WorkspaceSize = 64, 
            .Position = 8, 
            .Data = data 
        } 
    { 
        FillDecodeBuffer(); 
    } 
 
    TVector<float> TDecoder::DecodeAll(size_t sizeHint) { 
        TVector<float> result; 
        result.reserve(sizeHint); 
 
        while (Begin != End) { 
            result.insert(result.end(), Begin, End); 
            FillDecodeBuffer(); 
        } 
        return result; 
    } 
 
    size_t TDecoder::Decode(TArrayRef<float> dest) { 
        size_t count = 0; 
        while (count < dest.size()) { 
            if (dest.size() - count < size_t(End - Begin)) { 
                const auto size = dest.size() - count; 
                std::copy(Begin, Begin + size, dest.data() + count); 
                Begin += size; 
                return dest.size(); 
            } else { 
                std::copy(Begin, End, dest.data() + count); 
                count += End - Begin; 
                FillDecodeBuffer(); 
                if (Begin == End) { 
                    break; 
                } 
            } 
        } 
        return count; 
    } 
 
    size_t TDecoder::Skip(size_t count) { 
        size_t skippedCount = 0; 
        while (skippedCount < count) { 
            if (count - skippedCount < size_t(End - Begin)) { 
                const auto size = count - skippedCount; 
                Begin += size; 
                return count; 
            } else { 
                skippedCount += End - Begin; 
                FillDecodeBuffer(); 
                if (Begin == End) { 
                    break; 
                } 
            } 
        } 
        return skippedCount; 
    } 
 
    bool TDecoder::HasMore() const { 
        return Begin != End; 
    } 
 
    void TDecoder::FillDecodeBuffer() { 
        Begin = DecodeBuffer.data(); 
        End = DecodeBuffer.data(); 
 
        if (HitEos) { 
            return; 
        } 
 
        // This helps to keep most of the variables in the registers. 
        float* end = End; 
        TState state = State; 
 
        // It is faster to just zero all the memory here in one go 
        // and then avoid inner loop when writing zeros. There we 
        // can just increment end pointer. 
        std::fill(DecodeBuffer.begin(), DecodeBuffer.end(), 0.0f); 
 
        // Make sure that inside the loop we always have space to put 64 zeros and one other 
        // value. 
        float* cap = DecodeBuffer.data() + DecodeBuffer.size() - 64 - 1; 
 
        while (end < cap) { 
            if (state.Workspace % 2 == 1) { 
                // Decode zeros 
                // There we can just scan whole state.Workspace for ones because it contains 
                // zeros outside of the WorkspaceSize bits. 
                const auto negWorkspace = ~state.Workspace; 
                const int zeroCount = negWorkspace ? CountTrailingZeroBits(negWorkspace) : 64; 
                end += zeroCount; 
                state.SkipBits(zeroCount); 
                continue; 
            } 
            if (state.PeekBits(4) == 0x0e) { 
                *end++ = 1.0f; 
                state.SkipBits(4); 
                continue; 
            } 
            const auto& entry = huffInfo.DecodeLookup[state.PeekBits(6)]; 
            const auto code = ui32((state.NextBitsUnmasked(entry.TotalLength) >> entry.PrefLength) & entry.Mask) + entry.CodeBase; 
            if (Y_UNLIKELY(code == EOS)) { 
                HitEos = true; 
                break; 
            } 
            *end++ = BitCast<float>(code); 
        } 
 
        End = end; 
        State = state; 
    } 
 
    ui64 TDecoder::TState::PeekBits(int count) { 
        if (WorkspaceSize > count) { 
            return Workspace & Mask64(count); 
        } else { 
            if (Y_UNLIKELY(Position + 8 > Data.size())) { 
                ThrowInvalidOffset(Data.size(), Position); 
            } 
            return (Workspace | (ReadUnaligned<ui64>(Data.data() + Position) << WorkspaceSize)) & Mask64(count); 
        } 
    } 
 
    ui64 TDecoder::TState::NextBitsUnmasked(int count) { 
        if (WorkspaceSize > count) { 
            const auto result = Workspace; 
            Workspace >>= count; 
            WorkspaceSize -= count; 
            return result; 
        } else { 
            if (Y_UNLIKELY(Position + 8 > Data.size())) { 
                ThrowInvalidOffset(Data.size(), Position); 
            } 
            ui64 result = Workspace; 
            Workspace = ReadUnaligned<ui64>(Data.data() + Position); 
            Position += 8; 
            result |= Workspace << WorkspaceSize; 
            Workspace >>= count - WorkspaceSize; 
            WorkspaceSize += 64 - count; 
            return result; 
        } 
    } 
 
    void TDecoder::TState::SkipBits(int count) { 
        if (WorkspaceSize > count) { 
            Workspace >>= count; 
            WorkspaceSize -= count; 
        } else { 
            if (Y_UNLIKELY(Position + 8 > Data.size())) { 
                ThrowInvalidOffset(Data.size(), Position); 
            } 
            Workspace = ReadUnaligned<ui64>(Data.data() + Position); 
            Position += 8; 
            Workspace >>= count - WorkspaceSize; 
            WorkspaceSize += 64 - count; 
        } 
    } 
 
    TVector<float> Decode(TStringBuf data, size_t sizeHint) { 
        return TDecoder(data).DecodeAll(sizeHint); 
    } 
} // namespace NCodecs::NFloatHuff