anil.recoil.org
Pure OCaml implementation of the Brotli compression algorithm
1(* Brotli compression implementation *)
2(* Supports quality levels 0-11 with context modeling, block splitting, and optimal parsing *)
3
4(* Re-export from LZ77 for backward compatibility *)
5let min_match = Lz77.min_match
6
7(* Number of literal contexts *)
8let num_literal_contexts = 64
9
10(* Insert length code tables *)
11let insert_length_n_bits = [|
12 0; 0; 0; 0; 0; 0; 1; 1; 2; 2; 3; 3; 4; 4; 5; 5; 6; 7; 8; 9; 10; 12; 14; 24
13|]
14
15let insert_length_offset = [|
16 0; 1; 2; 3; 4; 5; 6; 8; 10; 14; 18; 26; 34; 50; 66; 98; 130; 194; 322; 578; 1090; 2114; 6210; 22594
17|]
18
19(* Get insert length code *)
20let get_insert_code length =
21 let rec find i =
22 if i >= 23 then 23
23 else if length < insert_length_offset.(i + 1) then i
24 else find (i + 1)
25 in
26 find 0
27
28(* Get copy length code *)
29let get_copy_code length =
30 let copy_length_offset = [|
31 2; 3; 4; 5; 6; 7; 8; 9; 10; 12; 14; 18; 22; 30; 38; 54; 70; 102; 134; 198; 326; 582; 1094; 2118
32 |] in
33 let rec find i =
34 if i >= 23 then 23
35 else if length < copy_length_offset.(i + 1) then i
36 else find (i + 1)
37 in
38 find 0
39
40(* Command code lookup tables from RFC 7932 *)
41let insert_range_lut = [| 0; 0; 8; 8; 0; 16; 8; 16; 16 |]
42let copy_range_lut = [| 0; 8; 0; 8; 16; 0; 16; 8; 16 |]
43
44(* Build command code from insert_code and copy_code.
45 use_implicit_distance: true ONLY for distance code 0 (last distance)
46
47 Per RFC 7932, command codes have range_idx in bits 7-6:
48 - range_idx 0-1 (cmd_code 0-127): Distance code 0 is IMPLICIT (not read from stream)
49 The decoder automatically uses distance code 0 (last used distance).
50 - range_idx 2-8 (cmd_code 128+): Distance code is EXPLICIT (read from stream)
51 Short codes 0-15 and long codes >= 16 are all written explicitly.
52
53 IMPORTANT: Only dist_code=Some 0 can use implicit distance (range_idx 0-1).
54 For all other short codes (1-15), we must use explicit distance (range_idx >= 2).
55*)
56let get_command_code insert_code copy_code use_implicit_distance =
57 let found = ref None in
58
59 (* Only use range_idx 0-1 for implicit distance code 0 *)
60 if use_implicit_distance then begin
61 for r = 0 to 1 do
62 if !found = None then begin
63 let insert_base = insert_range_lut.(r) in
64 let copy_base = copy_range_lut.(r) in
65 let insert_delta = insert_code - insert_base in
66 let copy_delta = copy_code - copy_base in
67 if insert_delta >= 0 && insert_delta < 8 &&
68 copy_delta >= 0 && copy_delta < 8 then begin
69 let cmd_code = (r lsl 6) lor (insert_delta lsl 3) lor copy_delta in
70 found := Some cmd_code
71 end
72 end
73 done
74 end;
75
76 (* Use range_idx 2-8 for explicit distance (including short codes 0-15) *)
77 if !found = None then begin
78 for r = 2 to 8 do
79 if !found = None then begin
80 let adjusted_r = r - 2 in
81 let insert_base = insert_range_lut.(adjusted_r) in
82 let copy_base = copy_range_lut.(adjusted_r) in
83 let insert_delta = insert_code - insert_base in
84 let copy_delta = copy_code - copy_base in
85 if insert_delta >= 0 && insert_delta < 8 &&
86 copy_delta >= 0 && copy_delta < 8 then begin
87 let cmd_code = (r lsl 6) lor (insert_delta lsl 3) lor copy_delta in
88 found := Some cmd_code
89 end
90 end
91 done
92 end;
93
94 match !found with
95 | Some cmd_code -> cmd_code
96 | None ->
97 (* Fallback - shouldn't happen if LZ77 limits copy_len properly *)
98 let insert_delta = min insert_code 7 in
99 let copy_delta = min copy_code 7 in
100 (2 lsl 6) lor (insert_delta lsl 3) lor copy_delta
101
102(* Encode window bits *)
103let encode_window_bits bw =
104 Bit_writer.write_bits bw 1 1;
105 Bit_writer.write_bits bw 3 5 (* 22-bit window *)
106
107(* Write empty last block *)
108let write_empty_last_block bw =
109 Bit_writer.write_bits bw 1 1;
110 Bit_writer.write_bits bw 1 1
111
112(* Write meta-block header *)
113let write_meta_block_header bw length is_last is_uncompressed =
114 Bit_writer.write_bits bw 1 (if is_last then 1 else 0);
115 if is_last then
116 Bit_writer.write_bits bw 1 0; (* ISEMPTY = 0 for non-empty last block *)
117 let nibbles = if length - 1 < (1 lsl 16) then 4 else if length - 1 < (1 lsl 20) then 5 else 6 in
118 Bit_writer.write_bits bw 2 (nibbles - 4);
119 for i = 0 to nibbles - 1 do
120 Bit_writer.write_bits bw 4 (((length - 1) lsr (i * 4)) land 0xF)
121 done;
122 if not is_last then
123 Bit_writer.write_bits bw 1 (if is_uncompressed then 1 else 0)
124
125(* Write uncompressed block *)
126let write_uncompressed_block bw src src_pos length =
127 write_meta_block_header bw length false true;
128 Bit_writer.align_to_byte bw;
129 Bit_writer.copy_bytes bw ~src ~src_pos ~len:length
130
131(* Count bits needed to represent values 0 to n-1 (ceiling of log2(n)) *)
132let count_bits n =
133 if n <= 1 then 0
134 else
135 let rec count v b = if v = 0 then b else count (v lsr 1) (b + 1) in
136 count (n - 1) 0
137
138(* Write simple prefix code - 1 to 4 symbols *)
139let write_simple_prefix_code bw symbols alphabet_size =
140 let n = Array.length symbols in
141 Bit_writer.write_bits bw 2 1; (* HSKIP = 1 means simple code *)
142 Bit_writer.write_bits bw 2 (n - 1); (* NSYM - 1 *)
143 let bits = count_bits (alphabet_size - 1) in
144 for i = 0 to n - 1 do
145 Bit_writer.write_bits bw bits symbols.(i)
146 done;
147 if n = 4 then Bit_writer.write_bits bw 1 0
148
149(* Static Huffman code for code lengths *)
150let write_code_length_symbol bw len =
151 match len with
152 | 0 -> Bit_writer.write_bits bw 2 0
153 | 1 -> Bit_writer.write_bits bw 4 7
154 | 2 -> Bit_writer.write_bits bw 3 3
155 | 3 -> Bit_writer.write_bits bw 2 2
156 | 4 -> Bit_writer.write_bits bw 2 1
157 | 5 -> Bit_writer.write_bits bw 4 15
158 | _ -> Bit_writer.write_bits bw 2 0
159
160(* Build valid Huffman code lengths using Kraft inequality *)
161let build_valid_code_lengths freqs max_len =
162 let n = Array.length freqs in
163 let lengths = Array.make n 0 in
164 let symbols = ref [] in
165 for i = n - 1 downto 0 do
166 if freqs.(i) > 0 then
167 symbols := (freqs.(i), i) :: !symbols
168 done;
169 let num_symbols = List.length !symbols in
170 if num_symbols = 0 then lengths
171 else if num_symbols = 1 then begin
172 let (_, sym) = List.hd !symbols in
173 lengths.(sym) <- 1;
174 lengths
175 end
176 else begin
177 let sorted = List.sort (fun (f1, _) (f2, _) -> compare f2 f1) !symbols in
178 let bits_needed = count_bits num_symbols in
179 let base_len = min max_len (max bits_needed 1) in
180 let len_to_use = ref base_len in
181 while (1 lsl !len_to_use) < num_symbols && !len_to_use < max_len do
182 incr len_to_use
183 done;
184 let slots_used = ref num_symbols in
185 let total_slots = 1 lsl !len_to_use in
186 List.iter (fun (_, sym) ->
187 let extra_slots = total_slots - !slots_used in
188 if extra_slots > 0 && !len_to_use > 1 then begin
189 let shorter_len = !len_to_use - 1 in
190 let extra_needed = (1 lsl (!len_to_use - shorter_len)) - 1 in
191 if extra_slots >= extra_needed then begin
192 lengths.(sym) <- shorter_len;
193 slots_used := !slots_used + extra_needed
194 end else
195 lengths.(sym) <- !len_to_use
196 end else
197 lengths.(sym) <- !len_to_use
198 ) sorted;
199 lengths
200 end
201
202(* Build canonical Huffman codes from lengths *)
203let build_codes lengths =
204 let n = Array.length lengths in
205 let codes = Array.make n 0 in
206 let max_len = Array.fold_left max 0 lengths in
207 if max_len = 0 then codes
208 else begin
209 let bl_count = Array.make (max_len + 1) 0 in
210 Array.iter (fun l -> if l > 0 then bl_count.(l) <- bl_count.(l) + 1) lengths;
211 let next_code = Array.make (max_len + 1) 0 in
212 let code = ref 0 in
213 for bits = 1 to max_len do
214 code := (!code + bl_count.(bits - 1)) lsl 1;
215 next_code.(bits) <- !code
216 done;
217 for i = 0 to n - 1 do
218 let len = lengths.(i) in
219 if len > 0 then begin
220 codes.(i) <- next_code.(len);
221 next_code.(len) <- next_code.(len) + 1
222 end
223 done;
224 codes
225 end
226
227(* Reverse bits for canonical Huffman *)
228let reverse_bits v n =
229 let r = ref 0 in
230 let v = ref v in
231 for _ = 0 to n - 1 do
232 r := (!r lsl 1) lor (!v land 1);
233 v := !v lsr 1
234 done;
235 !r
236
237(* Write a Huffman symbol *)
238let write_symbol bw codes lengths sym =
239 let len = lengths.(sym) in
240 if len > 0 then
241 Bit_writer.write_bits bw len (reverse_bits codes.(sym) len)
242
243(* RLE encoding for code lengths *)
244let emit_zeros_rle symbols_ref extras_ref total_ref run_len =
245 if run_len < 3 then begin
246 for _ = 1 to run_len do
247 symbols_ref := 0 :: !symbols_ref;
248 extras_ref := 0 :: !extras_ref;
249 incr total_ref
250 done
251 end else begin
252 let reps = ref (run_len - 3) in
253 let rec build acc_codes acc_extras =
254 let e = !reps land 7 in
255 reps := !reps lsr 3;
256 if !reps = 0 then
257 (17 :: acc_codes, e :: acc_extras)
258 else begin
259 reps := !reps - 1;
260 build (17 :: acc_codes) (e :: acc_extras)
261 end
262 in
263 let (codes, extras) = build [] [] in
264 List.iter2 (fun c e ->
265 symbols_ref := c :: !symbols_ref;
266 extras_ref := e :: !extras_ref
267 ) codes extras;
268 total_ref := !total_ref + run_len
269 end
270
271let emit_nonzero_rle symbols_ref extras_ref total_ref run_len prev_value_ref value =
272 let to_write = ref run_len in
273 if !prev_value_ref <> value then begin
274 symbols_ref := value :: !symbols_ref;
275 extras_ref := 0 :: !extras_ref;
276 prev_value_ref := value;
277 decr to_write;
278 incr total_ref
279 end;
280 if !to_write < 3 then begin
281 for _ = 1 to !to_write do
282 symbols_ref := value :: !symbols_ref;
283 extras_ref := 0 :: !extras_ref
284 done;
285 total_ref := !total_ref + !to_write
286 end else begin
287 let reps = ref (!to_write - 3) in
288 let rec build acc_codes acc_extras =
289 let e = !reps land 3 in
290 reps := !reps lsr 2;
291 if !reps = 0 then
292 (16 :: acc_codes, e :: acc_extras)
293 else begin
294 reps := !reps - 1;
295 build (16 :: acc_codes) (e :: acc_extras)
296 end
297 in
298 let (codes, extras) = build [] [] in
299 List.iter2 (fun c e ->
300 symbols_ref := c :: !symbols_ref;
301 extras_ref := e :: !extras_ref
302 ) codes extras;
303 total_ref := !total_ref + !to_write
304 end
305
306let generate_rle_sequence lengths num_symbols =
307 let symbols = ref [] in
308 let extras = ref [] in
309 let prev_value = ref 8 in
310 let total = ref 0 in
311 let i = ref 0 in
312 while !i < num_symbols do
313 let value = if !i < Array.length lengths then lengths.(!i) else 0 in
314 let run_start = !i in
315 while !i < num_symbols &&
316 (if !i < Array.length lengths then lengths.(!i) else 0) = value do
317 incr i
318 done;
319 let run_len = !i - run_start in
320 if value = 0 then
321 emit_zeros_rle symbols extras total run_len
322 else
323 emit_nonzero_rle symbols extras total run_len prev_value value
324 done;
325 let syms = Array.of_list (List.rev !symbols) in
326 let exts = Array.of_list (List.rev !extras) in
327 (syms, exts)
328
329(* Write complex prefix code with RLE encoding *)
330let write_complex_prefix_code bw lengths alphabet_size =
331 let last_nonzero = ref (-1) in
332 for i = 0 to min (alphabet_size - 1) (Array.length lengths - 1) do
333 if lengths.(i) > 0 then last_nonzero := i
334 done;
335 let num_symbols = !last_nonzero + 1 in
336 let (rle_symbols, rle_extra) = generate_rle_sequence lengths num_symbols in
337 let cl_histogram = Array.make Constants.code_length_codes 0 in
338 Array.iter (fun sym -> cl_histogram.(sym) <- cl_histogram.(sym) + 1) rle_symbols;
339 let cl_depths = build_valid_code_lengths cl_histogram Constants.huffman_max_code_length_code_length in
340 let num_codes = ref 0 in
341 for i = 0 to Constants.code_length_codes - 1 do
342 if cl_histogram.(i) > 0 then incr num_codes
343 done;
344 let skip_some =
345 if cl_depths.(Constants.code_length_code_order.(0)) = 0 &&
346 cl_depths.(Constants.code_length_code_order.(1)) = 0 then
347 if cl_depths.(Constants.code_length_code_order.(2)) = 0 then 3
348 else 2
349 else 0
350 in
351 let codes_to_store = ref Constants.code_length_codes in
352 if !num_codes > 1 then begin
353 while !codes_to_store > 0 &&
354 cl_depths.(Constants.code_length_code_order.(!codes_to_store - 1)) = 0 do
355 decr codes_to_store
356 done
357 end;
358 Bit_writer.write_bits bw 2 skip_some;
359 let space = ref 32 in
360 for i = skip_some to !codes_to_store - 1 do
361 if !space > 0 then begin
362 let idx = Constants.code_length_code_order.(i) in
363 let l = cl_depths.(idx) in
364 write_code_length_symbol bw l;
365 if l <> 0 then
366 space := !space - (32 lsr l)
367 end
368 done;
369 let cl_codes = build_codes cl_depths in
370 for i = 0 to Array.length rle_symbols - 1 do
371 let sym = rle_symbols.(i) in
372 if !num_codes > 1 then
373 write_symbol bw cl_codes cl_depths sym;
374 if sym = 16 then
375 Bit_writer.write_bits bw 2 rle_extra.(i)
376 else if sym = 17 then
377 Bit_writer.write_bits bw 3 rle_extra.(i)
378 done
379
380(* Write Huffman code definition - choose simple or complex *)
381let write_huffman_code bw lengths alphabet_size =
382 let symbols = ref [] in
383 for i = 0 to min (alphabet_size - 1) (Array.length lengths - 1) do
384 if i < Array.length lengths && lengths.(i) > 0 then
385 symbols := (i, lengths.(i)) :: !symbols
386 done;
387 let sorted = List.sort (fun (s1, l1) (s2, l2) ->
388 let c = compare l1 l2 in
389 if c <> 0 then c else compare s1 s2
390 ) !symbols in
391 let symbols = Array.of_list (List.map fst sorted) in
392 let num_symbols = Array.length symbols in
393 if num_symbols = 0 then
394 write_simple_prefix_code bw [|0|] alphabet_size
395 else if num_symbols <= 4 then
396 write_simple_prefix_code bw symbols alphabet_size
397 else
398 write_complex_prefix_code bw lengths alphabet_size
399
400(* Count used symbols in frequency array *)
401let count_used_symbols freqs =
402 let count = ref 0 in
403 Array.iter (fun f -> if f > 0 then incr count) freqs;
404 !count
405
406(* Write context map using RLE and IMTF encoding *)
407(* Encode a variable length uint8 (matches decode_var_len_uint8 in decoder) *)
408let write_var_len_uint8 bw n =
409 if n = 0 then
410 Bit_writer.write_bits bw 1 0
411 else if n = 1 then begin
412 Bit_writer.write_bits bw 1 1;
413 Bit_writer.write_bits bw 3 0 (* nbits = 0 means value 1 *)
414 end else begin
415 Bit_writer.write_bits bw 1 1;
416 (* Find nbits such that (1 << nbits) <= n < (1 << (nbits + 1)) *)
417 let rec find_nbits nb =
418 if n < (1 lsl (nb + 1)) then nb
419 else find_nbits (nb + 1)
420 in
421 let nbits = find_nbits 1 in
422 Bit_writer.write_bits bw 3 nbits;
423 Bit_writer.write_bits bw nbits (n - (1 lsl nbits))
424 end
425
426let write_context_map bw context_map num_trees =
427 (* Write NTREES - 1 using variable length encoding *)
428 write_var_len_uint8 bw (num_trees - 1);
429
430 if num_trees > 1 then begin
431 (* Write RLEMAX flag: 0 = no RLE *)
432 Bit_writer.write_bits bw 1 0;
433
434 (* With RLEMAX=0, alphabet size is just num_trees, symbols are values directly *)
435 let map_len = Array.length context_map in
436 let freq = Array.make num_trees 0 in
437 for i = 0 to map_len - 1 do
438 freq.(context_map.(i)) <- freq.(context_map.(i)) + 1
439 done;
440
441 (* Build Huffman code for context map values *)
442 let lengths = build_valid_code_lengths freq 15 in
443 let codes = build_codes lengths in
444
445 (* Write the Huffman code for num_trees symbols *)
446 write_huffman_code bw lengths num_trees;
447
448 (* Write the context map values *)
449 let num_symbols = count_used_symbols freq in
450 for i = 0 to map_len - 1 do
451 if num_symbols > 1 then
452 write_symbol bw codes lengths context_map.(i)
453 done;
454
455 (* Write IMTF flag: 0 = no inverse move-to-front *)
456 Bit_writer.write_bits bw 1 0
457 end
458
459(* Copy length extra bits table *)
460let copy_length_n_bits = [|
461 0; 0; 0; 0; 0; 0; 0; 0; 1; 1; 2; 2; 3; 3; 4; 4; 5; 5; 6; 7; 8; 9; 10; 24
462|]
463
464let copy_length_offset = [|
465 2; 3; 4; 5; 6; 7; 8; 9; 10; 12; 14; 18; 22; 30; 38; 54; 70; 102; 134; 198; 326; 582; 1094; 2118
466|]
467
468(* Encode distance for NPOSTFIX=0, NDIRECT=0 *)
469let encode_distance distance =
470 if distance < 1 then
471 (16, 1, 0)
472 else begin
473 let d = distance - 1 in
474 let nbits = ref 1 in
475 let range_start = ref 0 in
476 while d >= !range_start + (1 lsl (!nbits + 1)) && !nbits < 24 do
477 range_start := !range_start + (1 lsl (!nbits + 1));
478 incr nbits
479 done;
480 let half_size = 1 lsl !nbits in
481 let d_in_range = d - !range_start in
482 let lcode = if d_in_range >= half_size then 1 else 0 in
483 let dc = 2 * (!nbits - 1) + lcode in
484 let code = 16 + dc in
485 let extra = d_in_range - (lcode * half_size) in
486 (code, !nbits, extra)
487 end
488
489(* Quality level for dictionary matching *)
490let current_quality = ref 1
491
492(* Write a compressed block with context modeling for quality >= 5 *)
493let write_compressed_block_with_context bw src _src_pos _src_len is_last context_mode context_map num_lit_trees num_dist_trees dist_context_map commands =
494 let num_distance_codes = 16 + 48 in
495
496 (* Count frequencies for context-aware literal encoding *)
497 let lit_freqs = Array.init num_lit_trees (fun _ -> Array.make 256 0) in
498 let cmd_freq = Array.make 704 0 in
499 (* Distance frequencies per tree *)
500 let dist_freqs = Array.init num_dist_trees (fun _ -> Array.make num_distance_codes 0) in
501
502 (* Track previous bytes for context calculation *)
503 let prev1 = ref 0 in
504 let prev2 = ref 0 in
505
506 (* Helper to get distance code value *)
507 let get_dist_code_val dist_code distance =
508 match dist_code with
509 | Some code -> code
510 | None ->
511 let dist_code_val, _, _ = encode_distance distance in
512 min dist_code_val (num_distance_codes - 1)
513 in
514
515 (* Count literals with context and build command/distance frequencies *)
516 List.iter (fun cmd ->
517 match cmd with
518 | Lz77.Literals { start; len } ->
519 for i = start to start + len - 1 do
520 let c = Char.code (Bytes.get src i) in
521 let ctx_id = Context.get_context context_mode ~prev_byte1:!prev1 ~prev_byte2:!prev2 in
522 let tree_id = context_map.(ctx_id) in
523 lit_freqs.(tree_id).(c) <- lit_freqs.(tree_id).(c) + 1;
524 prev2 := !prev1;
525 prev1 := c
526 done;
527 let insert_code = get_insert_code len in
528 let copy_code = 0 in
529 let cmd_code = get_command_code insert_code copy_code false in
530 cmd_freq.(cmd_code) <- cmd_freq.(cmd_code) + 1;
531 (* Literals command with copy_code=0 has copy_len=2, so dist context = 0 *)
532 let dist_tree = dist_context_map.(0) in
533 dist_freqs.(dist_tree).(0) <- dist_freqs.(dist_tree).(0) + 1
534 | Lz77.InsertCopy { lit_start; lit_len; copy_len; distance; dist_code } ->
535 for i = lit_start to lit_start + lit_len - 1 do
536 let c = Char.code (Bytes.get src i) in
537 let ctx_id = Context.get_context context_mode ~prev_byte1:!prev1 ~prev_byte2:!prev2 in
538 let tree_id = context_map.(ctx_id) in
539 lit_freqs.(tree_id).(c) <- lit_freqs.(tree_id).(c) + 1;
540 prev2 := !prev1;
541 prev1 := c
542 done;
543 let insert_code = get_insert_code lit_len in
544 let copy_code = get_copy_code copy_len in
545 let use_implicit = dist_code = Some 0 in
546 let cmd_code = get_command_code insert_code copy_code use_implicit in
547 let range_idx = cmd_code lsr 6 in
548 cmd_freq.(cmd_code) <- cmd_freq.(cmd_code) + 1;
549 if range_idx >= 2 then begin
550 let dist_ctx = Context.distance_context copy_len in
551 let dist_tree = dist_context_map.(dist_ctx) in
552 let code_val = get_dist_code_val dist_code distance in
553 dist_freqs.(dist_tree).(code_val) <- dist_freqs.(dist_tree).(code_val) + 1
554 end
555 ) commands;
556
557 (* Build Huffman codes for each literal tree *)
558 let lit_lengths_arr = Array.init num_lit_trees (fun i ->
559 build_valid_code_lengths lit_freqs.(i) 15
560 ) in
561 let lit_codes_arr = Array.init num_lit_trees (fun i ->
562 build_codes lit_lengths_arr.(i)
563 ) in
564 let cmd_lengths = build_valid_code_lengths cmd_freq 15 in
565 let cmd_codes = build_codes cmd_lengths in
566 (* Build Huffman codes for each distance tree *)
567 let dist_lengths_arr = Array.init num_dist_trees (fun i ->
568 build_valid_code_lengths dist_freqs.(i) 15
569 ) in
570 let dist_codes_arr = Array.init num_dist_trees (fun i ->
571 build_codes dist_lengths_arr.(i)
572 ) in
573
574 (* Calculate total uncompressed size *)
575 let total_len = List.fold_left (fun acc cmd ->
576 match cmd with
577 | Lz77.Literals { len; _ } -> acc + len
578 | Lz77.InsertCopy { lit_len; copy_len; _ } -> acc + lit_len + copy_len
579 ) 0 commands in
580
581 (* Write meta-block header *)
582 write_meta_block_header bw total_len is_last false;
583
584 (* Block type counts: 1 for each category *)
585 Bit_writer.write_bits bw 1 0; (* NBLTYPESL = 1 *)
586 Bit_writer.write_bits bw 1 0; (* NBLTYPESI = 1 *)
587 Bit_writer.write_bits bw 1 0; (* NBLTYPESD = 1 *)
588
589 (* Distance parameters: NPOSTFIX=0, NDIRECT=0 *)
590 Bit_writer.write_bits bw 2 0;
591 Bit_writer.write_bits bw 4 0;
592
593 (* Context mode for literal block type 0 *)
594 Bit_writer.write_bits bw 2 (Context.int_of_mode context_mode);
595
596 (* Literal context map *)
597 write_context_map bw context_map num_lit_trees;
598
599 (* Distance context map: 4 contexts per block type *)
600 write_context_map bw dist_context_map num_dist_trees;
601
602 (* Write Huffman codes for all literal trees *)
603 for i = 0 to num_lit_trees - 1 do
604 write_huffman_code bw lit_lengths_arr.(i) 256
605 done;
606 write_huffman_code bw cmd_lengths 704;
607 (* Write Huffman codes for all distance trees *)
608 for i = 0 to num_dist_trees - 1 do
609 write_huffman_code bw dist_lengths_arr.(i) num_distance_codes
610 done;
611
612 (* Write commands with context-aware literal and distance encoding *)
613 let num_cmd_symbols = count_used_symbols cmd_freq in
614 prev1 := 0;
615 prev2 := 0;
616
617 List.iter (fun cmd ->
618 match cmd with
619 | Lz77.Literals { start; len } ->
620 let insert_code = get_insert_code len in
621 let copy_code = 0 in
622 let cmd_code = get_command_code insert_code copy_code false in
623 if num_cmd_symbols > 1 then
624 write_symbol bw cmd_codes cmd_lengths cmd_code;
625 if insert_length_n_bits.(insert_code) > 0 then begin
626 let extra = len - insert_length_offset.(insert_code) in
627 Bit_writer.write_bits bw insert_length_n_bits.(insert_code) extra
628 end;
629 for i = start to start + len - 1 do
630 let c = Char.code (Bytes.get src i) in
631 let ctx_id = Context.get_context context_mode ~prev_byte1:!prev1 ~prev_byte2:!prev2 in
632 let tree_id = context_map.(ctx_id) in
633 let num_symbols = count_used_symbols lit_freqs.(tree_id) in
634 if num_symbols > 1 then
635 write_symbol bw lit_codes_arr.(tree_id) lit_lengths_arr.(tree_id) c;
636 prev2 := !prev1;
637 prev1 := c
638 done
639
640 | Lz77.InsertCopy { lit_start; lit_len; copy_len; distance; dist_code } ->
641 let insert_code = get_insert_code lit_len in
642 let copy_code = get_copy_code copy_len in
643 let use_implicit = dist_code = Some 0 in
644 let cmd_code = get_command_code insert_code copy_code use_implicit in
645 let range_idx = cmd_code lsr 6 in
646 if num_cmd_symbols > 1 then
647 write_symbol bw cmd_codes cmd_lengths cmd_code;
648 if insert_length_n_bits.(insert_code) > 0 then begin
649 let extra = lit_len - insert_length_offset.(insert_code) in
650 Bit_writer.write_bits bw insert_length_n_bits.(insert_code) extra
651 end;
652 if copy_length_n_bits.(copy_code) > 0 then begin
653 let extra = copy_len - copy_length_offset.(copy_code) in
654 Bit_writer.write_bits bw copy_length_n_bits.(copy_code) extra
655 end;
656 for i = lit_start to lit_start + lit_len - 1 do
657 let c = Char.code (Bytes.get src i) in
658 let ctx_id = Context.get_context context_mode ~prev_byte1:!prev1 ~prev_byte2:!prev2 in
659 let tree_id = context_map.(ctx_id) in
660 let num_symbols = count_used_symbols lit_freqs.(tree_id) in
661 if num_symbols > 1 then
662 write_symbol bw lit_codes_arr.(tree_id) lit_lengths_arr.(tree_id) c;
663 prev2 := !prev1;
664 prev1 := c
665 done;
666 if range_idx >= 2 then begin
667 let dist_ctx = Context.distance_context copy_len in
668 let dist_tree = dist_context_map.(dist_ctx) in
669 let num_dist_symbols = count_used_symbols dist_freqs.(dist_tree) in
670 match dist_code with
671 | Some code ->
672 if num_dist_symbols > 1 then
673 write_symbol bw dist_codes_arr.(dist_tree) dist_lengths_arr.(dist_tree) code
674 | None ->
675 let dist_code_val, nbits, extra = encode_distance distance in
676 if num_dist_symbols > 1 then
677 write_symbol bw dist_codes_arr.(dist_tree) dist_lengths_arr.(dist_tree) dist_code_val;
678 if nbits > 0 then
679 Bit_writer.write_bits bw nbits extra
680 end
681 ) commands
682
683(* Write a compressed block with LZ77 commands *)
684let write_compressed_block bw src src_pos src_len is_last =
685 (* Dictionary matching provides additional compression for text content *)
686 let use_dict = !current_quality >= 3 in
687 let quality = !current_quality in
688
689 (* Generate commands using LZ77 or optimal parsing *)
690 let commands =
691 if quality >= 10 then
692 (* Use optimal greedy parsing with lazy matching for quality 10-11 *)
693 Optimal.generate_commands ~quality src src_pos src_len
694 else
695 (* Standard LZ77 for lower quality levels *)
696 Lz77.generate_commands ~use_dict ~quality src src_pos src_len
697 in
698
699 (* Use context modeling for quality >= 5 *)
700 if quality >= 5 then begin
701 let context_mode = Block_split.choose_context_mode src src_pos src_len in
702 (* For quality >= 7 with enough data, use multiple literal trees *)
703 let (context_map, num_lit_trees) =
704 if quality >= 7 && src_len >= 1024 then begin
705 let max_trees = if quality >= 9 then 4 else 2 in
706 let (cmap, _histograms, ntrees) =
707 Block_split.build_literal_context_map context_mode src src_pos src_len max_trees
708 in
709 (cmap, ntrees)
710 end else
711 (Array.make 64 0, 1)
712 in
713 (* Distance context map: 4 contexts based on copy_length *)
714 (* For now, use single distance tree (infrastructure ready for multiple) *)
715 let dist_context_map = Array.make 4 0 in
716 let num_dist_trees = 1 in
717 write_compressed_block_with_context bw src src_pos src_len is_last
718 context_mode context_map num_lit_trees num_dist_trees dist_context_map commands
719 end else begin
720 (* Original simple encoding for quality < 5 *)
721
722 (* Count frequencies for all three alphabets *)
723 let lit_freq = Array.make 256 0 in
724 let cmd_freq = Array.make 704 0 in
725 let num_distance_codes = 16 + 48 in
726 let dist_freq = Array.make num_distance_codes 0 in
727
728 (* Count literals and build command/distance frequencies *)
729 List.iter (fun cmd ->
730 match cmd with
731 | Lz77.Literals { start; len } ->
732 for i = start to start + len - 1 do
733 let c = Char.code (Bytes.get src i) in
734 lit_freq.(c) <- lit_freq.(c) + 1
735 done;
736 let insert_code = get_insert_code len in
737 let copy_code = 0 in
738 let cmd_code = get_command_code insert_code copy_code false in
739 cmd_freq.(cmd_code) <- cmd_freq.(cmd_code) + 1;
740 (* range_idx for Literals command with copy_code=0 is >= 2, so we need distance *)
741 dist_freq.(0) <- dist_freq.(0) + 1
742 | Lz77.InsertCopy { lit_start; lit_len; copy_len; distance; dist_code } ->
743 for i = lit_start to lit_start + lit_len - 1 do
744 let c = Char.code (Bytes.get src i) in
745 lit_freq.(c) <- lit_freq.(c) + 1
746 done;
747 let insert_code = get_insert_code lit_len in
748 let copy_code = get_copy_code copy_len in
749 (* Only dist_code=Some 0 can use implicit distance (range_idx 0-1) *)
750 let use_implicit = dist_code = Some 0 in
751 let cmd_code = get_command_code insert_code copy_code use_implicit in
752 let range_idx = cmd_code lsr 6 in
753 cmd_freq.(cmd_code) <- cmd_freq.(cmd_code) + 1;
754 (* Count distance code if range_idx >= 2 (explicit distance) *)
755 if range_idx >= 2 then begin
756 match dist_code with
757 | Some code -> dist_freq.(code) <- dist_freq.(code) + 1
758 | None ->
759 let dist_code_val, _, _ = encode_distance distance in
760 if dist_code_val < num_distance_codes then
761 dist_freq.(dist_code_val) <- dist_freq.(dist_code_val) + 1
762 else
763 dist_freq.(num_distance_codes - 1) <- dist_freq.(num_distance_codes - 1) + 1
764 end
765 (* For range_idx 0-1, distance code 0 is implicit, don't count *)
766 ) commands;
767
768 (* Build Huffman codes *)
769 let lit_lengths = build_valid_code_lengths lit_freq 15 in
770 let lit_codes = build_codes lit_lengths in
771 let cmd_lengths = build_valid_code_lengths cmd_freq 15 in
772 let cmd_codes = build_codes cmd_lengths in
773 let dist_lengths = build_valid_code_lengths dist_freq 15 in
774 let dist_codes = build_codes dist_lengths in
775
776 (* Calculate total uncompressed size for meta-block header *)
777 let total_len = List.fold_left (fun acc cmd ->
778 match cmd with
779 | Lz77.Literals { len; _ } -> acc + len
780 | Lz77.InsertCopy { lit_len; copy_len; _ } -> acc + lit_len + copy_len
781 ) 0 commands in
782
783 (* Write meta-block header *)
784 write_meta_block_header bw total_len is_last false;
785
786 (* Block type counts: 1 for each category *)
787 Bit_writer.write_bits bw 1 0;
788 Bit_writer.write_bits bw 1 0;
789 Bit_writer.write_bits bw 1 0;
790
791 (* Distance parameters: NPOSTFIX=0, NDIRECT=0 *)
792 Bit_writer.write_bits bw 2 0;
793 Bit_writer.write_bits bw 4 0;
794
795 (* Context mode for literal block type 0: LSB6 = 0 *)
796 Bit_writer.write_bits bw 2 0;
797
798 (* Literal context map: NTREESL = 1 tree *)
799 Bit_writer.write_bits bw 1 0;
800
801 (* Distance context map: NTREESD = 1 tree *)
802 Bit_writer.write_bits bw 1 0;
803
804 (* Write Huffman codes *)
805 write_huffman_code bw lit_lengths 256;
806 write_huffman_code bw cmd_lengths 704;
807 write_huffman_code bw dist_lengths num_distance_codes;
808
809 (* Write commands *)
810 let num_lit_symbols = count_used_symbols lit_freq in
811 let num_cmd_symbols = count_used_symbols cmd_freq in
812 let num_dist_symbols = count_used_symbols dist_freq in
813
814 List.iter (fun cmd ->
815 match cmd with
816 | Lz77.Literals { start; len } ->
817 let insert_code = get_insert_code len in
818 let copy_code = 0 in
819 let cmd_code = get_command_code insert_code copy_code false in
820 if num_cmd_symbols > 1 then
821 write_symbol bw cmd_codes cmd_lengths cmd_code;
822 if insert_length_n_bits.(insert_code) > 0 then begin
823 let extra = len - insert_length_offset.(insert_code) in
824 Bit_writer.write_bits bw insert_length_n_bits.(insert_code) extra
825 end;
826 if num_lit_symbols > 1 then begin
827 for i = start to start + len - 1 do
828 let c = Char.code (Bytes.get src i) in
829 write_symbol bw lit_codes lit_lengths c
830 done
831 end
832
833 | Lz77.InsertCopy { lit_start; lit_len; copy_len; distance; dist_code } ->
834 let insert_code = get_insert_code lit_len in
835 let copy_code = get_copy_code copy_len in
836 (* Only dist_code=Some 0 can use implicit distance (range_idx 0-1) *)
837 let use_implicit = dist_code = Some 0 in
838 let cmd_code = get_command_code insert_code copy_code use_implicit in
839 let range_idx = cmd_code lsr 6 in
840 if num_cmd_symbols > 1 then
841 write_symbol bw cmd_codes cmd_lengths cmd_code;
842 if insert_length_n_bits.(insert_code) > 0 then begin
843 let extra = lit_len - insert_length_offset.(insert_code) in
844 Bit_writer.write_bits bw insert_length_n_bits.(insert_code) extra
845 end;
846 if copy_length_n_bits.(copy_code) > 0 then begin
847 let extra = copy_len - copy_length_offset.(copy_code) in
848 Bit_writer.write_bits bw copy_length_n_bits.(copy_code) extra
849 end;
850 if num_lit_symbols > 1 then begin
851 for i = lit_start to lit_start + lit_len - 1 do
852 let c = Char.code (Bytes.get src i) in
853 write_symbol bw lit_codes lit_lengths c
854 done
855 end;
856 (* Write distance code.
857 For range_idx 0-1 (command codes 0-127), the decoder uses implicit distance code 0
858 and does NOT read from the stream. For range_idx >= 2, we must write the distance code. *)
859 if range_idx >= 2 then begin
860 match dist_code with
861 | Some code ->
862 (* Short codes 0-15 - just write the code, no extra bits *)
863 if num_dist_symbols > 1 then
864 write_symbol bw dist_codes dist_lengths code
865 | None ->
866 let dist_code_val, nbits, extra = encode_distance distance in
867 if num_dist_symbols > 1 then
868 write_symbol bw dist_codes dist_lengths dist_code_val;
869 if nbits > 0 then
870 Bit_writer.write_bits bw nbits extra
871 end
872 (* For range_idx 0-1, distance code 0 is implicit, don't write anything *)
873 ) commands
874 end
875
876(* Write a compressed block with only literals *)
877let write_literals_only_block bw src src_pos src_len is_last =
878 write_meta_block_header bw src_len is_last false;
879 Bit_writer.write_bits bw 1 0;
880 Bit_writer.write_bits bw 1 0;
881 Bit_writer.write_bits bw 1 0;
882 Bit_writer.write_bits bw 2 0;
883 Bit_writer.write_bits bw 4 0;
884 Bit_writer.write_bits bw 2 0;
885 Bit_writer.write_bits bw 1 0;
886 Bit_writer.write_bits bw 1 0;
887
888 let lit_freq = Array.make 256 0 in
889 for i = src_pos to src_pos + src_len - 1 do
890 let c = Char.code (Bytes.get src i) in
891 lit_freq.(c) <- lit_freq.(c) + 1
892 done;
893 let num_lit_symbols = count_used_symbols lit_freq in
894 let lit_lengths = build_valid_code_lengths lit_freq 15 in
895 let lit_codes = build_codes lit_lengths in
896
897 let insert_code = get_insert_code src_len in
898 let copy_code = 0 in
899 let cmd_code = get_command_code insert_code copy_code false in
900 let cmd_freq = Array.make 704 0 in
901 cmd_freq.(cmd_code) <- 1;
902 let cmd_lengths = build_valid_code_lengths cmd_freq 15 in
903
904 let num_distance_codes = 16 + 48 in
905 let dist_freq = Array.make num_distance_codes 0 in
906 dist_freq.(0) <- 1;
907 let dist_lengths = build_valid_code_lengths dist_freq 15 in
908
909 write_huffman_code bw lit_lengths 256;
910 write_huffman_code bw cmd_lengths 704;
911 write_huffman_code bw dist_lengths num_distance_codes;
912
913 if insert_length_n_bits.(insert_code) > 0 then begin
914 let extra = src_len - insert_length_offset.(insert_code) in
915 Bit_writer.write_bits bw insert_length_n_bits.(insert_code) extra
916 end;
917
918 if num_lit_symbols > 1 then begin
919 for i = src_pos to src_pos + src_len - 1 do
920 let c = Char.code (Bytes.get src i) in
921 write_symbol bw lit_codes lit_lengths c
922 done
923 end
924
925(* Main compression function *)
926let compress_into ?(quality=1) ~src ~src_pos ~src_len ~dst ~dst_pos () =
927 current_quality := quality;
928 let bw = Bit_writer.create ~dst ~pos:dst_pos ~len:(Bytes.length dst - dst_pos) in
929 encode_window_bits bw;
930
931 if src_len = 0 then begin
932 write_empty_last_block bw;
933 Bit_writer.flush bw - dst_pos
934 end
935 else if quality = 0 || src_len < 16 then begin
936 write_uncompressed_block bw src src_pos src_len;
937 write_empty_last_block bw;
938 Bit_writer.flush bw - dst_pos
939 end
940 else begin
941 try
942 if quality >= 2 && src_len >= min_match then
943 write_compressed_block bw src src_pos src_len true
944 else
945 write_literals_only_block bw src src_pos src_len true;
946 Bit_writer.flush bw - dst_pos
947 with _ ->
948 let bw = Bit_writer.create ~dst ~pos:dst_pos ~len:(Bytes.length dst - dst_pos) in
949 encode_window_bits bw;
950 write_uncompressed_block bw src src_pos src_len;
951 write_empty_last_block bw;
952 Bit_writer.flush bw - dst_pos
953 end
954
955let max_compressed_length input_len =
956 input_len + input_len / 8 + 64
957
958(* Streaming encoder state *)
959type streaming_encoder = {
960 mutable quality : int;
961 mutable dst : bytes;
962 mutable dst_pos : int;
963 mutable header_written : bool;
964 mutable finished : bool;
965}
966
967let create_streaming_encoder ?(quality=1) ~dst ~dst_pos () =
968 { quality; dst; dst_pos; header_written = false; finished = false }
969
970(* Write a chunk of data to the streaming encoder *)
971let streaming_write encoder ~src ~src_pos ~src_len ~is_last =
972 if encoder.finished then
973 invalid_arg "streaming encoder already finished";
974
975 current_quality := encoder.quality;
976 let bw = Bit_writer.create ~dst:encoder.dst ~pos:encoder.dst_pos
977 ~len:(Bytes.length encoder.dst - encoder.dst_pos) in
978
979 (* Write header on first chunk *)
980 if not encoder.header_written then begin
981 encode_window_bits bw;
982 encoder.header_written <- true
983 end;
984
985 if src_len = 0 then begin
986 if is_last then begin
987 write_empty_last_block bw;
988 encoder.finished <- true
989 end
990 end
991 else if encoder.quality = 0 || src_len < 16 then begin
992 (* For low quality or small blocks, write uncompressed *)
993 if is_last then begin
994 write_uncompressed_block bw src src_pos src_len;
995 write_empty_last_block bw;
996 encoder.finished <- true
997 end else begin
998 (* Non-last uncompressed block *)
999 write_meta_block_header bw src_len false true;
1000 Bit_writer.align_to_byte bw;
1001 Bit_writer.copy_bytes bw ~src ~src_pos ~len:src_len
1002 end
1003 end
1004 else begin
1005 try
1006 if encoder.quality >= 2 && src_len >= min_match then
1007 write_compressed_block bw src src_pos src_len is_last
1008 else
1009 write_literals_only_block bw src src_pos src_len is_last;
1010 if is_last then encoder.finished <- true
1011 with _ ->
1012 (* Fallback to uncompressed *)
1013 if is_last then begin
1014 write_uncompressed_block bw src src_pos src_len;
1015 write_empty_last_block bw;
1016 encoder.finished <- true
1017 end else begin
1018 write_meta_block_header bw src_len false true;
1019 Bit_writer.align_to_byte bw;
1020 Bit_writer.copy_bytes bw ~src ~src_pos ~len:src_len
1021 end
1022 end;
1023
1024 let written = Bit_writer.flush bw - encoder.dst_pos in
1025 encoder.dst_pos <- encoder.dst_pos + written;
1026 written
1027
1028let streaming_finish encoder =
1029 if not encoder.finished then begin
1030 let result = streaming_write encoder ~src:(Bytes.create 0) ~src_pos:0 ~src_len:0 ~is_last:true in
1031 encoder.finished <- true;
1032 result
1033 end else 0
1034
1035let streaming_bytes_written encoder =
1036 encoder.dst_pos
1037
1038(* Re-export command type for Debug module *)
1039type command = Lz77.command =
1040 | InsertCopy of { lit_start: int; lit_len: int; copy_len: int; distance: int; dist_code: int option }
1041 | Literals of { start: int; len: int }
1042
1043let generate_commands src src_pos src_len =
1044 Lz77.generate_commands src src_pos src_len