Skip to content

Commit

Permalink
Add Vec::replace and MutRef::replace
Browse files Browse the repository at this point in the history
  • Loading branch information
FnControlOption authored and Luthaf committed Feb 24, 2024
1 parent 587fedd commit f41eecb
Show file tree
Hide file tree
Showing 4 changed files with 84 additions and 3 deletions.
1 change: 1 addition & 0 deletions rustfmt.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
disable_all_formatting = true
25 changes: 24 additions & 1 deletion soa-derive-internal/src/refs.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use proc_macro2::TokenStream;
use proc_macro2::{Ident, Span, TokenStream};
use quote::quote;

use crate::input::Input;
Expand Down Expand Up @@ -26,6 +26,11 @@ pub fn derive(input: &Input) -> TokenStream {
.map(|field| field.ident.clone().unwrap())
.collect::<Vec<_>>();

let fields_names_hygienic = input.fields.iter()
.enumerate()
.map(|(i, _)| Ident::new(&format!("___soa_derive_private_{}", i), Span::call_site()))
.collect::<Vec<_>>();

let ref_fields_types = input.map_fields_nested_or(
|_, field_type| {
let field_ptr_type = names::ref_name(field_type);
Expand Down Expand Up @@ -57,6 +62,11 @@ pub fn derive(input: &Input) -> TokenStream {
|ident, _| quote! { self.#ident.clone() },
).collect::<Vec<_>>();

let ref_replace = input.map_fields_nested_or(
|ident, _| quote! { self.#ident.replace(field) },
|ident, _| quote! { ::std::mem::replace(&mut *self.#ident, field) },
).collect::<Vec<_>>();

quote! {
/// A reference to a
#[doc = #doc_url]
Expand Down Expand Up @@ -146,6 +156,19 @@ pub fn derive(input: &Input) -> TokenStream {
#( #fields_names: #to_owned, )*
}
}

/// Similar to [`std::mem::replace()`](https://doc.rust-lang.org/std/mem/fn.replace.html).
pub fn replace(&mut self, val: #name) -> #name {
#(
let field = unsafe { ::std::ptr::read(&val.#fields_names) };
let #fields_names_hygienic = #ref_replace;
)*
// if val implements Drop, we don't want to run it here, only
// when the vec itself will be dropped
::std::mem::forget(val);

#name{#(#fields_names: #fields_names_hygienic),*}
}
}
}
}
27 changes: 25 additions & 2 deletions soa-derive-internal/src/vec.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
use proc_macro2::{Ident, Span};
use proc_macro2::{TokenStream};
use proc_macro2::{Ident, Span, TokenStream};
use quote::TokenStreamExt;
use quote::quote;

Expand Down Expand Up @@ -68,6 +67,11 @@ pub fn derive(input: &Input) -> TokenStream {
|ident, _| quote! { Vec::from_raw_parts(data.#ident, len, capacity) },
).collect::<Vec<_>>();

let vec_replace = input.map_fields_nested_or(
|ident, _| quote! { self.#ident.replace(index, field) },
|ident, _| quote! { ::std::mem::replace(&mut self.#ident[index], field) },
).collect::<Vec<_>>();

let mut generated = quote! {
/// An analog to `
#[doc = #vec_name_str]
Expand Down Expand Up @@ -210,6 +214,25 @@ pub fn derive(input: &Input) -> TokenStream {
::std::mem::forget(element);
}

/// Similar to [`std::mem::replace()`](https://doc.rust-lang.org/std/mem/fn.replace.html).
pub fn replace(&mut self, index: usize, element: #name) -> #name {
if index > self.len() {
panic!("index out of bounds: the len is {} but the index is {}", self.len(), index);
}

// similar to push, we can not use move and have to rely on ptr
// read/write
#(
let field = unsafe { ::std::ptr::read(&element.#fields_names) };
let #fields_names_hygienic = #vec_replace;
)*
// if value implements Drop, we don't want to run it here, only
// when the vec itself will be dropped.
::std::mem::forget(element);

#name{#(#fields_names: #fields_names_hygienic),*}
}

/// Similar to [`
#[doc = #vec_name_str]
/// ::remove()`](https://doc.rust-lang.org/std/vec/struct.Vec.html#method.remove).
Expand Down
34 changes: 34 additions & 0 deletions tests/replace.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
#![allow(clippy::float_cmp)]

mod particles;
use self::particles::{Particle, ParticleVec};

#[test]
fn replace_element() {
let mut soa = ParticleVec::new();
soa.push(Particle::new(String::from("Na"), 22.990));
soa.push(Particle::new(String::from("Zn"), 65.380));
soa.push(Particle::new(String::from("Cl"), 35.453));

let particle = soa.replace(1, Particle::new(String::from("Br"), 79.904));
assert_eq!(particle.name, "Zn");
assert_eq!(particle.mass, 65.380);

assert_eq!(soa.name[1], "Br");
assert_eq!(soa.mass[1], 79.904);
}

#[test]
fn replace_mutable_reference() {
let mut soa = ParticleVec::new();
soa.push(Particle::new(String::from("Na"), 22.990));
soa.push(Particle::new(String::from("Zn"), 65.380));
soa.push(Particle::new(String::from("Cl"), 35.453));

let particle = soa.index_mut(1).replace(Particle::new(String::from("Br"), 79.904));
assert_eq!(particle.name, "Zn");
assert_eq!(particle.mass, 65.380);

assert_eq!(soa.name[1], "Br");
assert_eq!(soa.mass[1], 79.904);
}

0 comments on commit f41eecb

Please sign in to comment.