colvec-derive 0.0.4

#[cfg(not(test))]
use proc_macro::TokenStream;
#[cfg(test)]
use proc_macro2::TokenStream;
use quote::quote;
use syn::DeriveInput;

#[cfg(not(test))]
#[proc_macro_derive(ColVec)]
pub fn colvec_derive(input:TokenStream)->TokenStream{
	let input:DeriveInput=syn::parse_macro_input!(input);
	colvec_derive_inner(input)
}

fn colvec_derive_inner(input:DeriveInput)->TokenStream{
	match input.data{
		syn::Data::Struct(syn::DataStruct{fields:syn::Fields::Named(fields_named),..})=>derive_struct(input.ident,input.vis,fields_named),
		_=>unimplemented!("Only structs are supported"),
	}
}

fn derive_struct(ident:syn::Ident,vis:syn::Visibility,fields:syn::FieldsNamed)->TokenStream{
	let colvec_ident_string=format!("{ident}ColVec");
	let colvec_ident=syn::Ident::new(&colvec_ident_string,ident.span());

	let fields_count=fields.named.len();
	#[cfg_attr(not(feature = "std"), expect(unused_mut))]
	let mut colvec: syn::ItemStruct = syn::parse_quote!{
		#vis struct #colvec_ident<A: ::colvec::alloc::Allocator>{
			buf: ::colvec::raw::RawColVec<#fields_count, #ident, A>,
			len: usize,
		}
	};

	#[cfg(feature = "std")]
	match colvec.generics.params.first_mut(){
		Some(syn::GenericParam::Type(type_param))=>{
			let eq_token=syn::Token![=](type_param.ident.span());
			let default=syn::parse_quote!{::colvec::alloc::Global};
			type_param.default=Some((eq_token,default));
		},
		// colvec expression always contains type param
		_ => unreachable!(),
	}

	#[cfg(feature = "std")]
	let global = quote! {
		impl #colvec_ident<::colvec::alloc::Global>{
			#[inline]
			#[must_use]
			pub const fn new() -> Self {
				Self { buf: ::colvec::raw::RawColVec::new_in(::colvec::alloc::Global), len: 0 }
			}
			#[inline]
			#[must_use]
			#[track_caller]
			pub fn with_capacity(capacity: usize) -> Self {
				Self::with_capacity_in(capacity, ::colvec::alloc::Global)
			}
			#[inline]
			pub unsafe fn from_raw_parts(ptr: *mut u8, length: usize, capacity: usize) -> Self {
				unsafe { Self::from_raw_parts_in(ptr, length, capacity, ::colvec::alloc::Global) }
			}
		}
	};

	// this trait smuggles information about the input type into RawColVec and RawColVecInner
	let fields_types=fields.named.iter().map(|field|field.ty.clone());
	let struct_info = quote! {
		impl ::colvec::raw::StructInfo<#fields_count> for #ident{
			const LAYOUT: ::core::alloc::Layout = unsafe {
				let size = Self::FIELDS.size();
				let align = align_of::<#ident>();
				::core::alloc::Layout::from_size_align_unchecked(size, align)
			};
			const FIELDS: ::colvec::fields::Fields<#fields_count> = ::colvec::fields::Fields::from_sizes([
				#(size_of::<#fields_types>()),*
			]);
		}
	};

	let field_indices=0..fields_count;
	let field_types=fields.named.iter().map(|field|field.ty.clone());
	let field_idents=fields.named.iter().map(|field|field.ident.as_ref().unwrap().clone());
	let impls = quote! {
		impl<A: ::colvec::alloc::Allocator> #colvec_ident<A>{
			#[inline]
			pub const fn new_in(alloc: A) -> Self {
				Self { buf: ::colvec::raw::RawColVec::new_in(alloc), len: 0 }
			}
			#[inline]
			#[track_caller]
			pub fn with_capacity_in(capacity: usize, alloc: A) -> Self {
				Self { buf: ::colvec::raw::RawColVec::with_capacity_in(capacity, alloc), len: 0 }
			}
			#[inline]
			pub unsafe fn from_raw_parts_in(ptr: *mut u8, length: usize, capacity: usize, alloc: A) -> Self {
				unsafe { Self { buf: ::colvec::raw::RawColVec::from_raw_parts_in(ptr, capacity, alloc), len: length } }
			}
			#[must_use = "losing the pointer will leak memory"]
			pub fn into_raw_parts(self) -> (*mut u8, usize, usize) {
				let mut me = ::core::mem::ManuallyDrop::new(self);
				(me.as_mut_ptr(), me.len(), me.capacity())
			}
			#[must_use = "losing the pointer will leak memory"]
			pub fn into_raw_parts_with_alloc(self) -> (*mut u8, usize, usize, A) {
				let mut me = ::core::mem::ManuallyDrop::new(self);
				let len = me.len();
				let capacity = me.capacity();
				let ptr = me.as_mut_ptr();
				let alloc = unsafe { ::core::ptr::read(me.allocator()) };
				(ptr, len, capacity, alloc)
			}
			#[inline]
			pub const fn capacity(&self) -> usize {
				self.buf.capacity()
			}
			#[track_caller]
			pub fn reserve(&mut self, additional: usize) {
				self.buf.reserve(self.len, additional);
			}
			#[inline]
			const fn as_ptr(&self) -> *const u8 {
				// We shadow the slice method of the same name to avoid going through
				// `deref`, which creates an intermediate reference.
				self.buf.ptr()
			}
			#[inline]
			const fn as_mut_ptr(&mut self) -> *mut u8 {
				// We shadow the slice method of the same name to avoid going through
				// `deref_mut`, which creates an intermediate reference.
				self.buf.ptr()
			}
			#[inline]
			pub fn allocator(&self) -> &A {
				self.buf.allocator()
			}
			#[inline]
			pub unsafe fn set_len(&mut self, new_len: usize) {
				debug_assert!(new_len <= self.capacity());

				self.len = new_len;
			}
			pub fn push(&mut self, value: #ident){
				// Inform codegen that the length does not change across grow_one().
				let len = self.len;
				// This will panic or abort if we would allocate > isize::MAX bytes
				// or if the length increment would overflow for zero-sized types.
				if len == self.buf.capacity() {
					self.buf.grow_one();
				}
				unsafe {
					#(
						let end = self.as_mut_ptr()
							.add(self.buf.capacity() * <#ident as ::colvec::raw::StructInfo<#fields_count>>::FIELDS.offset_of(#field_indices))
							.cast::<#field_types>()
							.add(len);
						::core::ptr::write(end, value.#field_idents);
					)*
				}
				self.len = len + 1;
			}
			#[inline]
			#[track_caller]
			pub fn append(&mut self, other: &mut Self) {
				unsafe {
					self.append_elements(other);
					other.set_len(0);
				}
			}

			/// Appends elements to `self` from other buffer.
			#[inline]
			#[track_caller]
			unsafe fn append_elements(&mut self, other: &Self) {
				let count = other.len();
				self.reserve(count);
				let len = self.len();
				unsafe {
					<#ident as ::colvec::raw::StructInfo<#fields_count>>::FIELDS.move_fields(
						other.as_ptr(),
						self.as_mut_ptr(),
						other.capacity(),
						self.capacity(),
						len,
						count,
					)
				}
				self.len += count;
			}
		   #[inline]
			pub const fn len(&self) -> usize {
				self.len
			}
		}
	};

	let field_slice_mut_fn_idents=fields.named.iter().map(|field|{
		let ident=field.ident.as_ref().unwrap();
		let slice_ident=format!("{ident}_slice_mut");
		syn::Ident::new(&slice_ident,ident.span())
	});
	let drop = quote! {
		impl<A: ::colvec::alloc::Allocator> Drop for #colvec_ident<A>{
			fn drop(&mut self){
				unsafe {
					#(
						::core::ptr::drop_in_place(self.#field_slice_mut_fn_idents());
					)*
				}
			}
		}
	};

	let field_indices=0..fields_count;
	let field_types1=fields.named.iter().map(|field|field.ty.clone());
	let field_types2=field_types1.clone();
	let field_idents1=fields.named.iter().map(|field|field.ident.as_ref().unwrap().clone());
	let field_idents2=field_idents1.clone();
	let field_slice_fn_idents1=fields.named.iter().map(|field|{
		let ident=field.ident.as_ref().unwrap();
		let slice_ident=format!("{ident}_slice");
		syn::Ident::new(&slice_ident,ident.span())
	});
	let field_slice_fn_idents2=field_slice_fn_idents1.clone();
	let field_slice_fn_idents3=field_slice_fn_idents1.clone();
	let field_slice_mut_fn_idents1=fields.named.iter().map(|field|{
		let ident=field.ident.as_ref().unwrap();
		let slice_ident=format!("{ident}_slice_mut");
		syn::Ident::new(&slice_ident,ident.span())
	});
	let field_slice_mut_fn_idents2=field_slice_mut_fn_idents1.clone();
	let clone = quote! {
		impl<A: ::colvec::alloc::Allocator + Clone> Clone for #colvec_ident<A>
			where
				#(#field_types1: Clone,)*
		{
			fn clone(&self)->Self{
				let alloc = self.allocator().clone();
				struct DropGuard<'a, A: ::colvec::alloc::Allocator> {
					colvec: &'a mut #colvec_ident<A>,
					num_init: usize,
				}
				impl<'a, A: ::colvec::alloc::Allocator> Drop for DropGuard<'a, A> {
					#[inline]
					fn drop(&mut self) {
						// SAFETY:
						// items were marked initialized in the loop below
						unsafe {
							self.colvec.set_len(self.num_init);
						}
					}
				}
				let mut colvec = Self::with_capacity_in(self.len(), alloc);
				let mut guard = DropGuard { colvec: &mut colvec, num_init: 0 };
				#(
					let #field_slice_fn_idents1 = self.#field_slice_fn_idents2();
				)*
				// Create the pointer once to satisfy miri
				let cap = guard.colvec.buf.capacity();
				let ptr = guard.colvec.as_mut_ptr();
				#(
					let #field_slice_mut_fn_idents1 = unsafe {
						::core::slice::from_raw_parts_mut(
							ptr.add(cap * <#ident as ::colvec::raw::StructInfo<#fields_count>>::FIELDS.offset_of(#field_indices))
								.cast::<::core::mem::MaybeUninit<#field_types2>>(),
							self.len
						)
					};
				)*
				for i in 0..self.len() {
					guard.num_init = i;
					// Clone all fields first. Ensures the cloned values are dropped if clone panics.
					#(
						let #field_idents1 = unsafe { #field_slice_fn_idents3.get_unchecked(i) }.clone();
					)*
					#(
						unsafe{ #field_slice_mut_fn_idents2.get_unchecked_mut(i).write(#field_idents2) };
					)*
				}
				::core::mem::forget(guard);
				// SAFETY:
				// the colvec was allocated and initialized above to at least this length.
				unsafe {
					colvec.set_len(self.len());
				}
				colvec
			}
		}
	};

	let field_indices=0..fields_count;
	let field_types=fields.named.iter().map(|field|field.ty.clone());
	let field_slice_fn_idents=fields.named.iter().map(|field|{
		let ident=field.ident.as_ref().unwrap();
		let slice_ident=format!("{ident}_slice");
		syn::Ident::new(&slice_ident,ident.span())
	});
	let field_slice_mut_fn_idents=fields.named.iter().map(|field|{
		let ident=field.ident.as_ref().unwrap();
		let slice_ident=format!("{ident}_slice_mut");
		syn::Ident::new(&slice_ident,ident.span())
	});
	let field_access = quote! {
		impl<A: ::colvec::alloc::Allocator> #colvec_ident<A>{
			#(
				#[inline]
				pub const fn #field_slice_fn_idents(&self) -> &[#field_types] {
					unsafe {
						::core::slice::from_raw_parts(
							self.as_ptr()
								.add(self.buf.capacity() * <#ident as ::colvec::raw::StructInfo<#fields_count>>::FIELDS.offset_of(#field_indices))
								.cast::<#field_types>(),
							self.len
						)
					}
				}
				#[inline]
				pub const fn #field_slice_mut_fn_idents(&mut self) -> &mut [#field_types] {
					unsafe {
						::core::slice::from_raw_parts_mut(
							self.as_mut_ptr()
								.add(self.buf.capacity() * <#ident as ::colvec::raw::StructInfo<#fields_count>>::FIELDS.offset_of(#field_indices))
								.cast::<#field_types>(),
							self.len
						)
					}
				}
			)*
		}
	};

	#[cfg_attr(not(feature = "std"), expect(unused_mut))]
	let mut output=quote! {
		#colvec

		#struct_info

		#drop
		#clone
		#impls
		#field_access
	};

	#[cfg(feature = "std")]
	output.extend(global);

	output.into()
}

#[cfg(test)]
mod tests {
	use super::*;
	use syn::parse_quote;

	#[test]
	fn snapshot_test1() {
		let test1:syn::ItemStruct = parse_quote! {
			pub struct Test{
				field0:u8,
				field1:Option<u8>,
				field2:i16,
				field3:u32,
			}
		};

		let output = colvec_derive_inner(test1.into());

		// pretend it outputs a file
		let as_file = syn::parse_file(&output.to_string()).unwrap();

		// format it in a pretty way
		let formatted = prettyplease::unparse(&as_file);

		// snapshot-test it
		insta::assert_snapshot!(formatted);
	}
}