Rust で Monad, Monad Transformer そして Free Monad をエミュレートする

こんにちは。基礎技術研究部の茂木です。今回は Rust の話をします。

11 月 3 日にリリースされた Rust v1.65.0 では Generic Associated Types(以下、GATs)が入りました。 これにより、Higher Kinded Types のエミュレートが比較的しやすくなります。 そうなると Monad を実装してみたくなるのが人情というものです。 実際いくつも(GATs が Stable になる以前から)先行事例が存在します[1,2,3]。

本記事ではまず[1,2]の手法に従って Monad Transformer を実装しつつ、[3]の手法に切り替えたのち Free Monad を実装します。

なお筆者は関数型プログラミング初心者です。その意味で本記事は[4]における教科書や論文ではないのでご了承ください。

また本記事では Monad 自体の解説はしませんので、「すごい Haskell たのしく学ぼう！」*1等をご参照ください。

まずは HKT, Functor, Monad trait を実装します(Applicative は本記事では省略します)。

trait HKT {
    type C;
    type H<T>: HKT<C = T>;
}

trait Functor: HKT {
    type H<T>: Functor<C = T>;
    fn fmap<F, B>(self, f: F) -> <Self as Functor>::H<B>
    where
        F: FnMut(Self::C) -> B;
}

trait Monad: Functor {
    type H<T>: Monad<C = T>;
    fn pure<A>(t: A) -> <Self as Monad>::H<A>;
    fn bind<F, C>(self, f: F) -> <Self as Monad>::H<C>
    where
        F: FnMut(Self::C) -> <Self as Monad>::H<C>;
}

Option 型を使って、いわゆる Maybe Monad 相当を実装してみます。

impl<A> HKT for Option<A> {
    type C = A;
    type H<T> = Option<T>;
}

impl<A> Functor for Option<A> {
    type H<T> = Option<T>;
    fn fmap<F, B>(self, f: F) -> Option<B>
    where
        F: FnMut(Self::C) -> B,
    {
        self.map(f)
    }
}

impl<A> Monad for Option<A> {
    type H<T> = Option<T>;
    fn pure<T>(t: T) -> Option<T> {
        Option::<T>::Some(t)
    }
    fn bind<F, C>(self, f: F) -> Option<C>
    where
        F: FnMut(Self::C) -> Option<C>,
    {
        self.and_then(f)
    }
}

次に Vec 型を用いていわゆる List Monad を実装してみます。

impl<A> HKT for Vec<A> {
    type C = A;
    type H<T> = Vec<T>;
}
impl<A> Functor for Vec<A> {
    type H<T> = Vec<T>;
    fn fmap<F, B>(self, f: F) -> Vec<B>
    where
        F: FnMut(Self::C) -> B,
    {
        self.into_iter().map(f).collect()
    }
}

impl<A> Monad for Vec<A> {
    type H<T> = Vec<T>;
    fn pure<T>(t: T) -> Vec<T> {
        vec![t]
    }
    fn bind<F, C>(self, f: F) -> Vec<C>
    where
        F: FnMut(Self::C) -> Vec<C>,
    {
        self.into_iter().flat_map(f).collect()
    }
}

動作は省略しますが、このように GATs を用いることで Monad をエミュレートできました。

こうなると、次は Monad Transformer をやってみたくなります。実際[2]でも実装されているのですが、そこで実装されているのは IdentityT なので少し物足りません。ここでは MaybeT を実装してみます。

#[derive(Debug, PartialEq)]
struct MaybeT<M, T>
where
    M: Monad<C = Option<T>>,
{
    mvalue: M,
}

impl<M: Monad<C = Option<A>>, A> HKT for MaybeT<M, A>
where
    M: Monad,
{
    type C = A;
    type H<T> = MaybeT<<M as Monad>::H<<<M as HKT>::C as HKT>::H<T>>, T>;
}

impl<M: Monad<C = Option<T>>, T> Functor for MaybeT<M, T> {
    type H<U> = MaybeT<<M as Monad>::H<<<M as HKT>::C as HKT>::H<U>>, U>;
    fn fmap<F, B>(self, mut f: F) -> MaybeT<<M as Monad>::H<<<M as HKT>::C as HKT>::H<B>>, B>
    where
        F: FnMut(Self::C) -> B,
    {
        MaybeT {
            mvalue: self.mvalue.bind(|value| match value {
                Some(v) => M::pure(Option::<B>::Some(f(v))),
                None => M::pure::<Option<B>>(None),
            }),
        }
    }
}

impl<M: Monad<C = Option<T>>, T> Monad for MaybeT<M, T> {
    type H<U> = MaybeT<<M as Monad>::H<<<M as HKT>::C as HKT>::H<U>>, U>;
    fn pure<A>(t: A) -> <Self as Monad>::H<A> {
        MaybeT {
            mvalue: M::pure(Option::<A>::Some(t)),
        }
    }
    fn bind<F, C>(self, mut f: F) -> <Self as Monad>::H<C>
    where
        F: FnMut(Self::C) -> <Self as Monad>::H<C>,
    {
        MaybeT {
            mvalue: self.mvalue.bind(|value| match value {
                Some(v) => f(v).mvalue,
                None => M::pure::<Option<C>>(None),
            }),
        }
    }
}

ここで type H<T> = MaybeT<<M as Monad>::H<<<M as HKT>::C as HKT>::H<T>>, T> について解説します。 H<T> では MaybeT<M, A> の A を T に変えたいのですが、実は MaybeT<M, T> では不十分です。というのは、M に A が含まれてしまっているからです。

実際、例えば M = Vec<Option<A>> としてみます。このとき MaybeT<M, T> は MaybeT<Vec<Option<A>>, T> です。しかしここで H<T> として欲しいのは MaybeT<Vec<Option<T>>, T> です。 Option を単体の Functor/Monad として使用できないため、Option<A> のような形で Functor/Monad をエミュレートしていることが仇となっています。

では、どうすればよいのでしょうか。

ここでは Vec<Option<T>> を作りたいわけですが、これは M::H<Option<T>> と書き換えられます。 ここで M::C は Option<A> であり、Option<A>::H<T> 、即ち M::C::H<T> は Option<T> です。 よって type H<T> = MaybeT<<M as Monad>::H<<<M as HKT>::C as HKT>::H<T>>, T> とすればよいことが分かります。

それでは MaybeT を使ってみましょう。その前に次の 2 つの関数を定義します。

impl<M: Monad<C = Option<T>>, T> MaybeT<M, T> {
    fn run_maybe_t(self) -> M {
        self.mvalue
    }
}

fn lift<M: Monad>(ma: M) -> MaybeT<<M as Functor>::H<Option<M::C>>, M::C>
where
    <M as Functor>::H<Option<M::C>>: Monad,
{
    MaybeT {
        mvalue: ma.fmap(Option::<M::C>::pure),
    }
}

こうして、以下が動きます。

    let pure = MaybeT::<Vec<Option<i8>>, i8>::pure;
    let x = pure(10)
        .fmap(|x| 2 * x)
        .bind(|x| pure(x + 2))
        .bind(|x| lift(vec![x, x]))
        .run_maybe_t();
    assert_eq!(x, vec![Some(22), Some(22)]);

これは Haskell でいえば次に相当します。

x :: [Maybe Int]
x = runMaybeT $ (return 10) <&> (2*) >>= return . (+2) >>= \z -> lift [z, z]

さて、Monad Transformer が実装できたので、こうなると Free Monad や Extensible Effects も実装してみたくなります。 しかし、残念ながらここまで記載してきたエミュレート手法では Free Monad を実装できませんでした。 これは Free Monad の再帰的な定義と、Option<A> を Option<T> にするような内側の型の置き換えの相性が悪いことに起因しているように思えます。

筆者の Rust 力が低いだけかもしれませんので、みなさんも是非挑戦してみてください。

そこでエミュレート手法を[3]のものに変更します。

trait HKT {
    type H<T>;
}

trait Functor: HKT {
    fn fmap<F, A, B>(f: F, fa: Self::H<A>) -> Self::H<B>
    where
        F: FnOnce(A) -> B;
}

trait Monad: Functor {
    fn pure<A>(t: A) -> Self::H<A>;
    fn bind<F, A, C>(f: F, ma: Self::H<A>) -> Self::H<C>
    where
        F: FnOnce(A) -> Self::H<C>;
}

例えば Maybe Monad は以下のようになります。

struct Maybe {}

impl HKT for Maybe {
    type H<T> = Option<T>;
}

impl Functor for Maybe {
    fn fmap<F, A, B>(f: F, fa: Option<A>) -> Option<B>
    where
        F: FnOnce(A) -> B,
    {
        fa.map(f)
    }
}

impl Monad for Maybe {
    fn pure<T>(t: T) -> Option<T> {
        Option::<T>::Some(t)
    }
    fn bind<F, A, C>(f: F, ma: Option<A>) -> Option<C>
    where
        F: FnOnce(A) -> Option<C>,
    {
        ma.and_then(f)
    }
}

これは以下のように使います。

let a = Some(1);
let b = Maybe::fmap(|value| value + 1, a);

Option そのものを Maybe Monad として扱えないため、最初の手法より「エミュレーション感」が強く、またメソッドチェーンしていけないといった不便さもあります*2。しかし、Maybe そのものから値の型パラメーターを消すことができました。省略しますが、List Monad も同様に実装できます。

それでは Free Monad を実装してみましょう。

enum Free<F: Functor, A> {
    Pure(A),
    Join(Box<F::H<Free<F, A>>>),
}

struct FreeMonad<F> {
    _type: PhantomData<F>,
}

impl<F: Functor> HKT for FreeMonad<F> {
    type H<T> = Free<F, T>;
}

impl<F: Functor> Functor for FreeMonad<F> {
    fn fmap<G, A, B>(f: G, fa: Self::H<A>) -> Self::H<B>
    where
        G: FnOnce(A) -> B,
    {
        match fa {
            Free::Pure(v) => Free::Pure(f(v)),
            Free::Join(value) => Free::Join(Box::new(F::fmap(|v| FreeMonad::fmap(f, v), *value))),
        }
    }
}

impl<F: Functor> Monad for FreeMonad<F> {
    fn pure<T>(t: T) -> Free<F, T> {
        Free::Pure(t)
    }
    fn bind<G, A, C>(f: G, ma: Self::H<A>) -> Self::H<C>
    where
        G: FnOnce(A) -> Self::H<C>,
    {
        match ma {
            Free::Pure(v) => f(v),
            Free::Join(value) => Free::Join(Box::new(F::fmap(|v| FreeMonad::bind(f, v), *value))),
        }
    }
}

Option と Maybe のように、Free と FreeMonad を分けるのが肝です。 また、FreeMonad::bind/fmap に f をそのまま渡すため FnOnce にしています(|a|f(a) を渡せば FnMut でもいけますがすぐに再帰制限に引っかかります)。 このようにして Free Monad が実装できました。Free Monad は Functor から Monad が作れます。 実際に Writer Monad を Free Monad を使って実装してみましょう。なお実装は[5]を参考にしています。

fn lift_f<F: Functor, A>(fa: F::H<A>) -> Free<F, A> {
    Free::Join(Box::new(F::fmap(|a| Free::Pure(a), fa)))
}

struct Tell<W, A> {
    a: A,
    w: W,
}
struct Writer<W> {
    _type: PhantomData<W>,
}

impl<W> HKT for Writer<W> {
    type H<T> = Tell<W, T>;
}

impl<W> Functor for Writer<W> {
    fn fmap<F, A, B>(f: F, fa: Self::H<A>) -> Self::H<B>
    where
        F: FnOnce(A) -> B,
    {
        Tell {
            a: f(fa.a),
            w: fa.w,
        }
    }
}

type WriterMonad<W> = FreeMonad<Writer<W>>;

fn tell<W>(w: W) -> Free<Writer<W>, ()> {
    lift_f(Tell { a: (), w })
}

それでは動かしてみましょう。そのために 2 つの関数を用意します。

fn run_as_vec<W, A>(free: Free<Writer<W>, A>) -> (Vec<W>, A) {
    fn go<W, A>(mut acc: Vec<W>, free: Free<Writer<W>, A>) -> (Vec<W>, A) {
        match free {
            Free::Pure(a) => (acc, a),
            Free::Join(value) => {
                let Tell { a, w } = *value;
                acc.push(w);
                go(acc, a)
            }
        }
    }
    go(vec![], free)
}

fn run_as_list<W, A>(free: Free<Writer<W>, A>) -> (LinkedList<W>, A) {
    fn go<W, A>(mut acc: LinkedList<W>, free: Free<Writer<W>, A>) -> (LinkedList<W>, A) {
        match free {
            Free::Pure(a) => (acc, a),
            Free::Join(value) => {
                let Tell { a, w } = *value;
                acc.push_back(w);
                go(acc, a)
            }
        }
    }
    go(LinkedList::new(), free)
}

そうして、型と値を出力する補助関数を定義しつつ main 関数を以下のようにします。

fn print<T: Debug>(v: &T) {
    println!("{}", std::any::type_name::<T>());
    println!("{:?}", v)
}

fn main() {
    let a = tell("foo");
    let b = WriterMonad::<&str>::bind(|_| tell("hoge"), a);
    let (vec, _) = run_as_vec(b);
    print(&vec);

    let a = tell("foo");
    let b = WriterMonad::<&str>::bind(|_| tell("hoge"), a);
    let (vec, _) = run_as_list(b);
    print(&vec);
}

これを動かした結果は以下となります。

alloc::vec::Vec<&str>
["foo", "hoge"]
alloc::collections::linked_list::LinkedList<&str>
["foo", "hoge"]

つまり、同じ b に対して、Vec と LinkedList の 2 つの別々の解釈ができたわけです。

おわりに

Rust で GATs を駆使することで Monad、Monad Transformer、さらには Free Monad までエミュレートしました。マクロを使った do 構文や Extensible Effects は読者の課題とします。

ＦＦＲＩセキュリティでは製品開発等に Rust も使っています*3。サイバーセキュリティに関する研究開発にご興味ある方は、是非こちらの採用ページをご覧ください。

参考文献

[1]Rustでもモナドは実装できるのか？（再）, https://blog-dry.com/entry/2020/12/25/130250 (Accessed 09 November 2022.)

[2]Monads and GATs in nightly Rust, https://www.fpcomplete.com/blog/monads-gats-nightly-rust/ (Accessed 09 November 2022.)

[3]free.rs, https://github.com/romac/free.rs (Accessed 09 November 2022.)

[4]技術記事の3類型: 初心者による技術記事執筆のすすめ, https://zenn.dev/uhyo/articles/technical-articles (Accessed 09 November 2022.)

[5]Extensible Effects in Dotty, https://www.slideshare.net/SanshiroYoshida/extensible-effects-in-dotty (Accessed 09 November 2022.)

*1:Miran Lipovaca 著『すごい Haskell たのしく学ぼう！』田中英行・村主崇行共訳、2012 年、オーム社

*2:文法的にはこちらの方が Haskell の Maybe Monad に近いですし、メソッドチェーンも代わりに演算子をオーバーロードする手があります。

*3:もちろん本記事のようなことはしません