Skip to main content

colvec_derive/
lib.rs

1#[cfg(not(test))]
2use proc_macro::TokenStream;
3#[cfg(test)]
4use proc_macro2::TokenStream;
5use quote::quote;
6use syn::DeriveInput;
7
8#[cfg(not(test))]
9#[proc_macro_derive(ColVec)]
10pub fn colvec_derive(input:TokenStream)->TokenStream{
11	let input:DeriveInput=syn::parse_macro_input!(input);
12	colvec_derive_inner(input)
13}
14
15fn colvec_derive_inner(input:DeriveInput)->TokenStream{
16	match input.data{
17		syn::Data::Struct(syn::DataStruct{fields:syn::Fields::Named(fields_named),..})=>derive_struct(input.ident,input.vis,fields_named),
18		_=>unimplemented!("Only structs are supported"),
19	}
20}
21
22fn derive_struct(ident:syn::Ident,vis:syn::Visibility,fields:syn::FieldsNamed)->TokenStream{
23	let colvec_ident_string=format!("{ident}ColVec");
24	let colvec_ident=syn::Ident::new(&colvec_ident_string,ident.span());
25
26	let fields_count=fields.named.len();
27	let mut colvec: syn::ItemStruct = syn::parse_quote!{
28		#vis struct #colvec_ident<A: ::colvec::alloc::Allocator>{
29			buf: ::colvec::raw::RawColVec<#fields_count, #ident, A>,
30			len: usize,
31		}
32	};
33
34	#[cfg(feature = "std")]
35	match colvec.generics.params.first_mut(){
36		Some(syn::GenericParam::Type(type_param))=>{
37			type_param.eq_token=Some(syn::Token![=](type_param.ident.span()));
38			type_param.default=Some(syn::parse_quote!{::colvec::alloc::Global});
39		},
40		// colvec expression always contains type param
41		_ => unreachable!(),
42	}
43
44	#[cfg(feature = "std")]
45	let global = quote! {
46		impl #colvec_ident<::colvec::alloc::Global>{
47			#[inline]
48			#[must_use]
49			pub const fn new() -> Self {
50				Self { buf: ::colvec::raw::RawColVec::new_in(::colvec::alloc::Global), len: 0 }
51			}
52			#[inline]
53			#[must_use]
54			#[track_caller]
55			pub fn with_capacity(capacity: usize) -> Self {
56				Self::with_capacity_in(capacity, ::colvec::alloc::Global)
57			}
58		}
59	};
60
61	// this trait smuggles information about the input type into RawColVec and RawColVecInner
62	let fields_types=fields.named.iter().map(|field|field.ty.clone());
63	let struct_info = quote! {
64		impl ::colvec::raw::StructInfo<#fields_count> for #ident{
65			const LAYOUT: ::core::alloc::Layout = unsafe {
66				let size = Self::FIELDS.size();
67				let align = align_of::<#ident>();
68				::core::alloc::Layout::from_size_align_unchecked(size, align)
69			};
70			const FIELDS: ::colvec::fields::Fields<#fields_count> = ::colvec::fields::Fields::from_sizes([
71				#(size_of::<#fields_types>()),*
72			]);
73		}
74	};
75
76	let field_indices=0..fields_count;
77	let field_types=fields.named.iter().map(|field|field.ty.clone());
78	let field_idents=fields.named.iter().map(|field|field.ident.as_ref().unwrap().clone());
79	let impls = quote! {
80		impl<A: ::colvec::alloc::Allocator> #colvec_ident<A>{
81			#[inline]
82			pub const fn new_in(alloc: A) -> Self {
83				Self { buf: ::colvec::raw::RawColVec::new_in(alloc), len: 0 }
84			}
85			#[inline]
86			#[track_caller]
87			pub fn with_capacity_in(capacity: usize, alloc: A) -> Self {
88				Self { buf: ::colvec::raw::RawColVec::with_capacity_in(capacity, alloc), len: 0 }
89			}
90			#[inline]
91			pub const fn capacity(&self) -> usize {
92				self.buf.capacity()
93			}
94			#[track_caller]
95			pub fn reserve(&mut self, additional: usize) {
96				self.buf.reserve(self.len, additional);
97			}
98			#[inline]
99			const fn as_ptr(&self) -> *const u8 {
100				// We shadow the slice method of the same name to avoid going through
101				// `deref`, which creates an intermediate reference.
102				self.buf.ptr()
103			}
104			#[inline]
105			const fn as_mut_ptr(&mut self) -> *mut u8 {
106				// We shadow the slice method of the same name to avoid going through
107				// `deref_mut`, which creates an intermediate reference.
108				self.buf.ptr()
109			}
110			#[inline]
111			pub unsafe fn set_len(&mut self, new_len: usize) {
112				debug_assert!(new_len <= self.capacity());
113
114				self.len = new_len;
115			}
116			pub fn push(&mut self, value: #ident){
117				// Inform codegen that the length does not change across grow_one().
118				let len = self.len;
119				// This will panic or abort if we would allocate > isize::MAX bytes
120				// or if the length increment would overflow for zero-sized types.
121				if len == self.buf.capacity() {
122					self.buf.grow_one();
123				}
124				unsafe {
125					#(
126						let end = self.as_mut_ptr()
127							.add(self.buf.capacity() * <#ident as ::colvec::raw::StructInfo<#fields_count>>::FIELDS.offset_of(#field_indices))
128							.cast::<#field_types>()
129							.add(len);
130						::core::ptr::write(end, value.#field_idents);
131					)*
132				}
133				self.len = len + 1;
134			}
135			#[inline]
136			#[track_caller]
137			pub fn append(&mut self, other: &mut Self) {
138				unsafe {
139					self.append_elements(other);
140					other.set_len(0);
141				}
142			}
143
144			/// Appends elements to `self` from other buffer.
145			#[inline]
146			#[track_caller]
147			unsafe fn append_elements(&mut self, other: &Self) {
148				let count = other.len();
149				self.reserve(count);
150				let len = self.len();
151				unsafe {
152					<#ident as ::colvec::raw::StructInfo<#fields_count>>::FIELDS.move_fields(
153						other.as_ptr(),
154						self.as_mut_ptr(),
155						other.capacity(),
156						self.capacity(),
157						len,
158						count,
159					)
160				}
161				self.len += count;
162			}
163		   #[inline]
164			pub const fn len(&self) -> usize {
165				self.len
166			}
167		}
168	};
169
170	let field_indices=0..fields_count;
171	let field_types=fields.named.iter().map(|field|field.ty.clone());
172	let field_slice_fn_idents=fields.named.iter().map(|field|{
173		let ident=field.ident.as_ref().unwrap();
174		let slice_ident=format!("{ident}_slice");
175		syn::Ident::new(&slice_ident,ident.span())
176	});
177	let field_slice_mut_fn_idents=fields.named.iter().map(|field|{
178		let ident=field.ident.as_ref().unwrap();
179		let slice_ident=format!("{ident}_slice_mut");
180		syn::Ident::new(&slice_ident,ident.span())
181	});
182	let field_access = quote! {
183		impl<A: ::colvec::alloc::Allocator> #colvec_ident<A>{
184			#(
185				#[inline]
186				pub const fn #field_slice_fn_idents(&self) -> &[#field_types] {
187					unsafe {
188						::core::slice::from_raw_parts(
189							self.as_ptr()
190								.add(self.buf.capacity() * <#ident as ::colvec::raw::StructInfo<#fields_count>>::FIELDS.offset_of(#field_indices))
191								.cast::<#field_types>(),
192							self.len
193						)
194					}
195				}
196				#[inline]
197				pub const fn #field_slice_mut_fn_idents(&mut self) -> &mut [#field_types] {
198					unsafe {
199						::core::slice::from_raw_parts_mut(
200							self.as_mut_ptr()
201								.add(self.buf.capacity() * <#ident as ::colvec::raw::StructInfo<#fields_count>>::FIELDS.offset_of(#field_indices))
202								.cast::<#field_types>(),
203							self.len
204						)
205					}
206				}
207			)*
208		}
209	};
210
211	let mut output=quote! {
212		#colvec
213
214		#struct_info
215
216		#impls
217		#field_access
218	};
219
220	#[cfg(feature = "std")]
221	output.extend(global);
222
223	output.into()
224}
225
226#[cfg(test)]
227mod tests {
228	use super::*;
229	use syn::parse_quote;
230
231	#[test]
232	fn snapshot_test1() {
233		let test1:syn::ItemStruct = parse_quote! {
234			pub struct Test{
235				field0:u8,
236				field1:Option<u8>,
237				field2:i16,
238				field3:u32,
239			}
240		};
241
242		let output = colvec_derive_inner(test1.into());
243
244		// pretend it outputs a file
245		let as_file = syn::parse_file(&output.to_string()).unwrap();
246
247		// format it in a pretty way
248		let formatted = prettyplease::unparse(&as_file);
249
250		// snapshot-test it
251		insta::assert_snapshot!(formatted);
252	}
253}