pretty good actuallycodex

1 files changed, 65 insertions, 1 deletions

diff --git a/ocaml/lib/bitstream.ml b/ocaml/lib/bitstream.ml
index 39bfd6a..e758c2a 100644
--- a/ocaml/lib/bitstream.ml
+++ b/ocaml/lib/bitstream.ml
@@ -67,6 +67,18 @@ let reader_create buf =
     len = Bytes.length buf * 8;
   }
 
+(* Lookup table for trailing zero counts within a byte. Value for 0 is 8 so
+   callers can detect the "no one-bit present" case. *)
+let trailing_zeros =
+  let tbl = Array.make 256 8 in
+  for i = 1 to 255 do
+    let rec count n value =
+      if value land 1 = 1 then n else count (n + 1) (value lsr 1)
+    in
+    tbl.(i) <- count 0 i
+  done;
+  tbl
+
 (** Read a single bit *)
 let read_bit r =
   if r.bit_pos >= r.len then
@@ -80,7 +92,36 @@ let read_bit r =
 (** Read multiple bits as a Z.t - optimized for bulk reads *)
 let read_bits r nbits =
   if nbits = 0 then Z.zero
-  else if nbits <= 64 && (r.bit_pos mod 8 = 0) && nbits mod 8 = 0 then begin
+  else if nbits > 4096 then begin
+    (* Bulk path: copy bytes then convert to Z. *)
+    let byte_len = (nbits + 7) / 8 in
+    let buf = Bytes.make byte_len '\x00' in
+    let bits_done = ref 0 in
+
+    while !bits_done < nbits do
+      if (!bits_done land 7) = 0 && (r.bit_pos land 7) = 0 then begin
+        let rem_bits = nbits - !bits_done in
+        let bytes_to_copy = rem_bits / 8 in
+        if bytes_to_copy > 0 then begin
+          Bytes.blit r.buf (r.bit_pos / 8) buf (!bits_done / 8) bytes_to_copy;
+          r.bit_pos <- r.bit_pos + (bytes_to_copy * 8);
+          bits_done := !bits_done + (bytes_to_copy * 8)
+        end
+      end;
+
+      if !bits_done < nbits then begin
+        if read_bit r then begin
+          let byte_idx = !bits_done / 8 in
+          let bit_idx = !bits_done mod 8 in
+          let existing = Bytes.get_uint8 buf byte_idx in
+          Bytes.set_uint8 buf byte_idx (existing lor (1 lsl bit_idx))
+        end;
+        incr bits_done
+      end
+    done;
+
+    Z.of_bits (Bytes.unsafe_to_string buf)
+  end else if nbits <= 64 && (r.bit_pos mod 8 = 0) && nbits mod 8 = 0 then begin
     (* Fast path: byte-aligned, <= 8 bytes *)
     let byte_pos = r.bit_pos / 8 in
     let num_bytes = nbits / 8 in
@@ -144,3 +185,26 @@ let reader_pos r = r.bit_pos
 
 (** Check if at end of stream *)
 let reader_at_end r = r.bit_pos >= r.len
+
+let count_zero_bits_until_one r =
+  let buf = r.buf in
+  let len_bits = r.len in
+  let rec scan count bit_pos =
+    if bit_pos >= len_bits then
+      raise (Invalid_argument "count_zero_bits_until_one: end of stream")
+    else begin
+      let byte_idx = bit_pos lsr 3 in
+      let bit_off = bit_pos land 7 in
+      let byte = Bytes.get_uint8 buf byte_idx in
+      let masked = byte lsr bit_off in
+      if masked <> 0 then begin
+        let tz = trailing_zeros.(masked land 0xff) in
+        let zeros = count + tz in
+        r.bit_pos <- bit_pos + tz + 1; (* skip zeros and the terminating 1 bit *)
+        zeros
+      end else
+        let remaining = 8 - bit_off in
+        scan (count + remaining) (bit_pos + remaining)
+    end
+  in
+  scan 0 r.bit_pos