arrow_avro/reader/
vlq.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::errors::AvroError;
19
20/// Decoder for zig-zag encoded variable length (VLW) integers
21///
22/// See also:
23/// <https://avro.apache.org/docs/1.11.1/specification/#primitive-types-1>
24/// <https://protobuf.dev/programming-guides/encoding/#varints>
25#[derive(Debug, Default)]
26pub struct VLQDecoder {
27    /// Scratch space for decoding VLQ integers
28    in_progress: u64,
29    shift: u32,
30}
31
32impl VLQDecoder {
33    /// Decode a signed long from `buf`
34    ///
35    /// Returns `Err` if more than 10 continuation bytes accumulate across calls without a
36    /// terminator, which would otherwise overflow the accumulated `i64`.
37    pub fn long(&mut self, buf: &mut &[u8]) -> Result<Option<i64>, AvroError> {
38        while let Some(byte) = buf.first().copied() {
39            if self.shift == 63 && byte >= 0x02 {
40                self.in_progress = 0;
41                self.shift = 0;
42                return Err(AvroError::ParseError(
43                    "Malformed Avro varint: too many continuation bytes".to_string(),
44                ));
45            }
46            *buf = &buf[1..];
47            self.in_progress |= ((byte & 0x7F) as u64) << self.shift;
48            self.shift += 7;
49            if byte & 0x80 == 0 {
50                let val = self.in_progress;
51                self.in_progress = 0;
52                self.shift = 0;
53                return Ok(Some((val >> 1) as i64 ^ -((val & 1) as i64)));
54            }
55        }
56        Ok(None)
57    }
58}
59
60/// Read a varint from `buf` returning the decoded `u64` and the number of bytes read
61#[inline]
62pub(crate) fn read_varint(buf: &[u8]) -> Option<(u64, usize)> {
63    let first = *buf.first()?;
64    if first < 0x80 {
65        return Some((first as u64, 1));
66    }
67
68    if let Some(array) = buf.get(..10) {
69        return read_varint_array(array.try_into().unwrap());
70    }
71
72    read_varint_slow(buf)
73}
74
75/// Based on
76/// - <https://github.com/tokio-rs/prost/blob/master/prost/src/encoding/varint.rs#L71>
77/// - <https://github.com/google/protobuf/blob/3.3.x/src/google/protobuf/io/coded_stream.cc#L365-L406>
78/// - <https://github.com/protocolbuffers/protobuf-go/blob/v1.27.1/encoding/protowire/wire.go#L358>
79#[inline]
80fn read_varint_array(buf: [u8; 10]) -> Option<(u64, usize)> {
81    let mut in_progress = 0_u64;
82    for (idx, b) in buf.into_iter().take(9).enumerate() {
83        in_progress += (b as u64) << (7 * idx);
84        if b < 0x80 {
85            return Some((in_progress, idx + 1));
86        }
87        in_progress -= 0x80 << (7 * idx);
88    }
89
90    let b = buf[9] as u64;
91    in_progress += b << (7 * 9);
92    (b < 0x02).then_some((in_progress, 10))
93}
94
95#[inline(never)]
96#[cold]
97fn read_varint_slow(buf: &[u8]) -> Option<(u64, usize)> {
98    let mut value = 0;
99    for (count, _byte) in buf.iter().take(10).enumerate() {
100        let byte = buf[count];
101        value |= u64::from(byte & 0x7F) << (count * 7);
102        if byte <= 0x7F {
103            // Check for u64::MAX overflow. See [`ConsumeVarint`][1] for details.
104            // [1]: https://github.com/protocolbuffers/protobuf-go/blob/v1.27.1/encoding/protowire/wire.go#L358
105            return (count != 9 || byte < 2).then_some((value, count + 1));
106        }
107    }
108
109    None
110}
111
112pub(crate) fn skip_varint(buf: &[u8]) -> Option<usize> {
113    if let Some(array) = buf.get(..10) {
114        return skip_varint_array(array.try_into().unwrap());
115    }
116    skip_varint_slow(buf)
117}
118
119fn skip_varint_array(buf: [u8; 10]) -> Option<usize> {
120    // Using buf.into_iter().enumerate() regresses performance by 1% on x86-64
121    #[allow(clippy::needless_range_loop)]
122    for idx in 0..9 {
123        if buf[idx] < 0x80 {
124            return Some(idx + 1);
125        }
126    }
127    (buf[9] < 0x02).then_some(10)
128}
129
130#[cold]
131fn skip_varint_slow(buf: &[u8]) -> Option<usize> {
132    debug_assert!(
133        buf.len() < 10,
134        "should be only called on buffers too short for the fast path"
135    );
136    for (idx, &byte) in buf.iter().enumerate() {
137        if byte < 0x80 {
138            return Some(idx + 1);
139        }
140    }
141    None
142}
143
144#[cfg(test)]
145mod tests {
146    use super::*;
147
148    fn encode_var(mut n: u64, dst: &mut [u8]) -> usize {
149        let mut i = 0;
150
151        while n >= 0x80 {
152            dst[i] = 0x80 | (n as u8);
153            i += 1;
154            n >>= 7;
155        }
156
157        dst[i] = n as u8;
158        i + 1
159    }
160
161    fn varint_test(a: u64) {
162        let mut buf = [0_u8; 10];
163        let len = encode_var(a, &mut buf);
164        assert_eq!(read_varint(&buf[..len]).unwrap(), (a, len));
165        assert_eq!(read_varint(&buf).unwrap(), (a, len));
166    }
167
168    #[test]
169    fn test_varint() {
170        varint_test(0);
171        varint_test(4395932);
172        varint_test(u64::MAX);
173
174        for _ in 0..1000 {
175            varint_test(rand::random());
176        }
177    }
178
179    fn zigzag_encode(n: i64) -> u64 {
180        ((n << 1) ^ (n >> 63)) as u64
181    }
182
183    fn long_test(n: i64) {
184        let mut buf = [0_u8; 10];
185        let len = encode_var(zigzag_encode(n), &mut buf);
186        let mut decoder = VLQDecoder::default();
187        let mut slice = &buf[..len];
188        assert_eq!(decoder.long(&mut slice).unwrap(), Some(n));
189        assert!(slice.is_empty());
190    }
191
192    #[test]
193    fn test_long_roundtrip() {
194        long_test(0);
195        long_test(1);
196        long_test(-1);
197        long_test(4395932);
198        long_test(-4395932);
199        long_test(i64::MAX);
200        long_test(i64::MIN);
201
202        for _ in 0..1000 {
203            long_test(rand::random());
204        }
205    }
206
207    #[test]
208    fn test_long_overlong_varint_returns_error() {
209        // The Avro file magic followed by 13 continuation bytes and a terminator, as
210        // reported against #10290: previously drove `self.shift` past 63 and panicked
211        // with "attempt to shift left with overflow" inside `<< self.shift`.
212        let overlong = [0xFFu8; 13];
213        let mut decoder = VLQDecoder::default();
214        let mut buf = &overlong[..];
215        assert!(decoder.long(&mut buf).is_err());
216
217        // The decoder's state must be reset after a malformed varint, so a subsequent
218        // call decodes a fresh varint cleanly rather than staying stuck mid-decode.
219        let mut fresh = &[0x02u8][..]; // zigzag-encoded 1
220        assert_eq!(decoder.long(&mut fresh).unwrap(), Some(1));
221    }
222}
arrow_avro/reader/vlq.rs

arrow_avro/reader/
vlq.rs
