heapdump_vmo/
stack_trace_compression.rs

1// Copyright 2023 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5/// Given the size of an uncompressed stack trace (expressed as the number of stack frames), returns
6/// an upper bound of its compressed size.
7pub const fn max_compressed_size(num_frames: usize) -> usize {
8    num_frames * VARINT_MAX_ENCODED_LEN
9}
10
11/// Compresses a stack trace into a preallocated buffer.
12///
13/// The destination buffer's size must be suitable for the input size (see `max_compressed_size` in
14/// this crate).
15///
16/// This function returns the number of bytes actually written into the destination buffer (i.e. the
17/// compressed size).
18pub fn compress_into(src: &[usize], dest: &mut [u8]) -> usize {
19    assert!(dest.len() >= max_compressed_size(src.len()), "dest buffer is not big enough");
20
21    let mut offset = 0;
22    let mut prev = 0;
23    for value in src {
24        let zigzag = zigzag_encode(value.wrapping_sub(prev));
25        offset += varint_encode(zigzag, &mut dest[offset..]);
26        prev = *value;
27    }
28
29    offset
30}
31
32/// Compresses a stack trace into a dynamically-allocated vector of bytes.
33pub fn compress(src: &[usize]) -> Vec<u8> {
34    let mut buf = vec![0; max_compressed_size(src.len())];
35    let compressed_size = compress_into(src, &mut buf);
36    buf.truncate(compressed_size);
37    buf
38}
39
40/// Uncompresses a stack trace.
41///
42/// This function assumes that the uncompressed pointers will fit in usize, which is always true if
43/// the stack trace was generated on the same architecture.
44pub fn uncompress(src: &[u8]) -> Result<Vec<usize>, crate::Error> {
45    let mut result = Vec::new();
46
47    let mut offset = 0;
48    let mut value = 0;
49    while offset != src.len() {
50        let (zigzag, num_bytes) =
51            varint_decode(&src[offset..]).ok_or(crate::Error::InvalidInput)?;
52        offset += num_bytes;
53        value = zigzag_decode(zigzag).wrapping_add(value);
54        result.push(value);
55    }
56
57    Ok(result)
58}
59
60const VARINT_SHIFT: u32 = 7;
61const VARINT_VALUE_MASK: u8 = 0x7f;
62const VARINT_CONT_BIT: u8 = 0x80;
63const VARINT_MAX_ENCODED_LEN: usize = usize::BITS.div_ceil(VARINT_SHIFT) as usize;
64
65/// Encodes a value into the given buffer, returning the number of bytes that were written.
66fn varint_encode(mut value: usize, dest: &mut [u8]) -> usize {
67    assert!(dest.len() >= VARINT_MAX_ENCODED_LEN, "dest buffer is not big enough");
68
69    let mut offset = 0;
70    loop {
71        dest[offset] = value as u8 & VARINT_VALUE_MASK;
72        value >>= VARINT_SHIFT;
73
74        if value != 0 {
75            dest[offset] |= VARINT_CONT_BIT;
76            offset += 1;
77        } else {
78            return offset + 1;
79        }
80    }
81}
82
83/// Tries to decode a value from the given buffer, returning the value and the number of bytes that
84/// were consumed.
85fn varint_decode(src: &[u8]) -> Option<(usize, usize)> {
86    let mut result = 0;
87    let mut offset = 0;
88    let mut shift = 0;
89    loop {
90        // Read the next byte or return None if we are at the end of the array.
91        let input = src.get(offset)?;
92
93        // Left-shift the value at its final position, then right-shift it back to validate that
94        // we didn't lose any bit due to the left-shift overflowing.
95        let value = (input & VARINT_VALUE_MASK) as usize;
96        let value_shifted = value.checked_shl(shift)?;
97        if (value_shifted >> shift) != value {
98            return None; // overflow detected
99        }
100
101        result |= value_shifted;
102        if (input & VARINT_CONT_BIT) != 0 {
103            shift += VARINT_SHIFT;
104            offset += 1;
105        } else {
106            return Some((result, offset + 1));
107        }
108    }
109}
110
111const ZIGZAG_VALUE_MASK: usize = !1;
112const ZIGZAG_SIGN_BIT: usize = 1;
113
114fn zigzag_encode(mut value: usize) -> usize {
115    value = value.rotate_left(1);
116    if (value & ZIGZAG_SIGN_BIT) != 0 {
117        value ^= ZIGZAG_VALUE_MASK;
118    }
119    value
120}
121
122fn zigzag_decode(mut value: usize) -> usize {
123    if (value & ZIGZAG_SIGN_BIT) != 0 {
124        value ^= ZIGZAG_VALUE_MASK;
125    }
126    value.rotate_right(1)
127}
128
129#[cfg(test)]
130mod tests {
131    use super::*;
132    use test_case::test_case;
133
134    // Arbitrary constants, used instead of real code addresses in the tests below:
135    const A: usize = 0x0000_008a_b3fe_9821;
136    const B: usize = 0x0000_812c_6a4a_682e;
137    const C: usize = 0x0048_b5a0_2d45_9e5a;
138    const D: usize = 0xcccc_f148_39a5_9c5a;
139    const F: usize = 0xffff_ffff_ffff_ffff; // all 1's
140    const Z: usize = 0x0000_0000_0000_0000; // all 0's
141
142    #[test_case(&[] ; "empty")]
143    #[test_case(&[A] ; "A")]
144    #[test_case(&[B] ; "B")]
145    #[test_case(&[C] ; "C")]
146    #[test_case(&[D] ; "D")]
147    #[test_case(&[A, B] ; "AB")]
148    #[test_case(&[A, A] ; "AA")]
149    #[test_case(&[A, A, B] ; "AAB")]
150    #[test_case(&[A, B, C, D] ; "ABCD")]
151    #[test_case(&[A, A, B, B, C, C, D, D] ; "AABBCCDD")]
152    #[test_case(&[D, D, C, C, B, B, A, A] ; "DDCCBBAA")]
153    #[test_case(&[F, F, A, F, F] ; "FFAFF")]
154    #[test_case(&[Z, Z, A, Z, Z] ; "ZZAFF")]
155    #[test_case(&[F, Z, F] ; "FZF")]
156    #[test_case(&[Z, F, Z] ; "ZFZ")]
157    fn test_compress_and_uncompress(stack_trace: &[usize]) {
158        let compressed_data = compress(stack_trace);
159        assert_eq!(stack_trace, &uncompress(&compressed_data).unwrap());
160    }
161
162    #[test_case(&[] ; "empty")]
163    #[test_case(&[0xff, 0xff] ; "ends with CONT bit set")]
164    #[test_case(&[0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x02] ;
165        "valid encoding but too big (u64::MAX + 1)")]
166    #[test_case(&[0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x01] ;
167        "valid encoding but too big (u64::MAX * 2^14)")]
168    fn test_varint_decode_bad(data: &[u8]) {
169        assert_eq!(varint_decode(data), None);
170    }
171
172    #[test_case(0 ; "0")]
173    #[test_case(1 ; "1")]
174    #[test_case(-1 ; "negative 1")]
175    #[test_case(1000 ; "1000")]
176    #[test_case(-1000 ; "negative 1000")]
177    #[test_case(isize::MIN ; "min")]
178    #[test_case(isize::MAX ; "max")]
179    fn test_zigzag(value: isize) {
180        let encoded_value = zigzag_encode(value as usize);
181        let decoded_value = zigzag_decode(encoded_value) as isize;
182        assert_eq!(decoded_value, value, "encoded value: {:x}", encoded_value);
183    }
184}
heapdump_vmo/stack_trace_compression.rs

heapdump_vmo/
stack_trace_compression.rs
