Something went wrong. Try again.
Bitlevel streams for OCaml
Something went wrong. Try again.
11 kB · 385 lines
OCaml
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386(** Bitstream - Bit-level I/O for binary formats.
Provides forward and backward bitstream reading and writing for parsing and generating binary formats that operate at the bit level.
Forward streams read/write from the start of a buffer towards the end. Backward streams read/write from the end of a buffer towards the start, which is required by some compression algorithms (FSE, ANS). *)
(** {1 Slice Type} *)
module Slice = struct type t = { bytes : bytes; first : int; length : int; }
let make bytes ~first ~length = { bytes; first; length }
let of_bytes ?first ?length bytes = let first = Option.value first ~default:0 in let length = Option.value length ~default:(Bytes.length bytes - first) in { bytes; first; length }
let to_bytes t = Bytes.sub t.bytes t.first t.length
let is_empty t = t.length = 0
let sub t ~first ~length = { bytes = t.bytes; first = t.first + first; length }end
(** {1 Exceptions} *)
exception End_of_stream(** Raised when attempting to read past the end of the stream. *)
exception Invalid_state of string(** Raised when an operation requires a specific state (e.g., byte alignment). *)
exception Corrupted_stream of string(** Raised when stream data is malformed (e.g., invalid padding marker). *)
(** {1 Forward Bitstream Reader} *)
module Forward_reader = struct type t = { src : bytes; start_pos : int; limit : int; mutable byte_pos : int; mutable bit_pos : int; (* 0-7, bits consumed in current byte *) }
let of_slice (slice : Slice.t) = { src = slice.bytes; start_pos = slice.first; limit = slice.first + slice.length; byte_pos = slice.first; bit_pos = 0 }
let of_bytes src = of_slice (Slice.of_bytes src)
let create src ~pos ~len = of_slice (Slice.make src ~first:pos ~length:len)
let[@inline] remaining t = (t.limit - t.byte_pos) * 8 - t.bit_pos
let[@inline] is_byte_aligned t = t.bit_pos = 0
let[@inline] read_bits t n = if n <= 0 then 0 else if n > 57 then invalid_arg "read_bits: n > 57" else begin let result = ref 0 in let bits_read = ref 0 in while !bits_read < n do if t.byte_pos >= t.limit then raise End_of_stream; let byte = Bytes.get_uint8 t.src t.byte_pos in let available = 8 - t.bit_pos in let to_read = min available (n - !bits_read) in let mask = (1 lsl to_read) - 1 in let bits = (byte lsr t.bit_pos) land mask in result := !result lor (bits lsl !bits_read); bits_read := !bits_read + to_read; t.bit_pos <- t.bit_pos + to_read; if t.bit_pos >= 8 then begin t.bit_pos <- 0; t.byte_pos <- t.byte_pos + 1 end done; !result end
let[@inline] read_byte t = if t.bit_pos <> 0 then raise (Invalid_state "read_byte: not byte aligned"); if t.byte_pos >= t.limit then raise End_of_stream; let b = Bytes.get_uint8 t.src t.byte_pos in t.byte_pos <- t.byte_pos + 1; b
let rewind_bits t n = let total_bits = (t.byte_pos - t.start_pos) * 8 + t.bit_pos in let new_total = total_bits - n in if new_total < 0 then raise End_of_stream; t.byte_pos <- t.start_pos + new_total / 8; t.bit_pos <- new_total mod 8
let align t = if t.bit_pos <> 0 then begin t.bit_pos <- 0; t.byte_pos <- t.byte_pos + 1 end
let byte_position t = if t.bit_pos <> 0 then raise (Invalid_state "byte_position: not byte aligned"); t.byte_pos
let get_slice t n : Slice.t = if t.bit_pos <> 0 then raise (Invalid_state "get_slice: not byte aligned"); if t.byte_pos + n > t.limit then raise End_of_stream; let result = Slice.make t.src ~first:t.byte_pos ~length:n in t.byte_pos <- t.byte_pos + n; result
let get_bytes t n = Slice.to_bytes (get_slice t n)
let to_slice t : Slice.t = if t.bit_pos <> 0 then raise (Invalid_state "to_slice: not byte aligned"); Slice.make t.src ~first:t.byte_pos ~length:(t.limit - t.byte_pos)
let advance t n = if t.bit_pos <> 0 then raise (Invalid_state "advance: not byte aligned"); if t.byte_pos + n > t.limit then raise End_of_stream; t.byte_pos <- t.byte_pos + n
let sub t n = if t.bit_pos <> 0 then raise (Invalid_state "sub: not byte aligned"); if t.byte_pos + n > t.limit then raise End_of_stream; let result = of_slice (Slice.make t.src ~first:t.byte_pos ~length:n) in t.byte_pos <- t.byte_pos + n; result
let remaining_bytes t = if t.bit_pos <> 0 then raise (Invalid_state "remaining_bytes: not byte aligned"); t.limit - t.byte_pos
let skip_bits t n = ignore (read_bits t n)end
(** {1 Backward Bitstream Reader}
Reads bits from the end of a buffer towards the start. The stream starts with a padding marker (highest 1-bit indicates start of data). *)
module Backward_reader = struct type t = { src : bytes; start_pos : int; mutable bit_offset : int; (* Bits remaining from end, decreasing *) }
let of_slice (slice : Slice.t) = if slice.length = 0 then raise End_of_stream; let last_byte_pos = slice.first + slice.length - 1 in let last_byte = Bytes.get_uint8 slice.bytes last_byte_pos in if last_byte = 0 then raise (Corrupted_stream "invalid padding marker"); (* Find the highest set bit - this is the padding marker *) let rec find_marker byte bit = if bit < 0 then 0 else if (byte land (1 lsl bit)) <> 0 then bit else find_marker byte (bit - 1) in let padding = 8 - find_marker last_byte 7 in let bit_offset = slice.length * 8 - padding in { src = slice.bytes; start_pos = slice.first; bit_offset }
let of_bytes src ~pos ~len = of_slice (Slice.make src ~first:pos ~length:len)
let[@inline] remaining t = t.bit_offset
let[@inline] is_empty t = t.bit_offset <= 0
let[@inline] read_bits t n = if n <= 0 then 0 else if n > 57 then invalid_arg "read_bits: n > 57" else begin t.bit_offset <- t.bit_offset - n; let actual_offset = max 0 t.bit_offset in let actual_bits = if t.bit_offset < 0 then n + t.bit_offset else n in if actual_bits <= 0 then 0 else begin let byte_offset = t.start_pos + (actual_offset / 8) in let bit_offset = actual_offset mod 8 in let result = ref 0 in let bits_read = ref 0 in let current_byte = ref byte_offset in let current_bit = ref bit_offset in while !bits_read < actual_bits do let byte = Bytes.get_uint8 t.src !current_byte in let available = 8 - !current_bit in let to_read = min available (actual_bits - !bits_read) in let mask = (1 lsl to_read) - 1 in let bits = (byte lsr !current_bit) land mask in result := !result lor (bits lsl !bits_read); bits_read := !bits_read + to_read; current_bit := !current_bit + to_read; if !current_bit >= 8 then begin current_bit := 0; incr current_byte end done; (* If we read past the beginning, shift the result *) if t.bit_offset < 0 then !result lsl (-t.bit_offset) else !result end end
let peek_bits t n = let saved_offset = t.bit_offset in let result = read_bits t n in t.bit_offset <- saved_offset; resultend
(** {1 Forward Bitstream Writer} *)
module Forward_writer = struct type t = { dst : bytes; start_pos : int; mutable byte_pos : int; mutable bit_pos : int; (* 0-7, bits written in current byte *) mutable current_byte : int; }
let of_slice (slice : Slice.t) = { dst = slice.bytes; start_pos = slice.first; byte_pos = slice.first; bit_pos = 0; current_byte = 0 }
let of_bytes dst = of_slice (Slice.of_bytes dst)
let create dst ~pos = of_slice (Slice.make dst ~first:pos ~length:(Bytes.length dst - pos))
let flush t = if t.bit_pos > 0 then begin Bytes.set_uint8 t.dst t.byte_pos t.current_byte; t.byte_pos <- t.byte_pos + 1; t.bit_pos <- 0; t.current_byte <- 0 end
let[@inline] write_bits t value n = if n <= 0 then () else if n > 57 then invalid_arg "write_bits: n > 57" else begin let value = ref value in let remaining = ref n in
while !remaining > 0 do let available = 8 - t.bit_pos in let to_write = min available !remaining in let mask = (1 lsl to_write) - 1 in t.current_byte <- t.current_byte lor ((!value land mask) lsl t.bit_pos); value := !value lsr to_write; remaining := !remaining - to_write; t.bit_pos <- t.bit_pos + to_write;
if t.bit_pos = 8 then begin Bytes.set_uint8 t.dst t.byte_pos t.current_byte; t.byte_pos <- t.byte_pos + 1; t.bit_pos <- 0; t.current_byte <- 0 end done end
let write_byte t value = if t.bit_pos <> 0 then flush t; Bytes.set_uint8 t.dst t.byte_pos value; t.byte_pos <- t.byte_pos + 1
let write_slice t (slice : Slice.t) = if t.bit_pos <> 0 then flush t; Bytes.blit slice.bytes slice.first t.dst t.byte_pos slice.length; t.byte_pos <- t.byte_pos + slice.length
let write_bytes t src = write_slice t (Slice.of_bytes src)
let byte_position t = if t.bit_pos > 0 then t.byte_pos + 1 else t.byte_pos
let finalize t = flush t; t.byte_pos - t.start_pos
let to_slice t : Slice.t = flush t; Slice.make t.dst ~first:t.start_pos ~length:(t.byte_pos - t.start_pos)end
(** {1 Backward Bitstream Writer}
Accumulates bits to be read backwards. Used for FSE and Huffman encoding. *)
module Backward_writer = struct type t = { mutable bits : int64; mutable num_bits : int; buffer : bytes; mutable buf_pos : int; }
let create size = { bits = 0L; num_bits = 0; buffer = Bytes.create size; buf_pos = size }
let[@inline] write_bits t value n = if n > 0 then begin t.bits <- Int64.logor t.bits (Int64.shift_left (Int64.of_int value) t.num_bits); t.num_bits <- t.num_bits + n end
let flush_bytes t = while t.num_bits >= 8 do t.buf_pos <- t.buf_pos - 1; Bytes.set_uint8 t.buffer t.buf_pos (Int64.to_int (Int64.logand t.bits 0xFFL)); t.bits <- Int64.shift_right_logical t.bits 8; t.num_bits <- t.num_bits - 8 done
let finalize_to_slice t : Slice.t = write_bits t 1 1; if t.num_bits mod 8 <> 0 then t.num_bits <- ((t.num_bits + 7) / 8) * 8; flush_bytes t; let len = Bytes.length t.buffer - t.buf_pos in (* Reverse bytes in place so marker ends up at the end *) for i = 0 to len / 2 - 1 do let j = t.buf_pos + i in let k = t.buf_pos + len - 1 - i in let tmp = Bytes.get t.buffer j in Bytes.set t.buffer j (Bytes.get t.buffer k); Bytes.set t.buffer k tmp done; Slice.make t.buffer ~first:t.buf_pos ~length:len
let finalize t = Slice.to_bytes (finalize_to_slice t)
let current_size t = Bytes.length t.buffer - t.buf_pos + (t.num_bits + 7) / 8end