【n 番煎じ】Rust で標準入力を楽にするマクロを書いてみた

前口上

Rust で標準入力を扱う際、競技プログラミングのように形式が決まりきっていても結構なしんどさになります。例えば一行で整数が空白区切りで渡されるとき、それを Vec<i32> で受けるとなると次のようになります。

let mut buffer = String::new();
std::io::stdin().read_line(&mut buffer).unwrap();
let v = buffer
    .split_whitespace()
    .map(|s| s.parse::<i32>().unwrap())
    .collect::<Vec<_>>();

__既に有名なものがある__ようですが、せっかくなのでマクロの練習をしてみようと思い立って実装してみました。

既にあるもの

勉強のため__中身は見ないようにして__実装したのですが（友人が使っているのを見たことがあるので影響は受けていますが）、実装そのものもかなり似たものができてしまいました。
特に .bytes().map(|u| u.unwrap() as char) の箇所とか .skip_while(...).take_while(...) とか丸かぶりです。

再発明した車輪

参考

Rust はリファレンスに例がくっついていていい感じで助かります。

本体

お気楽パーサを書いている感じで楽しいです。

use std::io::Read;

macro_rules! read {
    () => { read!('next) };
    ($t:tt) => { // let ss = read!('line).split(',').collect::<Vec<_>>();
        read_inner! {
            std::io::stdin()
                .lock()
                .bytes()
                .map(|u| u.unwrap() as char)
                .skip_while(|&c| c.is_whitespace() /* 先頭の空白は無視 */), $t
        }
    };
    ($t:tt; $length:expr) => { // let v = read!(i32; 100);
        read!([$t; $length])
    };
    ($t:tt $(, $rest_t:tt)* $(,)?) => { // let (i, j) = read!(i32, i64);
        read!(($t $(, $rest_t)* ,));
    };
    ($($($i:ident)+ : $t:tt $(,)? $(;)?)*) => { // read! { let i: i32; j: i32; }
        $($($i)+ = read!($t);)*
    };
}

macro_rules! read_inner {
    ($iter:expr, ['seqc; $length:expr]) => { // 連続する $length 文字
        $iter.take($length).collect::<Vec<_>>()
    };
    ($iter:expr, [$t:tt; $length:expr]) => { // vec
        (0..$length).map(|_| read_inner!($iter, $t)).collect::<Vec<_>>()
    };
    ($iter:expr, ()) => { // void (読み飛ばし)
        (read_inner!($iter, 'next), ()).1
    };
    ($iter:expr, ($t:tt ,)) => { // tuple (要素 1)
        (read_inner!($iter, $t) ,)
    };
    ($iter:expr, ($t:tt $(, $rest_t:tt)+ $(,)?)) => { // tuple
        (read_inner!($iter, $t) $(, read_inner!($iter, $rest_t))+)
    };
    ($iter:expr, 'next) => { // 次 (空白区切り)
        $iter.take_while(|&c| !c.is_whitespace()).collect::<String>()
    };
    ($iter:expr, 'line) => { // 次 (改行区切り)
        $iter.take_while(|&c| c != '\n' && c != '\r').collect::<String>()
    };
    ($iter:expr, char) => {
        $iter.next().unwrap()
    };
    ($iter:expr, $t:tt) => { // 本質
        read_inner!($iter, 'next).parse::<$t>().unwrap()
    };
}

使用例

// 次の空白文字までの文字列を読み取る (String)
let s = read!();

// 型を指定
let i = read!(i32);

// タプルや vec の場合の省略記法
let (n, k) = read!(i32, i64); // read!((i32, i64)); と同じ
let v = read![char; n];       // read!([char; n]);  と同じ

// 例のマクロのように (ちょっと違うけど) 宣言しながら
read! {
    let h1: usize;
    let w1: usize;
    let mut stage: [[i32; w1]; h1]; // mut も二重配列もいける
}

一部特殊なもの

楽しくなって作りました。read_inner のパターンを追加していけばいくらでも増やせるので無限に蛇足ができます。

read! {
    let  t: 'next; // 次の空白文字の手前まで (String)
    let  l: 'line; // 次の改行文字の手前まで (String)
    let _x: ();    // 'next 分の読み飛ばし
    let  q: ['seqc; k]; // 連続する k 文字 (Vec<char>)
}

本質

基本的に次のようなものになるように展開しているだけです。

io::stdin().bytes().map(|u| u.unwrap() as char)
// 標準入力を空白文字に当たっている間読み飛ばす
    .skip_while(|c| c.is_whitespace())
// 空白文字ではない間読み取る
    .take_while(|c| !c.is_whitespace())
// String にする
    .collect::<String>()
// T にパース
    .parse::<T>().unwrap();

着想

Haskell で標準入力をどうのこうのしたい場合には簡単パーサを書いているのですが、イメージ的にはほとんど同じような気がするので参考として載せておきます。

import           Control.Monad.State   (StateT (..), evalStateT)
import           Data.Char             (isSpace)
import           Data.Maybe            (fromJust)

import           Data.ByteString.Char8 (ByteString)
import qualified Data.ByteString.Char8 as BS
import           Data.Vector.Unboxed   (Vector)
import qualified Data.Vector.Unboxed   as V

rInt :: StateT ByteString Maybe Int
rInt = StateT $ BS.readInt . BS.dropWhile isSpace

input :: StateT ByteString Maybe (Int, Int, Vector Int, Vector (Int, Int))
input = do
  n <- rInt
  m <- rInt
  v <- V.replicateM n rInt
  u <- V.replicateM m ((,) <$> rInt <*> rInt)
  return (n, m, v, u)

main :: IO ()
main = do
  (n, m, v, u) <- fromJust . evalStateT input <$> BS.getContents
  return ()

最後に

この記事を書いていて思ったことですが、実用的には次のようにいくつかの関数を用意するだけで良い気がします。

fn read_l() -> String {
    use std::io::Read;
    std::io::stdin()
        .lock()
        .bytes()
        .map(|u| u.unwrap() as char)
        .skip_while(|&c| c.is_whitespace())
        .take_while(|&c| c != '\n' && c != '\r')
        .collect()
}

fn read_s() -> String {
    use std::io::Read;
    std::io::stdin()
        .lock()
        .bytes()
        .map(|u| u.unwrap() as char)
        .skip_while(|&c| c.is_whitespace())
        .take_while(|&c| !c.is_whitespace())
        .collect()
}

fn read_t<T: std::str::FromStr>() -> Result<T, T::Err> {
    read_s().parse()
}

fn read_v<T: std::str::FromStr>(n: usize) -> Result<Vec<T>, T::Err> {
    (0..n).map(|_| read_t()).collect()
}

fn main() {
    let v = read_v::<i32>(100).unwrap();
    // ...
}