/* SPDX-License-Identifier: BSD-2-Clause */
#include "bitstream_writer.h"

#include <string.h>

void bsw_init(struct bs_writer *bs, uint8_t *buf, size_t cap)
{
	bs->buf = buf;
	bs->cap = cap;
	bs->pos_bytes = 0;
	bs->pos_bit = 0;
	bs->overflow = false;
	if (buf && cap)
		buf[0] = 0;
}

void bsw_put_u(struct bs_writer *bs, uint32_t v, int n)
{
	int i;

	if (n <= 0 || n > 32)
		return;

	for (i = n - 1; i >= 0; i--) {
		uint8_t bit = (uint8_t) ((v >> i) & 1u);

		if (bs->pos_bytes >= bs->cap) {
			bs->overflow = true;
			return;
		}

		if (bs->pos_bit == 0)
			bs->buf[bs->pos_bytes] = 0;

		bs->buf[bs->pos_bytes] |= (uint8_t) (bit << (7 - bs->pos_bit));
		bs->pos_bit++;
		if (bs->pos_bit == 8) {
			bs->pos_bit = 0;
			bs->pos_bytes++;
		}
	}
}

/*
 * Exp-Golomb ue(v) — H.264 9.1 / 9.1.1.
 * For an unsigned value v:
 *   codeNum = v
 *   leadingZeroBits = floor(log2(codeNum + 1))
 *   code = leadingZeroBits zeros, then '1', then leadingZeroBits bits
 *          of (codeNum + 1 - 2^leadingZeroBits)
 *
 * Total length = 2 * leadingZeroBits + 1 bits.  For v = 0 the
 * code is just "1" (1 bit).  For v in [1,2] the code is 3 bits, etc.
 */
void bsw_put_ue(struct bs_writer *bs, uint32_t v)
{
	uint32_t code_num = v;
	uint32_t code_num_plus_1 = code_num + 1u;
	int leading_zeros = 0;
	uint32_t tmp;

	tmp = code_num_plus_1 >> 1;
	while (tmp) {
		leading_zeros++;
		tmp >>= 1;
	}

	/* leading_zeros zero bits */
	bsw_put_u(bs, 0u, leading_zeros);
	/* one '1' bit + the (leading_zeros) low bits of code_num_plus_1 */
	bsw_put_u(bs, code_num_plus_1, leading_zeros + 1);
}

/*
 * Exp-Golomb se(v) — H.264 9.1.1: signed mapping is interleaved:
 *   0 -> 0, 1 -> 1, -1 -> 2, 2 -> 3, -2 -> 4, ...
 * i.e. codeNum = 2 * |v| - (v > 0 ? 1 : 0).
 */
void bsw_put_se(struct bs_writer *bs, int32_t v)
{
	uint32_t code_num;

	if (v > 0)
		code_num = (uint32_t) (2 * v - 1);
	else
		code_num = (uint32_t) (-2 * v);
	bsw_put_ue(bs, code_num);
}

void bsw_align_rbsp(struct bs_writer *bs)
{
	/* rbsp_stop_one_bit + zero-fill to byte boundary */
	bsw_put_u(bs, 1u, 1);
	while (bs->pos_bit != 0)
		bsw_put_u(bs, 0u, 1);
}

size_t bsw_bytes(const struct bs_writer *bs)
{
	return bs->pos_bytes;
}