1 files changed, 51 insertions, 7 deletions
diff --git a/ocaml/lib/bitstream.ml b/ocaml/lib/bitstream.ml
index 8c1ef5b..39bfd6a 100644
--- a/ocaml/lib/bitstream.ml
+++ b/ocaml/lib/bitstream.ml
@@ -77,14 +77,58 @@ let read_bit r =
   r.bit_pos <- r.bit_pos + 1;
   (byte_val lsr bit_off) land 1 = 1
 
-(** Read multiple bits as a Z.t *)
+(** Read multiple bits as a Z.t - optimized for bulk reads *)
 let read_bits r nbits =
-  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
+  if nbits = 0 then Z.zero
+  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
+    r.bit_pos <- r.bit_pos + nbits;
+
+    let result = ref Z.zero in
+    for i = 0 to num_bytes - 1 do
+      let byte_val = Z.of_int (Bytes.get_uint8 r.buf (byte_pos + i)) in
+      result := Z.logor !result (Z.shift_left byte_val (i * 8))
+    done;
+    !result
+  end else if nbits >= 8 then begin
+    (* Mixed path: read whole bytes + remaining bits *)
+    let result = ref Z.zero in
+    let bits_read = ref 0 in
+
+    (* Read as many whole bytes as possible *)
+    while !bits_read + 8 <= nbits && (r.bit_pos mod 8 <> 0 || !bits_read = 0) do
+      if read_bit r then
+        result := Z.logor !result (Z.shift_left Z.one !bits_read);
+      bits_read := !bits_read + 1
+    done;
+
+    (* Now read whole bytes efficiently if byte-aligned *)
+    while !bits_read + 8 <= nbits && (r.bit_pos mod 8 = 0) do
+      let byte_pos = r.bit_pos / 8 in
+      let byte_val = Z.of_int (Bytes.get_uint8 r.buf byte_pos) in
+      result := Z.logor !result (Z.shift_left byte_val !bits_read);
+      r.bit_pos <- r.bit_pos + 8;
+      bits_read := !bits_read + 8
+    done;
+
+    (* Read remaining bits *)
+    while !bits_read < nbits do
+      if read_bit r then
+        result := Z.logor !result (Z.shift_left Z.one !bits_read);
+      bits_read := !bits_read + 1
+    done;
+    !result
+  end else begin
+    (* Small reads: use original bit-by-bit approach *)
+    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
 
 (** Peek at a bit without advancing *)
 let peek_bit r =