1 files changed, 118 insertions, 0 deletions
diff --git a/ocaml/lib/bitstream.ml b/ocaml/lib/bitstream.ml
new file mode 100644
index 0000000..748e204
--- /dev/null
+++ b/ocaml/lib/bitstream.ml
@@ -0,0 +1,118 @@
+type writer = {
+  buf_ref : Bytes.t Stdlib.ref;
+  mutable bit_pos : int;
+}
+
+type reader = {
+  buf: bytes;
+  mutable bit_pos: int;
+  len: int;
+}
+
+let writer_create () =
+  { buf_ref = ref (Bytes.make 1024 '\x00'); bit_pos = 0 }
+
+let writer_ensure (w : writer) bits_needed =
+  let bytes_needed = (w.bit_pos + bits_needed + 7) / 8 in
+  let buf = !(w.buf_ref) in
+  if bytes_needed > Bytes.length buf then begin
+    let new_size = max (bytes_needed * 2) (Bytes.length buf * 2) in
+    let new_buf = Bytes.make new_size '\x00' in
+    Bytes.blit buf 0 new_buf 0 (Bytes.length buf);
+    w.buf_ref := new_buf
+  end
+
+let write_bit (w : writer) bit =
+  writer_ensure w 1;
+  let byte_pos = w.bit_pos / 8 in
+  let bit_off = w.bit_pos mod 8 in
+  if bit then begin
+    let buf = !(w.buf_ref) in
+    let old_byte = Bytes.get_uint8 buf byte_pos in
+    Bytes.set_uint8 buf byte_pos (old_byte lor (1 lsl bit_off))
+  end;
+  w.bit_pos <- w.bit_pos + 1
+
+let write_bits (w : writer) value nbits =
+  writer_ensure w nbits;
+  for i = 0 to nbits - 1 do
+    let bit = Z.testbit value i in
+    write_bit w bit
+  done
+
+let writer_to_bytes (w : writer) =
+  let byte_len = (w.bit_pos + 7) / 8 in
+  Bytes.sub !(w.buf_ref) 0 byte_len
+
+let writer_pos (w : writer) = w.bit_pos
+
+let reader_create buf =
+  { buf; bit_pos = 0; len = Bytes.length buf * 8 }
+
+let read_bit r =
+  if r.bit_pos >= r.len then
+    raise (Invalid_argument "read_bit: end of stream");
+  let byte_pos = r.bit_pos / 8 in
+  let bit_off = r.bit_pos mod 8 in
+  let byte_val = Bytes.get_uint8 r.buf byte_pos in
+  r.bit_pos <- r.bit_pos + 1;
+  (byte_val lsr bit_off) land 1 = 1
+
+let read_bits r nbits =
+  if nbits = 0 then Z.zero
+  else if nbits >= 64 then begin
+    (* For large reads, align to byte boundary then use fast path *)
+    let bit_offset = r.bit_pos mod 8 in
+    let result = ref Z.zero in
+    let bits_read = ref 0 in
+
+    (* Read initial unaligned bits to reach byte boundary *)
+    if bit_offset <> 0 then begin
+      let align_bits = 8 - bit_offset in
+      let to_read = min align_bits nbits in
+      for i = 0 to to_read - 1 do
+        if read_bit r then
+          result := Z.logor !result (Z.shift_left Z.one i)
+      done;
+      bits_read := to_read
+    end;
+
+    (* Now we're byte-aligned, read full bytes directly *)
+    if !bits_read < nbits then begin
+      let remaining = nbits - !bits_read in
+      let full_bytes = remaining / 8 in
+
+      if full_bytes > 0 then begin
+        let byte_pos = r.bit_pos / 8 in
+        let bytes_data = Bytes.sub r.buf byte_pos full_bytes in
+        let bytes_value = Z.of_bits (Bytes.to_string bytes_data) in
+        result := Z.logor !result (Z.shift_left bytes_value !bits_read);
+        r.bit_pos <- r.bit_pos + (full_bytes * 8);
+        bits_read := !bits_read + (full_bytes * 8)
+      end;
+
+      (* Read final partial byte *)
+      let final_bits = nbits - !bits_read in
+      for i = 0 to final_bits - 1 do
+        if read_bit r then
+          result := Z.logor !result (Z.shift_left Z.one (!bits_read + i))
+      done
+    end;
+    !result
+  end else begin
+    (* Small read: bit by bit is fine *)
+    let result = ref Z.zero in
+    for i = 0 to nbits - 1 do
+      if read_bit r then
+        result := Z.logor !result (Z.shift_left Z.one i)
+    done;
+    !result
+  end
+
+let reader_pos r = r.bit_pos
+
+let count_zero_bits_until_one r =
+  let rec loop count =
+    if read_bit r then count else loop (count + 1)
+  in
+  loop 0