ydb/library/cpp/compproto/bit.h

#pragma once

#include <util/generic/array_ref.h>
#include <util/generic/vector.h>
#include <util/stream/output.h>
#include <util/stream/input.h>

#include "huff.h"
#include "compressor.h"
#include "metainfo.h"

namespace NCompProto {
    struct TBitBuffer {
        TVector<ui8> Out;
        ui64 Position;
        ui64 Size;
        ui8 Counter;
        TBitBuffer() {
            Clear();
        }

        static ui64 Read(const ui8* out, ui64 position, size_t size) {
            const ui8* dst = out + (size_t(position >> 3));
            ui64 outCode = *reinterpret_cast<const ui64*>(dst);
            ui8 shift = position & 7;
            return ((1ULL << size) - 1) & (outCode >> shift);
        }

        ui64 Read(ui64 position, size_t size) {
            return Read(&Out[0], position, size);
        }

        void Code(ui64 value, size_t size) {
            if (++Counter == 0) {
                Junk(257 * 64);
            }
            Position = Code(value, size, Position);
        }
        ui64 Code(ui64 value, size_t size, ui64 position) {
            ui8* dst = &Out[size_t(position >> 3)];
            ui64& outCode = *(ui64*)dst;
            ui8 shift = position & 7;
            ui64 mask = ((1ULL << size) - 1) << shift;
            outCode = ((value << shift) & mask) | (outCode & ~mask);
            return position + size;
        }
        void Junk(size_t junk = 1024) {
            size_t need = size_t(Position >> 3);
            if (Out.size() * 8 < Position + junk) {
                Out.resize(((need + junk) * 3) / 2);
                Size = Out.size();
            }
        }
        void Insert(ui64 value, ui64 position, size_t size) {
            ui64 newVal = Read(position, 56);
            position = Code(value, size, position);
            value = newVal;

            while (position < Position + 64) {
                newVal = Read(position + 56 - size, 56);
                position = Code(value, 56, position);
                value = newVal;
            }

            Position += size;
        }

        size_t ByteLength() const {
            return (Position + 7) / 8;
        }

        void ResetPosition() {
            Position = 0;
            Size = 0;
        }

        void Clear() {
            Counter = 0;
            Position = 0;
            Size = 0;
            Out.clear();
            Junk();
        }

        TArrayRef<const char> AsDataRegion() const {
            return TArrayRef<const char>(reinterpret_cast<const char*>(Out.data()), ByteLength());
        }
    };

    struct THuff {
        TCoder Coder;
        ui64 Position;
        THuff() {
            Position = (ui64)(-1);
        }
        void Load(IInputStream& stream) {
            TString name;
            Coder.InitDefault();
            Coder.Entries.clear();
            while (1) {
                stream >> name;
                if (name == "end")
                    return;
                else if (name == "entry") {
                    TCoderEntry entry;
                    ui32 value;
                    stream >> value;
                    entry.MinValue = value;
                    stream >> value;
                    entry.Prefix = value;
                    stream >> value;
                    entry.PrefixBits = value;
                    stream >> value;
                    entry.AllBits = value;
                    Coder.Entries.push_back(entry);
                }
            }
        }

        void Save(IOutputStream& stream, TString offset) {
            TString step = "    ";
            for (size_t i = 0; i < Coder.Entries.size(); ++i) {
                stream << offset << step << "entry ";
                stream << (ui32)Coder.Entries[i].MinValue << " ";
                stream << (ui32)Coder.Entries[i].Prefix << " ";
                stream << (ui32)Coder.Entries[i].PrefixBits << " ";
                stream << (ui32)Coder.Entries[i].AllBits << " ";
                stream << Endl;
            }
            stream << offset << "end" << Endl;
        }

        void BeginElement(TBitBuffer& out) {
            Position = out.Position;
        }
        void Add(ui32 value, TBitBuffer& out) {
            size_t val = 0;
            ui64 code = Coder.Code(value, val);
            out.Code(code, val);
        }
        void AddDelayed(ui32 value, TBitBuffer& out) {
            size_t val = 0;
            ui64 code = Coder.Code(value, val);
            if (Position == (ui64)(-1)) {
                ythrow yexception() << "Position == (ui64)(-1)";
            }
            out.Insert(code, Position, val);
            out.Junk();
        }
    };

    struct THist {
        TAccum Accum;
        THist() {
        }

        void Load(IInputStream& stream) {
            TString name;
            while (1) {
                stream >> name;
                if (name == "end")
                    return;
            }
        }
        // TODO: why not const TString& ???
        void Save(IOutputStream& /*stream*/, TString /*offset*/) {
        }

        void Add(ui32 value, TEmpty& /*empty*/) {
            Accum.Add(value);
        }
        void AddDelayed(ui32 value, TEmpty& /*empty*/) {
            Accum.Add(value);
        }
        void BeginElement(TEmpty& /*empty*/) {
        }
    };

    struct THistToHuff {
        static THistToHuff Instance() {
            return THistToHuff();
        }
        TEmpty Build() const {
            return TEmpty();
        }
        void Build(THuff& info, const THist& hist) const {
            Analyze(hist.Accum, info.Coder.Entries);
            info.Coder.Normalize();
        }
    };

    struct IDecompressor {
        // sequentially decompresses whole structure according to metainfo, starts at position offset
        virtual void Decompress(const TMetaInfo<TTable>* table, const ui8* codes, ui64& offset) = 0;
        // decompresses one record of outer repeated structure starting at position offset
        virtual void DecompressOne(const TMetaInfo<TTable>* table, const ui8* codes, ui64& offset, ui32 prevIndex = -1) = 0;
        virtual ~IDecompressor() = default;
    };

    template <class X>
    struct TMetaIterator: public IDecompressor {
        X Self;
        TMetaIterator() {
            Self.Parent = this;
        }

    private:
        inline void DecompressSingle(ui32 repeatedIndex, const TMetaInfo<TTable>* table, const ui8* codes, ui64& offset) {
            Self.BeginElement(repeatedIndex);
            ui32 mask = table->Mask.Decompress(codes, offset);
            size_t index = 0;
            ui32 scalarMask = table->ScalarMask;
            while (mask || scalarMask) {
                if (mask & 1) {
                    if (scalarMask & 1) {
                        ui32 val = table->Scalar[index].Decompress(codes, offset);
                        Self.SetScalar(index, val);
                    } else {
                        Self.GetDecompressor(index).Decompress(table->Repeated[index].Get(), codes, offset);
                    }
                } else if ((scalarMask & 1) && table->Default[index].Type == TScalarDefaultValue::Fixed) {
                    Self.SetScalar(index, table->Default[index].Value);
                }
                scalarMask = scalarMask >> 1;
                mask = mask >> 1;
                ++index;
            }
            Self.EndElement();
        }

        inline void DecompressSingleScalarsOnly(ui32 repeatedIndex, const TMetaInfo<TTable>* table, const ui8* codes, ui64& offset) {
            Self.BeginElement(repeatedIndex);
            ui32 mask = table->Mask.Decompress(codes, offset);
            ui32 scalarMask = table->ScalarMask;
            size_t index = 0;
            while (scalarMask) {
                if (mask & 1) {
                    ui32 val = table->Scalar[index].Decompress(codes, offset);
                    Self.SetScalar(index, val);
                } else if (table->Default[index].Type == TScalarDefaultValue::Fixed) {
                    Self.SetScalar(index, table->Default[index].Value);
                }
                mask = mask >> 1;
                scalarMask = scalarMask >> 1;
                ++index;
            }
            Self.EndElement();
        }

    public:
        void Decompress(const TMetaInfo<TTable>* table, const ui8* codes, ui64& offset) override {
            ui64 locOffset = offset;
            ui32 count = table->Count.Decompress(codes, locOffset);
            ui32 repeatedIndex = (ui32)(-1);
            Self.BeginSelf(count, table->Id);
            if (table->RepeatedMask) {
                for (ui32 i = 0; i < count; ++i) {
                    repeatedIndex += table->Index.Decompress(codes, locOffset);
                    DecompressSingle(repeatedIndex, table, codes, locOffset);
                }
            } else {
                for (ui32 i = 0; i < count; ++i) {
                    repeatedIndex += table->Index.Decompress(codes, locOffset);
                    DecompressSingleScalarsOnly(repeatedIndex, table, codes, locOffset);
                }
            }
            offset = locOffset;
            Self.EndSelf();
        }

        // XXX: iterator needed?
        //
        // Following two functions serves the purpose of decompressing outer repeated-structure(such structure is mandatory now).
        // They can work for inner repeated structures too, if you supply correct table, codes and offset parameters.
        void DecompressOne(const TMetaInfo<TTable>* table, const ui8* codes, ui64& offset, ui32 prevIndex = -1) override {
            table->Index.Decompress(codes, offset);
            DecompressSingle(prevIndex, table, codes, offset);
        }

        ui32 DecompressCount(const TMetaInfo<TTable>* table, const ui8* codes, ui64& offset) {
            return table->Count.Decompress(codes, offset);
        }
    };

    template <class X>
    struct TParentHold {
        TMetaIterator<X>* Parent;
        TParentHold()
            : Parent(nullptr)
        {
        }
    };

    struct TEmptyDecompressor: public TParentHold<TEmptyDecompressor> {
        void BeginSelf(ui32 /*count*/, ui32 /*id*/) {
        }
        void EndSelf() {
        }
        void BeginElement(ui32 /*element*/) {
        }
        void EndElement() {
        }
        void SetScalar(size_t /*index*/, ui32 /*val*/) {
        }
        TMetaIterator<TEmptyDecompressor>& GetDecompressor(size_t index) {
            Y_UNUSED(index);
            return *Parent;
        }
    };

    inline TMetaIterator<TEmptyDecompressor>& GetEmptyDecompressor() {
        static TMetaIterator<TEmptyDecompressor> empty;
        return empty;
    }

    struct THuffToTable {
        static THuffToTable Instance() {
            return THuffToTable();
        }
        void Build(TTable& table, const THuff& huff) const {
            // can happen if a field was never serialized across whole structure
            if (huff.Coder.Entries.empty())
                return;
            for (ui32 i = 0; i < 64; ++i) {
                const TCoderEntry& entry = huff.Coder.GetEntry(i, table.Id[i]);
                table.CodeBase[i] = entry.MinValue;
                table.PrefLength[i] = entry.PrefixBits;
                table.CodeMask[i] = (1ULL << (entry.AllBits - entry.PrefixBits)) - 1ULL;
                table.Length[i] = entry.AllBits;
            }
        }
    };

    struct THuffToTableWithDecompressor: private THuffToTable {
        TSimpleSharedPtr<IDecompressor> Decompressor;
        THuffToTableWithDecompressor(TSimpleSharedPtr<IDecompressor> decompressor)
            : Decompressor(decompressor)
        {
        }
        TSimpleSharedPtr<IDecompressor> Build() const {
            return Decompressor;
        }
        void Build(TTable& table, const THuff& huff) const {
            THuffToTable::Build(table, huff);
        }
    };

}