src/bisection_tree/either.rs


use std::iter::Chain;
use std::sync::Arc;

use crate::types::*;
use crate::mapping::{Apply, Differentiable};
use crate::iter::{Mappable, MapF, MapZ};
use crate::sets::Cube;
use crate::loc::Loc;

use super::support::*;
use super::aggregator::*;

/// A structure for storing two [`SupportGenerator`]s summed/chain together.
///
/// This is needed to work with sums of different types of [`Support`]s.
#[derive(Debug,Clone)]
pub struct BothGenerators<A, B>(
    pub(super) Arc<A>,
    pub(super) Arc<B>,
);

/// A structure for a [`Support`] that can be either `A` or `B`.
///
/// This is needed to work with sums of different types of [`Support`]s.
#[derive(Debug,Clone)]
pub enum EitherSupport<A, B> {
    Left(A),
    Right(B),
}

impl<A, B, F : Num> HasScalarField for EitherSupport<A, B>
where A : HasScalarField<Field = F>,
      B : HasScalarField<Field = F> {
    type Field = F;
}

// We need type alias bounds to access associate types.
#[allow(type_alias_bounds)]
type BothAllDataIter<
    'a,
    G1 : SupportGenerator<N>,
    G2 : SupportGenerator<N, RealField = G1::RealField>,
    const N : usize
> = Chain<
    MapF<G1::AllDataIter<'a>, (usize, EitherSupport<G1::SupportType, G2::SupportType>)>,
    MapZ<G2::AllDataIter<'a>, usize, (usize, EitherSupport<G1::SupportType, G2::SupportType>)>,
>;

impl<G1, G2> BothGenerators<G1, G2> {
    /// Helper for [`all_left_data`].
    #[inline]
    fn map_left<F : Float, const N : usize>((d, support) : (G1::Id, G1::SupportType))
    -> (usize, EitherSupport<G1::SupportType, G2::SupportType>)
    where G1 : SupportGenerator<N, Id=usize, RealField =F>,
          G2 : SupportGenerator<N, Id=usize, RealField = F> {

        let id : usize = d.into();
        (id.into(), EitherSupport::Left(support))
    }

    /// Helper for [`all_right_data`].
    #[inline]
    fn map_right<F : Float, const N : usize>(n0 : &usize, (d, support) : (G2::Id, G2::SupportType))
    -> (usize, EitherSupport<G1::SupportType, G2::SupportType>)
    where G1 : SupportGenerator<N, Id=usize, RealField = F>,
          G2 : SupportGenerator<N, Id=usize, RealField = F> {

        let id : usize = d.into();
        ((n0+id).into(), EitherSupport::Right(support))
    }

    /// Calls [`SupportGenerator::all_data`] on the “left” support generator.
    ///
    /// Converts both the id and the [`Support`] into a form that corresponds to `BothGenerators`.
    #[inline]
    pub(super) fn all_left_data<F : Float, const N : usize>(&self)
    -> MapF<G1::AllDataIter<'_>, (usize, EitherSupport<G1::SupportType, G2::SupportType>)>
    where G1 : SupportGenerator<N, Id=usize, RealField = F>,
          G2 : SupportGenerator<N, Id=usize, RealField = F> {
        self.0.all_data().mapF(Self::map_left)
    }

    /// Calls [`SupportGenerator::all_data`] on the “right” support generator.
    ///
    /// Converts both the id and the [`Support`] into a form that corresponds to `BothGenerators`.
    #[inline]
    pub(super) fn all_right_data<F : Float, const N : usize>(&self)
    -> MapZ<G2::AllDataIter<'_>, usize, (usize, EitherSupport<G1::SupportType, G2::SupportType>)>
    where G1 : SupportGenerator<N, Id=usize, RealField=F>,
          G2 : SupportGenerator<N, Id=usize, RealField=F> {
        let n0 = self.0.support_count();
        self.1.all_data().mapZ(n0, Self::map_right)
    }
}

impl<F : Float, G1, G2> HasScalarField for BothGenerators<G1, G2>
where G1 : HasScalarField<Field=F>,
      G2 : HasScalarField<Field=F> {
    type Field = F;
}

impl<F : Float, G1, G2, const N : usize>
SupportGenerator<N>
for BothGenerators<G1, G2>
where G1 : SupportGenerator<N, Id=usize> + HasRealField<RealField=F>,
      G2 : SupportGenerator<N, Id=usize> + HasRealField<RealField=F>,
      Self : HasRealField<RealField = F> {

    type Id = usize;
    type SupportType = EitherSupport<G1::SupportType, G2::SupportType>;
    type AllDataIter<'a> = BothAllDataIter<'a, G1, G2, N> where G1 : 'a, G2 : 'a;

    #[inline]
    fn support_for(&self,  id : Self::Id)
    -> Self::SupportType {
        let n0 = self.0.support_count();
        if id < n0 {
            EitherSupport::Left(self.0.support_for(id.into()))
        } else {
            EitherSupport::Right(self.1.support_for((id-n0).into()))
        }
    }

    #[inline]
    fn support_count(&self) -> usize {
        self.0.support_count() + self.1.support_count()
    }

    #[inline]
    fn all_data(&self) -> Self::AllDataIter<'_> {
        self.all_left_data().chain(self.all_right_data())
    }
}

impl<F: Float, S1, S2, const N : usize> Support<N> for EitherSupport<S1, S2>
where S1 : Support<N, RealField = F>,
      S2 : Support<N, RealField = F> {

    #[inline]
    fn support_hint(&self) -> Cube<F,N> {
        match self {
            EitherSupport::Left(ref a) => a.support_hint(),
            EitherSupport::Right(ref b) => b.support_hint(),
        }
    }

    #[inline]
    fn in_support(&self, x : &Loc<F,N>) -> bool {
        match self {
            EitherSupport::Left(ref a) => a.in_support(x),
            EitherSupport::Right(ref b) => b.in_support(x),
        }
    }

    #[inline]
    fn bisection_hint(&self, cube : &Cube<F, N>) -> [Option<F>; N] {
        match self {
            EitherSupport::Left(ref a) => a.bisection_hint(cube),
            EitherSupport::Right(ref b) => b.bisection_hint(cube),
        }
    }
}

impl<F : Float, A, S1, S2, const N : usize> LocalAnalysis<A, Cube<F, N>>
for EitherSupport<S1, S2>
where A : Aggregator,
      S1 : LocalAnalysis<A, Cube<F, N>>,
      S2 : LocalAnalysis<A, Cube<F, N>> {

    #[inline]
    fn local_analysis(&self, cube : &Cube<F, N>) -> A {
        match self {
            EitherSupport::Left(ref a) => a.local_analysis(cube),
            EitherSupport::Right(ref b) => b.local_analysis(cube),
        }
    }
}

impl<A, S1, S2> GlobalAnalysis<A> for EitherSupport<S1, S2>
where A : Aggregator,
      S1 : GlobalAnalysis<A>,
      S2 : GlobalAnalysis<A>, {

    #[inline]
    fn global_analysis(&self) -> A {
        match self {
            EitherSupport::Left(ref a) => a.global_analysis(),
            EitherSupport::Right(ref b) => b.global_analysis(),
        }
    }
}

impl<F, S1, S2, X> Apply<X> for EitherSupport<S1, S2>
where S1 : Apply<X, Output=F>,
      S2 : Apply<X, Output=F> {
    type Output = F;
    #[inline]
    fn apply(&self, x : X) -> F {
        match self {
            EitherSupport::Left(ref a) => a.apply(x),
            EitherSupport::Right(ref b) => b.apply(x),
        }
    }
}

impl<F : Float, S1, S2, X> Differentiable<X> for EitherSupport<S1, S2>
where S1 : Differentiable<X, Derivative=F, RealField=F>,
      S2 : Differentiable<X, Derivative=F, RealField=F> {
    type Derivative = F;
    #[inline]
    fn differential(&self, x : X) -> F {
        match self {
            EitherSupport::Left(ref a) => a.differential(x),
            EitherSupport::Right(ref b) => b.differential(x),
        }
    }
}

macro_rules! make_either_scalarop_rhs {
    ($trait:ident, $fn:ident, $trait_assign:ident, $fn_assign:ident) => {
        impl<F : Float, G1, G2>
        std::ops::$trait_assign<F>
        for BothGenerators<G1, G2>
        where G1 : std::ops::$trait_assign<F> + Clone,
              G2 : std::ops::$trait_assign<F> + Clone, {
            #[inline]
            fn $fn_assign(&mut self, t : F) {
                Arc::make_mut(&mut self.0).$fn_assign(t);
                Arc::make_mut(&mut self.1).$fn_assign(t);
            }
        }

        impl<'a, F : Float, G1, G2>
        std::ops::$trait<F>
        for &'a BothGenerators<G1, G2>
        where &'a G1 : std::ops::$trait<F,Output=G1>,
              &'a G2 : std::ops::$trait<F,Output=G2> {
            type Output = BothGenerators<G1, G2>;
            #[inline]
            fn $fn(self, t : F) -> BothGenerators<G1, G2> {
                BothGenerators(Arc::new(self.0.$fn(t)),
                               Arc::new(self.1.$fn(t)))
            }
        }
    }
}

make_either_scalarop_rhs!(Mul, mul, MulAssign, mul_assign);
make_either_scalarop_rhs!(Div, div, DivAssign, div_assign);

impl<G1, G2> std::ops::Neg for BothGenerators<G1, G2>
where G1 : std::ops::Neg + Clone,
      G2 : std::ops::Neg + Clone, {
    type Output = BothGenerators<G1::Output, G2::Output>;
    #[inline]
    fn neg(self) -> Self::Output {
        BothGenerators(Arc::new(Arc::unwrap_or_clone(self.0).neg()),
                       Arc::new(Arc::unwrap_or_clone(self.1).neg()))
    }
}
/*
impl<'a, G1, G2> std::ops::Neg for &'a BothGenerators<G1, G2>
where &'a G1 : std::ops::Neg, &'a G2 : std::ops::Neg, {
    type Output = BothGenerators<<&'a G1 as std::ops::Neg>::Output,
                                 <&'a G2 as std::ops::Neg>::Output>;
    fn neg(self) -> Self::Output {
        BothGenerators(self.0.neg(), self.1.neg())
    }
}
*/