Rust で競技プログラミング用のテストマクロと AtCoder のテストコード生成ツールを作る

ここのところ競技プログラミングが流行ってきて(?) AtCoder の参加者が増えているような気がします。
ということで、今回は AtCoderに Rust で参加する時に作成したテストマクロなどを紹介します。

モチベーション

AtCoderに参加していると、問題を解く以外のところにかける労力はできるだけ減らしたくなりますね。
具体的には次のようなことがめんどくさい。

  • 出力結果のテストがめんどくさい
  • ターミナルにいちいちサンプルの入力、コピペするのがだるい
  • めんどいので提出時にソースコードを全部コピペしたい (テストコード含んだまま出したい)
  • Rust にはテストアトリビュートがせっかく用意されてるのでどうせなら cargo test したい

というわけで、👇のように書いた時に cargo test だけで全てのテストが終わるようなマクロが欲しいです。

test! {
    // "input" => "output" みたいな感じの入出力のペアを何個でも書ける
}

test! マクロ内に記述するコードを毎回書くのもかなり面倒なので、ついでに test! マクロ内に記述する入出力の全てのサンプルを AtCoder のページからワンクリックでコピペできるようなブラウザのアドオンも作ります。

入力マクロを利用する

まず最初に、標準入力を受け取るマクロを作ります。
これについては、tanakhさんが作成した素晴らしい input! マクロを利用させていただきます。

solve関数を作る

問題の入出力を文字列のまま引数と戻り値にするsolve関数を作ります。

fn solve(src: &str) -> String {
    input! {
        source = src,
        n: usize    // 問題に合わせて記述
    }
    // 問題を解くコードを記述
}

この関数を使うとき、main関数は👇のように書けます。

fn main() {
    println!("{}", solve(&stdin!()));
}

テストマクロを作る

いよいよsolve関数を利用して test! マクロを作成します。

macro_rules! test {
    ($($input:expr => $output:expr),* $(,)*) => {
        #[test]
        fn solve_test() {
            $(
                assert_eq!(solve($input), $output);
            )*
        }
    };
}

このコードを使ってABC133のA問題を例に👇のようにテストコードを記述することができます。

// abc133_a
test! {
    "4 2 9" => "8",
    "4 2 7" => "7",
    "4 2 8" => "8"
}

複数行ある時とかは raw string literal なども検討してみましょう。(r"hogehoge"みたいに書くやつ)
ちなみにこのマクロが展開されるとこうなります。

#[test]
fn solve_test() {
    assert_eq!(solve("4 2 9"), "8");
    assert_eq!(solve("4 2 7"), "7");
    assert_eq!(solve("4 2 8"), "8");
}

cargo test でテストする

先程のテストコードを、solve関数を完成させて実行してみましょう。

solve関数

// abc133_a
fn solve(src: &str) -> String {
    input! {
        source = src,
        n: i32,    // 問題に合わせて記述
        a: i32,
        b: i32
    }
    // 問題を解くコード
    std::cmp::min(n * a, b).to_string()
}

実行結果

(´・_・`)👉 cargo test
    Finished debug [unoptimized + debuginfo] target(s) in 0.0 secs
     Running target/debug/deps/a-19521ac5975cd6cf

running 1 test
test solve_test ... ok

test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured

テストが通りました!
状況に応じて自分でコーナーケースを追加したりしつつテストが通ったらそのまま全部コピペして提出しましょう。

AtCoder用にテストコード生成ツールを作る

このままだけだと、test! マクロの中に記述するテストコードがまだ面倒です。
したがって、まとめてコピーするブラウザ拡張機能を作ります。(私の場合はFireFox)

const main = () => {
  const samplePairs = getSamplePairs();
  if (samplePairs.length < 1) {
    return;
  }
  // 日本語と英語それぞれに同じサンプル要素があるので半分だけ使う
  const halfSamplePairs = samplePairs.slice(0, samplePairs.length / 2);
  const testCodeText = toCodeFromSamplePairs(halfSamplePairs);
  console.log(testCodeText);
  generateElements(testCodeText);
};

// 入出力サンプルペア要素を [[inputElement, outputElement], ...] の形で取得
const getSamplePairs = () => {
  let samplePairs = [];
  let num = 0;
  while ((input = document.getElementById(`pre-sample${num++}`))) {
    output = document.getElementById(`pre-sample${num++}`);
    samplePairs.push([input, output]);
  }
  return samplePairs;
};

// 入出力サンプルペアをテストコードに変換
const toCodeFromSamplePair = samplePair => {
  const [input, output] = samplePair;
  return `r"${input.innerText}" => r"${output.innerText.trim()}"`;
};

// 入出力サンプルペアのリストからテストコードに変換
const toCodeFromSamplePairs = samplePairs => {
  return samplePairs.map(toCodeFromSamplePair).join(',\n\n');
};

// コピーボタンとコピー内容(テストコード)の要素を生成する
const generateElements = testCodeText => {
  const button = document.createElement('button');
  button.innerText = 'Copy sample test code';
  button.addEventListener('click', () => {
    document.getSelection().selectAllChildren(testCodeElement);
    document.execCommand('copy');
    const p = document.createElement('p');
    p.innerText = 'Copied!!!';
    p.style = 'color: blue;';
    wrapper.appendChild(p);
  });

  const testCodeElement = document.createElement('pre');
  testCodeElement.innerText = testCodeText;

  const wrapper = document.createElement('div');
  wrapper.appendChild(button);
  wrapper.appendChild(testCodeElement);

  const root = document.getElementById('task-statement');
  root.appendChild(wrapper);
};

main();

作りました。
いろいろと雑ですが許してください。

manifest.json

こんな感じで manifest.json を書いてatcoderページでだけ機能するようにします。

{
  "manifest_version": 2,
  "name": "AtCoder Sample Test-code",
  "version": "1.0",
  "description": "Copy AtCoder test code of sample case for Rust",
  "content_scripts": [
    {
      "matches": ["*://*.atcoder.jp/*"],
      "js": ["atcoder_sample_testcode.js"]
    }
  ]
}

また、アドオンをインストールする際に署名が必要になったりするのでこのへんとかこのへん を参考にしました。
AMOでも非公開で署名できるっぽいのでそれを使わせていただきました。

使ってみた様子

あまりかっこよくはありませんが問題ページにこんな感じで表示され、ボタンを押すと表示されてるテストコードがクリップボードにコピーされます。
複数行の入力にも対応するために、raw string literal を使っています。

copy sample test code at AtCoder

これで割と楽にAtCoderのテストができるようになりました。

最終的なコード

まとまったコードを置いておきます。
スニペットなどを作る時などの参考にしてもらえたら嬉しいです。

// Input macro
macro_rules! input {
    (source = $s:expr, $($r:tt)*) => {
        let mut iter = $s.split_whitespace();
        input_inner!{iter, $($r)*}
    };
    ($($r:tt)*) => {
        let 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! stdin {
    () => {{
        use std::io::Read;
        let mut s = String::new();
        std::io::stdin().read_to_string(&mut s).unwrap();
        s
    }};
}

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")
    };
}

// Test Code
macro_rules! test {
    ($($input:expr => $output:expr),* $(,)*) => {
        #[test]
        fn solve_test() {
            $(
                assert_eq!(solve($input), $output);
            )*
        }
    };
}

test! {
    // テストコードをここにコピペする
}

fn main() {
    println!("{}", solve(&stdin!()));
}

fn solve(src: &str) -> String {
    input! {
        source = src,
        n: i32,    // 問題に合わせて記述
    }
    // 問題を解くコード
}

おたより発掘

関連記事

モダンな開発用ターミナル環境のためのツール紹介

デザインも頼めるシステム開発会社をお探しならBPS株式会社までどうぞ 開発エンジニア積極採用中です! Ruby on Rails の開発なら実績豊富なBPS

この記事の著者

jhonda

2017年入社。 いろんな言語を触ったり開発環境をいじくり回すのが好きです。 健康管理が趣味です。

jhondaの書いた記事

BPSアドベントカレンダー

週刊Railsウォッチ