Source file `binary.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
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
(* Binary encoding/decoding functions for ClickHouse protocol *)

open Varint

(* Helper to read exactly n bytes *)
let really_input_bytes ic n =
  let buf = Bytes.create n in
  really_input ic buf 0 n;
  buf

(* Basic types - Little Endian *)
let[@inline] read_uint8 ic = input_byte ic
let[@inline] write_uint8 oc v = output_byte oc v

let[@inline] read_int32_le ic =
  let a = input_byte ic in
  let b = input_byte ic in
  let c = input_byte ic in
  let d = input_byte ic in
  Int32.logor (Int32.of_int a)
    (Int32.logor
       (Int32.shift_left (Int32.of_int b) 8)
       (Int32.logor (Int32.shift_left (Int32.of_int c) 16) (Int32.shift_left (Int32.of_int d) 24)))

let[@inline] write_int32_le oc v =
  let a = Int32.to_int (Int32.logand v 0xFFl) in
  let b = Int32.to_int (Int32.logand (Int32.shift_right_logical v 8) 0xFFl) in
  let c = Int32.to_int (Int32.logand (Int32.shift_right_logical v 16) 0xFFl) in
  let d = Int32.to_int (Int32.logand (Int32.shift_right_logical v 24) 0xFFl) in
  output_byte oc a;
  output_byte oc b;
  output_byte oc c;
  output_byte oc d

let[@inline] read_uint64_le ic =
  let low = read_int32_le ic in
  let high = read_int32_le ic in
  Int64.logor
    (Int64.logand (Int64.of_int32 low) 0xFFFFFFFFL)
    (Int64.shift_left (Int64.of_int32 high) 32)

let[@inline] write_uint64_le oc v =
  let low32 = Int64.to_int32 (Int64.logand v 0xFFFFFFFFL) in
  let high32 = Int64.to_int32 (Int64.shift_right_logical v 32) in
  write_int32_le oc low32;
  write_int32_le oc high32

let[@inline] read_float64_le ic =
  let bits = read_uint64_le ic in
  Int64.float_of_bits bits

let[@inline] write_float64_le oc v =
  let bits = Int64.bits_of_float v in
  write_uint64_le oc bits

(* String reading/writing *)
let read_str ic =
  let len = read_varint_int ic in
  if len = 0 then "" else really_input_string ic len

let write_str oc s =
  let len = String.length s in
  write_varint_int oc len;
  output_string oc s

(* Bytes-reader variants *)
let really_input_bytes_br br n =
  let buf = Bytes.create n in
  Buffered_reader.really_input br buf 0 n;
  buf

(* keep read_uint8_br as alias *)

let[@inline] read_int32_le_br br =
  let a = Buffered_reader.input_byte br in
  let b = Buffered_reader.input_byte br in
  let c = Buffered_reader.input_byte br in
  let d = Buffered_reader.input_byte br in
  Int32.logor (Int32.of_int a)
    (Int32.logor
       (Int32.shift_left (Int32.of_int b) 8)
       (Int32.logor (Int32.shift_left (Int32.of_int c) 16) (Int32.shift_left (Int32.of_int d) 24)))

let[@inline] read_uint64_le_br br =
  let low = read_int32_le_br br in
  let high = read_int32_le_br br in
  Int64.logor
    (Int64.logand (Int64.of_int32 low) 0xFFFFFFFFL)
    (Int64.shift_left (Int64.of_int32 high) 32)

let[@inline] read_int64_le_br br =
  let low = read_int32_le_br br in
  let high = read_int32_le_br br in
  Int64.logor (Int64.of_int32 low) (Int64.shift_left (Int64.of_int32 high) 32)

let[@inline] read_float64_le_br br =
  let bits = read_uint64_le_br br in
  Int64.float_of_bits bits

let read_varint_int_br br : int =
  let rec loop shift acc =
    let b = Buffered_reader.input_byte br in
    let v = acc lor ((b land 0x7f) lsl shift) in
    if b land 0x80 <> 0 then loop (shift + 7) v else v
  in
  loop 0 0

let read_str_br br =
  let len = read_varint_int_br br in
  if len = 0 then "" else Buffered_reader.really_input_string br len

(* Bytes helpers for compression module *)
let[@inline] bytes_get_int32_le buf offset =
  let a = Char.code (Bytes.get buf offset) in
  let b = Char.code (Bytes.get buf (offset + 1)) in
  let c = Char.code (Bytes.get buf (offset + 2)) in
  let d = Char.code (Bytes.get buf (offset + 3)) in
  Int32.logor (Int32.of_int a)
    (Int32.logor
       (Int32.shift_left (Int32.of_int b) 8)
       (Int32.logor (Int32.shift_left (Int32.of_int c) 16) (Int32.shift_left (Int32.of_int d) 24)))

let[@inline] bytes_set_int32_le buf offset v =
  let a = Int32.to_int (Int32.logand v 0xFFl) in
  let b = Int32.to_int (Int32.logand (Int32.shift_right_logical v 8) 0xFFl) in
  let c = Int32.to_int (Int32.logand (Int32.shift_right_logical v 16) 0xFFl) in
  let d = Int32.to_int (Int32.logand (Int32.shift_right_logical v 24) 0xFFl) in
  Bytes.set buf offset (Char.chr a);
  Bytes.set buf (offset + 1) (Char.chr b);
  Bytes.set buf (offset + 2) (Char.chr c);
  Bytes.set buf (offset + 3) (Char.chr d)

let[@inline] bytes_get_int64_le buf offset =
  let low = bytes_get_int32_le buf offset in
  let high = bytes_get_int32_le buf (offset + 4) in
  Int64.logor
    (Int64.logand (Int64.of_int32 low) 0xFFFFFFFFL)
    (Int64.shift_left (Int64.logand (Int64.of_int32 high) 0xFFFFFFFFL) 32)

let[@inline] bytes_set_int64_le buf offset v =
  let low32 = Int64.to_int32 (Int64.logand v 0xFFFFFFFFL) in
  let high32 = Int64.to_int32 (Int64.shift_right_logical v 32) in
  bytes_set_int32_le buf offset low32;
  bytes_set_int32_le buf (offset + 4) high32

let read_int16_le ic =
  let a = input_byte ic in
  let b = input_byte ic in
  let v = a lor (b lsl 8) in
  if v land 0x8000 <> 0 then Int32.of_int (v - 0x10000) else Int32.of_int v

let[@inline] read_int16_le_br br =
  let a = Buffered_reader.input_byte br in
  let b = Buffered_reader.input_byte br in
  let v = a lor (b lsl 8) in
  if v land 0x8000 <> 0 then Int32.of_int (v - 0x10000) else Int32.of_int v

let[@inline] read_uint8_br br = Buffered_reader.input_byte br

let read_int64_le ic =
  let low = read_int32_le ic in
  let high = read_int32_le ic in
  Int64.logor (Int64.of_int32 low) (Int64.shift_left (Int64.of_int32 high) 32)