1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158

//! This crate provides a `recycle` extension method for Vec.
//! It's intended to change the type of the Vec while "recycling"
//! the underlying allocation. This is a trick that is useful especially
//! when storing data with short lifetimes in Vec:
//! ```
//! # use std::error::Error;
//! # use crate::recycle_vec::VecExt;
//! #
//! # struct Stream(bool);
//! #
//! # impl Stream {
//! #     fn new() -> Self {
//! #         Stream(false)
//! #     }
//! #
//! #     fn next(&mut self) -> Option<&[u8]> {
//! #         if self.0 {
//! #             None
//! #         } else {
//! #             self.0 = true;
//! #             Some(&b"foo"[..])
//! #         }
//! #     }
//! # }
//! #
//! # fn process(input: &[Object<'_>]) -> Result<(), Box<dyn Error>> {
//! #     for obj in input {
//! #         let _ = obj.reference;
//! #     }
//! #     Ok(())
//! # }
//! #
//! # struct Object<'a> {
//! #     reference: &'a [u8],
//! # }
//! #
//! # fn deserialize<'a>(input: &'a [u8], output: &mut Vec<Object<'a>>) -> Result<(), Box<dyn Error>> {
//! #     output.push(Object { reference: input });
//! #     Ok(())
//! # }
//! #
//! # fn main() -> Result<(), Box<dyn Error>> {
//! # let mut stream = Stream::new();
//! let mut objects: Vec<Object<'static>> = Vec::new();    // Any lifetime goes here
//!
//! while let Some(byte_chunk) = stream.next() {           // `byte_chunk` lifetime starts
//!     let mut temp: Vec<Object<'_>> = objects.recycle(); // `temp` lifetime starts
//!
//!     // Zero-copy parsing; deserialized `Object`s have references to `byte_chunk`
//!     deserialize(byte_chunk, &mut temp)?;
//!     process(&temp)?;
//!
//!     objects = temp.recycle();                          // `temp` lifetime ends
//! }                                                      // `byte_chunk` lifetime ends
//! # Ok(())
//! # }
//! ```
//! # Notes about safety
//! This crate uses internally `unsafe` to achieve it's functionality.
//! However, it provides a safe interface. To achieve safety, it does
//! the following precautions:
//! 1. It truncates the `Vec` to zero length, dropping all the values.
//! This ensures that no values of arbitrary types are transmuted
//! accidentally.
//! 2. It checks that the sizes and alignments of the source and target
//! types match. This ensures that the underlying block of memory backing
//! `Vec` is compatible layout-wise. The sizes and alignments are checked
//! statically, so if the compile will fail in case of a mismatch.
//! 3. It creates a new `Vec` value using `from_raw_parts`, instead of
//! transmuting, an operation whose soundness would be questionable.

#![no_std]

extern crate alloc;

use alloc::vec::Vec;

struct AssertSameLayout<A, B>(core::marker::PhantomData<(A, B)>);
impl<A, B> AssertSameLayout<A, B> {
    const OK: () = assert!(
        core::mem::size_of::<A>() == core::mem::size_of::<B>() && core::mem::align_of::<A>() == core::mem::align_of::<B>(),
        "types must have identical size and alignment"
    );
}

/// A trait that provides an API for recycling Vec's internal buffers
pub trait VecExt<T> {
    /// Allows re-interpreting the type of a Vec to reuse the allocation.
    /// The vector is emptied and any values contained in it will be dropped.
    /// The target type must have the same size and alignment as the source type.
    /// This API doesn't transmute any values of T to U, because it makes sure
    /// to empty the vector before any unsafe operations.
    fn recycle<U>(self) -> Vec<U>;
}

impl<T> VecExt<T> for Vec<T> {
    fn recycle<U>(mut self) -> Vec<U> {
        self.clear();

        () = AssertSameLayout::<T, U>::OK;

        let cap = self.capacity();
        let ptr = self.as_mut_ptr() as *mut U;
        core::mem::forget(self);
        unsafe { Vec::from_raw_parts(ptr, 0, cap) }
    }
}

/// Tests that `recycle` successfully re-interprets the type to have different lifetime from the original
#[test]
fn test_recycle_lifetime() {
    use crate::alloc::string::ToString;
    let s_1 = "foo".to_string();
    let mut buf = Vec::with_capacity(100);
    {
        let mut buf2 = buf;
        let s_2 = "foo".to_string();
        buf2.push(s_2.as_str());

        assert_eq!(buf2.len(), 1);
        assert_eq!(buf2.capacity(), 100);

        buf = buf2.recycle();
    }
    buf.push(s_1.as_str());
}

/// Tests that `recycle` successfully re-interprets the type itself
#[test]
fn test_recycle_type() {
    use crate::alloc::string::ToString;
    let s = "foo".to_string();
    let mut buf = Vec::with_capacity(100);
    {
        let mut buf2 = buf.recycle();

        let mut i = Vec::new();
        i.push(1);
        i.push(2);
        i.push(3);

        buf2.push(i.as_slice());

        assert_eq!(buf2.len(), 1);
        assert_eq!(buf2.capacity(), 100);

        buf = buf2.recycle();
    }
    buf.push(s.as_str());
}

#[test]
fn test_layout_assert() {
    let t = trybuild::TestCases::new();
    t.pass("tests/force_build.rs");
    t.compile_fail("tests/recycle_incompatible_size.rs");
    t.compile_fail("tests/recycle_incompatible_alignment.rs");
}