Source file containers_leb128.ml

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
(* adapted from ocaml-protoc from code by c-cube *)

module Byte_slice = CCByte_slice
module Byte_buffer = CCByte_buffer

module Decode = struct
  let skip (sl : Byte_slice.t) off : int =
    let shift = ref 0 in
    let continue = ref true in

    let off = ref off in
    let n_consumed = ref 0 in

    while !continue do
      if sl.len <= 0 then invalid_arg "out of bound";
      incr n_consumed;
      let b = Char.code (Bytes.get sl.bs !off) in
      let cur = b land 0x7f in
      if cur <> b then (
        (* at least one byte follows this one *)
        incr off;
        shift := !shift + 7
      ) else if !shift < 63 || b land 0x7f <= 1 then
        continue := false
      else
        invalid_arg "leb128 varint is too long"
    done;

    !n_consumed

  let u64 (sl : Byte_slice.t) (off : int) : int64 * int =
    let shift = ref 0 in
    let res = ref 0L in
    let continue = ref true in

    let off = ref off in
    let n_consumed = ref 0 in

    while !continue do
      if sl.len <= 0 then invalid_arg "out of bound";
      incr n_consumed;
      let b = Char.code (Bytes.get sl.bs !off) in
      let cur = b land 0x7f in
      if cur <> b then (
        (* at least one byte follows this one *)
        (res := Int64.(logor !res (shift_left (of_int cur) !shift)));
        incr off;
        shift := !shift + 7
      ) else if !shift < 63 || b land 0x7f <= 1 then (
        (res := Int64.(logor !res (shift_left (of_int b) !shift)));
        continue := false
      ) else
        invalid_arg "leb128 varint is too long"
    done;

    !res, !n_consumed

  let[@inline] uint_truncate sl off =
    let v, n_consumed = u64 sl off in
    Int64.to_int v, n_consumed

  let[@inline] decode_zigzag (v : int64) : int64 =
    Int64.(logxor (shift_right v 1) (neg (logand v Int64.one)))

  let[@inline] i64 sl off : int64 * int =
    let v, n_consumed = u64 sl off in
    decode_zigzag v, n_consumed

  let[@inline] int_truncate sl off =
    let v, n_consumed = u64 sl off in
    Int64.to_int (decode_zigzag v), n_consumed
end

module Encode = struct
  let[@inline] encode_zigzag (i : int64) : int64 =
    Int64.(logxor (shift_left i 1) (shift_right i 63))

  external varint_size : (int64[@unboxed]) -> int
    = "caml_cc_leb128_varint_size_byte" "caml_cc_leb128_varint_size"
  [@@noalloc]
  (** Compute how many bytes this int would occupy as varint *)

  external varint_slice : bytes -> (int[@untagged]) -> (int64[@unboxed]) -> unit
    = "caml_cc_leb128_varint_byte" "caml_cc_leb128_varint"
  [@@noalloc]
  (** Write this int as varint into the given slice *)

  let[@inline] u64 (buf : Byte_buffer.t) (i : int64) =
    let n = varint_size i in
    Byte_buffer.ensure_free buf n;
    assert (buf.len + n <= Bytes.length buf.bs);
    varint_slice buf.bs buf.len i;
    buf.len <- buf.len + n

  let[@inline] i64 buf i : unit = u64 buf (encode_zigzag i)
  let[@inline] uint buf i : unit = u64 buf (Int64.of_int i)
  let[@inline] int buf i : unit = u64 buf (encode_zigzag (Int64.of_int i))
end