宣言的マクロでfor内包表記を部分的に実装する

Rustの宣言的マクロ(declarative macro)でScalaのfor内包表記(for comprehension)を部分的に実装してみた。Scalaの

for {
  x <- xs
  y <- ys
} yield doSomething(x, y)

をRustで

for_vec! {
  x <- xs,
  y <- ys;
  do_something(x, y)
}

のように書けるようにする。実装したのはVecとOptionに対して使えるマクロ。複数の型を混在させたりガード式を書ける機能は実装してない。

元のコードは「なっとく!関数型プログラミング」("Grokking Functional Programming")の例の引用。

Vec

macro_rules! for_vec {
    ($v:ident <- $iterable:expr; $yield:expr) => {
        $iterable.iter().map(|$v| $yield).collect::<Vec<_>>()
    };

    ($x:ident <- $xs:expr, $($y:ident <- $ys:expr),+; $yield:expr) => {
        $xs.iter().flat_map(|$x| for_vec! { $($y <- $ys),+ ; $yield }).collect::<Vec<_>>()
    }
}

Vecに対するfor内包表記をRustのflat_mapmapを使った記述に変換する。マクロ利用時はx <- xsのような式をカンマ区切りで書くことで、リストからの要素取り出しを表現する。

マクロのマッチャーでは、xは識別子:identで、またxsは任意の式を使えるように:exprでマッチする。最後に、セミコロンのあとに書いた式をyield用の式としてマッチする。マッチで取り出したプログラム上の要素がメタ変数($v$iterable$yieldなど)に入る。

マッチしたら、アームでメタ変数とflat_mapmapを使った変換後のプログラムを書く。

x <- xsのような式の列挙をうまく扱うにはマクロの繰り返しと再帰を使う。x <- xsのような式が3つ以上あるとき、2個目以降の式は$($y:ident <- $ys:expr),+でマッチする。再帰的にflat_mapを適用し、再帰のベースケースでmapを適用するようにマクロを書く。

使用例:

#[derive(Debug)]
struct Book {
    title: String,
    authors: Vec<String>,
}

#[derive(Debug)]
struct Movie {
    title: String,
}

fn book_adaptations(author: &String) -> Vec<Movie> {
    if author == "Tolkien" {
        vec![
            Movie {
                title: "An Unexpected Journey".to_string(),
            },
            Movie {
                title: "The Desolation of Smaug".to_string(),
            },
        ]
    } else {
        vec![]
    }
}

fn main() {
  let books = vec![
      Book {
          title: "FP in Scala".to_string(),
          authors: vec!["Chiusano".to_string(), "Bjarnason".to_string()],
      },
      Book {
          title: "The Hobbit".to_string(),
          authors: vec!["Tolkien".to_string()],
      },
  ];
  
  let recommendations = for_vec! {
      book <- books,
      author <- book.authors,
      movie <- book_adaptations(author);
      format!("You may like {}, because you linked {}'s {}", movie.title, author, book.title)
  };

  println!("{:#?}", recommendations);
}
// stdout
[
    "You may like An Unexpected Journey, because you linked Tolkien's The Hobbit",
    "You may like The Desolation of Smaug, because you linked Tolkien's The Hobbit",
]

Option

macro_rules! for_option {
    ($v:ident <- $option:expr; $yield:expr) => {
        $option.map(|$v| $yield)
    };

    ($x:ident <- $xopt:expr, $($y:ident <- $yopt:expr),+; $yield:expr) => {
        $xopt.and_then(|$x| for_option! { $($y <- $yopt),+ ; $yield })
    }
}

マクロの構成方法はVecと同じ。RustのOptionだとOption::and_thenがflat mapにあたるメソッドなので、マクロでそれを使った形式に変換する。

利用例:

#[derive(Debug)]
struct Event {
    name: String,
    start: isize,
    end: isize,
}

fn parse(name: String, start: isize, end: isize) -> Option<Event> {
    for_option! {
        valid_name <- validate_name(&name),
        valid_end <- validate_end(end),
        valid_start <- validate_start(start, end);
        Event { name: valid_name, start: valid_start, end: valid_end }
    }
}

fn validate_name(name: &String) -> Option<String> {
    if !name.is_empty() {
        Some(name.to_owned())
    } else {
        None
    }
}

fn validate_end(end: isize) -> Option<isize> {
    if end < 3000 {
        Some(end)
    } else {
        None
    }
}

fn validate_start(start: isize, end: isize) -> Option<isize> {
    if start <= end {
        Some(start)
    } else {
        None
    }
}

fn main() {
    let apollo = parse("Apollo Program".to_string(), 1961, 1972);
    println!("{:?}", apollo);
    let ww2 = parse("".to_string(), 1939, 1945);
    println!("{:?}", ww2);
}
// stdout
Some(Event { name: "Apollo Program", start: 1961, end: 1972 })
None

感想

DSLを定義するには手軽、がんばりすぎると可読性難