feat: (bugged) #[derive(MultilinearMap)] ` ( (#3)

* feat: (bugged) `#[derive(MultilinearMap)]`

Author: Autoparallel

Cargo.lock

@@ -3,7 +3,9 @@ name = "tensor"
 version = "0.1.0"
 edition = "2021"
 
-# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+[dependencies]
+tensor_macros = { path = "macros/" }
+
 
 [dev-dependencies]
 # tracing-subscriber = { version = "0.3.18", default-features = false, features = [

Cargo.toml

@@ -0,0 +1,12 @@
+[package]
+name = "tensor_macros"
+version = "0.1.0"
+edition = "2021"
+
+[lib]
+proc-macro = true
+
+[dependencies]
+syn = { version = "2.0.60", features = ["full"] }
+quote = "1.0.36"
+proc-macro2 = "1.0.81"

macros/Cargo.toml

@@ -0,0 +1,72 @@
+use proc_macro::TokenStream;
+use quote::quote;
+use syn::{parse_macro_input, DeriveInput, GenericParam, Ident};
+
+#[proc_macro_derive(MultilinearMap)]
+pub fn multilinear_map_derive(input: TokenStream) -> TokenStream {
+    let ast = parse_macro_input!(input as DeriveInput);
+    let struct_name = &ast.ident;
+
+    let generics = &ast.generics;
+    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
+
+    let const_generics: Vec<_> = generics
+        .params
+        .iter()
+        .filter_map(|param| match param {
+            GenericParam::Const(const_param) => Some(&const_param.ident),
+            _ => None,
+        })
+        .collect();
+
+    let input_params = const_generics.iter().enumerate().map(|(i, ident)| {
+        let param_name = Ident::new(&format!("v_{}", i), ident.span());
+        quote! { #param_name: V<#ident, F> }
+    });
+
+    let loop_indices: Vec<_> = (0..const_generics.len())
+        .map(|i| Ident::new(&format!("i_{}", i), proc_macro2::Span::call_site()))
+        .collect();
+
+    let component_product = loop_indices.iter().zip(0..).map(|(index, i)| {
+        let param_name = Ident::new(&format!("v_{}", i), index.span());
+        quote! { * #param_name.0[#index] }
+    });
+
+    // Add the calculation to the innermost loop
+    let coefficient_access =
+        loop_indices
+            .iter()
+            .fold(quote! { self.coefficients }, |acc, index| {
+                quote! { #acc.0[#index] }
+            });
+
+    let mut loop_nest = quote! {
+        sum += #coefficient_access #(#component_product)*;
+    };
+
+    for (index, ident) in loop_indices.iter().rev().zip(const_generics.iter().rev()) {
+        loop_nest = quote! {
+            for #index in 0..#ident {
+                #loop_nest
+            }
+        };
+    }
+
+    loop_nest = quote! {
+        #loop_nest
+
+    };
+
+    let expanded = quote! {
+        impl #impl_generics #struct_name #ty_generics #where_clause {
+            pub fn multilinear_map(&self, #(#input_params),*) -> F {
+                let mut sum = F::default();
+                #loop_nest
+                sum
+            }
+        }
+    };
+
+    TokenStream::from(expanded)
+}

macros/src/lib.rs

@@ -1,25 +1,29 @@
 use super::*;
 
+// TODO: Could probably just assign a valence to the tensors and use N0, N1, N2,
+// etc. as dims
+
 #[macro_export]
 macro_rules! tensor {
-    ($name:ident, $($const:ident),+) => {
-        pub struct $name<$(const $const: usize),+, F>
-        where F: Default +  Copy,
+    ($name:ident, $($consts:ident),+) => {
+        #[derive(tensor_macros::MultilinearMap)]
+        pub struct $name<$(const $consts: usize),+, F>
+        where F: Default + Copy + AddAssign + Mul<F, Output = F>,
         {
-            pub coefficients: coeff_builder!($($const),+; F),
+            pub coefficients: coeff_builder!($($consts),+; F),
         }
 
-        impl<$(const $const: usize),+, F: Default + Copy> Default for  $name<$($const),+, F> {
+        impl<$(const $consts: usize),+, F: Default + Copy + AddAssign + Mul<F, Output = F>> Default for  $name<$($consts),+, F> {
             fn default() -> Self {
-                let coefficients = <def_builder!($($const),+; F)>::default();
+                let coefficients = <def_builder!($($consts),+; F)>::default();
                 $name { coefficients }
             }
 
         }
 
-        impl<$(const $const: usize),+, F> Debug for $name<$($const),+, F>
+        impl<$(const $consts: usize),+, F> Debug for $name<$($consts),+, F>
         where
-            F: Default + Copy + Debug,
+            F: Default + Copy + Debug + AddAssign + Mul<F, Output = F>,
         {
             fn fmt(&self, f: &mut Formatter<'_>) -> Result {
                 f.debug_struct(stringify!($name))
@@ -28,34 +32,31 @@ macro_rules! tensor {
             }
         }
 
-        impl<$(const $const: usize),+, F> Add for $name<$($const),+, F>
+        impl<$(const $consts: usize),+, F> Add for $name<$($consts),+, F>
         where
-            F: Add<Output = F> + Copy + Default,
+            F: Add<Output = F> + Copy + Default + AddAssign + Mul<F, Output = F>,
         {
             type Output = Self;
 
             fn add(self, other: Self) -> Self::Output {
                 let mut result = Self::default();
-                add_tensors!(result.coefficients, self.coefficients, other.coefficients; $($const),+);
+                add_tensors!(result.coefficients, self.coefficients, other.coefficients; $($consts),+);
                 result
             }
         }
 
-        impl<$(const $const: usize),+, F> Mul<F> for $name<$($const),+, F>
+        impl<$(const $consts: usize),+, F> Mul<F> for $name<$($consts),+, F>
         where
-            F: Mul<Output = F> + Copy + Default,
+            F: Mul<Output = F> + Copy + Default + AddAssign,
         {
             type Output = Self;
 
             fn mul(self, scalar: F) -> Self::Output {
                 let mut result = Self::default();
-                scalar_mul_tensor!(result.coefficients, self.coefficients, scalar; $($const),+);
+                scalar_mul_tensor!(result.coefficients, self.coefficients, scalar; $($consts),+);
                 result
             }
         }
-
-
-
     }
 }
 
@@ -103,11 +104,14 @@ macro_rules! scalar_mul_tensor {
     };
 }
 
+tensor!(TensorTester, M, N, P);
+
 #[cfg(test)]
 mod tests {
 
     use super::*;
     tensor!(Tensor2, M, N);
+
     tensor!(Tensor3, M, N, P);
 
     use log::{debug, info};
@@ -161,4 +165,41 @@ mod tests {
         let tensor2 = tensor1 * scalar;
         info!("output: {:?}", tensor2.coefficients);
     }
+
+    #[test]
+    fn multilinear_map() {
+        log();
+        //           / 1    0     0 \
+        // tensor =  \ 0    1     0 /
+        let mut tensor = Tensor2::<2, 3, f64>::default();
+        tensor.coefficients.0[0].0[0] = 1.0;
+        tensor.coefficients.0[1].0[1] = 1.0;
+        debug!("tensor: {:?}", tensor);
+
+        //        / -1 \
+        // v_0 =  \  1 /
+        let mut v_0 = V::default();
+        v_0.0[0] = -1.0;
+        v_0.0[1] = 1.0;
+        debug!("v_0: {:?}", v_0);
+
+        //        / 1 \
+        //       |  2  |
+        // v_1 =  \ 3 /
+        let mut v_1 = V::default();
+        v_1.0[0] = 1.0;
+        v_1.0[1] = 2.0;
+        v_1.0[2] = 3.0;
+        debug!("v_1: {:?}", v_1);
+
+        //                      / 1 \
+        // tensor.map(_,v_1) =  \ 2 /
+        //
+        // then the next is:
+        //                                     / 1 \
+        // tensor.map(v_0, v_1) = < -1    1 >  \ 2 /   = -1 + 2 = 1
+        let output = tensor.multilinear_map(v_0, v_1);
+        info!("output: {:?}", output);
+        assert_eq!(output, 1.0);
+    }
 }