Refactoring to be able to implement joint stereo

6 files changed, 133 insertions, 85 deletions

diff --git a/src/atrac/atrac3.h b/src/atrac/atrac3.h
index 1dfb9ee..5dcdb9e 100644
--- a/src/atrac/atrac3.h
+++ b/src/atrac/atrac3.h
@@ -234,6 +234,12 @@ public:
         const std::vector<TGainPoint>& GetGainPoints(uint32_t i) const {
             return Info[i];
         }
+        void Reset()
+        {
+            for (auto& v : Info) {
+                v.clear();
+            }
+        }
     };
 
     struct TTonalVal {
@@ -244,16 +250,18 @@ public:
 };
 
 struct TAtrac3EncoderSettings {
-    explicit TAtrac3EncoderSettings(uint32_t bitrate, bool noGainControll, bool noTonalComponents)
+    explicit TAtrac3EncoderSettings(uint32_t bitrate, bool noGainControll, bool noTonalComponents, uint8_t sourceChannels)
         : ConteinerParams(TAtrac3Data::GetContainerParamsForBitrate(bitrate))
         , NoGainControll(noGainControll)
         , NoTonalComponents(noTonalComponents)
+        , SourceChannels(sourceChannels)
     {
         std::cout << bitrate << " " << ConteinerParams->Bitrate << std::endl;
     }
     const TContainerParams* ConteinerParams;
     const bool NoGainControll;
     const bool NoTonalComponents;
+    const uint8_t SourceChannels;
 };
 
 } // namespace NAtrac3
diff --git a/src/atrac/atrac3_bitstream.cpp b/src/atrac/atrac3_bitstream.cpp
index d76bae2..79ca9b5 100644
--- a/src/atrac/atrac3_bitstream.cpp
+++ b/src/atrac/atrac3_bitstream.cpp
@@ -215,10 +215,10 @@ void TAtrac3BitStreamWriter::EncodeSpecs(const vector<TScaledBlock>& scaledBlock
     }
 }
 
-uint8_t TAtrac3BitStreamWriter::GroupTonalComponents(const std::vector<TTonalComponent>& tonalComponents,
+uint8_t TAtrac3BitStreamWriter::GroupTonalComponents(const std::vector<TTonalBlock>& tonalComponents,
                                                      TTonalComponentsSubGroup groups[64])
 {
-    for (const TTonalComponent& tc : tonalComponents) {
+    for (const TTonalBlock& tc : tonalComponents) {
         assert(tc.ScaledBlock.Values.size() < 8);
         assert(tc.ScaledBlock.Values.size() > 0);
         assert(tc.QuantIdx >1);
@@ -253,7 +253,7 @@ uint8_t TAtrac3BitStreamWriter::GroupTonalComponents(const std::vector<TTonalCom
     return tcsgn;
 }
 
-uint16_t TAtrac3BitStreamWriter::EncodeTonalComponents(const std::vector<TTonalComponent>& tonalComponents,
+uint16_t TAtrac3BitStreamWriter::EncodeTonalComponents(const std::vector<TTonalBlock>& tonalComponents,
                                                        NBitStream::TBitStream* bitStream, uint8_t numQmfBand)
 {
     const uint16_t bitsUsed = bitStream->GetSizeInBits();
@@ -388,54 +388,70 @@ vector<uint32_t> TAtrac3BitStreamWriter::CalcBitsAllocation(const std::vector<TS
 }
 
 
-void TAtrac3BitStreamWriter::WriteSoundUnit(const TAtrac3Data::SubbandInfo& subbandInfo,
-                                            const std::vector<TTonalComponent>& tonalComponents,
-                                            const vector<TScaledBlock>& scaledBlocks)
+void TAtrac3BitStreamWriter::WriteSoundUnit(const vector<TSingleChannelElement>& singleChannelElements)
 {
-    NBitStream::TBitStream bitStream;
-    if (Params.Js) {
-        //TODO
-    } else {
-        bitStream.Write(0x28, 6); //0x28 - id
-    }
-    const uint8_t numQmfBand = subbandInfo.GetQmfNum();
-    bitStream.Write(numQmfBand - 1, 2);
-
-    //write gain info
-    for (uint32_t band = 0; band < numQmfBand; ++band) {
-        const vector<TAtrac3Data::SubbandInfo::TGainPoint>& GainPoints = subbandInfo.GetGainPoints(band);
-        assert(GainPoints.size() < TAtrac3Data::SubbandInfo::MaxGainPointsNum);
-        bitStream.Write(GainPoints.size(), 3);
-        int s = 0;
-        for (const TAtrac3Data::SubbandInfo::TGainPoint& point : GainPoints) {
-            bitStream.Write(point.Level, 4);
-            bitStream.Write(point.Location, 5);
-            s++;
-            assert(s < 8);
+
+    assert(singleChannelElements.size() == 1 || singleChannelElements.size() == 2);
+    NBitStream::TBitStream bitStreams[2];
+    uint16_t usedBits[2];
+    for (uint32_t channel = 0; channel < singleChannelElements.size(); channel++) {
+        const TSingleChannelElement& sce = singleChannelElements[channel];
+        const TAtrac3Data::SubbandInfo& subbandInfo = sce.SubbandInfo;
+        const std::vector<TTonalBlock>& tonalComponents = sce.TonalBlocks;
+
+        NBitStream::TBitStream* bitStream = &bitStreams[channel];
+
+        if (Params.Js) {
+            //TODO
+            abort();
+        } else {
+            bitStream->Write(0x28, 6); //0x28 - id
+        }
+        const uint8_t numQmfBand = subbandInfo.GetQmfNum();
+        bitStream->Write(numQmfBand - 1, 2);
+
+        //write gain info
+        for (uint32_t band = 0; band < numQmfBand; ++band) {
+            const vector<TAtrac3Data::SubbandInfo::TGainPoint>& GainPoints = subbandInfo.GetGainPoints(band);
+            assert(GainPoints.size() < TAtrac3Data::SubbandInfo::MaxGainPointsNum);
+            bitStream->Write(GainPoints.size(), 3);
+            int s = 0;
+            for (const TAtrac3Data::SubbandInfo::TGainPoint& point : GainPoints) {
+                bitStream->Write(point.Level, 4);
+                bitStream->Write(point.Location, 5);
+                s++;
+                assert(s < 8);
+            }
         }
+        const uint16_t bitsUsedByGainInfo = bitStream->GetSizeInBits() - 8;
+        const uint16_t bitsUsedByTonal = EncodeTonalComponents(tonalComponents, bitStream, numQmfBand);
+        usedBits[channel] = bitsUsedByGainInfo + bitsUsedByTonal;
     }
-    const uint16_t bitsUsedByGainInfo = bitStream.GetSizeInBits() - 8;
-    const uint16_t bitsUsedByTonal = EncodeTonalComponents(tonalComponents, &bitStream, numQmfBand);
-    //spec
-    EncodeSpecs(scaledBlocks, &bitStream, bitsUsedByTonal + bitsUsedByGainInfo);
-
-    if (!Container)
-        abort();
-    if (OutBuffer.empty()) {
-        std::vector<char> channel = bitStream.GetBytes();
-        assert(channel.size() <= Params.FrameSz/2);
-        channel.resize(Params.FrameSz/2);
-        OutBuffer.insert(OutBuffer.end(), channel.begin(), channel.end());
-    } else {
-        std::vector<char> channel = bitStream.GetBytes();
 
-        assert(channel.size() <= Params.FrameSz/2);
-        channel.resize(Params.FrameSz/2);
-        OutBuffer.insert(OutBuffer.end(), channel.begin(), channel.end());
-        Container->WriteFrame(OutBuffer);
-        OutBuffer.clear();
+    for (uint32_t channel = 0; channel < singleChannelElements.size(); channel++) {
+        const TSingleChannelElement& sce = singleChannelElements[channel];
+        const vector<TScaledBlock>& scaledBlocks = sce.ScaledBlocks;
+        NBitStream::TBitStream* bitStream = &bitStreams[channel];
+        //spec
+        EncodeSpecs(scaledBlocks, bitStream, usedBits[channel]);
+
+        if (!Container)
+            abort();
+        std::vector<char> channelData = bitStream->GetBytes();
+        assert(channelData.size() <= (size_t)Params.FrameSz >> 1);
+        channelData.resize(Params.FrameSz >> 1);
+        OutBuffer.insert(OutBuffer.end(), channelData.begin(), channelData.end());
     }
 
+    //No mone mode for atrac3, just make duplicate of first channel
+    if (singleChannelElements.size() == 1 && !Params.Js) {
+        int sz = OutBuffer.size();
+        assert(sz == Params.FrameSz >> 1);
+        OutBuffer.resize(sz << 1);
+        std::copy_n(OutBuffer.begin(), sz, OutBuffer.begin() + sz);
+    }
+    Container->WriteFrame(OutBuffer);
+    OutBuffer.clear();
 }
 
 } // namespace NAtrac3
diff --git a/src/atrac/atrac3_bitstream.h b/src/atrac/atrac3_bitstream.h
index 30a2dff..9177bb2 100644
--- a/src/atrac/atrac3_bitstream.h
+++ b/src/atrac/atrac3_bitstream.h
@@ -29,8 +29,8 @@
 namespace NAtracDEnc {
 namespace NAtrac3 {
 
-struct TTonalComponent {
-    TTonalComponent(const TAtrac3Data::TTonalVal* valPtr, uint8_t quantIdx, const TScaledBlock& scaledBlock)
+struct TTonalBlock {
+    TTonalBlock(const TAtrac3Data::TTonalVal* valPtr, uint8_t quantIdx, const TScaledBlock& scaledBlock)
         : ValPtr(valPtr)
         , QuantIdx(quantIdx)
         , ScaledBlock(scaledBlock)
@@ -43,7 +43,7 @@ struct TTonalComponent {
 class TAtrac3BitStreamWriter : public virtual TAtrac3Data {
     struct TTonalComponentsSubGroup {
         std::vector<uint8_t> SubGroupMap;
-        std::vector<const TTonalComponent*> SubGroupPtr;
+        std::vector<const TTonalBlock*> SubGroupPtr;
     };
     TOma* Container;
     const TContainerParams Params;
@@ -68,10 +68,10 @@ class TAtrac3BitStreamWriter : public virtual TAtrac3Data {
     void EncodeSpecs(const std::vector<TScaledBlock>& scaledBlocks, NBitStream::TBitStream* bitStream,
                      uint16_t bitsUsed);
 
-    uint8_t GroupTonalComponents(const std::vector<TTonalComponent>& tonalComponents,
+    uint8_t GroupTonalComponents(const std::vector<TTonalBlock>& tonalComponents,
                                  TTonalComponentsSubGroup groups[64]);
 
-    uint16_t EncodeTonalComponents(const std::vector<TTonalComponent>& tonalComponents,
+    uint16_t EncodeTonalComponents(const std::vector<TTonalBlock>& tonalComponents,
                                    NBitStream::TBitStream* bitStream, uint8_t numQmfBand);
 public:
     TAtrac3BitStreamWriter(TOma* container, const TContainerParams& params) //no mono mode for atrac3
@@ -80,9 +80,14 @@ public:
     {
 
     }
-    void WriteSoundUnit(const TAtrac3Data::SubbandInfo& subbandInfo,
-                        const std::vector<TTonalComponent>& tonalComponents,
-                        const std::vector<TScaledBlock>& scaledBlocks);
+
+    struct TSingleChannelElement {
+        TAtrac3Data::SubbandInfo SubbandInfo;
+        std::vector<TTonalBlock> TonalBlocks;
+        std::vector<TScaledBlock> ScaledBlocks;
+    };
+
+    void WriteSoundUnit(const std::vector<TSingleChannelElement>& singleChannelElements);
 };
 
 } // namespace NAtrac3
diff --git a/src/atrac3denc.cpp b/src/atrac3denc.cpp
index 3ca00bc..1a394f0 100644
--- a/src/atrac3denc.cpp
+++ b/src/atrac3denc.cpp
@@ -94,6 +94,7 @@ TAtrac3Processor::TAtrac3Processor(TCompressedIOPtr&& oma, TAtrac3EncoderSetting
     : Oma(std::move(oma))
     , Params(std::move(encoderSettings))
     , TransientDetectors(2 * 4, TTransientDetector(8, 256)) //2 - channels, 4 - bands
+    , SingleChannelElements(Params.SourceChannels)
 {}
 
 TAtrac3Processor::~TAtrac3Processor()
@@ -164,9 +165,8 @@ TAtrac3Data::TTonalComponents TAtrac3Processor::ExtractTonalComponents(TFloat* s
     return res;
 }
 
-std::vector<TTonalComponent> TAtrac3Processor::MapTonalComponents(const TTonalComponents& tonalComponents)
+void TAtrac3Processor::MapTonalComponents(const TTonalComponents& tonalComponents, vector<TTonalBlock>* componentMap)
 {
-    vector<TTonalComponent> componentMap;
     for (uint16_t i = 0; i < tonalComponents.size();) {
         const uint16_t startPos = i;
         uint16_t curPos;
@@ -179,9 +179,8 @@ std::vector<TTonalComponent> TAtrac3Processor::MapTonalComponents(const TTonalCo
         for (uint8_t j = 0; j < len; ++j)
             tmp[j] = tonalComponents[startPos + j].Val;
         const TScaledBlock& scaledBlock = Scaler.Scale(tmp, len);
-        componentMap.push_back({&tonalComponents[startPos], 7, scaledBlock});
+        componentMap->push_back({&tonalComponents[startPos], 7, scaledBlock});
     }
-    return componentMap;
 }
 
 
@@ -240,11 +239,11 @@ TAtrac3Processor::TTransientParam TAtrac3Processor::CalcTransientParam(const std
     return {attackLocation, attackRelation, releaseLocation, releaseRelation};
 }
 
-TAtrac3Data::SubbandInfo TAtrac3Processor::CreateSubbandInfo(TFloat* in[4],
-                                                             uint32_t channel,
-                                                             TTransientDetector* transientDetector)
+void TAtrac3Processor::CreateSubbandInfo(TFloat* in[4],
+                                         uint32_t channel,
+                                         TTransientDetector* transientDetector,
+                                         TAtrac3Data::SubbandInfo* subbandInfo)
 {
-    TAtrac3Data::SubbandInfo siCur;
     for (int band = 0; band < 4; ++band) {
 
         TFloat invBuf[256];
@@ -308,11 +307,10 @@ TAtrac3Data::SubbandInfo TAtrac3Processor::CreateSubbandInfo(TFloat* in[4],
         }
 
         if (hasTransient) {
-            siCur.AddSubbandCurve(band, std::move(curve));
+            subbandInfo->AddSubbandCurve(band, std::move(curve));
         }
 
     }
-    return siCur;
 }
 
 
@@ -326,24 +324,35 @@ TPCMEngine<TFloat>::TProcessLambda TAtrac3Processor::GetEncodeLambda()
 
     TAtrac3BitStreamWriter* bitStreamWriter = new TAtrac3BitStreamWriter(omaptr, *Params.ConteinerParams);
     return [this, bitStreamWriter](TFloat* data, const TPCMEngine<TFloat>::ProcessMeta& meta) {
-        for (uint32_t channel=0; channel < 2; channel++) {
+        using TSce = TAtrac3BitStreamWriter::TSingleChannelElement;
+
+        // TTonalBlock has pointer to the TTonalVal so TTonalComponents must be avaliable
+        // TODO: this code should be rewritten
+        TTonalComponents tonals[2];
+
+        assert(SingleChannelElements.size() == meta.Channels);
+        for (uint32_t channel = 0; channel < SingleChannelElements.size(); channel++) {
             vector<TFloat> specs(1024);
             TFloat src[NumSamples];
 
             for (size_t i = 0; i < NumSamples; ++i) {
-                src[i] = data[meta.Channels == 1 ? i : (i * 2 + channel)] / 4.0; //no mono mode in atrac3. //TODO we can TFloat frame after encoding
+                src[i] = data[i * meta.Channels  + channel] / 4.0;
             }
 
             TFloat* p[4] = {&PcmBuffer[channel][0][0], &PcmBuffer[channel][1][0], &PcmBuffer[channel][2][0], &PcmBuffer[channel][3][0]};
             SplitFilterBank[channel].Split(&src[0], p);
             
-            TAtrac3Data::SubbandInfo siCur = Params.NoGainControll ?
-                TAtrac3Data::SubbandInfo() : CreateSubbandInfo(p, channel, &TransientDetectors[channel*4]); //4 detectors per band
+            TSce* sce = &SingleChannelElements[channel];
+
+            sce->SubbandInfo.Reset();
+            if (!Params.NoGainControll) {
+                CreateSubbandInfo(p, channel, &TransientDetectors[channel*4], &sce->SubbandInfo); //4 detectors per band
+            }
 
             TFloat* maxOverlapLevels = PrevPeak[channel];
 
-            Mdct(specs.data(), p, maxOverlapLevels, MakeGainModulatorArray(siCur));
-            TTonalComponents tonals = Params.NoTonalComponents ? 
+            Mdct(specs.data(), p, maxOverlapLevels, MakeGainModulatorArray(sce->SubbandInfo));
+            tonals[channel] = Params.NoTonalComponents ?
                     TAtrac3Data::TTonalComponents() : ExtractTonalComponents(specs.data(), [](const TFloat* spec, uint16_t len) {
                 std::vector<TFloat> magnitude(len);
                 for (uint16_t i = 0; i < len; ++i) {
@@ -358,11 +367,15 @@ TPCMEngine<TFloat>::TProcessLambda TAtrac3Processor::GetEncodeLambda()
                 return NAN;
             });
 
-            const std::vector<TTonalComponent>& components = MapTonalComponents(tonals);
+            sce->TonalBlocks.clear();
+            MapTonalComponents(tonals[channel], &sce->TonalBlocks);
 
             //TBlockSize for ATRAC3 - 4 subband, all are long (no short window)
-            bitStreamWriter->WriteSoundUnit(siCur, components, Scaler.ScaleFrame(specs, TBlockSize())); 
+            sce->ScaledBlocks = std::move(Scaler.ScaleFrame(specs, TBlockSize()));
+
         }
+
+        bitStreamWriter->WriteSoundUnit(SingleChannelElements);
     };
 }
 
diff --git a/src/atrac3denc.h b/src/atrac3denc.h
index 1f936c6..c77c4ab 100644
--- a/src/atrac3denc.h
+++ b/src/atrac3denc.h
@@ -75,6 +75,7 @@ class TAtrac3Processor : public IProcessor<TFloat>, public TAtrac3MDCT {
     Atrac3SplitFilterBank<TFloat> SplitFilterBank[2];
     TScaler<TAtrac3Data> Scaler;
     std::vector<TTransientDetector> TransientDetectors;
+    std::vector<NAtrac3::TAtrac3BitStreamWriter::TSingleChannelElement> SingleChannelElements;
     typedef std::array<uint8_t, NumSpecs> TonalComponentMask;
 public:
     struct TTransientParam {
@@ -90,11 +91,13 @@ public:
 #endif
     TFloat LimitRel(TFloat x);
     TTransientParam CalcTransientParam(const std::vector<TFloat>& gain, TFloat lastMax);
-    TAtrac3Data::SubbandInfo CreateSubbandInfo(TFloat* in[4], uint32_t channel, TTransientDetector* transientDetector);
+    void CreateSubbandInfo(TFloat* in[4], uint32_t channel,
+                           TTransientDetector* transientDetector,
+                           TAtrac3Data::SubbandInfo* subbandInfo);
     TonalComponentMask AnalyzeTonalComponent(TFloat* specs);
     TTonalComponents ExtractTonalComponents(TFloat* specs, TTonalDetector fn);
 
-    std::vector<NAtrac3::TTonalComponent> MapTonalComponents(const TTonalComponents& tonalComponents);
+    void MapTonalComponents(const TTonalComponents& tonalComponents, std::vector<NAtrac3::TTonalBlock>* componentMap);
 public:
     TAtrac3Processor(TCompressedIOPtr&& oma, NAtrac3::TAtrac3EncoderSettings&& encoderSettings);
     ~TAtrac3Processor();
diff --git a/src/main.cpp b/src/main.cpp
index 58f78cd..b74dea1 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -112,6 +112,13 @@ static void CheckInputFormat(const TWav* p)
         throw std::runtime_error("unsupported sample rate");
 }
 
+static TWavPtr OpenWavFile(const string& inFile)
+{
+    TWav* wavPtr = new TWav(inFile);
+    CheckInputFormat(wavPtr);
+    return TWavPtr(wavPtr);
+}
+
 static void PrepareAtrac1Encoder(const string& inFile,
                                  const string& outFile, 
                                  const bool noStdOut, 
@@ -173,20 +180,15 @@ static void PrepareAtrac3Encoder(const string& inFile,
                                  const bool noStdOut,
                                  NAtrac3::TAtrac3EncoderSettings&& encoderSettings,
                                  uint64_t* totalSamples,
-                                 TWavPtr* wavIO,
+                                 const TWavPtr& wavIO,
                                  TPcmEnginePtr* pcmEngine,
                                  TAtracProcessorPtr* atracProcessor)
 {
     std::cout << "WARNING: ATRAC3 is uncompleted, result will be not good )))" << std::endl;
     if (!noStdOut)
         std::cout << "bitrate " << encoderSettings.ConteinerParams->Bitrate << std::endl;
-    {
-        TWav* wavPtr = new TWav(inFile);
-        CheckInputFormat(wavPtr);
-        wavIO->reset(wavPtr);
-    }
-    const int numChannels = (*wavIO)->GetChannelNum();
-    *totalSamples = (*wavIO)->GetTotalSamples();
+    const int numChannels = encoderSettings.SourceChannels;
+    *totalSamples = wavIO->GetTotalSamples();
     TCompressedIOPtr omaIO = TCompressedIOPtr(new TOma(outFile,
                                                        "test",
                                                        numChannels,
@@ -194,7 +196,7 @@ static void PrepareAtrac3Encoder(const string& inFile,
                                                        encoderSettings.ConteinerParams->FrameSz));
     pcmEngine->reset(new TPCMEngine<TFloat>(4096,
                                             numChannels,
-                                            TPCMEngine<TFloat>::TReaderPtr((*wavIO)->GetPCMReader<TFloat>())));
+                                            TPCMEngine<TFloat>::TReaderPtr(wavIO->GetPCMReader<TFloat>())));
     atracProcessor->reset(new TAtrac3Processor(std::move(omaIO), std::move(encoderSettings)));
 }
 
@@ -333,9 +335,10 @@ int main(int argc, char* const* argv)
             case (E_ENCODE | E_ATRAC3):
             {
                 using NAtrac3::TAtrac3Data;
-                NAtrac3::TAtrac3EncoderSettings encoderSettings(bitrate * 1024, noGainControl, noTonalComponents);
+                wavIO = OpenWavFile(inFile);
+                NAtrac3::TAtrac3EncoderSettings encoderSettings(bitrate * 1024, noGainControl, noTonalComponents, wavIO->GetChannelNum());
                 PrepareAtrac3Encoder(inFile, outFile, noStdOut, std::move(encoderSettings),
-                &totalSamples, &wavIO, &pcmEngine, &atracProcessor);
+                &totalSamples, wavIO, &pcmEngine, &atracProcessor);
                 pcmFrameSz = TAtrac3Data::NumSamples;;
             }
             break;