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    // ============ VAR_UINT32 (TypeId = 12) ============
204
205    #[inline(always)]
206    pub fn write_var_u32(&mut self, value: u32) {
207        self._write_var_u32(value)
208    }
209
210    #[inline(always)]
211    fn _write_var_u32(&mut self, value: u32) {
212        if value < 0x80 {
213            self.bf.push(value as u8);
214        } else if value < 0x4000 {
215            // 2 bytes
216            let u1 = ((value as u8) & 0x7F) | 0x80;
217            let u2 = (value >> 7) as u8;
218            self.write_u16(((u2 as u16) << 8) | u1 as u16);
219        } else if value < 0x200000 {
220            // 3 bytes
221            let u1 = ((value as u8) & 0x7F) | 0x80;
222            let u2 = (((value >> 7) as u8) & 0x7F) | 0x80;
223            let u3 = (value >> 14) as u8;
224            self.write_u16(((u2 as u16) << 8) | u1 as u16);
225            self.bf.push(u3);
226        } else if value < 0x10000000 {
227            // 4 bytes
228            let u1 = ((value as u8) & 0x7F) | 0x80;
229            let u2 = (((value >> 7) as u8) & 0x7F) | 0x80;
230            let u3 = (((value >> 14) as u8) & 0x7F) | 0x80;
231            let u4 = (value >> 21) as u8;
232            self.write_u32(
233                ((u4 as u32) << 24) | ((u3 as u32) << 16) | ((u2 as u32) << 8) | u1 as u32,
234            );
235        } else {
236            // 5 bytes
237            let u1 = ((value as u8) & 0x7F) | 0x80;
238            let u2 = (((value >> 7) as u8) & 0x7F) | 0x80;
239            let u3 = (((value >> 14) as u8) & 0x7F) | 0x80;
240            let u4 = (((value >> 21) as u8) & 0x7F) | 0x80;
241            let u5 = (value >> 28) as u8;
242            self.write_u32(
243                ((u4 as u32) << 24) | ((u3 as u32) << 16) | ((u2 as u32) << 8) | u1 as u32,
244            );
245            self.bf.push(u5);
246        }
247    }
248
249    // ============ UINT64 (TypeId = 13) ============
250
251    #[inline(always)]
252    pub fn write_u64(&mut self, value: u64) {
253        #[cfg(target_endian = "little")]
254        {
255            let bytes = unsafe { &*(&value as *const u64 as *const [u8; 8]) };
256            self.bf.extend_from_slice(bytes);
257        }
258        #[cfg(target_endian = "big")]
259        {
260            self.bf.extend_from_slice(&value.to_le_bytes());
261        }
262    }
263
264    // ============ VAR_UINT64 (TypeId = 14) ============
265
266    #[inline(always)]
267    pub fn write_var_u64(&mut self, value: u64) {
268        self._write_var_u64(value);
269    }
270
271    #[inline(always)]
272    fn _write_var_u64(&mut self, value: u64) {
273        if value < 0x80 {
274            self.bf.push(value as u8);
275        } else if value < 0x4000 {
276            let u1 = ((value as u8) & 0x7F) | 0x80;
277            let u2 = (value >> 7) as u8;
278            self.write_u16(((u2 as u16) << 8) | u1 as u16);
279        } else if value < 0x200000 {
280            let u1 = ((value as u8) & 0x7F) | 0x80;
281            let u2 = (((value >> 7) as u8) & 0x7F) | 0x80;
282            let u3 = (value >> 14) as u8;
283            self.write_u16(((u2 as u16) << 8) | u1 as u16);
284            self.bf.push(u3);
285        } else if value < 0x10000000 {
286            let u1 = ((value as u8) & 0x7F) | 0x80;
287            let u2 = (((value >> 7) as u8) & 0x7F) | 0x80;
288            let u3 = (((value >> 14) as u8) & 0x7F) | 0x80;
289            let u4 = (value >> 21) as u8;
290            self.write_u32(
291                ((u4 as u32) << 24) | ((u3 as u32) << 16) | ((u2 as u32) << 8) | u1 as u32,
292            );
293        } else if value < 0x800000000 {
294            let u1 = ((value as u8) & 0x7F) | 0x80;
295            let u2 = (((value >> 7) as u8) & 0x7F) | 0x80;
296            let u3 = (((value >> 14) as u8) & 0x7F) | 0x80;
297            let u4 = (((value >> 21) as u8) & 0x7F) | 0x80;
298            let u5 = (value >> 28) as u8;
299            self.write_u32(
300                ((u4 as u32) << 24) | ((u3 as u32) << 16) | ((u2 as u32) << 8) | u1 as u32,
301            );
302            self.bf.push(u5);
303        } else if value < 0x40000000000 {
304            let u1 = ((value as u8) & 0x7F) | 0x80;
305            let u2 = (((value >> 7) as u8) & 0x7F) | 0x80;
306            let u3 = (((value >> 14) as u8) & 0x7F) | 0x80;
307            let u4 = (((value >> 21) as u8) & 0x7F) | 0x80;
308            let u5 = (((value >> 28) as u8) & 0x7F) | 0x80;
309            let u6 = (value >> 35) as u8;
310            self.write_u32(
311                ((u4 as u32) << 24) | ((u3 as u32) << 16) | ((u2 as u32) << 8) | u1 as u32,
312            );
313            self.write_u16(((u6 as u16) << 8) | u5 as u16);
314        } else if value < 0x2000000000000 {
315            let u1 = ((value as u8) & 0x7F) | 0x80;
316            let u2 = (((value >> 7) as u8) & 0x7F) | 0x80;
317            let u3 = (((value >> 14) as u8) & 0x7F) | 0x80;
318            let u4 = (((value >> 21) as u8) & 0x7F) | 0x80;
319            let u5 = (((value >> 28) as u8) & 0x7F) | 0x80;
320            let u6 = (((value >> 35) as u8) & 0x7F) | 0x80;
321            let u7 = (value >> 42) as u8;
322            self.write_u32(
323                ((u4 as u32) << 24) | ((u3 as u32) << 16) | ((u2 as u32) << 8) | u1 as u32,
324            );
325            self.write_u16(((u6 as u16) << 8) | u5 as u16);
326            self.bf.push(u7);
327        } else if value < 0x100000000000000 {
328            let u1 = ((value as u8) & 0x7F) | 0x80;
329            let u2 = (((value >> 7) as u8) & 0x7F) | 0x80;
330            let u3 = (((value >> 14) as u8) & 0x7F) | 0x80;
331            let u4 = (((value >> 21) as u8) & 0x7F) | 0x80;
332            let u5 = (((value >> 28) as u8) & 0x7F) | 0x80;
333            let u6 = (((value >> 35) as u8) & 0x7F) | 0x80;
334            let u7 = (((value >> 42) as u8) & 0x7F) | 0x80;
335            let u8 = (value >> 49) as u8;
336            self.write_u64(
337                (u8 as u64) << 56
338                    | (u7 as u64) << 48
339                    | (u6 as u64) << 40
340                    | (u5 as u64) << 32
341                    | (u4 as u64) << 24
342                    | (u3 as u64) << 16
343                    | (u2 as u64) << 8
344                    | (u1 as u64),
345            );
346        } else {
347            let u1 = ((value as u8) & 0x7F) | 0x80;
348            let u2 = (((value >> 7) as u8) & 0x7F) | 0x80;
349            let u3 = (((value >> 14) as u8) & 0x7F) | 0x80;
350            let u4 = (((value >> 21) as u8) & 0x7F) | 0x80;
351            let u5 = (((value >> 28) as u8) & 0x7F) | 0x80;
352            let u6 = (((value >> 35) as u8) & 0x7F) | 0x80;
353            let u7 = (((value >> 42) as u8) & 0x7F) | 0x80;
354            let u8 = (((value >> 49) as u8) & 0x7F) | 0x80;
355            let u9 = (value >> 56) as u8;
356            self.write_u64(
357                (u8 as u64) << 56
358                    | (u7 as u64) << 48
359                    | (u6 as u64) << 40
360                    | (u5 as u64) << 32
361                    | (u4 as u64) << 24
362                    | (u3 as u64) << 16
363                    | (u2 as u64) << 8
364                    | (u1 as u64),
365            );
366            self.bf.push(u9);
367        }
368    }
369
370    // ============ TAGGED_UINT64 (TypeId = 15) ============
371
372    /// Write unsigned long using fory Tagged(Small long as int) encoding.
373    /// If value is in [0, 0x7fffffff], encode as 4 bytes: `((value as u32) << 1)`.
374    /// Otherwise write as 9 bytes: `0b1 | little-endian 8 bytes u64`.
375    #[inline(always)]
376    pub fn write_tagged_u64(&mut self, value: u64) {
377        if value <= i32::MAX as u64 {
378            // Fits in 31 bits, encode as 4 bytes with bit 0 = 0
379            let v = (value as u32) << 1;
380            self.write_u32(v);
381        } else {
382            // Write flag byte (0b1) followed by 8-byte u64
383            self.bf.push(0b1);
384            self.write_u64(value);
385        }
386    }
387
388    // ============ FLOAT32 (TypeId = 17) ============
389
390    #[inline(always)]
391    pub fn write_f32(&mut self, value: f32) {
392        #[cfg(target_endian = "little")]
393        {
394            let bytes = unsafe { &*(&value as *const f32 as *const [u8; 4]) };
395            self.bf.extend_from_slice(bytes);
396        }
397        #[cfg(target_endian = "big")]
398        {
399            self.bf.extend_from_slice(&value.to_bits().to_le_bytes());
400        }
401    }
402
403    // ============ FLOAT16 (TypeId = 16) ============
404    #[inline(always)]
405    pub fn write_f16(&mut self, value: float16) {
406        self.write_u16(value.to_bits());
407    }
408
409    // ============ BFLOAT16 (TypeId = 18) ============
410    #[inline(always)]
411    pub fn write_bf16(&mut self, value: bfloat16) {
412        self.write_u16(value.to_bits());
413    }
414
415    // ============ FLOAT64 (TypeId = 18) ============
416
417    #[inline(always)]
418    pub fn write_f64(&mut self, value: f64) {
419        #[cfg(target_endian = "little")]
420        {
421            let bytes = unsafe { &*(&value as *const f64 as *const [u8; 8]) };
422            self.bf.extend_from_slice(bytes);
423        }
424        #[cfg(target_endian = "big")]
425        {
426            self.bf.extend_from_slice(&value.to_bits().to_le_bytes());
427        }
428    }
429
430    // ============ STRING (TypeId = 19) ============
431
432    #[inline(always)]
433    pub fn write_utf8_string(&mut self, s: &str) {
434        let bytes = s.as_bytes();
435        let len = bytes.len();
436        self.bf.reserve(len);
437        self.bf.extend_from_slice(bytes);
438    }
439
440    // ============ Rust-specific types (i128, u128, isize, usize) ============
441
442    #[inline(always)]
443    pub fn write_i128(&mut self, value: i128) {
444        self.write_u128(value as u128);
445    }
446
447    #[inline(always)]
448    pub fn write_u128(&mut self, value: u128) {
449        #[cfg(target_endian = "little")]
450        {
451            let bytes = unsafe { &*(&value as *const u128 as *const [u8; 16]) };
452            self.bf.extend_from_slice(bytes);
453        }
454        #[cfg(target_endian = "big")]
455        {
456            self.bf.extend_from_slice(&value.to_le_bytes());
457        }
458    }
459
460    #[inline(always)]
461    pub fn write_isize(&mut self, value: isize) {
462        const SIZE: usize = std::mem::size_of::<isize>();
463        match SIZE {
464            2 => self.write_i16(value as i16),
465            4 => self.write_var_i32(value as i32),
466            8 => self.write_var_i64(value as i64),
467            _ => unreachable!("unsupported isize size"),
468        }
469    }
470
471    #[inline(always)]
472    pub fn write_usize(&mut self, value: usize) {
473        const SIZE: usize = std::mem::size_of::<usize>();
474        match SIZE {
475            2 => self.write_u16(value as u16),
476            4 => self.write_var_u32(value as u32),
477            8 => self.write_var_u64(value as u64),
478            _ => unreachable!("unsupported usize size"),
479        }
480    }
481
482    // ============ Other helper methods ============
483
484    #[inline(always)]
485    pub fn write_var_u36_small(&mut self, value: u64) {
486        assert!(
487            value < (1u64 << 36),
488            "value too large for 36-bit variable-length integer"
489        );
490        if value < 0x80 {
491            self.bf.push(value as u8);
492        } else if value < 0x4000 {
493            let b0 = ((value & 0x7F) as u8) | 0x80;
494            let b1 = (value >> 7) as u8;
495            let combined = ((b1 as u16) << 8) | (b0 as u16);
496            self.write_u16(combined);
497        } else if value < 0x200000 {
498            let b0 = (value & 0x7F) | 0x80;
499            let b1 = ((value >> 7) & 0x7F) | 0x80;
500            let b2 = value >> 14;
501            let combined = b0 | (b1 << 8) | (b2 << 16);
502            self.write_u32(combined as u32);
503        } else if value < 0x10000000 {
504            let b0 = (value & 0x7F) | 0x80;
505            let b1 = ((value >> 7) & 0x7F) | 0x80;
506            let b2 = ((value >> 14) & 0x7F) | 0x80;
507            let b3 = value >> 21;
508            let combined = b0 | (b1 << 8) | (b2 << 16) | (b3 << 24);
509            self.write_u32(combined as u32);
510        } else {
511            let b0 = (value & 0x7F) | 0x80;
512            let b1 = ((value >> 7) & 0x7F) | 0x80;
513            let b2 = ((value >> 14) & 0x7F) | 0x80;
514            let b3 = ((value >> 21) & 0x7F) | 0x80;
515            let b4 = value >> 28;
516            let combined = b0 | (b1 << 8) | (b2 << 16) | (b3 << 24) | (b4 << 32);
517            self.write_u64(combined);
518        }
519    }
520}
521
522#[derive(Default)]
523#[allow(clippy::needless_lifetimes)]
524pub struct Reader<'a> {
525    pub(crate) bf: &'a [u8],
526    pub(crate) cursor: usize,
527}
528
529#[allow(clippy::needless_lifetimes)]
530impl<'a> Reader<'a> {
531    // ============ Utility methods ============
532
533    #[inline(always)]
534    pub fn new(bf: &[u8]) -> Reader<'_> {
535        Reader { bf, cursor: 0 }
536    }
537
538    #[inline(always)]
539    pub(crate) fn move_next(&mut self, additional: usize) {
540        self.cursor += additional;
541    }
542
543    #[inline(always)]
544    pub(crate) fn move_back(&mut self, additional: usize) {
545        self.cursor -= additional;
546    }
547
548    #[inline(always)]
549    pub fn sub_slice(&self, start: usize, end: usize) -> Result<&[u8], Error> {
550        // Allow start == bf.len() when end == bf.len() to support empty slices at buffer end
551        if start > self.bf.len() || end > self.bf.len() || end < start {
552            Err(Error::buffer_out_of_bound(
553                start,
554                self.bf.len(),
555                self.bf.len(),
556            ))
557        } else {
558            Ok(&self.bf[start..end])
559        }
560    }
561
562    #[inline(always)]
563    pub fn slice_after_cursor(&self) -> &[u8] {
564        &self.bf[self.cursor..]
565    }
566
567    #[inline(always)]
568    pub fn get_cursor(&self) -> usize {
569        self.cursor
570    }
571
572    #[inline(always)]
573    fn value_at(&self, index: usize) -> Result<u8, Error> {
574        match self.bf.get(index) {
575            None => Err(Error::buffer_out_of_bound(
576                index,
577                self.bf.len(),
578                self.bf.len(),
579            )),
580            Some(v) => Ok(*v),
581        }
582    }
583
584    #[inline(always)]
585    pub(crate) fn check_bound(&self, n: usize) -> Result<(), Error> {
586        let end = self
587            .cursor
588            .checked_add(n)
589            .ok_or_else(|| Error::buffer_out_of_bound(self.cursor, n, self.bf.len()))?;
590        if end > self.bf.len() {
591            Err(Error::buffer_out_of_bound(self.cursor, n, self.bf.len()))
592        } else {
593            Ok(())
594        }
595    }
596
597    #[inline(always)]
598    fn read_u8_uncheck(&mut self) -> u8 {
599        let result = unsafe { self.bf.get_unchecked(self.cursor) };
600        self.move_next(1);
601        *result
602    }
603
604    #[inline(always)]
605    pub fn skip(&mut self, len: usize) -> Result<(), Error> {
606        self.check_bound(len)?;
607        self.move_next(len);
608        Ok(())
609    }
610
611    #[inline(always)]
612    pub fn read_bytes(&mut self, len: usize) -> Result<&[u8], Error> {
613        self.check_bound(len)?;
614        let result = &self.bf[self.cursor..self.cursor + len];
615        self.move_next(len);
616        Ok(result)
617    }
618
619    #[inline(always)]
620    pub fn reset_cursor_to_here(&self) -> impl FnOnce(&mut Self) {
621        let raw_cursor = self.cursor;
622        move |this: &mut Self| {
623            this.cursor = raw_cursor;
624        }
625    }
626
627    pub fn set_cursor(&mut self, cursor: usize) {
628        self.cursor = cursor;
629    }
630
631    // ============ BOOL (TypeId = 1) ============
632
633    #[inline(always)]
634    pub fn read_bool(&mut self) -> Result<bool, Error> {
635        Ok(self.read_u8()? != 0)
636    }
637
638    // ============ INT8 (TypeId = 2) ============
639
640    #[inline(always)]
641    pub fn read_i8(&mut self) -> Result<i8, Error> {
642        Ok(self.read_u8()? as i8)
643    }
644
645    // ============ INT16 (TypeId = 3) ============
646
647    #[inline(always)]
648    pub fn read_i16(&mut self) -> Result<i16, Error> {
649        Ok(self.read_u16()? as i16)
650    }
651
652    // ============ INT32 (TypeId = 4) ============
653
654    #[inline(always)]
655    pub fn read_i32(&mut self) -> Result<i32, Error> {
656        Ok(self.read_u32()? as i32)
657    }
658
659    // ============ VARINT32 (TypeId = 5) ============
660
661    #[inline(always)]
662    pub fn read_var_i32(&mut self) -> Result<i32, Error> {
663        let encoded = self.read_var_u32()?;
664        Ok(((encoded >> 1) as i32) ^ -((encoded & 1) as i32))
665    }
666
667    // ============ INT64 (TypeId = 6) ============
668
669    #[inline(always)]
670    pub fn read_i64(&mut self) -> Result<i64, Error> {
671        Ok(self.read_u64()? as i64)
672    }
673
674    // ============ VARINT64 (TypeId = 7) ============
675
676    #[inline(always)]
677    pub fn read_var_i64(&mut self) -> Result<i64, Error> {
678        let encoded = self.read_var_u64()?;
679        Ok(((encoded >> 1) as i64) ^ -((encoded & 1) as i64))
680    }
681
682    // ============ TAGGED_INT64 (TypeId = 8) ============
683
684    /// Read signed fory Tagged(Small long as int) encoded i64.
685    /// If bit 0 of the first 4 bytes is 0, return the value >> 1 (arithmetic shift).
686    /// Otherwise, skip the flag byte and read 8 bytes as i64.
687    #[inline(always)]
688    pub fn read_tagged_i64(&mut self) -> Result<i64, Error> {
689        self.check_bound(4)?;
690        let i = LittleEndian::read_i32(&self.bf[self.cursor..]);
691        if (i & 0b1) != 0b1 {
692            // Bit 0 is 0, small value encoded in 4 bytes
693            self.cursor += 4;
694            Ok((i >> 1) as i64) // arithmetic right shift preserves sign
695        } else {
696            // Bit 0 is 1, big value: skip flag byte and read 8 bytes
697            self.check_bound(9)?;
698            self.cursor += 1;
699            let value = LittleEndian::read_i64(&self.bf[self.cursor..]);
700            self.cursor += 8;
701            Ok(value)
702        }
703    }
704
705    // ============ UINT8 (TypeId = 9) ============
706
707    #[inline(always)]
708    pub fn peek_u8(&mut self) -> Result<u8, Error> {
709        let result = self.value_at(self.cursor)?;
710        Ok(result)
711    }
712
713    #[inline(always)]
714    pub fn read_u8(&mut self) -> Result<u8, Error> {
715        let result = self.value_at(self.cursor)?;
716        self.move_next(1);
717        Ok(result)
718    }
719
720    // ============ UINT16 (TypeId = 10) ============
721
722    #[inline(always)]
723    pub fn read_u16(&mut self) -> Result<u16, Error> {
724        self.check_bound(2)?;
725        let result = LittleEndian::read_u16(&self.bf[self.cursor..self.cursor + 2]);
726        self.cursor += 2;
727        Ok(result)
728    }
729
730    // ============ UINT32 (TypeId = 11) ============
731
732    #[inline(always)]
733    pub fn read_u32(&mut self) -> Result<u32, Error> {
734        self.check_bound(4)?;
735        let result = LittleEndian::read_u32(&self.bf[self.cursor..self.cursor + 4]);
736        self.cursor += 4;
737        Ok(result)
738    }
739
740    // ============ VAR_UINT32 (TypeId = 12) ============
741
742    #[inline(always)]
743    pub fn read_var_u32(&mut self) -> Result<u32, Error> {
744        let b0 = self.value_at(self.cursor)? as u32;
745        if b0 < 0x80 {
746            self.move_next(1);
747            return Ok(b0);
748        }
749
750        let b1 = self.value_at(self.cursor + 1)? as u32;
751        let mut encoded = (b0 & 0x7F) | ((b1 & 0x7F) << 7);
752        if b1 < 0x80 {
753            self.move_next(2);
754            return Ok(encoded);
755        }
756
757        let b2 = self.value_at(self.cursor + 2)? as u32;
758        encoded |= (b2 & 0x7F) << 14;
759        if b2 < 0x80 {
760            self.move_next(3);
761            return Ok(encoded);
762        }
763
764        let b3 = self.value_at(self.cursor + 3)? as u32;
765        encoded |= (b3 & 0x7F) << 21;
766        if b3 < 0x80 {
767            self.move_next(4);
768            return Ok(encoded);
769        }
770
771        let b4 = self.value_at(self.cursor + 4)? as u32;
772        encoded |= b4 << 28;
773        self.move_next(5);
774        Ok(encoded)
775    }
776
777    // ============ UINT64 (TypeId = 13) ============
778
779    #[inline(always)]
780    pub fn read_u64(&mut self) -> Result<u64, Error> {
781        self.check_bound(8)?;
782        let result = LittleEndian::read_u64(&self.bf[self.cursor..self.cursor + 8]);
783        self.cursor += 8;
784        Ok(result)
785    }
786
787    // ============ VAR_UINT64 (TypeId = 14) ============
788
789    #[inline(always)]
790    pub fn read_var_u64(&mut self) -> Result<u64, Error> {
791        let b0 = self.value_at(self.cursor)? as u64;
792        if b0 < 0x80 {
793            self.move_next(1);
794            return Ok(b0);
795        }
796
797        let b1 = self.value_at(self.cursor + 1)? as u64;
798        let mut result = (b0 & 0x7F) | ((b1 & 0x7F) << 7);
799        if b1 < 0x80 {
800            self.move_next(2);
801            return Ok(result);
802        }
803
804        let b2 = self.value_at(self.cursor + 2)? as u64;
805        result |= (b2 & 0x7F) << 14;
806        if b2 < 0x80 {
807            self.move_next(3);
808            return Ok(result);
809        }
810
811        let b3 = self.value_at(self.cursor + 3)? as u64;
812        result |= (b3 & 0x7F) << 21;
813        if b3 < 0x80 {
814            self.move_next(4);
815            return Ok(result);
816        }
817
818        let b4 = self.value_at(self.cursor + 4)? as u64;
819        result |= (b4 & 0x7F) << 28;
820        if b4 < 0x80 {
821            self.move_next(5);
822            return Ok(result);
823        }
824
825        let b5 = self.value_at(self.cursor + 5)? as u64;
826        result |= (b5 & 0x7F) << 35;
827        if b5 < 0x80 {
828            self.move_next(6);
829            return Ok(result);
830        }
831
832        let b6 = self.value_at(self.cursor + 6)? as u64;
833        result |= (b6 & 0x7F) << 42;
834        if b6 < 0x80 {
835            self.move_next(7);
836            return Ok(result);
837        }
838
839        let b7 = self.value_at(self.cursor + 7)? as u64;
840        result |= (b7 & 0x7F) << 49;
841        if b7 < 0x80 {
842            self.move_next(8);
843            return Ok(result);
844        }
845
846        let b8 = self.value_at(self.cursor + 8)? as u64;
847        result |= (b8 & 0xFF) << 56;
848        self.move_next(9);
849        Ok(result)
850    }
851
852    // ============ TAGGED_UINT64 (TypeId = 15) ============
853
854    /// Read unsigned fory Tagged(Small long as int) encoded u64.
855    /// If bit 0 of the first 4 bytes is 0, return the value >> 1.
856    /// Otherwise, skip the flag byte and read 8 bytes as u64.
857    #[inline(always)]
858    pub fn read_tagged_u64(&mut self) -> Result<u64, Error> {
859        self.check_bound(4)?;
860        let i = LittleEndian::read_u32(&self.bf[self.cursor..]);
861        if (i & 0b1) != 0b1 {
862            // Bit 0 is 0, small value encoded in 4 bytes
863            self.cursor += 4;
864            Ok((i >> 1) as u64)
865        } else {
866            // Bit 0 is 1, big value: skip flag byte and read 8 bytes
867            self.check_bound(9)?;
868            self.cursor += 1;
869            let value = LittleEndian::read_u64(&self.bf[self.cursor..]);
870            self.cursor += 8;
871            Ok(value)
872        }
873    }
874
875    // ============ FLOAT32 (TypeId = 17) ============
876
877    #[inline(always)]
878    pub fn read_f32(&mut self) -> Result<f32, Error> {
879        self.check_bound(4)?;
880        let result = LittleEndian::read_f32(&self.bf[self.cursor..self.cursor + 4]);
881        self.cursor += 4;
882        Ok(result)
883    }
884
885    // ============ FLOAT64 (TypeId = 18) ============
886    #[inline(always)]
887    pub fn read_f16(&mut self) -> Result<float16, Error> {
888        self.check_bound(2)?;
889        let bits = LittleEndian::read_u16(&self.bf[self.cursor..self.cursor + 2]);
890        self.cursor += 2;
891        Ok(float16::from_bits(bits))
892    }
893
894    #[inline(always)]
895    pub fn read_bf16(&mut self) -> Result<bfloat16, Error> {
896        self.check_bound(2)?;
897        let bits = LittleEndian::read_u16(&self.bf[self.cursor..self.cursor + 2]);
898        self.cursor += 2;
899        Ok(bfloat16::from_bits(bits))
900    }
901
902    pub fn read_f64(&mut self) -> Result<f64, Error> {
903        self.check_bound(8)?;
904        let result = LittleEndian::read_f64(&self.bf[self.cursor..self.cursor + 8]);
905        self.cursor += 8;
906        Ok(result)
907    }
908
909    // ============ STRING (TypeId = 19) ============
910
911    #[inline(always)]
912    pub fn read_latin1_string(&mut self, len: usize) -> Result<String, Error> {
913        self.check_bound(len)?;
914        if len < SIMD_THRESHOLD {
915            // Fast path for small buffers
916            unsafe {
917                let src = self.sub_slice(self.cursor, self.cursor + len)?;
918
919                // Check if all bytes are ASCII (< 0x80)
920                let is_ascii = src.iter().all(|&b| b < 0x80);
921
922                if is_ascii {
923                    // ASCII fast path: Latin1 == UTF-8, direct copy
924                    let mut vec = Vec::with_capacity(len);
925                    let dst = vec.as_mut_ptr();
926                    std::ptr::copy_nonoverlapping(src.as_ptr(), dst, len);
927                    vec.set_len(len);
928                    self.move_next(len);
929                    Ok(String::from_utf8_unchecked(vec))
930                } else {
931                    // Contains Latin1 bytes (0x80-0xFF): must convert to UTF-8
932                    let mut out: Vec<u8> = Vec::with_capacity(len * 2);
933                    let out_ptr = out.as_mut_ptr();
934                    let mut out_len = 0;
935
936                    for &b in src {
937                        if b < 0x80 {
938                            *out_ptr.add(out_len) = b;
939                            out_len += 1;
940                        } else {
941                            // Latin1 -> UTF-8 encoding
942                            *out_ptr.add(out_len) = 0xC0 | (b >> 6);
943                            *out_ptr.add(out_len + 1) = 0x80 | (b & 0x3F);
944                            out_len += 2;
945                        }
946                    }
947
948                    out.set_len(out_len);
949                    self.move_next(len);
950                    Ok(String::from_utf8_unchecked(out))
951                }
952            }
953        } else {
954            // Use SIMD for larger strings where the overhead is amortized
955            read_latin1_simd(self, len)
956        }
957    }
958
959    #[inline(always)]
960    pub fn read_utf8_string(&mut self, len: usize) -> Result<String, Error> {
961        self.check_bound(len)?;
962        let src = &self.bf[self.cursor..self.cursor + len];
963        // Rust is the only runtime that checks UTF-8 string payloads by default; other runtimes
964        // preserve their platform replacement behavior for invalid byte sequences.
965        let string =
966            std::str::from_utf8(src).map_err(|_| Error::encoding_error("invalid UTF-8 string"))?;
967        let string = string.to_owned();
968        self.move_next(len);
969        Ok(string)
970    }
971
972    #[inline(always)]
973    pub fn read_utf8_string_unchecked(&mut self, len: usize) -> Result<String, Error> {
974        self.check_bound(len)?;
975        // don't use simd for memory copy, copy_non_overlapping is faster
976        unsafe {
977            let mut vec = Vec::with_capacity(len);
978            let src = self.bf.as_ptr().add(self.cursor);
979            let dst = vec.as_mut_ptr();
980            // Use fastest possible copy - copy_nonoverlapping compiles to memcpy
981            std::ptr::copy_nonoverlapping(src, dst, len);
982            vec.set_len(len);
983            self.move_next(len);
984            Ok(String::from_utf8_unchecked(vec))
985        }
986    }
987
988    #[inline(always)]
989    pub fn read_utf16_string(&mut self, len: usize) -> Result<String, Error> {
990        self.check_bound(len)?;
991        let slice = self.sub_slice(self.cursor, self.cursor + len)?;
992        let units: Vec<u16> = slice
993            .chunks_exact(2)
994            .map(|c| u16::from_le_bytes([c[0], c[1]]))
995            .collect();
996        self.move_next(len);
997        Ok(String::from_utf16_lossy(&units))
998    }
999
1000    // ============ Rust-specific types (i128, u128, isize, usize) ============
1001
1002    #[inline(always)]
1003    pub fn read_i128(&mut self) -> Result<i128, Error> {
1004        Ok(self.read_u128()? as i128)
1005    }
1006
1007    #[inline(always)]
1008    pub fn read_u128(&mut self) -> Result<u128, Error> {
1009        self.check_bound(16)?;
1010        let result = LittleEndian::read_u128(&self.bf[self.cursor..self.cursor + 16]);
1011        self.cursor += 16;
1012        Ok(result)
1013    }
1014
1015    #[inline(always)]
1016    pub fn read_isize(&mut self) -> Result<isize, Error> {
1017        const SIZE: usize = std::mem::size_of::<isize>();
1018        match SIZE {
1019            2 => Ok(self.read_i16()? as isize),
1020            4 => Ok(self.read_var_i32()? as isize),
1021            8 => Ok(self.read_var_i64()? as isize),
1022            _ => unreachable!("unsupported isize size"),
1023        }
1024    }
1025
1026    #[inline(always)]
1027    pub fn read_usize(&mut self) -> Result<usize, Error> {
1028        const SIZE: usize = std::mem::size_of::<usize>();
1029        match SIZE {
1030            2 => Ok(self.read_u16()? as usize),
1031            4 => Ok(self.read_var_u32()? as usize),
1032            8 => Ok(self.read_var_u64()? as usize),
1033            _ => unreachable!("unsupported usize size"),
1034        }
1035    }
1036
1037    // ============ Other helper methods ============
1038
1039    #[inline(always)]
1040    pub fn read_var_u36_small(&mut self) -> Result<u64, Error> {
1041        // Keep this API panic-free even if cursor is externally set past buffer end.
1042        self.check_bound(0)?;
1043        let start = self.cursor;
1044        let slice = self.slice_after_cursor();
1045
1046        if slice.len() >= 8 {
1047            // here already check bound
1048            let bulk = self.read_u64()?;
1049            let mut result = bulk & 0x7F;
1050            let mut read_idx = start;
1051
1052            if (bulk & 0x80) != 0 {
1053                read_idx += 1;
1054                result |= (bulk >> 1) & 0x3F80;
1055                if (bulk & 0x8000) != 0 {
1056                    read_idx += 1;
1057                    result |= (bulk >> 2) & 0x1FC000;
1058                    if (bulk & 0x800000) != 0 {
1059                        read_idx += 1;
1060                        result |= (bulk >> 3) & 0xFE00000;
1061                        if (bulk & 0x80000000) != 0 {
1062                            read_idx += 1;
1063                            result |= (bulk >> 4) & 0xFF0000000;
1064                        }
1065                    }
1066                }
1067            }
1068            self.cursor = read_idx + 1;
1069            return Ok(result);
1070        }
1071
1072        let mut result = 0u64;
1073        let mut shift = 0;
1074        while self.cursor < self.bf.len() {
1075            let b = self.read_u8_uncheck();
1076            result |= ((b & 0x7F) as u64) << shift;
1077            if (b & 0x80) == 0 {
1078                break;
1079            }
1080            shift += 7;
1081            if shift >= 36 {
1082                return Err(Error::encode_error("var_u36_small overflow"));
1083            }
1084        }
1085        Ok(result)
1086    }
1087}
1088
1089#[allow(clippy::needless_lifetimes)]
1090unsafe impl<'a> Send for Reader<'a> {}
1091#[allow(clippy::needless_lifetimes)]
1092unsafe impl<'a> Sync for Reader<'a> {}
fory_core/buffer.rs

fory_core/
buffer.rs
