Merge pull request #916 from huili2/dynamic_dec_au_size

enable dynamic AU size in decoder

codec/decoder/core/inc/decoder_context.h

@@ -229,6 +229,7 @@ typedef struct TagWelsDecoderContext {
   SVlcTable			sVlcTable;		 // vlc table
 
   SBitStringAux		sBs;
+  int32_t iMaxBsBufferSizeInByte; //actual memory size for BS buffer
 
   /* Global memory external */
 

codec/decoder/core/inc/decoder_core.h

@@ -48,6 +48,28 @@
 #include "codec_def.h"
 
 namespace WelsDec {
+/*
+ * InitBsBuffer
+ * Memory allocation for Bitstream Buffer
+ * return:
+ *	0 - success; otherwise returned error_no defined in error_no.h.
+ */
+int32_t InitBsBuffer (PWelsDecoderContext pCtx);
+
+/*
+ * ExpandBsBuffer
+ * Expand current BS buffer and copy its content
+ * new buffer size will consider input size as a reference
+ * return:
+ *	0 - success; otherwise returned error_no defined in error_no.h.
+ */
+int32_t ExpandBsBuffer (PWelsDecoderContext pCtx, const int32_t kiSrcLen);
+
+/*
+ * CheckBsBuffer
+ * Check if current buffer size is enough
+ */
+int32_t CheckBsBuffer (PWelsDecoderContext pCtx, const int32_t kiSrcLen);
 
 /*
  * WelsInitMemory

codec/decoder/core/inc/wels_const.h

@@ -91,8 +91,9 @@
 #define LAYER_NUM_EXCHANGEABLE	1
 
 #define MAX_NAL_UNIT_NUM_IN_AU	32	// predefined maximal number of NAL Units in an access unit
-#define MAX_ACCESS_UNIT_CAPACITY	1048576	// Maximal AU capacity in bytes: (1<<20) = 1024 KB predefined
-#define BS_BUFFER_SIZE  (MAX_ACCESS_UNIT_CAPACITY * 3) //for delay case, keep three AU size to prevent buffer overwrite
+#define MIN_ACCESS_UNIT_CAPACITY	1048576	// Min AU capacity in bytes: (1<<20) = 1024 KB predefined
+#define MAX_BUFFERED_NUM 3 //mamixum stored number of AU|packet to prevent overwrite
+#define MAX_ACCESS_UNIT_CAPACITY 7077888 //Maximum AU size in bytes for level 5.2 for single frame
 #define MAX_MACROBLOCK_CAPACITY 5000 //Maximal legal MB capacity, 15000 bits is enough
 
 #endif//WELS_CONSTANCE_H__

codec/decoder/core/src/decoder_core.cpp

@@ -297,6 +297,57 @@ bool FillDefaultSliceHeaderExt (PSliceHeaderExt pShExt, PNalUnitHeaderExt pNalEx
   return true;
 }
 
+int32_t InitBsBuffer (PWelsDecoderContext pCtx) {
+  if (pCtx == NULL)
+    return ERR_INFO_INVALID_PTR;
+
+  pCtx->iMaxBsBufferSizeInByte = MIN_ACCESS_UNIT_CAPACITY * MAX_BUFFERED_NUM;
+  if ((pCtx->sRawData.pHead = static_cast<uint8_t*> (WelsMalloc (pCtx->iMaxBsBufferSizeInByte,
+                              "pCtx->sRawData.pHead"))) == NULL) {
+    return ERR_INFO_OUT_OF_MEMORY;
+  }
+  pCtx->sRawData.pStartPos = pCtx->sRawData.pCurPos = pCtx->sRawData.pHead;
+  pCtx->sRawData.pEnd = pCtx->sRawData.pHead + pCtx->iMaxBsBufferSizeInByte;
+  return ERR_NONE;
+}
+
+int32_t ExpandBsBuffer (PWelsDecoderContext pCtx, const int kiSrcLen) {
+  if (pCtx == NULL)
+    return ERR_INFO_INVALID_PTR;
+  int32_t iExpandStepShift = 1;
+  int32_t iNewBuffLen = WELS_MAX ((kiSrcLen * MAX_BUFFERED_NUM), (pCtx->iMaxBsBufferSizeInByte << iExpandStepShift));
+  //allocate new bs buffer
+  uint8_t* pNewBsBuff = static_cast<uint8_t*> (WelsMalloc (iNewBuffLen, "pCtx->sRawData.pHead"));
+  if (pNewBsBuff == NULL)
+    return ERR_INFO_OUT_OF_MEMORY;
+
+  //Copy current buffer status to new buffer
+  memcpy (pNewBsBuff, pCtx->sRawData.pHead, pCtx->iMaxBsBufferSizeInByte);
+  pCtx->iMaxBsBufferSizeInByte = iNewBuffLen;
+  pCtx->sRawData.pStartPos = pNewBsBuff + (pCtx->sRawData.pStartPos - pCtx->sRawData.pHead);
+  pCtx->sRawData.pCurPos = pNewBsBuff + (pCtx->sRawData.pCurPos - pCtx->sRawData.pHead);
+  pCtx->sRawData.pEnd = pNewBsBuff + iNewBuffLen;
+  WelsFree (pCtx->sRawData.pHead, "pCtx->sRawData.pHead");
+  pCtx->sRawData.pHead = pNewBsBuff;
+  return ERR_NONE;
+}
+
+int32_t CheckBsBuffer (PWelsDecoderContext pCtx, const int32_t kiSrcLen) {
+  if (kiSrcLen > MAX_ACCESS_UNIT_CAPACITY) { //exceeds max allowed data
+    WelsLog (pCtx, WELS_LOG_WARNING, "Max AU size exceeded. Allowed size = %d, current size = %d", MAX_ACCESS_UNIT_CAPACITY,
+             kiSrcLen);
+    pCtx->iErrorCode |= dsBitstreamError;
+    return ERR_INFO_INVALID_ACCESS;
+  } else if (kiSrcLen > pCtx->iMaxBsBufferSizeInByte /
+             MAX_BUFFERED_NUM) { //may lead to buffer overwrite, prevent it by expanding buffer
+    if (ExpandBsBuffer (pCtx, kiSrcLen)) {
+      return ERR_INFO_OUT_OF_MEMORY;
+    }
+  }
+
+  return ERR_NONE;
+}
+
 /*
  * WelsInitMemory
  * Memory request for new introduced data
@@ -313,13 +364,8 @@ int32_t WelsInitMemory (PWelsDecoderContext pCtx) {
   if (MemInitNalList (&pCtx->pAccessUnitList, MAX_NAL_UNIT_NUM_IN_AU) != 0)
     return ERR_INFO_OUT_OF_MEMORY;
 
-  if ((pCtx->sRawData.pHead = static_cast<uint8_t*> (WelsMalloc (BS_BUFFER_SIZE,
-                              "pCtx->sRawData->pHead"))) == NULL) {
+  if (InitBsBuffer (pCtx) != 0)
     return ERR_INFO_OUT_OF_MEMORY;
-  }
-  pCtx->sRawData.pStartPos               =
-    pCtx->sRawData.pCurPos                 = pCtx->sRawData.pHead;
-  pCtx->sRawData.pEnd                     = pCtx->sRawData.pHead + BS_BUFFER_SIZE;
 
   pCtx->uiTargetDqId			= (uint8_t) - 1;
   pCtx->bEndOfStreamFlag	= false;

codec/decoder/plus/src/welsDecoderExt.cpp

@@ -49,6 +49,7 @@
 
 //#include "macros.h"
 #include "decoder.h"
+#include "decoder_core.h"
 
 extern "C" {
 #include "decoder_core.h"
@@ -322,11 +323,7 @@ DECODING_STATE CWelsDecoder::DecodeFrame2 (const unsigned char* kpSrc,
     const int kiSrcLen,
     unsigned char** ppDst,
     SBufferInfo* pDstInfo) {
-  if (kiSrcLen > MAX_ACCESS_UNIT_CAPACITY - MAX_MACROBLOCK_CAPACITY) {//prevent from residual reading overflow
-    m_pDecContext->iErrorCode |= dsOutOfMemory;
-    IWelsTrace::WelsVTrace (m_pTrace, IWelsTrace::WELS_LOG_INFO,
-                            "max AU size exceeded. Allowed size = %d, current size = %d",
-                            MAX_ACCESS_UNIT_CAPACITY - MAX_MACROBLOCK_CAPACITY, kiSrcLen);
+  if (CheckBsBuffer (m_pDecContext, kiSrcLen)) {
     return dsOutOfMemory;
   }
   if (kiSrcLen > 0 && kpSrc != NULL) {