ArithmeticEncoder.h

#pragma once


#include <libutl/Encoder.h>
#include <libutl/Vector.h>


UTL_NS_BEGIN;

// ArithContext ////////////////////////////////////////////////////////////////////////////////////


class ArithContext : public Object
{
    UTL_CLASS_DECL(ArithContext, Object);

public:
    ArithContext(uint_t numSymbols, uint_t inc, uint_t incLimit)
    {
        _F = nullptr;
        set(numSymbols, inc, incLimit);
    }

    virtual void copy(const Object& rhs);

    uint_t
    freq(uint_t symbol) const
    {
        uint_t low, high;
        interval(&low, &high, symbol);
        return high - low;
    }

    void interval(uint_t* low, uint_t* high, uint_t symbol) const;

    uint_t
    numSymbols() const
    {
        return _numSymbols;
    }

    void set(uint_t numSymbols, uint_t inc, uint_t incLimit);

    uint_t symbol(uint_t target);

    uint_t
    totFreq() const
    {
        return _totFreq;
    }

    void update(uint_t symbol);

private:
    void
    init()
    {
        _F = nullptr;
        set(0, 0, 0);
    }
    void
    deInit()
    {
        delete[] _F;
    }
    void halve();

private:
    uint_t _numSymbols;
    uint_t _size;
    uint_t _mid;
    uint_t* _F;
    uint_t _inc;
    uint_t _incLimit;
    uint_t _totFreq;
};


#undef _L
#undef _R

// ArithmeticEncoder ///////////////////////////////////////////////////////////////////////////////


class ArithmeticEncoder : public Encoder
{
    UTL_CLASS_DECL(ArithmeticEncoder, Encoder);

public:
    ArithmeticEncoder(uint_t mode,
                      Stream* stream = nullptr,
                      bool owner = true,
                      ArithContext* ctx = nullptr,
                      bool eob = false)
    {
        _ctx = nullptr;
        start(mode, stream, owner, ctx, eob);
    }

    virtual size_t decode(byte_t* block, size_t num);

    inline uint_t decode(ArithContext* ctx);

    virtual size_t encode(const byte_t* block, size_t num);

    inline void encode(uint_t symbol, ArithContext* ctx);

    void
    setContext(ArithContext* ctx)
    {
        delete _ctx;
        _ctx = ctx;
    }

    void start(uint_t mode,
               Stream* stream,
               bool owner = true,
               ArithContext* ctx = nullptr,
               bool eob = false);

protected:
    virtual void clear();
    virtual void finishEncoding();

private:
    void
    init()
    {
        _ctx = nullptr;
        clearSelf();
    }
    void
    deInit()
    {
        close();
    }
    void clearSelf();
    inline void putBit(bool b);

private:
    uint_t _half, _qtr;
    uint_t _L, _R;
    uint_t _followBits;
    ArithContext* _ctx;
    uint_t _eob;
};


void
ArithmeticEncoder::putBit(bool b)
{
    _stream->putBit(b);
    while (_followBits > 0)
    {
        _stream->putBit(!b);
        _followBits--;
    }
}


void
ArithmeticEncoder::encode(uint_t symbol, ArithContext* ctx)
{
    uint_t l, h, t, r, rl;

    ctx->interval(&l, &h, symbol);
    t = ctx->totFreq();
    r = _R / t;
    rl = r * l;

    // adjust _L, _R
    _L += rl;
    if (h < t)
    {
        _R = r * h - rl;
    }
    else
    {
        _R -= rl;
    }

    // write bits as necessary
    while (_R <= _qtr)
    {
        if (_L + _R <= _half)
        {
            putBit(0);
        }
        else if (_L >= _half)
        {
            putBit(1);
            _L -= _half;
        }
        else
        {
            _followBits++;
            _L -= _qtr;
        }
        _L <<= 1;
        _R <<= 1;
    }

    ctx->update(symbol);
}


uint_t
ArithmeticEncoder::decode(ArithContext* ctx)
{
    uint_t l, h, t, r, rl, target, symbol;
    bool bit;

    t = ctx->totFreq();
    r = _R / t;

    // find l, h, symbol
    target = min(t - 1, _L / r);
    symbol = ctx->symbol(target);
    ctx->interval(&l, &h, symbol);
    rl = r * l;

    // adjust _L, _R
    _L -= rl;
    if (h < t)
    {
        _R = r * h - rl;
    }
    else
    {
        _R -= rl;
    }

    // read bits as necessary
    while (_R <= _qtr)
    {
        _R <<= 1;
        bit = _stream->getBit();
        _L = (_L << 1) | (bit ? 1 : 0);
    }

    ctx->update(symbol);
    return symbol;
}


UTL_NS_END;