vector! and dvector! macros

This commit is contained in:
Andreas Longva 2021-04-30 14:09:36 +02:00
parent d56db1a079
commit 07d41e457b
2 changed files with 94 additions and 4 deletions

View File

@ -131,3 +131,62 @@ pub fn dmatrix(stream: TokenStream) -> TokenStream {
proc_macro::TokenStream::from(output)
}
struct Vector {
elements: Vec<Expr>,
}
impl Vector {
fn to_array_tokens(&self) -> TokenStream2 {
let mut data = TokenStream2::new();
data.append_separated(&self.elements, Punct::new(',', Spacing::Alone));
TokenStream2::from(TokenTree::Group(Group::new(Delimiter::Bracket, data)))
}
fn len(&self) -> usize {
self.elements.len()
}
}
impl Parse for Vector {
fn parse(input: ParseStream) -> Result<Self> {
// The syntax of a vector is just the syntax of a single matrix row
if input.is_empty() {
Ok(Self {
elements: Vec::new()
})
} else {
let elements = MatrixRowSyntax::parse_separated_nonempty(input)?.into_iter().collect();
Ok(Self {
elements
})
}
}
}
#[proc_macro]
pub fn vector(stream: TokenStream) -> TokenStream {
let vector = parse_macro_input!(stream as Vector);
let len = vector.len();
let array_tokens = vector.to_array_tokens();
let output = quote! {
nalgebra::SVector::<_, #len>
::from_array_storage(nalgebra::ArrayStorage([#array_tokens]))
};
proc_macro::TokenStream::from(output)
}
#[proc_macro]
pub fn dvector(stream: TokenStream) -> TokenStream {
let vector = parse_macro_input!(stream as Vector);
let len = vector.len();
let array_tokens = vector.to_array_tokens();
let output = quote! {
nalgebra::DVector::<_>
::from_vec_storage(nalgebra::VecStorage::new(
nalgebra::Dynamic::new(#len),
nalgebra::Const::<1>,
vec!#array_tokens))
};
proc_macro::TokenStream::from(output)
}

View File

@ -1,12 +1,12 @@
use nalgebra_macros::{dmatrix, matrix};
use nalgebra::{DMatrix, SMatrix, Matrix3x2, Matrix1x2, Matrix1x3, Matrix1x4, Matrix2x1, Matrix2, Matrix2x3, Matrix2x4, Matrix3x1, Matrix3, Matrix3x4, Matrix4x1, Matrix4x2, Matrix4x3, Matrix4};
use nalgebra_macros::{dmatrix, dvector, matrix, vector};
use nalgebra::{DMatrix, DVector, SMatrix, Matrix3x2, Matrix1x2, Matrix1x3, Matrix1x4, Matrix2x1, Matrix2, Matrix2x3, Matrix2x4, Matrix3x1, Matrix3, Matrix3x4, Matrix4x1, Matrix4x2, Matrix4x3, Matrix4, Vector1, Vector2, Vector3, Vector4, Vector5, SVector, Vector6};
fn same_type<T>(_: &T, _: &T) {}
fn check_statically_same_type<T>(_: &T, _: &T) {}
/// Wrapper for `assert_eq` that also asserts that the types are the same
macro_rules! assert_eq_and_type {
($left:expr, $right:expr $(,)?) => {
same_type(&$left, &$right);
check_statically_same_type(&$left, &$right);
assert_eq!($left, $right);
};
}
@ -84,3 +84,34 @@ fn dmatrix_small_dims_exhaustive() {
assert_eq_and_type!(dmatrix![1, 2, 3, 4; 5, 6, 7, 8; 9, 10, 11, 12; 13, 14, 15, 16],
DMatrix::from_row_slice(4, 4, &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]));
}
#[test]
fn vector_small_dims_exhaustive() {
assert_eq_and_type!(vector![], SVector::<i32, 0>::zeros());
assert_eq_and_type!(vector![1], Vector1::<i32>::new(1));
assert_eq_and_type!(vector![1, 2], Vector2::new(1, 2));
assert_eq_and_type!(vector![1, 2, 3], Vector3::new(1, 2, 3));
assert_eq_and_type!(vector![1, 2, 3, 4], Vector4::new(1, 2, 3, 4));
assert_eq_and_type!(vector![1, 2, 3, 4, 5], Vector5::new(1, 2, 3, 4, 5));
assert_eq_and_type!(vector![1, 2, 3, 4, 5, 6], Vector6::new(1, 2, 3, 4, 5, 6));
}
#[test]
fn vector_const_fn() {
// Ensure that vector! can be used in const contexts
const _: SVector<i32, 0> = vector![];
const _: Vector1<i32> = vector![1];
const _: Vector2<i32> = vector![1, 2];
const _: Vector6<i32> = vector![1, 2, 3, 4, 5, 6];
}
#[test]
fn dvector_small_dims_exhaustive() {
assert_eq_and_type!(dvector![], DVector::<i32>::zeros(0));
assert_eq_and_type!(dvector![1], DVector::from_column_slice(&[1]));
assert_eq_and_type!(dvector![1, 2], DVector::from_column_slice(&[1, 2]));
assert_eq_and_type!(dvector![1, 2, 3], DVector::from_column_slice(&[1, 2, 3]));
assert_eq_and_type!(dvector![1, 2, 3, 4], DVector::from_column_slice(&[1, 2, 3, 4]));
assert_eq_and_type!(dvector![1, 2, 3, 4, 5], DVector::from_column_slice(&[1, 2, 3, 4, 5]));
assert_eq_and_type!(dvector![1, 2, 3, 4, 5, 6], DVector::from_column_slice(&[1, 2, 3, 4, 5, 6]));
}