Rustで競技プログラミングの入力をスッキリ記述するマクロ

モチベーション

Rustで標準入力から望みのデータを入力するのは意外と面倒で、標準ライブラリにあるものを使って素直に書こうとすると、結構コード量がかさみます。例えば、文字列を一行読み込むのに、

let s = {
    let mut s = String::new(); // バッファを確保
    std::io::stdin().read_line(&mut s).unwrap(); // 一行読む。失敗を無視
    s.trim_right().to_owned() // 改行コードが末尾にくっついてくるので削る
};

こんな風に書く必要があって、ここからさらに、この行に整数が2個含まれているなら、

let (n, m) = {
    let ws = s.split_whitespace(); // 空白区切りの単語に分解する
    let n: i32 = ws.next().unwrap().parse().unwrap(); // イテレータから値を取り出して整数に
    let m: i32 = ws.next().unwrap().parse().unwrap();
    (n, m)
};

などと書く必要があります。が、とてもじゃないけど、のんびりこんなコードを書いているわけにはいきません。

行ずつ読むAPIだと、このように毎度単語に分解してparse()してというのがあまりにも面倒なので、次のコードのような、C++のcinやJavaのScanner的な空白区切りで読むトークナイザーを作って頑張っていたのですが、それでもいまいち思ったようにすっきり書けている気がしませんでした。

let content = read_string(); // 標準入力から全部Stringに読み込む関数
let mut sc = Scanner::new(&content); // Scannerを作る
let n: i32 = sc.next(); // 次のトークンを読み込む（FromStrな型へのparseを行う）
let m: i32 = sc.next();

このぐらいだとそれなりにはすっきりはしていますが、例えばn行m列の行列を読み込むような場合だと、

let n: usize = sc.next();
let m: usize = sc.next();
let mat: Vec<Vec<i32>> = (0..n).map(|_| (0..m)(|_| sc.next()).collect::<Vec<_>()).collect();

のような定型文を書いていたりして、なんだかなあという感じでした。

また、RustのString型は、char（UTF32）の配列ではなく、UTF8でエンコードされたバイト列でコンテンツを保持しており、ランダムアクセスができないので、ランダムアクセスが必要になるアルゴリズムで使う場合には、Vec<char>へ変換する必要がありました。

let v: Vec<char> = sc.next::<String>().chars().collect();

競技プログラミングでは、文字の二次元配列で二次元データを与えてくる問題がそれなりにあるので、そういうのだと、ただの入力で割とタイプが増えてしまって、やっぱり、もうちょっと何とかならないものかと常々考えていました。

本来であれば、問題文に書いてあるような入力の仕様を書くだけで、それに対するパーザーは勝手に出来上がってしかるべきだし、それ以上のことを書かないといけないのは何かがおかしいという気持ちをもとに、それっぽいものを作ってみたら、それなりに便利なものができた気がするので、ここに書いてみることにします。

input! マクロ

構文上のノイズが極限まで少なくなるように、また、実行効率を犠牲にしないように（無駄なメモリ確保を行わないなど）という設計思想の元、マクロを書いてみました。

まず最初に挙げたような、整数二つを読み込んで変数にバインドするような例であれば、次のように書けます。

fn main() {
    input!{
        n: usize,
        m: usize,
    }
    // solver code...
}

競技プログラミングによくありがちな、長さnの配列を読み込ませるものであれば、次のように書けます。

// n
// v_1 v_2 ... v_n

// こういうやつを読み込むコード
input!{
    n: usize,
    v: [i32; n],
}

vは、Vec<i32>の値になります。

配列はネストできるので、n×m行列であれば、

input!{
    n: usize,
    m: usize,
    mat: [[i32; m]; n],
}

こうなります。

ちょっと面倒な文字列をVec<char>として読み込むのも、

input!{
    cs: chars, // トークンを Vec<char> として読む
}

このように書けます。よくある、二次元テキストとして盤面が与えられるようなものなら、次のように書けます。

// 3 4
// ####
// #..#
// .#.#

// こういうやつを Vec<Vec<char>> として読み込むコード
input!{
    h: usize,
    w: usize,
    board: [chars; h],
}

AtCoderなどでデータを入力していて地味に面倒な、インデックスが1-originとして与えられるときに0-originへの変換なども、自動でできるようにしてみました。

// n m
// f_1 t_1
// f_2 t_2
// ...
// f_m t_m

// こういうグラフを読み込むコード
input!{
    n: usize, // ノード数
    m: usize, // 枝数
    edges: [(usize1, usize1); m], // usize1 は 1-origin to 0-origin変換を行う
}

割とノイズが少なくて理想的な記述ができるようになったと思います。

実装

というわけで実装を置いておきます。こういう典型的なケースがうまくすっきり入力できないぞとかありましたら教えていただけると嬉しいです。

macro_rules! input {
    (source = $s:expr, $($r:tt)*) => {
        let mut iter = $s.split_whitespace();
        input_inner!{iter, $($r)*}
    };
    ($($r:tt)*) => {
        let mut s = {
            use std::io::Read;
            let mut s = String::new();
            std::io::stdin().read_to_string(&mut s).unwrap();
            s
        };
        let mut iter = s.split_whitespace();
        input_inner!{iter, $($r)*}
    };
}

macro_rules! input_inner {
    ($iter:expr) => {};
    ($iter:expr, ) => {};

    ($iter:expr, $var:ident : $t:tt $($r:tt)*) => {
        let $var = read_value!($iter, $t);
        input_inner!{$iter $($r)*}
    };
}

macro_rules! read_value {
    ($iter:expr, ( $($t:tt),* )) => {
        ( $(read_value!($iter, $t)),* )
    };

    ($iter:expr, [ $t:tt ; $len:expr ]) => {
        (0..$len).map(|_| read_value!($iter, $t)).collect::<Vec<_>>()
    };

    ($iter:expr, chars) => {
        read_value!($iter, String).chars().collect::<Vec<char>>()
    };

    ($iter:expr, usize1) => {
        read_value!($iter, usize) - 1
    };

    ($iter:expr, $t:ty) => {
        $iter.next().unwrap().parse::<$t>().expect("Parse error")
    };
}

追記: 遅延入力

Rustで書くからには、どうしてもゼロオーバーヘッドで書きたくて、不要なメモリの動的確保やらを排除したかったのですが、その結果として上のコードでは、入力全体を一気に読み込んで、それから解析するというコードになっていました。

それでもサブミットする分には問題無いのですが、手元でテストする場合には、毎回Ctrl+DでEOFを入力しなくてはいけなくて幾分面倒です。なので、トークン毎にメモリ確保を行う代わりに、入力が与えられるたびに（コンソールだと大抵は行ごとに）値が得られるような実装を作ってみました。

毎回charにデコードした上でStringを作るので、前のバージョンよりも6倍ほど遅いですが、c++のcin(cin.tie(0)とsync_with_stdio(false)あり)と比べるても2倍遅い程度で収まっており（むしろ前のバージョンがC++より3倍ぐらい速かった）、6倍遅い状況でも、整数1000万個で2秒程度だったので、競技プログラミングで良くある、上限数十万程度の整数を読み込むのであれば、0.1秒もかからないと思います。なので、現実的には殆どの場合ではこちらで問題がないのではないかと思います。

ただ、やっぱり速度と利便性を両立できないものかと思ってしまいますし、Rustってのはそういうのをやるために作られた言語だと思いますので、どなたか上の実装と同じぐらい速い実装を思いつかれた方がいらっしゃいましたら、ご連絡くださると大変嬉しいです。

macro_rules! input {
    (source = $s:expr, $($r:tt)*) => {
        let mut iter = $s.split_whitespace();
        let mut next = || { iter.next().unwrap() };
        input_inner!{next, $($r)*}
    };
    ($($r:tt)*) => {
        let stdin = std::io::stdin();
        let mut bytes = std::io::Read::bytes(std::io::BufReader::new(stdin.lock()));
        let mut next = move || -> String{
            bytes
                .by_ref()
                .map(|r|r.unwrap() as char)
                .skip_while(|c|c.is_whitespace())
                .take_while(|c|!c.is_whitespace())
                .collect()
        };
        input_inner!{next, $($r)*}
    };
}

macro_rules! input_inner {
    ($next:expr) => {};
    ($next:expr, ) => {};

    ($next:expr, $var:ident : $t:tt $($r:tt)*) => {
        let $var = read_value!($next, $t);
        input_inner!{$next $($r)*}
    };
}

macro_rules! read_value {
    ($next:expr, ( $($t:tt),* )) => {
        ( $(read_value!($next, $t)),* )
    };

    ($next:expr, [ $t:tt ; $len:expr ]) => {
        (0..$len).map(|_| read_value!($next, $t)).collect::<Vec<_>>()
    };

    ($next:expr, chars) => {
        read_value!($next, String).chars().collect::<Vec<char>>()
    };

    ($next:expr, usize1) => {
        read_value!($next, usize) - 1
    };

    ($next:expr, $t:ty) => {
        $next().parse::<$t>().expect("Parse error")
    };
}