[x265] [PATCH] Modify TEncEntropy structure to support multiple color space formats

[ashok at multicorewareinc.com]

# HG changeset patch
# User ashok at multicorewareinc.com
# Date 1388754397 -19800
#      Fri Jan 03 18:36:37 2014 +0530
# Node ID be8c07a80a181f9a20d51abd01b90ac9fe94f4ea
# Parent  e4ab306e0347e5a65c52be2d0845a6cf592c5713
Modify TEncEntropy structure to support multiple color space formats

diff -r e4ab306e0347 -r be8c07a80a18 source/Lib/TLibEncoder/TEncEntropy.cpp
--- a/source/Lib/TLibEncoder/TEncEntropy.cpp	Fri Jan 03 18:35:01 2014 +0530
+++ b/source/Lib/TLibEncoder/TEncEntropy.cpp	Fri Jan 03 18:36:37 2014 +0530
@@ -225,7 +225,8 @@
     {
         m_bakAbsPartIdxCU = absPartIdx;
     }
-    if (log2TrafoSize == 2)
+
+    if ((log2TrafoSize == 2) && !(cu->getChromaFormat() == CHROMA_444))
     {
         uint32_t partNum = cu->getPic()->getNumPartInCU() >> ((depth - 1) << 1);
         if ((absPartIdx % partNum) == 0)
@@ -275,7 +276,15 @@
 
     const uint32_t trDepthCurr = depth - cu->getDepth(absPartIdx);
     const bool bFirstCbfOfCU = trDepthCurr == 0;
-    if (bFirstCbfOfCU || log2TrafoSize > 2)
+
+    bool mCodeAll = true;
+    const uint32_t numPels = (width >> cu->getHorzChromaShift()) * (height >> cu->getHorzChromaShift());
+    if(numPels < (MIN_TU_SIZE * MIN_TU_SIZE))
+    {
+        mCodeAll = false;
+    }
+
+    if (bFirstCbfOfCU || mCodeAll)
     {
         if (bFirstCbfOfCU || cu->getCbf(absPartIdx, TEXT_CHROMA_U, trDepthCurr - 1))
         {
@@ -286,7 +295,7 @@
             m_entropyCoderIf->codeQtCbf(cu, absPartIdx, TEXT_CHROMA_V, trDepthCurr);
         }
     }
-    else if (log2TrafoSize == 2)
+    else
     {
         assert(cu->getCbf(absPartIdx, TEXT_CHROMA_U, trDepthCurr) == cu->getCbf(absPartIdx, TEXT_CHROMA_U, trDepthCurr - 1));
         assert(cu->getCbf(absPartIdx, TEXT_CHROMA_V, trDepthCurr) == cu->getCbf(absPartIdx, TEXT_CHROMA_V, trDepthCurr - 1));
@@ -306,17 +315,17 @@
 
         absPartIdx += partNum;
         offsetLuma += size;
-        offsetChroma += (size >> 2);
+        offsetChroma += (size >> (cu->getHorzChromaShift() + cu->getVertChromaShift()));
         xEncodeTransform(cu, offsetLuma, offsetChroma, absPartIdx, depth, width, height, trIdx, bCodeDQP);
 
         absPartIdx += partNum;
         offsetLuma += size;
-        offsetChroma += (size >> 2);
+        offsetChroma += (size >> (cu->getHorzChromaShift() + cu->getVertChromaShift()));
         xEncodeTransform(cu, offsetLuma, offsetChroma, absPartIdx, depth, width, height, trIdx, bCodeDQP);
 
         absPartIdx += partNum;
         offsetLuma += size;
-        offsetChroma += (size >> 2);
+        offsetChroma += (size >> (cu->getHorzChromaShift() + cu->getVertChromaShift()));
         xEncodeTransform(cu, offsetLuma, offsetChroma, absPartIdx, depth, width, height, trIdx, bCodeDQP);
     }
     else
@@ -358,20 +367,8 @@
         {
             m_entropyCoderIf->codeCoeffNxN(cu, (cu->getCoeffY() + offsetLuma), absPartIdx, width, height, depth, TEXT_LUMA);
         }
-        if (log2TrafoSize > 2)
-        {
-            int trWidth = width >> 1;
-            int trHeight = height >> 1;
-            if (cbfU)
-            {
-                m_entropyCoderIf->codeCoeffNxN(cu, (cu->getCoeffCb() + offsetChroma), absPartIdx, trWidth, trHeight, depth, TEXT_CHROMA_U);
-            }
-            if (cbfV)
-            {
-                m_entropyCoderIf->codeCoeffNxN(cu, (cu->getCoeffCr() + offsetChroma), absPartIdx, trWidth, trHeight, depth, TEXT_CHROMA_V);
-            }
-        }
-        else
+
+        if ((log2TrafoSize == 2) && !(cu->getChromaFormat() == CHROMA_444))
         {
             uint32_t partNum = cu->getPic()->getNumPartInCU() >> ((depth - 1) << 1);
             if ((absPartIdx % partNum) == (partNum - 1))
@@ -386,6 +383,19 @@
                 }
             }
         }
+        else
+        {
+            int trWidth  = width >> cu->getHorzChromaShift();
+            int trHeight = height >> cu->getVertChromaShift();
+            if (cbfU)
+            {
+                m_entropyCoderIf->codeCoeffNxN(cu, (cu->getCoeffCb() + offsetChroma), absPartIdx, trWidth, trHeight, depth, TEXT_CHROMA_U);
+            }
+            if (cbfV)
+            {
+                m_entropyCoderIf->codeCoeffNxN(cu, (cu->getCoeffCr() + offsetChroma), absPartIdx, trWidth, trHeight, depth, TEXT_CHROMA_V);
+            }
+        }
     }
 }
 
@@ -412,10 +422,22 @@
     {
         absPartIdx = 0;
     }
+
     if (cu->isIntra(absPartIdx)) // If it is Intra mode, encode intra prediction mode.
     {
         encodeIntraDirModeLuma(cu, absPartIdx, true);
-        encodeIntraDirModeChroma(cu, absPartIdx, bRD);
+        if (cu->getChromaFormat() != CHROMA_400)
+        {
+            encodeIntraDirModeChroma(cu, absPartIdx, bRD);
+
+            if ((cu->getChromaFormat() == CHROMA_444) && (cu->getPartitionSize(absPartIdx) == SIZE_NxN))
+            {
+                uint32_t partOffset = (cu->getPic()->getNumPartInCU() >> (cu->getDepth(absPartIdx) << 1)) >> 2;
+                encodeIntraDirModeChroma(cu, absPartIdx + partOffset, bRD);
+                encodeIntraDirModeChroma(cu, absPartIdx + partOffset*2, bRD);
+                encodeIntraDirModeChroma(cu, absPartIdx + partOffset*3, bRD);
+            }
+        }
     }
     else                        // if it is Inter mode, encode motion vector and reference index
     {
@@ -572,7 +594,7 @@
 {
     uint32_t minCoeffSize = cu->getPic()->getMinCUWidth() * cu->getPic()->getMinCUHeight();
     uint32_t lumaOffset   = minCoeffSize * absPartIdx;
-    uint32_t chromaOffset = lumaOffset >> 2;
+    uint32_t chromaOffset = lumaOffset >> (cu->getHorzChromaShift() + cu->getVertChromaShift());
 
     if (cu->isIntra(absPartIdx))
     {