libavcenc: update bitstream context state only in successful writes

During entropy error, even though the bitstream buffer is not written
with additional bytes, the bitstream context state is getting updated.
As the number of bits left in current word is updated, the get_num_bits
used for computing header and texture bits can become negative causing
overflow errors

Change-Id: I2f990071a9935b2ee328dbd3915dfbefccbab4c5

3 files changed, 18 insertions, 15 deletions

diff --git a/encoder/ih264e_bitstream.c b/encoder/ih264e_bitstream.c
index 9f73f69..ef8dc3d 100644
--- a/encoder/ih264e_bitstream.c
+++ b/encoder/ih264e_bitstream.c
@@ -182,7 +182,6 @@ IH264E_ERROR_T ih264e_put_bits(bitstrm_t *ps_bitstrm,
         /* 4. insert remaining bits of code starting from msb of cur word   */
         /* 5. update bitsleft in current word and stream buffer offset      */
         /********************************************************************/
-        IH264E_ERROR_T status = IH264E_SUCCESS;
         WORD32  i, rem_bits = (code_len - bits_left_in_cw);
 
         /* insert parital code corresponding to bits left in cur word */
@@ -193,7 +192,8 @@ IH264E_ERROR_T ih264e_put_bits(bitstrm_t *ps_bitstrm,
             /* flush the bits in cur word byte by byte and copy to stream */
             UWORD8   u1_next_byte = (u4_cur_word >> (i-8)) & 0xFF;
 
-            status |= ih264e_put_byte_epb(ps_bitstrm, u1_next_byte);
+            IH264E_ERROR_T status = ih264e_put_byte_epb(ps_bitstrm, u1_next_byte);
+            if (status != IH264E_SUCCESS) return status;
         }
 
         /* insert the remaining bits from code val into current word */
@@ -202,7 +202,7 @@ IH264E_ERROR_T ih264e_put_bits(bitstrm_t *ps_bitstrm,
         /* update the state variables and return success */
         ps_bitstrm->u4_cur_word         = u4_cur_word;
         ps_bitstrm->i4_bits_left_in_cw  = WORD_SIZE - rem_bits;
-        return (status);
+        return (IH264E_SUCCESS);
 
     }
 }
@@ -262,7 +262,6 @@ IH264E_ERROR_T ih264e_put_rbsp_trailing_bits(bitstrm_t *ps_bitstrm)
     UWORD32 u4_cur_word = ps_bitstrm->u4_cur_word;
     WORD32  bits_left_in_cw = ps_bitstrm->i4_bits_left_in_cw;
     WORD32  bytes_left_in_cw = (bits_left_in_cw - 1) >> 3;
-    IH264E_ERROR_T status = IH264E_SUCCESS;
 
     /* insert a 1 at the end of current word and flush all the bits */
     u4_cur_word |= (1 << (bits_left_in_cw - 1));
@@ -275,7 +274,8 @@ IH264E_ERROR_T ih264e_put_rbsp_trailing_bits(bitstrm_t *ps_bitstrm)
         /* flush the bits in cur word byte by byte  and copy to stream */
         UWORD8   u1_next_byte = (u4_cur_word >> (i-8)) & 0xFF;
 
-        status |= ih264e_put_byte_epb(ps_bitstrm, u1_next_byte);
+        IH264E_ERROR_T status = ih264e_put_byte_epb(ps_bitstrm, u1_next_byte);
+        if (status != IH264E_SUCCESS) return status;
     }
 
     /* Default init values for scratch variables of bitstream context */
@@ -283,7 +283,7 @@ IH264E_ERROR_T ih264e_put_rbsp_trailing_bits(bitstrm_t *ps_bitstrm)
     ps_bitstrm->i4_bits_left_in_cw  = WORD_SIZE;
     ps_bitstrm->i4_zero_bytes_run   = 0;
 
-    return (status);
+    return (IH264E_SUCCESS);
 }
 
 /**
diff --git a/encoder/ih264e_bitstream.h b/encoder/ih264e_bitstream.h
index 4f592f3..58bfc5f 100644
--- a/encoder/ih264e_bitstream.h
+++ b/encoder/ih264e_bitstream.h
@@ -265,10 +265,11 @@ static inline IH264E_ERROR_T ih264e_put_byte_epb(bitstrm_t *ps_bitstrm, UWORD8 b
        UWORD8 u1_next_byte = (ps_bitstrm->u4_cur_word >> (i - 8)) & 0xFF;      \
        err |= ih264e_put_byte_epb(ps_bitstrm, u1_next_byte);                   \
     }                                                                          \
-    ps_bitstrm->u4_cur_word = 0;                                               \
-    ps_bitstrm->i4_bits_left_in_cw = WORD_SIZE;                                \
-}                                                                              \
-
+    if (err == IH264E_SUCCESS) {                                               \
+        ps_bitstrm->u4_cur_word = 0;                                           \
+        ps_bitstrm->i4_bits_left_in_cw = WORD_SIZE;                            \
+    }                                                                          \
+}
 
 
 /*****************************************************************************/
diff --git a/encoder/ih264e_cabac.c b/encoder/ih264e_cabac.c
index 2d91058..fb54475 100644
--- a/encoder/ih264e_cabac.c
+++ b/encoder/ih264e_cabac.c
@@ -315,11 +315,13 @@ IH264E_ERROR_T ih264e_cabac_flush(cabac_ctxt_t *ps_cabac_ctxt)
         last_byte &= 0xFF;
         status |= ih264e_put_byte_epb(ps_stream, last_byte);
 
-        /* update the state variables and return success */
-        ps_stream->i4_zero_bytes_run = 0;
-        /* Default init values for scratch variables of bitstream context */
-        ps_stream->u4_cur_word = 0;
-        ps_stream->i4_bits_left_in_cw = WORD_SIZE;
+        if (status == IH264E_SUCCESS) {
+            /* update the state variables and return success */
+            ps_stream->i4_zero_bytes_run = 0;
+            /* Default init values for scratch variables of bitstream context */
+            ps_stream->u4_cur_word = 0;
+            ps_stream->i4_bits_left_in_cw = WORD_SIZE;
+        }
 
     }
     return status;