Rustで内包表記 (List comprehension)

追伸

すごいHaskellの人が、きれいなHaskell風の内包表記のcrateを公開してくれているので、これをおすすめしたい。

comprehension - Rust

内包表記マクロ

素のRustには内包表記は用意されていないが、マクロを利用してPython風の内包表記を実現したcrateは幾つか書かれている。

cute - Rust
mapcomp - Rust

Rustには比較的強力で書きやすいハイジニックなマクロが用意されていて、構文はかなりの自由度で操作できる。ここでは、Haskell風の内包表記が欲しかったので書いてみた。

macro_rules! lc {
    [$x:expr; $a:ident <- $as:expr, $( $b:ident <- $bs:expr ),+]
        => { $as.flat_map(|$a| lc![$x; $( $b <- $bs ),+]).collect::<Vec<_>>() };
    [$x:expr; $a:ident <- $as:expr]
        => { $as.map(|$a| $x).collect::<Vec<_>>() };
    [$x:expr; $a:ident <- $as:expr, $( $b:ident <- $bs:expr ),+; $( $e:expr ),+]
        => { $as.flat_map(|$a| lc![$x; $( $b <- $bs ),+; $( $e ),+]).collect::<Vec<_>>() };
    [$x:expr; $a:ident <- $as:expr; $( $e:expr ),+]
        => { $as.flat_map(|$a| lc![$x; $( $e ),+]).collect::<Vec<_>>() };
    [$x:expr; $c:expr, $( $e:expr ),+]
        => { if $c { lc![$x; $( $e ),+] } else { vec![] } };
    [$x:expr; $c:expr] => { if $c { vec![$x] } else { vec![] } };
    [$x:expr] => { vec![$x] };
}

Schemeのsyntax-rulesを思わせる、書きやすいマクロだと思う。0回以上の繰り返しの量化子*の挙動が微妙だった (よく理解出来ていない) ので+のみを使った。

lc![x*x; x <- 1..20; x%2 == 1]
=> [1, 9, 25, 49, 81, 121, 169, 225, 289, 361]

// ラマヌジャンのタクシー数を見つける
lc![(a,b,c,d,a*a*a + b*b*b); a <- 1..20, b <- a..20, c <- 1..20, d <- c..20;
    a < c, a*a*a + b*b*b == c*c*c + d*d*d]
=> [(1, 12, 9, 10, 1729), (2, 16, 9, 15, 4104)]

Vecのiter()を使いたい

Vecのiter()をジェネレータ部分に使う時は、参照が渡るので、関数適用や比較演算などでderefやclone()が必要になることがある。そこで、最初にすべてcloned()相当の処理を行ってしまうことも考えた。効率は良くない気がする。

macro_rules! lc2 {
    [$x:expr; $a:ident <- $as:expr, $( $b:ident <- $bs:expr ),+]
        => { $as.map(|x| x.clone()).flat_map(|$a| lc2![$x; $( $b <- $bs ),+]).collect::<Vec<_>>() };
    [$x:expr; $a:ident <- $as:expr]
        => { $as.map(|x| x.clone()).map(|$a| $x).collect::<Vec<_>>() };
    [$x:expr; $a:ident <- $as:expr, $( $b:ident <- $bs:expr ),+; $( $e:expr ),+]
        => { $as.map(|x| x.clone()).flat_map(|$a| lc2![$x; $( $b <- $bs ),+; $( $e ),+]).collect::<Vec<_>>() };
    [$x:expr; $a:ident <- $as:expr; $( $e:expr ),+]
        => { $as.map(|x| x.clone()).flat_map(|$a| lc2![$x; $( $e ),+]).collect::<Vec<_>>() };
    [$x:expr; $c:expr, $( $e:expr ),+]
        => { if $c { lc2![$x; $( $e ),+] } else { vec![] } };
    [$x:expr; $c:expr] => { if $c { vec![$x] } else { vec![] } };
    [$x:expr] => { vec![$x] };
}

使ってみた例。lcではiter().cloned()するか、y.clone()または*yなどが必要な部分がある。

let v0: Vec<f64> = vec![0.1,0.2,0.3,0.5,0.8];
let v1: Vec<f64> = vec![1.41421,1.73205];
lc2![x.powf(y); x <- v0.iter(), y <- v1.iter(); x*5.0 > y]
=> [0.18219634046785238, 0.37521515374493114, 0.3010239124332578, 0.7293716698543952, 0.6794335871063408]

参考文献

マクロの書き方については、この辺りが参考になった。

Macros - The Rust Programming Language
Rustのマクロを覚える - Qiita