1/*
2 * Copyright (C) 2010 The Android Open Source Project
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *      http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17#include "ABitReader.h"
18
19#include <media/stagefright/foundation/ADebug.h>
20
21namespace android {
22
23ABitReader::ABitReader(const uint8_t *data, size_t size)
24    : mData(data),
25      mSize(size),
26      mReservoir(0),
27      mNumBitsLeft(0),
28      mOverRead(false) {
29}
30
31ABitReader::~ABitReader() {
32}
33
34bool ABitReader::fillReservoir() {
35    if (mSize == 0) {
36        mOverRead = true;
37        return false;
38    }
39
40    mReservoir = 0;
41    size_t i;
42    for (i = 0; mSize > 0 && i < 4; ++i) {
43        mReservoir = (mReservoir << 8) | *mData;
44
45        ++mData;
46        --mSize;
47    }
48
49    mNumBitsLeft = 8 * i;
50    mReservoir <<= 32 - mNumBitsLeft;
51    return true;
52}
53
54uint32_t ABitReader::getBits(size_t n) {
55    uint32_t ret;
56    CHECK(getBitsGraceful(n, &ret));
57    return ret;
58}
59
60uint32_t ABitReader::getBitsWithFallback(size_t n, uint32_t fallback) {
61    uint32_t ret = fallback;
62    (void)getBitsGraceful(n, &ret);
63    return ret;
64}
65
66bool ABitReader::getBitsGraceful(size_t n, uint32_t *out) {
67    if (n > 32) {
68        return false;
69    }
70
71    uint32_t result = 0;
72    while (n > 0) {
73        if (mNumBitsLeft == 0) {
74            if (!fillReservoir()) {
75                return false;
76            }
77        }
78
79        size_t m = n;
80        if (m > mNumBitsLeft) {
81            m = mNumBitsLeft;
82        }
83
84        result = (result << m) | (mReservoir >> (32 - m));
85        mReservoir <<= m;
86        mNumBitsLeft -= m;
87
88        n -= m;
89    }
90
91    *out = result;
92    return true;
93}
94
95bool ABitReader::skipBits(size_t n) {
96    uint32_t dummy;
97    while (n > 32) {
98        if (!getBitsGraceful(32, &dummy)) {
99            return false;
100        }
101        n -= 32;
102    }
103
104    if (n > 0) {
105        return getBitsGraceful(n, &dummy);
106    }
107    return true;
108}
109
110void ABitReader::putBits(uint32_t x, size_t n) {
111    if (mOverRead) {
112        return;
113    }
114
115    CHECK_LE(n, 32u);
116
117    while (mNumBitsLeft + n > 32) {
118        mNumBitsLeft -= 8;
119        --mData;
120        ++mSize;
121    }
122
123    mReservoir = (mReservoir >> n) | (x << (32 - n));
124    mNumBitsLeft += n;
125}
126
127size_t ABitReader::numBitsLeft() const {
128    return mSize * 8 + mNumBitsLeft;
129}
130
131const uint8_t *ABitReader::data() const {
132    return mData - (mNumBitsLeft + 7) / 8;
133}
134
135NALBitReader::NALBitReader(const uint8_t *data, size_t size)
136    : ABitReader(data, size),
137      mNumZeros(0) {
138}
139
140bool NALBitReader::atLeastNumBitsLeft(size_t n) const {
141    // check against raw size and reservoir bits first
142    size_t numBits = numBitsLeft();
143    if (n > numBits) {
144        return false;
145    }
146
147    ssize_t numBitsRemaining = (ssize_t)n - (ssize_t)mNumBitsLeft;
148
149    size_t size = mSize;
150    const uint8_t *data = mData;
151    int32_t numZeros = mNumZeros;
152    while (size > 0 && numBitsRemaining > 0) {
153        bool isEmulationPreventionByte = (numZeros >= 2 && *data == 3);
154
155        if (*data == 0) {
156            ++numZeros;
157        } else {
158            numZeros = 0;
159        }
160
161        if (!isEmulationPreventionByte) {
162            numBitsRemaining -= 8;
163        }
164
165        ++data;
166        --size;
167    }
168
169    return (numBitsRemaining <= 0);
170}
171
172bool NALBitReader::fillReservoir() {
173    if (mSize == 0) {
174        mOverRead = true;
175        return false;
176    }
177
178    mReservoir = 0;
179    size_t i = 0;
180    while (mSize > 0 && i < 4) {
181        bool isEmulationPreventionByte = (mNumZeros >= 2 && *mData == 3);
182
183        if (*mData == 0) {
184            ++mNumZeros;
185        } else {
186            mNumZeros = 0;
187        }
188
189        // skip emulation_prevention_three_byte
190        if (!isEmulationPreventionByte) {
191            mReservoir = (mReservoir << 8) | *mData;
192            ++i;
193        }
194
195        ++mData;
196        --mSize;
197    }
198
199    mNumBitsLeft = 8 * i;
200    mReservoir <<= 32 - mNumBitsLeft;
201    return true;
202}
203
204}  // namespace android
205