ans/bitstream.c at master · static-code-generators/ans · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
/**
 * @file
 * Bitstream lib
 */

#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <stdarg.h>
#include "bitstream.h"

static int fpeek(FILE *stream)
{
    int c;
    c = fgetc(stream);
    ungetc(c, stream);
    return c;
}

struct bitstream * initbitstream(FILE *fp, enum bs_type type, ...)
{
    struct bitstream *bs = malloc(sizeof(struct bitstream));
    bs->fp = fp;
    bs->buf = 0;
    bs->bufoffset = 0;
    bs->type = type;
    bs->last = 0;
    va_list ap;
    va_start(ap, type);
    if (type == BS_READ) {
        bs->padding = va_arg(ap, int);
    }
    va_end(ap);
    return bs;
}

void write_bit(struct bitstream *bs, char c)
{
    assert(c == '0' || c == '1');
    assert(bs->type == BS_WRITE);
    assert(bs->bufoffset < 8);
    /* Write bit to buffer */
    bs->buf = (bs->buf << 1) | ((int) (c - '0'));
    bs->bufoffset += 1;
    /* If buffer full, write it to file */
    if (bs->bufoffset == 8) {
        fwrite(&bs->buf, sizeof(uint8_t), 1, bs->fp);
        bs->buf = 0;
        bs->bufoffset = 0;
    }
}

int read_bit(struct bitstream *bs)
{
    assert(bs->type == BS_READ);
    /* If buffer empty, read from file */
    if (bs->bufoffset == 0) {
        fread(&bs->buf, sizeof(uint8_t), 1, bs->fp);
        bs->bufoffset = 8;
        if (fpeek(bs->fp) == EOF) {
            bs->last = 1;
        }
    }
    /* Check if last byte */
    if (bs->last) {
        /* Don't read padding bits */
        if (bs->bufoffset <= bs->padding) {
            return -1;
        }
    }
    int bit = bs->buf >> (bs->bufoffset - 1);
    bs->buf &= ~(1 << (bs->bufoffset - 1));
    bs->bufoffset -= 1;
    assert(bit == 1 || bit == 0);
    return bit;
}

uint8_t read_byte(struct bitstream *bs)
{
    /* Read next byte from file */
    fread(&bs->buf, sizeof(uint8_t), 1, bs->fp);
    bs->bufoffset = 0;
    if (fpeek(bs->fp) == EOF) {
        bs->last = 1;
    }
    return bs->buf;
}

uint64_t read_nbits(struct bitstream *bs, uint8_t n)
{
    assert(n <= 64);
    uint64_t ret = 0;
    int bit;
    while (n--) {
        bit = read_bit(bs);
        if (bit == -1) {
            return 0;
        }
        ret <<= 1;
        ret |= bit;
    }
    return ret;
}

void write_bitstring(struct bitstream *bs, char *str)
{
    while (*str) {
        write_bit(bs, *str);
        str++;
    }
}

int closebitstream(struct bitstream *bs)
{
    if (bs->type == BS_WRITE) {
        int padding = 0;
        if (bs->bufoffset > 0) {
            padding = 8 - bs->bufoffset;
            bs->buf <<= padding;
            fwrite(&bs->buf, sizeof(uint8_t), 1, bs->fp);
        }
        free(bs);
        return padding;
    } else {
        free(bs);
        return 0;
    }
}

#ifdef RUN_MAIN
int main(int argc, char **argv)
{
    FILE *fp = fopen(argv[1], "rb");
    struct bitstream *bs = initbitstream(fp, BS_READ, 0);
    uint8_t byte;
    while (bs->last != 1) {
        byte = read_byte(bs);
        printf("%02x ", byte);
    }
    return 0;
}
#endif