Skip to main content

fory_core/
buffer.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use crate::error::Error;
19use crate::types::bfloat16::bfloat16;
20use crate::types::float16::float16;
21use crate::util::buffer_rw_string::read_latin1_simd;
22use byteorder::{ByteOrder, LittleEndian};
23use std::cmp::max;
24
25/// Threshold for using SIMD optimizations in string operations.
26/// For buffers smaller than this, direct copy is faster than SIMD setup overhead.
27const SIMD_THRESHOLD: usize = 128;
28
29pub struct Writer<'a> {
30    pub(crate) bf: &'a mut Vec<u8>,
31}
32impl<'a> Writer<'a> {
33    // ============ Utility methods ============
34
35    #[inline(always)]
36    pub fn from_buffer(bf: &'a mut Vec<u8>) -> Writer<'a> {
37        Writer { bf }
38    }
39
40    #[inline(always)]
41    pub fn dump(&self) -> Vec<u8> {
42        self.bf.clone()
43    }
44
45    #[inline(always)]
46    pub fn reset(&mut self) {
47        self.bf.clear();
48    }
49
50    #[inline(always)]
51    pub fn len(&self) -> usize {
52        self.bf.len()
53    }
54
55    #[inline(always)]
56    pub fn is_empty(&self) -> bool {
57        self.bf.is_empty()
58    }
59
60    #[inline(always)]
61    pub fn reserve(&mut self, additional: usize) {
62        if self.bf.capacity() - self.len() < additional {
63            self.bf.reserve(max(additional * 2, self.bf.capacity()));
64        }
65    }
66
67    #[inline(always)]
68    pub fn skip(&mut self, len: usize) {
69        self.bf.resize(self.bf.len() + len, 0);
70    }
71
72    #[inline(always)]
73    pub fn set_bytes(&mut self, offset: usize, data: &[u8]) {
74        self.bf
75            .get_mut(offset..offset + data.len())
76            .unwrap()
77            .copy_from_slice(data);
78    }
79
80    #[inline(always)]
81    pub fn write_bytes(&mut self, v: &[u8]) -> usize {
82        self.bf.extend_from_slice(v);
83        v.len()
84    }
85
86    #[inline(always)]
87    pub(crate) unsafe fn write_bytes_from_ptr(&mut self, ptr: *const u8, len: usize) {
88        let offset = self.bf.len();
89        self.bf.reserve(len);
90        std::ptr::copy_nonoverlapping(ptr, self.bf.as_mut_ptr().add(offset), len);
91        self.bf.set_len(offset + len);
92    }
93
94    // ============ BOOL (TypeId = 1) ============
95
96    #[inline(always)]
97    pub fn write_bool(&mut self, value: bool) {
98        self.bf.push(if value { 1 } else { 0 });
99    }
100
101    // ============ INT8 (TypeId = 2) ============
102
103    #[inline(always)]
104    pub fn write_i8(&mut self, value: i8) {
105        self.bf.push(value as u8);
106    }
107
108    // ============ INT16 (TypeId = 3) ============
109
110    #[inline(always)]
111    pub fn write_i16(&mut self, value: i16) {
112        self.write_u16(value as u16);
113    }
114
115    // ============ INT32 (TypeId = 4) ============
116
117    #[inline(always)]
118    pub fn write_i32(&mut self, value: i32) {
119        self.write_u32(value as u32);
120    }
121
122    // ============ VARINT32 (TypeId = 5) ============
123
124    #[inline(always)]
125    pub fn write_var_i32(&mut self, value: i32) {
126        let zigzag = ((value as u32) << 1) ^ ((value >> 31) as u32);
127        self._write_var_u32(zigzag)
128    }
129
130    // ============ INT64 (TypeId = 6) ============
131
132    #[inline(always)]
133    pub fn write_i64(&mut self, value: i64) {
134        self.write_u64(value as u64);
135    }
136
137    // ============ VARINT64 (TypeId = 7) ============
138
139    #[inline(always)]
140    pub fn write_var_i64(&mut self, value: i64) {
141        let zigzag = ((value as u64) << 1) ^ ((value >> 63) as u64);
142        self._write_var_u64(zigzag);
143    }
144
145    // ============ TAGGED_INT64 (TypeId = 8) ============
146
147    /// Write signed long using fory Tagged(Small long as int) encoding.
148    /// If value is in [0xc0000000, 0x3fffffff] (i.e., [-1073741824, 1073741823]),
149    /// encode as 4 bytes: `((value as i32) << 1)`.
150    /// Otherwise write as 9 bytes: `0b1 | little-endian 8 bytes i64`.
151    #[inline(always)]
152    pub fn write_tagged_i64(&mut self, value: i64) {
153        const HALF_MIN_INT_VALUE: i64 = i32::MIN as i64 / 2; // -1073741824
154        const HALF_MAX_INT_VALUE: i64 = i32::MAX as i64 / 2; // 1073741823
155        if (HALF_MIN_INT_VALUE..=HALF_MAX_INT_VALUE).contains(&value) {
156            // Fits in 31 bits (with sign), encode as 4 bytes with bit 0 = 0
157            let v = (value as i32) << 1;
158            self.write_i32(v);
159        } else {
160            // Write flag byte (0b1) followed by 8-byte i64
161            self.bf.push(0b1);
162            self.write_i64(value);
163        }
164    }
165
166    // ============ UINT8 (TypeId = 9) ============
167
168    #[inline(always)]
169    pub fn write_u8(&mut self, value: u8) {
170        self.bf.push(value);
171    }
172
173    // ============ UINT16 (TypeId = 10) ============
174
175    #[inline(always)]
176    pub fn write_u16(&mut self, value: u16) {
177        #[cfg(target_endian = "little")]
178        {
179            let bytes = unsafe { &*(&value as *const u16 as *const [u8; 2]) };
180            self.bf.extend_from_slice(bytes);
181        }
182        #[cfg(target_endian = "big")]
183        {
184            self.bf.extend_from_slice(&value.to_le_bytes());
185        }
186    }
187
188    // ============ UINT32 (TypeId = 11) ============
189
190    #[inline(always)]
191    pub fn write_u32(&mut self, value: u32) {
192        #[cfg(target_endian = "little")]
193        {
194            let bytes = unsafe { &*(&value as *const u32 as *const [u8; 4]) };
195            self.bf.extend_from_slice(bytes);
196        }
197        #[cfg(target_endian = "big")]
198        {
199            self.bf.extend_from_slice(&value.to_le_bytes());
200        }
201    }
202
203    #[inline(always)]
204    fn write_u24(&mut self, value: u32) {
205        let offset = self.bf.len();
206        self.bf.reserve(4);
207        // The four-byte store stays inside reserved capacity; only the three wire bytes become
208        // initialized vector contents.
209        unsafe {
210            self.bf
211                .as_mut_ptr()
212                .add(offset)
213                .cast::<u32>()
214                .write_unaligned(value.to_le());
215            self.bf.set_len(offset + 3);
216        }
217    }
218
219    #[inline(always)]
220    fn write_u40(&mut self, value: u64) {
221        let offset = self.bf.len();
222        self.bf.reserve(8);
223        // The eight-byte store stays inside reserved capacity; only the five wire bytes become
224        // initialized vector contents.
225        unsafe {
226            self.bf
227                .as_mut_ptr()
228                .add(offset)
229                .cast::<u64>()
230                .write_unaligned(value.to_le());
231            self.bf.set_len(offset + 5);
232        }
233    }
234
235    // ============ VAR_UINT32 (TypeId = 12) ============
236
237    #[inline(always)]
238    pub fn write_var_u32(&mut self, value: u32) {
239        self._write_var_u32(value)
240    }
241
242    #[inline(always)]
243    fn _write_var_u32(&mut self, value: u32) {
244        if value < 0x80 {
245            self.bf.push(value as u8);
246        } else if value < 0x4000 {
247            // 2 bytes
248            let u1 = ((value as u8) & 0x7F) | 0x80;
249            let u2 = (value >> 7) as u8;
250            self.write_u16(((u2 as u16) << 8) | u1 as u16);
251        } else if value < 0x200000 {
252            // 3 bytes
253            let u1 = ((value as u8) & 0x7F) | 0x80;
254            let u2 = (((value >> 7) as u8) & 0x7F) | 0x80;
255            let u3 = (value >> 14) as u8;
256            self.write_u24(((u3 as u32) << 16) | ((u2 as u32) << 8) | u1 as u32);
257        } else if value < 0x10000000 {
258            // 4 bytes
259            let u1 = ((value as u8) & 0x7F) | 0x80;
260            let u2 = (((value >> 7) as u8) & 0x7F) | 0x80;
261            let u3 = (((value >> 14) as u8) & 0x7F) | 0x80;
262            let u4 = (value >> 21) as u8;
263            self.write_u32(
264                ((u4 as u32) << 24) | ((u3 as u32) << 16) | ((u2 as u32) << 8) | u1 as u32,
265            );
266        } else {
267            // 5 bytes
268            let u1 = ((value as u8) & 0x7F) | 0x80;
269            let u2 = (((value >> 7) as u8) & 0x7F) | 0x80;
270            let u3 = (((value >> 14) as u8) & 0x7F) | 0x80;
271            let u4 = (((value >> 21) as u8) & 0x7F) | 0x80;
272            let u5 = (value >> 28) as u8;
273            self.write_u40(
274                ((u5 as u64) << 32)
275                    | ((u4 as u64) << 24)
276                    | ((u3 as u64) << 16)
277                    | ((u2 as u64) << 8)
278                    | u1 as u64,
279            );
280        }
281    }
282
283    // ============ UINT64 (TypeId = 13) ============
284
285    #[inline(always)]
286    pub fn write_u64(&mut self, value: u64) {
287        #[cfg(target_endian = "little")]
288        {
289            let bytes = unsafe { &*(&value as *const u64 as *const [u8; 8]) };
290            self.bf.extend_from_slice(bytes);
291        }
292        #[cfg(target_endian = "big")]
293        {
294            self.bf.extend_from_slice(&value.to_le_bytes());
295        }
296    }
297
298    // ============ VAR_UINT64 (TypeId = 14) ============
299
300    #[inline(always)]
301    pub fn write_var_u64(&mut self, value: u64) {
302        self._write_var_u64(value);
303    }
304
305    #[inline(always)]
306    fn _write_var_u64(&mut self, value: u64) {
307        if value < 0x80 {
308            self.bf.push(value as u8);
309        } else if value < 0x4000 {
310            let u1 = ((value as u8) & 0x7F) | 0x80;
311            let u2 = (value >> 7) as u8;
312            self.write_u16(((u2 as u16) << 8) | u1 as u16);
313        } else if value < 0x200000 {
314            let u1 = ((value as u8) & 0x7F) | 0x80;
315            let u2 = (((value >> 7) as u8) & 0x7F) | 0x80;
316            let u3 = (value >> 14) as u8;
317            self.write_u24(((u3 as u32) << 16) | ((u2 as u32) << 8) | u1 as u32);
318        } else if value < 0x10000000 {
319            let u1 = ((value as u8) & 0x7F) | 0x80;
320            let u2 = (((value >> 7) as u8) & 0x7F) | 0x80;
321            let u3 = (((value >> 14) as u8) & 0x7F) | 0x80;
322            let u4 = (value >> 21) as u8;
323            self.write_u32(
324                ((u4 as u32) << 24) | ((u3 as u32) << 16) | ((u2 as u32) << 8) | u1 as u32,
325            );
326        } else if value < 0x800000000 {
327            let u1 = ((value as u8) & 0x7F) | 0x80;
328            let u2 = (((value >> 7) as u8) & 0x7F) | 0x80;
329            let u3 = (((value >> 14) as u8) & 0x7F) | 0x80;
330            let u4 = (((value >> 21) as u8) & 0x7F) | 0x80;
331            let u5 = (value >> 28) as u8;
332            self.write_u40(
333                ((u5 as u64) << 32)
334                    | ((u4 as u64) << 24)
335                    | ((u3 as u64) << 16)
336                    | ((u2 as u64) << 8)
337                    | u1 as u64,
338            );
339        } else if value < 0x40000000000 {
340            let u1 = ((value as u8) & 0x7F) | 0x80;
341            let u2 = (((value >> 7) as u8) & 0x7F) | 0x80;
342            let u3 = (((value >> 14) as u8) & 0x7F) | 0x80;
343            let u4 = (((value >> 21) as u8) & 0x7F) | 0x80;
344            let u5 = (((value >> 28) as u8) & 0x7F) | 0x80;
345            let u6 = (value >> 35) as u8;
346            self.write_u32(
347                ((u4 as u32) << 24) | ((u3 as u32) << 16) | ((u2 as u32) << 8) | u1 as u32,
348            );
349            self.write_u16(((u6 as u16) << 8) | u5 as u16);
350        } else if value < 0x2000000000000 {
351            let u1 = ((value as u8) & 0x7F) | 0x80;
352            let u2 = (((value >> 7) as u8) & 0x7F) | 0x80;
353            let u3 = (((value >> 14) as u8) & 0x7F) | 0x80;
354            let u4 = (((value >> 21) as u8) & 0x7F) | 0x80;
355            let u5 = (((value >> 28) as u8) & 0x7F) | 0x80;
356            let u6 = (((value >> 35) as u8) & 0x7F) | 0x80;
357            let u7 = (value >> 42) as u8;
358            self.write_u32(
359                ((u4 as u32) << 24) | ((u3 as u32) << 16) | ((u2 as u32) << 8) | u1 as u32,
360            );
361            self.write_u16(((u6 as u16) << 8) | u5 as u16);
362            self.bf.push(u7);
363        } else if value < 0x100000000000000 {
364            let u1 = ((value as u8) & 0x7F) | 0x80;
365            let u2 = (((value >> 7) as u8) & 0x7F) | 0x80;
366            let u3 = (((value >> 14) as u8) & 0x7F) | 0x80;
367            let u4 = (((value >> 21) as u8) & 0x7F) | 0x80;
368            let u5 = (((value >> 28) as u8) & 0x7F) | 0x80;
369            let u6 = (((value >> 35) as u8) & 0x7F) | 0x80;
370            let u7 = (((value >> 42) as u8) & 0x7F) | 0x80;
371            let u8 = (value >> 49) as u8;
372            self.write_u64(
373                (u8 as u64) << 56
374                    | (u7 as u64) << 48
375                    | (u6 as u64) << 40
376                    | (u5 as u64) << 32
377                    | (u4 as u64) << 24
378                    | (u3 as u64) << 16
379                    | (u2 as u64) << 8
380                    | (u1 as u64),
381            );
382        } else {
383            let u1 = ((value as u8) & 0x7F) | 0x80;
384            let u2 = (((value >> 7) as u8) & 0x7F) | 0x80;
385            let u3 = (((value >> 14) as u8) & 0x7F) | 0x80;
386            let u4 = (((value >> 21) as u8) & 0x7F) | 0x80;
387            let u5 = (((value >> 28) as u8) & 0x7F) | 0x80;
388            let u6 = (((value >> 35) as u8) & 0x7F) | 0x80;
389            let u7 = (((value >> 42) as u8) & 0x7F) | 0x80;
390            let u8 = (((value >> 49) as u8) & 0x7F) | 0x80;
391            let u9 = (value >> 56) as u8;
392            self.write_u64(
393                (u8 as u64) << 56
394                    | (u7 as u64) << 48
395                    | (u6 as u64) << 40
396                    | (u5 as u64) << 32
397                    | (u4 as u64) << 24
398                    | (u3 as u64) << 16
399                    | (u2 as u64) << 8
400                    | (u1 as u64),
401            );
402            self.bf.push(u9);
403        }
404    }
405
406    // ============ TAGGED_UINT64 (TypeId = 15) ============
407
408    /// Write unsigned long using fory Tagged(Small long as int) encoding.
409    /// If value is in [0, 0x7fffffff], encode as 4 bytes: `((value as u32) << 1)`.
410    /// Otherwise write as 9 bytes: `0b1 | little-endian 8 bytes u64`.
411    #[inline(always)]
412    pub fn write_tagged_u64(&mut self, value: u64) {
413        if value <= i32::MAX as u64 {
414            // Fits in 31 bits, encode as 4 bytes with bit 0 = 0
415            let v = (value as u32) << 1;
416            self.write_u32(v);
417        } else {
418            // Write flag byte (0b1) followed by 8-byte u64
419            self.bf.push(0b1);
420            self.write_u64(value);
421        }
422    }
423
424    // ============ FLOAT32 (TypeId = 17) ============
425
426    #[inline(always)]
427    pub fn write_f32(&mut self, value: f32) {
428        #[cfg(target_endian = "little")]
429        {
430            let bytes = unsafe { &*(&value as *const f32 as *const [u8; 4]) };
431            self.bf.extend_from_slice(bytes);
432        }
433        #[cfg(target_endian = "big")]
434        {
435            self.bf.extend_from_slice(&value.to_bits().to_le_bytes());
436        }
437    }
438
439    // ============ FLOAT16 (TypeId = 16) ============
440    #[inline(always)]
441    pub fn write_f16(&mut self, value: float16) {
442        self.write_u16(value.to_bits());
443    }
444
445    // ============ BFLOAT16 (TypeId = 18) ============
446    #[inline(always)]
447    pub fn write_bf16(&mut self, value: bfloat16) {
448        self.write_u16(value.to_bits());
449    }
450
451    // ============ FLOAT64 (TypeId = 18) ============
452
453    #[inline(always)]
454    pub fn write_f64(&mut self, value: f64) {
455        #[cfg(target_endian = "little")]
456        {
457            let bytes = unsafe { &*(&value as *const f64 as *const [u8; 8]) };
458            self.bf.extend_from_slice(bytes);
459        }
460        #[cfg(target_endian = "big")]
461        {
462            self.bf.extend_from_slice(&value.to_bits().to_le_bytes());
463        }
464    }
465
466    // ============ STRING (TypeId = 19) ============
467
468    #[inline(always)]
469    pub fn write_utf8_string(&mut self, s: &str) {
470        let bytes = s.as_bytes();
471        let len = bytes.len();
472        self.bf.reserve(len);
473        self.bf.extend_from_slice(bytes);
474    }
475
476    // ============ Rust-specific types (i128, u128, isize, usize) ============
477
478    #[inline(always)]
479    pub fn write_i128(&mut self, value: i128) {
480        self.write_u128(value as u128);
481    }
482
483    #[inline(always)]
484    pub fn write_u128(&mut self, value: u128) {
485        #[cfg(target_endian = "little")]
486        {
487            let bytes = unsafe { &*(&value as *const u128 as *const [u8; 16]) };
488            self.bf.extend_from_slice(bytes);
489        }
490        #[cfg(target_endian = "big")]
491        {
492            self.bf.extend_from_slice(&value.to_le_bytes());
493        }
494    }
495
496    #[inline(always)]
497    pub fn write_isize(&mut self, value: isize) {
498        const SIZE: usize = std::mem::size_of::<isize>();
499        match SIZE {
500            2 => self.write_i16(value as i16),
501            4 => self.write_var_i32(value as i32),
502            8 => self.write_var_i64(value as i64),
503            _ => unreachable!("unsupported isize size"),
504        }
505    }
506
507    #[inline(always)]
508    pub fn write_usize(&mut self, value: usize) {
509        const SIZE: usize = std::mem::size_of::<usize>();
510        match SIZE {
511            2 => self.write_u16(value as u16),
512            4 => self.write_var_u32(value as u32),
513            8 => self.write_var_u64(value as u64),
514            _ => unreachable!("unsupported usize size"),
515        }
516    }
517
518    // ============ Other helper methods ============
519
520    #[inline(always)]
521    pub fn write_var_u36_small(&mut self, value: u64) {
522        assert!(
523            value < (1u64 << 36),
524            "value too large for 36-bit variable-length integer"
525        );
526        if value < 0x80 {
527            self.bf.push(value as u8);
528        } else if value < 0x4000 {
529            let b0 = ((value as u8) & 0x7f) | 0x80;
530            let b1 = (value >> 7) as u8;
531            self.write_u16(((b1 as u16) << 8) | b0 as u16);
532        } else if value < 0x200000 {
533            let b0 = ((value as u8) & 0x7f) | 0x80;
534            let b1 = (((value >> 7) as u8) & 0x7f) | 0x80;
535            let b2 = (value >> 14) as u8;
536            self.write_u24(((b2 as u32) << 16) | ((b1 as u32) << 8) | b0 as u32);
537        } else if value < 0x10000000 {
538            let b0 = ((value as u8) & 0x7f) | 0x80;
539            let b1 = (((value >> 7) as u8) & 0x7f) | 0x80;
540            let b2 = (((value >> 14) as u8) & 0x7f) | 0x80;
541            let b3 = (value >> 21) as u8;
542            self.write_u32(
543                ((b3 as u32) << 24) | ((b2 as u32) << 16) | ((b1 as u32) << 8) | b0 as u32,
544            );
545        } else if value < (1u64 << 35) {
546            let b0 = ((value as u8) & 0x7f) | 0x80;
547            let b1 = (((value >> 7) as u8) & 0x7f) | 0x80;
548            let b2 = (((value >> 14) as u8) & 0x7f) | 0x80;
549            let b3 = (((value >> 21) as u8) & 0x7f) | 0x80;
550            let b4 = (value >> 28) as u8;
551            self.write_u40(
552                ((b4 as u64) << 32)
553                    | ((b3 as u64) << 24)
554                    | ((b2 as u64) << 16)
555                    | ((b1 as u64) << 8)
556                    | b0 as u64,
557            );
558        } else {
559            // Standard seven-bit varuint framing needs a sixth byte for bit 35.
560            let b0 = ((value as u8) & 0x7f) | 0x80;
561            let b1 = (((value >> 7) as u8) & 0x7f) | 0x80;
562            let b2 = (((value >> 14) as u8) & 0x7f) | 0x80;
563            let b3 = (((value >> 21) as u8) & 0x7f) | 0x80;
564            let b4 = (((value >> 28) as u8) & 0x7f) | 0x80;
565            let b5 = (value >> 35) as u8;
566            self.write_u32(
567                ((b3 as u32) << 24) | ((b2 as u32) << 16) | ((b1 as u32) << 8) | b0 as u32,
568            );
569            self.write_u16(((b5 as u16) << 8) | b4 as u16);
570        }
571    }
572}
573
574#[derive(Default)]
575#[allow(clippy::needless_lifetimes)]
576pub struct Reader<'a> {
577    pub(crate) bf: &'a [u8],
578    pub(crate) cursor: usize,
579}
580
581#[allow(clippy::needless_lifetimes)]
582impl<'a> Reader<'a> {
583    // ============ Utility methods ============
584
585    #[inline(always)]
586    pub fn new(bf: &[u8]) -> Reader<'_> {
587        Reader { bf, cursor: 0 }
588    }
589
590    #[inline(always)]
591    pub(crate) fn move_next(&mut self, additional: usize) {
592        self.cursor += additional;
593    }
594
595    #[inline(always)]
596    pub(crate) fn move_back(&mut self, additional: usize) {
597        self.cursor -= additional;
598    }
599
600    #[inline(always)]
601    pub fn sub_slice(&self, start: usize, end: usize) -> Result<&[u8], Error> {
602        // Allow start == bf.len() when end == bf.len() to support empty slices at buffer end
603        if start > self.bf.len() || end > self.bf.len() || end < start {
604            Err(Error::buffer_out_of_bound(
605                start,
606                self.bf.len(),
607                self.bf.len(),
608            ))
609        } else {
610            Ok(&self.bf[start..end])
611        }
612    }
613
614    #[inline(always)]
615    pub fn slice_after_cursor(&self) -> &[u8] {
616        &self.bf[self.cursor..]
617    }
618
619    #[inline(always)]
620    pub fn get_cursor(&self) -> usize {
621        self.cursor
622    }
623
624    #[inline(always)]
625    fn value_at(&self, index: usize) -> Result<u8, Error> {
626        match self.bf.get(index) {
627            None => Err(Error::buffer_out_of_bound(
628                index,
629                self.bf.len(),
630                self.bf.len(),
631            )),
632            Some(v) => Ok(*v),
633        }
634    }
635
636    #[inline(always)]
637    pub(crate) fn check_bound(&self, n: usize) -> Result<(), Error> {
638        let end = self
639            .cursor
640            .checked_add(n)
641            .ok_or_else(|| Error::buffer_out_of_bound(self.cursor, n, self.bf.len()))?;
642        if end > self.bf.len() {
643            Err(Error::buffer_out_of_bound(self.cursor, n, self.bf.len()))
644        } else {
645            Ok(())
646        }
647    }
648
649    #[inline(always)]
650    pub fn skip(&mut self, len: usize) -> Result<(), Error> {
651        self.check_bound(len)?;
652        self.move_next(len);
653        Ok(())
654    }
655
656    #[inline(always)]
657    pub fn read_bytes(&mut self, len: usize) -> Result<&[u8], Error> {
658        self.check_bound(len)?;
659        let result = &self.bf[self.cursor..self.cursor + len];
660        self.move_next(len);
661        Ok(result)
662    }
663
664    #[inline(always)]
665    pub fn reset_cursor_to_here(&self) -> impl FnOnce(&mut Self) {
666        let raw_cursor = self.cursor;
667        move |this: &mut Self| {
668            this.cursor = raw_cursor;
669        }
670    }
671
672    pub fn set_cursor(&mut self, cursor: usize) {
673        self.cursor = cursor;
674    }
675
676    // ============ BOOL (TypeId = 1) ============
677
678    #[inline(always)]
679    pub fn read_bool(&mut self) -> Result<bool, Error> {
680        Ok(self.read_u8()? != 0)
681    }
682
683    // ============ INT8 (TypeId = 2) ============
684
685    #[inline(always)]
686    pub fn read_i8(&mut self) -> Result<i8, Error> {
687        Ok(self.read_u8()? as i8)
688    }
689
690    // ============ INT16 (TypeId = 3) ============
691
692    #[inline(always)]
693    pub fn read_i16(&mut self) -> Result<i16, Error> {
694        Ok(self.read_u16()? as i16)
695    }
696
697    // ============ INT32 (TypeId = 4) ============
698
699    #[inline(always)]
700    pub fn read_i32(&mut self) -> Result<i32, Error> {
701        Ok(self.read_u32()? as i32)
702    }
703
704    // ============ VARINT32 (TypeId = 5) ============
705
706    #[inline(always)]
707    pub fn read_var_i32(&mut self) -> Result<i32, Error> {
708        let encoded = self.read_var_u32()?;
709        Ok(((encoded >> 1) as i32) ^ -((encoded & 1) as i32))
710    }
711
712    // ============ INT64 (TypeId = 6) ============
713
714    #[inline(always)]
715    pub fn read_i64(&mut self) -> Result<i64, Error> {
716        Ok(self.read_u64()? as i64)
717    }
718
719    // ============ VARINT64 (TypeId = 7) ============
720
721    #[inline(always)]
722    pub fn read_var_i64(&mut self) -> Result<i64, Error> {
723        let encoded = self.read_var_u64()?;
724        Ok(((encoded >> 1) as i64) ^ -((encoded & 1) as i64))
725    }
726
727    // ============ TAGGED_INT64 (TypeId = 8) ============
728
729    /// Read signed fory Tagged(Small long as int) encoded i64.
730    /// If bit 0 of the first 4 bytes is 0, return the value >> 1 (arithmetic shift).
731    /// Otherwise, skip the flag byte and read 8 bytes as i64.
732    #[inline(always)]
733    pub fn read_tagged_i64(&mut self) -> Result<i64, Error> {
734        self.check_bound(4)?;
735        let i = LittleEndian::read_i32(&self.bf[self.cursor..]);
736        if (i & 0b1) != 0b1 {
737            // Bit 0 is 0, small value encoded in 4 bytes
738            self.cursor += 4;
739            Ok((i >> 1) as i64) // arithmetic right shift preserves sign
740        } else {
741            // Bit 0 is 1, big value: skip flag byte and read 8 bytes
742            self.check_bound(9)?;
743            self.cursor += 1;
744            let value = LittleEndian::read_i64(&self.bf[self.cursor..]);
745            self.cursor += 8;
746            Ok(value)
747        }
748    }
749
750    // ============ UINT8 (TypeId = 9) ============
751
752    #[inline(always)]
753    pub fn peek_u8(&mut self) -> Result<u8, Error> {
754        let result = self.value_at(self.cursor)?;
755        Ok(result)
756    }
757
758    #[inline(always)]
759    pub fn read_u8(&mut self) -> Result<u8, Error> {
760        let result = self.value_at(self.cursor)?;
761        self.move_next(1);
762        Ok(result)
763    }
764
765    // ============ UINT16 (TypeId = 10) ============
766
767    #[inline(always)]
768    pub fn read_u16(&mut self) -> Result<u16, Error> {
769        self.check_bound(2)?;
770        let result = LittleEndian::read_u16(&self.bf[self.cursor..self.cursor + 2]);
771        self.cursor += 2;
772        Ok(result)
773    }
774
775    // ============ UINT32 (TypeId = 11) ============
776
777    #[inline(always)]
778    pub fn read_u32(&mut self) -> Result<u32, Error> {
779        self.check_bound(4)?;
780        let result = LittleEndian::read_u32(&self.bf[self.cursor..self.cursor + 4]);
781        self.cursor += 4;
782        Ok(result)
783    }
784
785    // ============ VAR_UINT32 (TypeId = 12) ============
786
787    #[inline(always)]
788    pub fn read_var_u32(&mut self) -> Result<u32, Error> {
789        let b0 = self.value_at(self.cursor)? as u32;
790        if b0 < 0x80 {
791            self.move_next(1);
792            return Ok(b0);
793        }
794
795        let b1 = self.value_at(self.cursor + 1)? as u32;
796        let mut encoded = (b0 & 0x7F) | ((b1 & 0x7F) << 7);
797        if b1 < 0x80 {
798            self.move_next(2);
799            return Ok(encoded);
800        }
801
802        let b2 = self.value_at(self.cursor + 2)? as u32;
803        encoded |= (b2 & 0x7F) << 14;
804        if b2 < 0x80 {
805            self.move_next(3);
806            return Ok(encoded);
807        }
808
809        let b3 = self.value_at(self.cursor + 3)? as u32;
810        encoded |= (b3 & 0x7F) << 21;
811        if b3 < 0x80 {
812            self.move_next(4);
813            return Ok(encoded);
814        }
815
816        let b4 = self.value_at(self.cursor + 4)? as u32;
817        encoded |= b4 << 28;
818        self.move_next(5);
819        Ok(encoded)
820    }
821
822    // ============ UINT64 (TypeId = 13) ============
823
824    #[inline(always)]
825    pub fn read_u64(&mut self) -> Result<u64, Error> {
826        self.check_bound(8)?;
827        let result = LittleEndian::read_u64(&self.bf[self.cursor..self.cursor + 8]);
828        self.cursor += 8;
829        Ok(result)
830    }
831
832    // ============ VAR_UINT64 (TypeId = 14) ============
833
834    #[inline(always)]
835    pub fn read_var_u64(&mut self) -> Result<u64, Error> {
836        let b0 = self.value_at(self.cursor)? as u64;
837        if b0 < 0x80 {
838            self.move_next(1);
839            return Ok(b0);
840        }
841
842        let b1 = self.value_at(self.cursor + 1)? as u64;
843        let mut result = (b0 & 0x7F) | ((b1 & 0x7F) << 7);
844        if b1 < 0x80 {
845            self.move_next(2);
846            return Ok(result);
847        }
848
849        let b2 = self.value_at(self.cursor + 2)? as u64;
850        result |= (b2 & 0x7F) << 14;
851        if b2 < 0x80 {
852            self.move_next(3);
853            return Ok(result);
854        }
855
856        let b3 = self.value_at(self.cursor + 3)? as u64;
857        result |= (b3 & 0x7F) << 21;
858        if b3 < 0x80 {
859            self.move_next(4);
860            return Ok(result);
861        }
862
863        let b4 = self.value_at(self.cursor + 4)? as u64;
864        result |= (b4 & 0x7F) << 28;
865        if b4 < 0x80 {
866            self.move_next(5);
867            return Ok(result);
868        }
869
870        let b5 = self.value_at(self.cursor + 5)? as u64;
871        result |= (b5 & 0x7F) << 35;
872        if b5 < 0x80 {
873            self.move_next(6);
874            return Ok(result);
875        }
876
877        let b6 = self.value_at(self.cursor + 6)? as u64;
878        result |= (b6 & 0x7F) << 42;
879        if b6 < 0x80 {
880            self.move_next(7);
881            return Ok(result);
882        }
883
884        let b7 = self.value_at(self.cursor + 7)? as u64;
885        result |= (b7 & 0x7F) << 49;
886        if b7 < 0x80 {
887            self.move_next(8);
888            return Ok(result);
889        }
890
891        let b8 = self.value_at(self.cursor + 8)? as u64;
892        result |= (b8 & 0xFF) << 56;
893        self.move_next(9);
894        Ok(result)
895    }
896
897    // ============ TAGGED_UINT64 (TypeId = 15) ============
898
899    /// Read unsigned fory Tagged(Small long as int) encoded u64.
900    /// If bit 0 of the first 4 bytes is 0, return the value >> 1.
901    /// Otherwise, skip the flag byte and read 8 bytes as u64.
902    #[inline(always)]
903    pub fn read_tagged_u64(&mut self) -> Result<u64, Error> {
904        self.check_bound(4)?;
905        let i = LittleEndian::read_u32(&self.bf[self.cursor..]);
906        if (i & 0b1) != 0b1 {
907            // Bit 0 is 0, small value encoded in 4 bytes
908            self.cursor += 4;
909            Ok((i >> 1) as u64)
910        } else {
911            // Bit 0 is 1, big value: skip flag byte and read 8 bytes
912            self.check_bound(9)?;
913            self.cursor += 1;
914            let value = LittleEndian::read_u64(&self.bf[self.cursor..]);
915            self.cursor += 8;
916            Ok(value)
917        }
918    }
919
920    // ============ FLOAT32 (TypeId = 17) ============
921
922    #[inline(always)]
923    pub fn read_f32(&mut self) -> Result<f32, Error> {
924        self.check_bound(4)?;
925        let result = LittleEndian::read_f32(&self.bf[self.cursor..self.cursor + 4]);
926        self.cursor += 4;
927        Ok(result)
928    }
929
930    // ============ FLOAT64 (TypeId = 18) ============
931    #[inline(always)]
932    pub fn read_f16(&mut self) -> Result<float16, Error> {
933        self.check_bound(2)?;
934        let bits = LittleEndian::read_u16(&self.bf[self.cursor..self.cursor + 2]);
935        self.cursor += 2;
936        Ok(float16::from_bits(bits))
937    }
938
939    #[inline(always)]
940    pub fn read_bf16(&mut self) -> Result<bfloat16, Error> {
941        self.check_bound(2)?;
942        let bits = LittleEndian::read_u16(&self.bf[self.cursor..self.cursor + 2]);
943        self.cursor += 2;
944        Ok(bfloat16::from_bits(bits))
945    }
946
947    pub fn read_f64(&mut self) -> Result<f64, Error> {
948        self.check_bound(8)?;
949        let result = LittleEndian::read_f64(&self.bf[self.cursor..self.cursor + 8]);
950        self.cursor += 8;
951        Ok(result)
952    }
953
954    // ============ STRING (TypeId = 19) ============
955
956    #[inline(always)]
957    pub fn read_latin1_string(&mut self, len: usize) -> Result<String, Error> {
958        self.check_bound(len)?;
959        if len < SIMD_THRESHOLD {
960            // Fast path for small buffers
961            unsafe {
962                let src = self.sub_slice(self.cursor, self.cursor + len)?;
963
964                // Check if all bytes are ASCII (< 0x80)
965                let is_ascii = src.iter().all(|&b| b < 0x80);
966
967                if is_ascii {
968                    // ASCII fast path: Latin1 == UTF-8, direct copy
969                    let mut vec = Vec::with_capacity(len);
970                    let dst = vec.as_mut_ptr();
971                    std::ptr::copy_nonoverlapping(src.as_ptr(), dst, len);
972                    vec.set_len(len);
973                    self.move_next(len);
974                    Ok(String::from_utf8_unchecked(vec))
975                } else {
976                    // Contains Latin1 bytes (0x80-0xFF): must convert to UTF-8
977                    let mut out: Vec<u8> = Vec::with_capacity(len * 2);
978                    let out_ptr = out.as_mut_ptr();
979                    let mut out_len = 0;
980
981                    for &b in src {
982                        if b < 0x80 {
983                            *out_ptr.add(out_len) = b;
984                            out_len += 1;
985                        } else {
986                            // Latin1 -> UTF-8 encoding
987                            *out_ptr.add(out_len) = 0xC0 | (b >> 6);
988                            *out_ptr.add(out_len + 1) = 0x80 | (b & 0x3F);
989                            out_len += 2;
990                        }
991                    }
992
993                    out.set_len(out_len);
994                    self.move_next(len);
995                    Ok(String::from_utf8_unchecked(out))
996                }
997            }
998        } else {
999            // Use SIMD for larger strings where the overhead is amortized
1000            read_latin1_simd(self, len)
1001        }
1002    }
1003
1004    #[inline(always)]
1005    pub fn read_utf8_string(&mut self, len: usize) -> Result<String, Error> {
1006        self.check_bound(len)?;
1007        let src = &self.bf[self.cursor..self.cursor + len];
1008        // Rust is the only runtime that checks UTF-8 string payloads by default; other runtimes
1009        // preserve their platform replacement behavior for invalid byte sequences.
1010        let string =
1011            std::str::from_utf8(src).map_err(|_| Error::encoding_error("invalid UTF-8 string"))?;
1012        let string = string.to_owned();
1013        self.move_next(len);
1014        Ok(string)
1015    }
1016
1017    /// Reads bytes without validating UTF-8.
1018    ///
1019    /// # Safety
1020    ///
1021    /// The next `len` bytes must be valid UTF-8. Violating this requirement creates an invalid
1022    /// [`String`] and breaks its required invariant.
1023    ///
1024    /// ```compile_fail
1025    /// use fory_core::buffer::Reader;
1026    ///
1027    /// let mut reader = Reader::new(b"valid");
1028    /// let _ = reader.read_utf8_string_unchecked(5);
1029    /// ```
1030    #[inline(always)]
1031    pub unsafe fn read_utf8_string_unchecked(&mut self, len: usize) -> Result<String, Error> {
1032        self.check_bound(len)?;
1033        let mut vec = Vec::with_capacity(len);
1034        let src = unsafe { self.bf.as_ptr().add(self.cursor) };
1035        let dst = vec.as_mut_ptr();
1036        unsafe {
1037            std::ptr::copy_nonoverlapping(src, dst, len);
1038            vec.set_len(len);
1039        }
1040        self.move_next(len);
1041        Ok(unsafe { String::from_utf8_unchecked(vec) })
1042    }
1043
1044    #[inline(always)]
1045    pub fn read_utf16_string(&mut self, len: usize) -> Result<String, Error> {
1046        self.check_bound(len)?;
1047        let slice = self.sub_slice(self.cursor, self.cursor + len)?;
1048        let units: Vec<u16> = slice
1049            .chunks_exact(2)
1050            .map(|c| u16::from_le_bytes([c[0], c[1]]))
1051            .collect();
1052        self.move_next(len);
1053        Ok(String::from_utf16_lossy(&units))
1054    }
1055
1056    // ============ Rust-specific types (i128, u128, isize, usize) ============
1057
1058    #[inline(always)]
1059    pub fn read_i128(&mut self) -> Result<i128, Error> {
1060        Ok(self.read_u128()? as i128)
1061    }
1062
1063    #[inline(always)]
1064    pub fn read_u128(&mut self) -> Result<u128, Error> {
1065        self.check_bound(16)?;
1066        let result = LittleEndian::read_u128(&self.bf[self.cursor..self.cursor + 16]);
1067        self.cursor += 16;
1068        Ok(result)
1069    }
1070
1071    #[inline(always)]
1072    pub fn read_isize(&mut self) -> Result<isize, Error> {
1073        const SIZE: usize = std::mem::size_of::<isize>();
1074        match SIZE {
1075            2 => Ok(self.read_i16()? as isize),
1076            4 => Ok(self.read_var_i32()? as isize),
1077            8 => Ok(self.read_var_i64()? as isize),
1078            _ => unreachable!("unsupported isize size"),
1079        }
1080    }
1081
1082    #[inline(always)]
1083    pub fn read_usize(&mut self) -> Result<usize, Error> {
1084        const SIZE: usize = std::mem::size_of::<usize>();
1085        match SIZE {
1086            2 => Ok(self.read_u16()? as usize),
1087            4 => Ok(self.read_var_u32()? as usize),
1088            8 => Ok(self.read_var_u64()? as usize),
1089            _ => unreachable!("unsupported usize size"),
1090        }
1091    }
1092
1093    // ============ Other helper methods ============
1094
1095    #[inline(always)]
1096    pub fn read_var_u36_small(&mut self) -> Result<u64, Error> {
1097        // Keep this API panic-free even if cursor is externally set past buffer end.
1098        self.check_bound(0)?;
1099        let start = self.cursor;
1100        let slice = self.slice_after_cursor();
1101
1102        if slice.len() >= 8 {
1103            // Decode speculatively without advancing so malformed input leaves the cursor intact.
1104            let bulk = LittleEndian::read_u64(&slice[..8]);
1105            let mut result = bulk & 0x7F;
1106            let mut read_idx = start;
1107
1108            if (bulk & 0x80) != 0 {
1109                read_idx += 1;
1110                result |= (bulk >> 1) & 0x3F80;
1111                if (bulk & 0x8000) != 0 {
1112                    read_idx += 1;
1113                    result |= (bulk >> 2) & 0x1FC000;
1114                    if (bulk & 0x800000) != 0 {
1115                        read_idx += 1;
1116                        result |= (bulk >> 3) & 0xFE00000;
1117                        if (bulk & 0x80000000) != 0 {
1118                            read_idx += 1;
1119                            result |= (bulk >> 4) & 0x7F0000000;
1120                            if (bulk & 0x8000000000) != 0 {
1121                                let sixth = ((bulk >> 40) & 0xFF) as u8;
1122                                // Only bit 35 belongs to a 36-bit value; continuation or higher
1123                                // payload bits would extend the value beyond the wire type.
1124                                if sixth > 1 {
1125                                    return Err(Error::invalid_data("var_u36_small overflow"));
1126                                }
1127                                read_idx += 1;
1128                                result |= (sixth as u64) << 35;
1129                            }
1130                        }
1131                    }
1132                }
1133            }
1134            self.cursor = read_idx + 1;
1135            return Ok(result);
1136        }
1137
1138        let mut result = 0u64;
1139        for index in 0..5 {
1140            let b = self.value_at(start + index)?;
1141            result |= ((b & 0x7F) as u64) << (index * 7);
1142            if (b & 0x80) == 0 {
1143                self.cursor = start + index + 1;
1144                return Ok(result);
1145            }
1146        }
1147
1148        let sixth = self.value_at(start + 5)?;
1149        // The sixth group may contain only bit 35 and must terminate the varuint.
1150        if sixth > 1 {
1151            return Err(Error::invalid_data("var_u36_small overflow"));
1152        }
1153        result |= (sixth as u64) << 35;
1154        self.cursor = start + 6;
1155        Ok(result)
1156    }
1157}
1158
1159#[allow(clippy::needless_lifetimes)]
1160unsafe impl<'a> Send for Reader<'a> {}
1161#[allow(clippy::needless_lifetimes)]
1162unsafe impl<'a> Sync for Reader<'a> {}
1163
1164#[cfg(test)]
1165mod tests {
1166    use super::{Reader, Writer};
1167
1168    #[test]
1169    fn varuint_boundary_roundtrip() {
1170        let cases = [
1171            (0x3fff_u32, 2),
1172            (0x4000, 3),
1173            (0x1f_ffff, 3),
1174            (0x20_0000, 4),
1175            (0x0fff_ffff, 4),
1176            (0x1000_0000, 5),
1177            (u32::MAX, 5),
1178        ];
1179        for (value, expected_len) in cases {
1180            let mut buffer = Vec::new();
1181            let mut writer = Writer::from_buffer(&mut buffer);
1182            writer.write_var_u32(value);
1183            assert_eq!(writer.len(), expected_len);
1184            let bytes = writer.dump();
1185            assert_eq!(Reader::new(&bytes).read_var_u32().unwrap(), value);
1186            let mut padded = bytes;
1187            padded.extend_from_slice(&[0; 8]);
1188            let mut reader = Reader::new(&padded);
1189            assert_eq!(reader.read_var_u32().unwrap(), value);
1190            assert_eq!(reader.get_cursor(), expected_len);
1191
1192            let mut buffer = Vec::new();
1193            let mut writer = Writer::from_buffer(&mut buffer);
1194            writer.write_var_u64(u64::from(value));
1195            assert_eq!(writer.len(), expected_len);
1196            let bytes = writer.dump();
1197            assert_eq!(
1198                Reader::new(&bytes).read_var_u64().unwrap(),
1199                u64::from(value)
1200            );
1201        }
1202    }
1203}