[AT3P] Adjust number of quant units during encoding.

3 files changed, 41 insertions, 1 deletions

diff --git a/src/lib/bs_encode/encode.cpp b/src/lib/bs_encode/encode.cpp
index 4c17ce9..230ae0f 100644
--- a/src/lib/bs_encode/encode.cpp
+++ b/src/lib/bs_encode/encode.cpp
@@ -32,6 +32,7 @@ public:
     void DoSubmit(size_t gotBits) noexcept;
     bool DoCheck(size_t gotBits) const noexcept;
     void DoRun(void* frameData, TBitStream& bs);
+    uint32_t DoGetCurGlobalConsumption() const noexcept;
 private:
     std::vector<IBitStreamPartEncoder::TPtr> Encoders;
     size_t CurEncPos;
@@ -109,6 +110,9 @@ void TBitStreamEncoder::TImpl::DoRun(void* frameData, TBitStream& bs)
             }
             if (status == IBitStreamPartEncoder::EStatus::Repeat) {
                 cont = true;
+                for (size_t i = 0; i < CurEncPos; i++) {
+                    Encoders[i]->Reset();
+                }
                 RepeatEncPos = 0;
                 break;
             }
@@ -120,6 +124,15 @@ void TBitStreamEncoder::TImpl::DoRun(void* frameData, TBitStream& bs)
     }
 }
 
+uint32_t TBitStreamEncoder::TImpl::DoGetCurGlobalConsumption() const noexcept
+{
+    uint32_t consumption = 0;
+    for (size_t i = 0; i < CurEncPos; i++) {
+        consumption += Encoders[i]->GetConsumption();
+    }
+    return consumption;
+}
+
 /////
 
 TBitStreamEncoder::TBitStreamEncoder(std::vector<IBitStreamPartEncoder::TPtr>&& encoders)
@@ -161,4 +174,11 @@ bool TBitAllocHandler::Check(size_t gotBits) const noexcept
     const TBitStreamEncoder::TImpl* self = static_cast<const TBitStreamEncoder::TImpl*>(this);
     return self->DoCheck(gotBits);
 }
+
+uint32_t TBitAllocHandler::GetCurGlobalConsumption() const noexcept
+{
+    const TBitStreamEncoder::TImpl* self = static_cast<const TBitStreamEncoder::TImpl*>(this);
+    return self->DoGetCurGlobalConsumption();
+}
+
 }
diff --git a/src/lib/bs_encode/encode.h b/src/lib/bs_encode/encode.h
index 9671e10..3302d10 100644
--- a/src/lib/bs_encode/encode.h
+++ b/src/lib/bs_encode/encode.h
@@ -34,6 +34,9 @@ public:
     float Continue() noexcept;
     bool Check(size_t gitBits) const noexcept;
     void Submit(size_t gotBits) noexcept;
+
+    // Returns consumption of all previous encoded parts (except part from this method called)
+    uint32_t GetCurGlobalConsumption() const noexcept;
 };
 
 class IBitStreamPartEncoder {
@@ -47,7 +50,8 @@ public:
     virtual ~IBitStreamPartEncoder() = default;
     virtual EStatus Encode(void* frameData, TBitAllocHandler& ba) = 0;
     virtual void Dump(NBitStream::TBitStream& bs) = 0;
-
+    virtual void Reset() noexcept {};
+    virtual uint32_t GetConsumption() const noexcept = 0;
 };
 
 class TBitStreamEncoder {
diff --git a/src/lib/bs_encode/encode_ut.cpp b/src/lib/bs_encode/encode_ut.cpp
index e9293a2..39f0ff1 100644
--- a/src/lib/bs_encode/encode_ut.cpp
+++ b/src/lib/bs_encode/encode_ut.cpp
@@ -47,6 +47,10 @@ public:
     void Dump(NBitStream::TBitStream& bs) override {
          EXPECT_EQ(EncCalls, ExpCalls);
     }
+
+    uint32_t GetConsumption() const noexcept override {
+        return 0;
+    }
 private:
     const size_t ExpCalls;
     size_t EncCalls = 0;
@@ -74,6 +78,10 @@ public:
              bs.Write(1, 1);
          }
     }
+
+    uint32_t GetConsumption() const noexcept override {
+        return 1 * Bits;
+    }
 private:
     const size_t ExpCalls;
     size_t EncCalls = 0;
@@ -94,6 +102,10 @@ public:
     void Dump(NBitStream::TBitStream& bs) override {
          EXPECT_EQ(EncCalls, ExpCalls);
     }
+
+    uint32_t GetConsumption() const noexcept override {
+        return 0;
+    }
 private:
     const size_t ExpCalls;
     size_t EncCalls = 0;
@@ -115,6 +127,10 @@ public:
     void Dump(NBitStream::TBitStream& bs) override {
          EXPECT_EQ(EncCalls, 1);
     }
+
+    uint32_t GetConsumption() const noexcept override {
+        return 0;
+    }
 private:
     size_t EncCalls = 0;
 };