aoj-checker.rb

require 'json'
require 'net/http'
require 'open3'
require 'optparse'
require 'pathname'
require 'rbconfig'
require 'tempfile'
require 'uri'

TEST_CASE_MAX = 120

def main
  opts = parse_args(ARGV)
  unless opts.include? :command
    exit_err("You have to specify command!\n Usage: ruby aoj-checker.rb -P 0511 -C ./a.out")
  end
  timeout = opts.include?(:timeout) ? opts[:timeout].to_f : 5.0
  if opts.include? :problem
    d = Downloader.new opts[:problem]
    n = run_with_downloader d, opts[:specs], opts[:command], timeout
    success n
  elsif opts.include? :input
    unless opts.include? :output
      exit_err 'You have to specify both intput and output files.
                Usage: ruby aoj-checker.rb -C ./a.out -I in1.text -O out1.txt'
    end
    in_str = File.open(opts[:input]) {|file| file.read}
    run_1test in_str, opts[:command], timeout, opts[:output]
    success 1
  else
    exit_err 'You have to specify problem id or test case files.
              Usage: ruby aoj-checker.rb -C ./a.out -P 0511'
  end
end

def success(n)
  puts "Congrats! You passed %i tests!" % n
end

def run_1test(in_str, commands, timeout, out_file, test_spec=nil)
  res = exec_command in_str, commands, timeout
  res.check commands
  res = exec_diff commands, res.out, out_file
  res.check_for_diff test_spec
end

def run_with_downloader(downloader, specs, commands, timeout)
  spec_iter = specs.nil? ? 1..TEST_CASE_MAX : specs
  dir = Pathname.new 'aoj-' + downloader.id + '-tests'
  dir.mkdir unless dir.exist?
  suc = 0
  spec_iter.each do |spec|
    in_file = dir.join('in' + spec.to_s + '.txt')
    out_file = dir.join('out' + spec.to_s + '.txt')
    in_str = if in_file.exist? and out_file.exist?
      in_file.open { |f| f.read }
    else
      res = downloader.get spec
      break unless res.ok? suc > 0, !specs.nil?
      File.open(in_file.to_path, 'w') { |f| f.write(res.in) }
      File.open(out_file.to_path, 'w') { |f| f.write(res.out) }
      res.in
    end
    run_1test in_str, commands, timeout, out_file.to_path, test_spec=spec
    puts 'Passed test case %i' % spec
    suc += 1
  end
  suc
end

class DlResult
  attr_reader :in, :out, :err
  def initialize(i, o, err)
    @in, @out, @err = [i, o, err]
  end
  def ok?(done, specified)
    if @err
      exit_err 'Failed to get a specified test case: %s' % @err if specified
      exit_err 'Failed to get a test case: %s' % @err unless done
      false
    else
      true
    end
  end
end

# see http://developers.u-aizu.ac.jp/index for detail
class Downloader
  PREFIX = 'https://judgedat.u-aizu.ac.jp/testcases/'
  attr_reader :id
  
  def initialize(id)
    @id = id
  end
  
  def get(spec)
    uri = PREFIX + @id.to_s + '/' + spec.to_s
    self.parse Net::HTTP.get URI.parse(uri)
  end
  
  def parse(s)
    res = JSON.parse(s)
    if res.instance_of? Array and res[0].include? 'id'
      return DlResult.new nil, nil, res[0]['message']
    elsif res.include? 'problemId'
      return DlResult.new res['in'], res['out'], nil
    else
      exit_err 'Unexpected json: ' + s
    end
  end
end

def parse_args(argv)
  opt = OptionParser.new
  args = {}
  opt.on('-P', '-p', '--problem=PROBLEM_ID', desc = '
         Specifies the problem id.') {|x| args[:problem] = x }
  opt.on('-I', '-i', '--input=INPUT_FILE', desc ='
         Specifies the input file you want to use.
         Ignored when used with -P.') {|x| args[:input] = x }
  opt.on('-O', '-o', '--output=OUTPUT_FILE', desc = '
         Specifies the correct output file..
         Ignored when used with -P, or used without -I.') {|x| args[:output] = x }
  opt.on('-C', '-c', '--command=COMMAND', desc = '
         Specifies the command to exec your program.
         E.g. -C ./a.out') {|x| args[:command] = x.split(' ').map{|s| s.strip} }
  opt.on('-S', '-s', '--specs=SPEC1, SPEC2,..', desc = '
         Specifies the testcase id you want to check(e.g. -S "2, 3, 18")
         Only enabled when used with -P') do |x|
    args[:specs] = x.split(',').map{|s| s.split(' ')}.flatten
  end
  opt.on('-T', '-t', '--timeout=TIMEOUT', desc = '
         Set timeout to your command.
         E.g. ruby aoj-checker.rb -C ./a.out -T 2.5') {|x| args[:timeout] = x}
  opt.parse! argv
  args
end

def exit_err(s)
  puts 'Error!'
  puts s
  exit!
end

class CmdResult
  attr_accessor :out, :err, :stat, :user_op, :answer
  
  def initialize 
    @out, @err, @stat, @user_op, @answer = [nil, nil, nil, nil, nil]
  end

  def set(oes)
    @out, @err, @stat = oes
  end
  
  def check(commands)
    return if @stat.success?
    exit_err "command %s exited with code %i\nstderr: %s\nstdout:%s" \
             % [commands.join(' '), @stat.exitstatus, @err, @out]
  end
  
  def check_for_diff(test_spec=nil)
    return if @stat.success?
    exit_str = 'diff failed '
    exit_str += 'for %i ' % test_spec if test_spec
    exit_str += "\nstderr:%s \nstdout:%s" % [@err, @out]
    puts exit_str + 'Your output: %s Answer: %s' % [@user_op, @answer]
    exit!
  end
end

def exec_diff(commands, in_str, out_file, test_spec=nil)
  diff_cmd ||= (
    host_os = RbConfig::CONFIG['host_os']
    case host_os
    # TODO: maybe msys or mingw has diff?
    when /mswin|msys|mingw|cygwin|bccwin|wince|emc/
      ['fc', '/N']
    else
      ['diff', '-u']
    end
  )
  name = commands.last.split('.')[0] + '-output.txt'
  res = CmdResult.new
  res.user_op = name
  res.answer = out_file
  File.open(name, 'w') do |in_file|
    in_file << in_str
    in_file.close
    res.set Open3.capture3(*diff_cmd, out_file, in_file.path)
  end
  res
end

def exec_command(in_str, commands, timeout, test_spec=nil)
  in_r,  in_w  = IO.pipe
  out_r, out_w = IO.pipe
  err_r, err_w = IO.pipe
  in_w.sync = true
  pid = nil
  res = CmdResult.new

  out_reader = nil
  err_reader = nil
  wait_thr = nil

  begin
    Timeout.timeout(timeout) do
      pid = spawn(*commands, :in => in_r, :out => out_w, :err => err_w)
      wait_thr = Process.detach(pid)
      in_r.close; out_w.close; err_w.close
      out_reader = Thread.new { out_r.read }
      err_reader = Thread.new { err_r.read }
      in_w.write in_str
      in_w.close
    end
  rescue Timeout::Error
    Process.kill(:TERM, pid)
    err_str = commands.join(' ') + ' timeouted'
    err_str += ' for testcase ' + test_spec.to_s if test_spec
    exit_err err_str
  ensure
    res.stat = wait_thr.value if wait_thr
    res.out = out_reader.value if out_reader
    res.err = err_reader.value if err_reader
    out_r.close unless out_r.closed?
    err_r.close unless err_r.closed?
  end
  res
end

main