nmr-schedule 0.2.1

//! Contains various modifiers and filters that can be applied to generators.

mod basic;
mod iterate;
mod psf_polisher;
mod tm_filter;

pub use basic::*;
pub use iterate::*;
use ndarray::Dimension;
pub use psf_polisher::*;
pub use tm_filter::*;

use alloc::borrow::ToOwned;

use crate::{
    generators::{Generator, Trace},
    Schedule,
};

/// Generators returned by modifiers
pub mod generators {
    pub use super::basic::generators::*;
    pub use super::iterate::generators::*;
    pub use super::psf_polisher::generators::*;
    pub use super::tm_filter::generators::*;
}

/// Post-processes existing generators.
///
/// Modifiers get access to the original generator, so they may call the underlying generator repeatedly to get schedules they would like.
pub trait Modifier<Dim: Dimension> {
    /// The new generator after applying the modifier
    type Output<T: Generator<Dim>>: Generator<Dim>;

    /// Apply the modifier to a schedule generator
    fn modify<T: Generator<Dim>>(self, generator: T) -> Self::Output<T>;
}

/// Post-process an existing schedule without requiring access to the generator.
///
/// Filters do not get access to the underlying generator and operate only on a schedule. This allows filters to be applied to user-supplied schedules unlike modifiers.
pub trait Filter<Dim: Dimension> {
    /// Apply the filter to a schedule
    fn filter(&self, sched: Schedule<Dim>) -> Schedule<Dim> {
        self.filter_with_iter_and_trace(sched, 0).into_sched()
    }

    /// Apply the filter to a schedule while providing a trace
    fn filter_with_iter_and_trace(&self, sched: Schedule<Dim>, iteration: u64) -> Trace<Dim>;

    /// Filter a schedule from an existing trace while pushing a new trace onto it
    fn filter_from_iter_and_trace(&self, mut trace: Trace<Dim>, iteration: u64) -> Trace<Dim> {
        let mut trace_after = self.filter_with_iter_and_trace(trace.sched().to_owned(), iteration);

        trace.stack.append(&mut trace_after.stack);

        trace
    }
}

/// Creates boilerplate (including a builder extension trait) for modifiers
///
/// Example:
/// ```
/// # use nmr_schedule::{*, generators::*, modifiers::*, pdf::*};
/// # use ndarray::*;
/// # pub struct MagicifyGenerator<T, F>(T, F);
/// # impl<T, F> Generator<Ix1> for MagicifyGenerator<T, F> {
/// #     fn _generate(&self, count: usize, dims: Ix1, iteration: u64) -> Trace<Ix1> {
/// #         todo!()
/// #     }
/// # }
/// # fn typecheck() {
/// #     Quantiles::new(linear).magicify(12, |s| 5.);
/// # }
///
/// modifier!(<F: Fn(&Schedule<Ix1>) -> f64> Magicify<Ix1>,
///     MagicifyBuilder,
///     r"Apply magic to the schedule",
///     magicify,
///     magic_amt: usize,
///     magic_fn: F
/// );
///
/// impl<F: Fn(&Schedule<Ix1>) -> f64> Modifier<Ix1> for Magicify<F> {
///     type Output<T: Generator<Ix1>> = MagicifyGenerator<T, F>;
///
///     fn modify<T: Generator<Ix1>>(self, generator: T) -> Self::Output<T> {
///         // ...
///         # todo!()
///     }
/// }
/// ```
///
/// - The parameter in angle brackets is an optional list generic parameters with bounds
/// - The second parameter is the name of the modifier. The macro will create a tuple-struct with that name and all the parameters.
/// - The third supplies documentation to the modifier and to the builder method.
/// - The fourth is the name of the builder trait that will be automatically implemented on all generators
/// - The fifth is the name of the builder method that will be able to be called on generators (for example `.tm_filter()`).
/// - The rest are arguments to be passed into the modifier's `new` function, the builder method, and added as fields to the modifier.
#[macro_export]
macro_rules! modifier {
    ($(<$($name:ident : $bounds:path),+>)? $modifier:ident <$dim:ident>, $trait_name:ident, $doc:expr, $builder_name: ident, $($arg:ident : $ty:ty),*) => {
        #[doc=$doc]
        #[doc=concat!("\n\nCan be used as the `.", stringify!($builder_name), "()` method.")]
        #[derive(Debug)]
        #[allow(missing_copy_implementations)]
        pub struct $modifier $(<$($name: $bounds),+>)? ($($ty),*);

        impl $(<$($name: $bounds),+>)? $modifier $(<$($name),+>)? {
            #[doc=concat!("Create a new instance of the modifier.")]
            #[allow(clippy::new_without_default)]
            pub const fn new($($arg : $ty),*) -> $modifier $(<$($name),+>)? {
                $modifier($($arg),*)
            }
        }

        #[doc=concat!("Builder trait for the [`", stringify!($modifier), "`] modifier.")]
        pub trait $trait_name: Generator<$dim> {
            #[doc=$doc]
            fn $builder_name<$($($name: $bounds),+)?>(self, $($arg : $ty),*) -> <$modifier $(<$($name),+>)? as Modifier<$dim>>::Output<Self> where Self: Sized +  {
                self.then($modifier::new($($arg),*))
            }
        }

        impl<T: Generator<$dim>> $trait_name for T {}
    };
    ($dim:ident $(<$($name:ident : $bounds:path),+>)? $modifier:ident, $trait_name:ident, $doc:expr, $builder_name:ident, $($arg:ident : $ty:ty),*) => {
        #[doc=$doc]
        #[doc=concat!("\n\nCan be used as the `.", stringify!($builder_name), "()` method.")]
        #[derive(Debug)]
        #[allow(missing_copy_implementations)]
        pub struct $modifier <$dim: Dimension, $($($name: $bounds),+)?> ($($ty),*, PhantomData<$dim>);

        impl <$dim: Dimension, $($($name: $bounds),+)?> $modifier <$dim $(, $($name),+)?> {
            #[doc=concat!("Create a new instance of the modifier.")]
            #[allow(clippy::new_without_default)]
            pub const fn new($($arg : $ty),*) -> $modifier <$dim $(, $($name),+)?> {
                $modifier($($arg),*, PhantomData)
            }
        }

        #[doc=concat!("Builder trait for the [`", stringify!($modifier), "`] modifier.")]
        pub trait $trait_name<$dim: Dimension>: Generator<$dim> {
            #[doc=$doc]
            fn $builder_name<$($($name: $bounds),+)?>(self, $($arg : $ty),*) -> <$modifier <$dim $(, $($name),+)?> as Modifier<$dim>>::Output<Self> where Self: Sized +  {
                self.then($modifier::new($($arg),*))
            }
        }

        impl<Dim: Dimension, T: Generator<Dim>> $trait_name<Dim> for T {}
    };
}