higher kinded types in rust with functor, applicative, and monad traits + implementation for custom maybe type

hkt.rs

//! Higher-kinded types in Rust
//!
//! Functor, Applicative, and Monad traits
//!
//! Example implementations for custom Maybe type

// -- Trait Definitions -- //

/// A constructor for a type of a higher kind with one type parameter
pub trait TypeClass {
    type Type<A>;
}

/// A higher-kinded type with one type parameter
pub trait Kinded<A> {
    type Kind: TypeClass<Type<A> = Self>;
}

/// Utility for applying a type constructor to a type with one type parameter
pub type Apply<F, A> = <F as TypeClass>::Type<A>;

/// A higher-order abstraction that encapsulates values within a context and allows mapping
/// functions over these values while preserving the structure of the context
pub trait Functor<A>: Kinded<A> {
    /// Applies a transformation `A -> B` to a value of type `F<A>`
    fn fmap<B, M: FnMut(A) -> B>(self, f: M) -> Apply<Self::Kind, B>;
}

/// A higher-order abstraction that enables applying functions wrapped in a context to values
/// in the same context, allowing for parallel composition of independent operations
pub trait Applicative<A>: Functor<A> {
    /// Lifts a value into the applicative context
    fn pure(b: A) -> Apply<Self::Kind, A>;

    /// Applies functions contained in an applicative context to values in this applicative context
    ///
    /// NOTE: applies from right to left
    fn apply<B, F: FnMut(A) -> B>(self, ff: Apply<Self::Kind, F>) -> Apply<Self::Kind, B>;
}

/// A higher-order abstraction that wraps values and provides operations for sequential composition
/// while preserving context, allowing chainable transformations with controlled side effects
pub trait Monad<A>: Applicative<A> {
    /// Applies a function to the value in this monad, allowing chaining of computations
    /// that return values wrapped in the same context and flattening the resulting structure
    fn bind<B, F: FnMut(A) -> Apply<Self::Kind, B>>(self, f: F) -> Apply<Self::Kind, B>;
}

// -- Implementations -- //

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
pub enum Maybe<A> {
    #[default]
    Nothing,
    Just(A),
}

impl<A> Maybe<A> {
    pub fn is_nothing(&self) -> bool {
        match self {
            Maybe::Nothing => true,
            Maybe::Just(_) => false,
        }
    }

    pub fn is_just(&self) -> bool {
        match self {
            Maybe::Nothing => false,
            Maybe::Just(_) => true,
        }
    }
}

pub struct MaybeKind;

impl TypeClass for MaybeKind {
    type Type<A> = Maybe<A>;
}

impl<A> Kinded<A> for Maybe<A> {
    type Kind = MaybeKind;
}

impl<A> Functor<A> for Maybe<A> {
    fn fmap<B, F: FnMut(A) -> B>(self, mut f: F) -> Apply<MaybeKind, B> {
        match self {
            Maybe::Just(x) => Maybe::Just(f(x)),
            Maybe::Nothing => Maybe::Nothing,
        }
    }
}

impl<A> Applicative<A> for Maybe<A> {
    fn pure(a: A) -> Apply<MaybeKind, A> {
        Maybe::Just(a)
    }

    fn apply<B, F: FnMut(A) -> B>(self, ff: Apply<MaybeKind, F>) -> Apply<MaybeKind, B> {
        match (self, ff) {
            (Maybe::Just(a), Maybe::Just(mut f)) => Maybe::Just(f(a)),
            _ => Maybe::Nothing,
        }
    }
}

impl<A> Monad<A> for Maybe<A> {
    fn bind<B, F: FnMut(A) -> Apply<MaybeKind, B>>(self, mut f: F) -> Apply<MaybeKind, B> {
        match self {
            Maybe::Just(a) => f(a),
            Maybe::Nothing => Maybe::Nothing,
        }
    }
}

// -- Functions -- //

/// Applies a transformation `A -> B` to a value of type `F<A>`
pub fn fmap<A, B, FA: Functor<A>, F: FnMut(A) -> B>(g: F, f: FA) -> Apply<FA::Kind, B> {
    f.fmap(g)
}

/// Lifts a value into the applicative context
pub fn pure<FA: Applicative<A>, A>(a: A) -> Apply<FA::Kind, A> {
    FA::pure(a)
}

/// Applies a function from an applicative context to a value in an applicative context
///
/// NOTE: applies from left to right
pub fn apply<A, B, F, FA>(ff: Apply<FA::Kind, F>, fa: FA) -> Apply<FA::Kind, B>
where
    F: FnMut(A) -> B,
    FA: Applicative<A>,
{
    fa.apply::<B, F>(ff)
}

/// Identity function, returns the value it was given
pub fn id<A>(a: A) -> A {
    a
}

/// Function composition, from right to left
pub trait Composable<A, B> {
    /// Composes two functions from right to left
    fn compose<C, F: Fn(C) -> A>(self, f: F) -> impl Fn(C) -> B;
}

impl<A, B, ThisFn: Fn(A) -> B> Composable<A, B> for ThisFn {
    fn compose<C, FF: Fn(C) -> A>(self, f: FF) -> impl Fn(C) -> B {
        move |c| self(f(c))
    }
}

// -- Tests -- //

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn functor_identity() {
        let x = Maybe::pure(69);

        let left = x.fmap(id);
        let right = id(x);

        assert_eq!(left, right);
    }

    #[test]
    fn functor_composition() {
        let x = Maybe::pure(69);

        let f = |x: i32| x * 11;
        let g = |x: i32| x / 3;

        let left = x.fmap(f.compose(g));
        let right = x.fmap(g).fmap(f);

        assert_eq!(left, right);
    }

    #[test]
    fn applicative_identity() {
        let x = Maybe::pure(69);

        let left = apply(Maybe::pure(id), x);
        let right = x;

        assert_eq!(left, right);
    }

    #[test]
    fn applicative_composition() {
        let times_11 = |x: i32| x * 11;
        let div_3 = |x: i32| x / 3;

        let u = Maybe::pure(times_11);
        let v = Maybe::pure(div_3);
        let w = Maybe::pure(69);

        // manually compose the functions bc generic rust fn composition is.... something
        let f = |x| times_11(div_3(x));

        let pure_f = Maybe::pure(f);

        let left = w.apply(pure_f);
        let right = w.apply(v).apply(u);

        assert_eq!(left, right);
    }

    #[test]
    fn applicative_homomorphism() {
        let x = 69;
        let f = |x: i32| (x * 11) / 3;

        let pure_x = Maybe::pure(x);
        let pure_f = Maybe::pure(f);

        let left = pure_x.apply(pure_f);
        let right = Maybe::pure(f(x));

        assert_eq!(left, right);
    }

    #[test]
    fn applicative_interchange() {
        const Y: i32 = 69;

        fn calc(x: i32) -> i32 {
            (x * 11) / 3
        }

        fn apply_to_y(f: impl Fn(i32) -> i32) -> i32 {
            f(Y)
        }

        let u = Maybe::pure(calc);
        let pure_y = Maybe::pure(Y);
        let pure_apply_to_y = Maybe::pure(apply_to_y);

        let left = pure_y.apply(u);
        let right = u.apply(pure_apply_to_y);

        assert_eq!(left, right);
    }

    #[test]
    fn functor_satisfies_applicative() {
        let x = Maybe::pure(69);
        let f = |x: i32| (x * 11) / 3;

        let pure_f = Maybe::pure(f);

        let left = x.fmap(f);
        let right = x.apply(pure_f);

        assert_eq!(left, right);
    }

    #[test]
    fn monad_left_identity() {
        let x: i32 = 69;
        let f = |x: i32| Maybe::pure((x * 11) / 3);

        let left = Maybe::pure(x).bind(f);
        let right = f(x);

        assert_eq!(left, right);
    }

    #[test]
    fn monad_right_identity() {
        let m = Maybe::pure(69);

        let left = m.bind(Maybe::pure);
        let right = m;

        assert_eq!(left, right);
    }

    #[test]
    fn monad_associativity() {
        let m = Maybe::pure(69);
        let k = |x: i32| Maybe::pure(x * 11);
        let h = |x: i32| Maybe::pure(x / 3);

        let left = m.bind(|x| k(x).bind(h));
        let right = m.bind(k).bind(h);

        assert_eq!(left, right);
    }

    #[test]
    fn monad_applicative_relationship() {
        let f = |x: i32| (x * 11) / 3;
        let m1 = Maybe::pure(f);
        let m2 = Maybe::pure(69);

        let left = m2.apply(m1);
        let right = m1.bind(|x1| m2.bind(|x2| Maybe::pure(x1(x2))));

        assert_eq!(left, right);
    }
}