Rust monads with binding Fn

playground.rs

use std::ops::Mul;

pub mod hkt {
    pub trait HktAt<T> {
        type Out: HasHkt<Hkt=Self, Arg=T>;
    }
    
    pub trait HasHkt {
        type Arg;
        type Hkt: HktAt<Self::Arg, Out=Self>;
    }
    
    pub type Apply<Hkt, T> where Hkt: HktAt<T> = Hkt::Out;
    
    pub trait Monad<'a>: HasHkt + Sized + 'a {
        fn bind<U, F>(self, f: F) -> Apply<Self::Hkt, U>
        where F: Fn(Self::Arg) -> Apply<Self::Hkt, U> + 'a,
              Apply<Self::Hkt, U>: Monad<'a>;
        
        fn ret(x: Self::Arg) -> Self;
        
        fn join(self) -> Self::Arg
        where <Self as HasHkt>::Arg: HasHkt<Hkt=Self::Hkt> + Monad<'a>,
              <Self as HasHkt>::Hkt: HktAt<<Self::Arg as HasHkt>::Arg, Out=Self::Arg>, {
            self.bind(|x| x)
        }
        
        fn then<U>(self, other: Apply<Self::Hkt, U>) -> Apply<Self::Hkt, U>
        where Self::Hkt: HktAt<U>,
              Apply<Self::Hkt, U>: Monad<'a> + Clone, {
            self.bind::<U, _>(move |_| other.clone())
        }
    }
}
use hkt::*;

pub mod func {
    use std::marker::PhantomData;
    use std::rc::Rc;
    use ::hkt::*;
    
    pub struct Func_<'a, R> {
        phantom: PhantomData<Fn(R) + 'a>,
    }
    
    #[derive(Clone)]
    pub struct Func<'a, R, T> {
        func: Rc<Fn(R) -> T + 'a>,
    }
    
    impl<'a, R, T> Func<'a, R, T> {
        pub fn new<F>(f: F) -> Func<'a, R, T> where F: Fn(R) -> T + 'a {
            Func{func: Rc::new(f)}
        }
        
        pub fn call(&self, r: R) -> T { (self.func)(r) }
    }
    
    impl<'a, T, R> HktAt<T> for Func_<'a, R> {
        type Out = Func<'a, R, T>;
    }
     
    impl<'a, R, T> HasHkt for Func<'a, R, T> {
        type Arg = T;
        type Hkt = Func_<'a, R>;
    }
    
    impl<'a, R, T> Monad<'a> for Func<'a, R, T>
    where R: Clone + 'a, T: Clone + 'a, {
        fn bind<U, F>(self, f: F) -> Func<'a, R, U>
        where F: Fn(T) -> Func<'a, R, U> + 'a, {
            Func::new(move |r: R| f(self.call(r.clone())).call(r))
        }
        
        fn ret(x: T) -> Func<'a, R, T> { Func::new(move |_| x.clone()) }
    }
}
use func::*;

pub mod state {
    use std::marker::PhantomData;
    use std::rc::Rc;
    use ::hkt::*;
    
    pub struct State_<'a, S> {
        phantom: PhantomData<Fn(S) -> S + 'a>,
    }
    
    #[derive(Clone)]
    pub struct State<'a, S, T> {
        func: Rc<Fn(S) -> (S, T) + 'a>,
    }
    
    impl<'a, S, T> State<'a, S, T> {
        pub fn new<F>(f: F) -> State<'a, S, T> where F: Fn(S) -> (S, T) + 'a {
            State{func: Rc::new(f)}
        }
        
        pub fn call(&self, r: S) -> (S, T) { (self.func)(r) }
    }
    
    impl<'a, S, T> HktAt<T> for State_<'a, S> {
        type Out = State<'a, S, T>;
    }
     
    impl<'a, S, T> HasHkt for State<'a, S, T> {
        type Arg = T;
        type Hkt = State_<'a, S>;
    }
    
    impl<'a, S, T> Monad<'a> for State<'a, S, T>
    where S: Clone + 'a, T: Clone + 'a, {
        fn bind<U, F>(self, f: F) -> State<'a, S, U>
        where F: Fn(T) -> State<'a, S, U> + 'a, {
            State::new(move |s1: S| {
                let (s2, t) = self.call(s1);
                f(t).call(s2)
            })
        }
        
        fn ret(x: T) -> State<'a, S, T> { 
            State::new(move |s| (s, x.clone()))
        }
    }
    
    pub fn put<'a, T: Clone + 'a>(t: T) -> State<'a, T, T> {
        State::new(move |_| (t.clone(), t.clone()))
    }
    
    pub fn get<'a, S: Clone + 'a>() -> State<'a, S, S> {
        State::new(move |s: S| (s.clone(), s))
    }
}
use state::*;

fn mul_state<'a, S, T>(t: T) -> State<'a, S, T>
where S: for<'b> Mul<&'b T, Output=S> + Clone + 'a,
      T: Clone + 'a, {
    let t2 = t.clone();
    get().bind(move |x| put(x*&t)).then(State::ret(t2))
}

fn main() {
    println!("{}",
        Func::ret(Func::ret(2)).join().then(Func::ret(4)).call(0)    
    );
    
    
    println!("{:?}",
        State::ret(2)
        .bind(mul_state)
        .then(State::ret(4))
        .bind(mul_state)
        .call(3)
    );
}