judge

#!/usr/bin/env python3

"""
	This program is free software: you can redistribute it and/or modify
	it under the terms of the GNU General Public License as published by
	the Free Software Foundation, either version 3 of the License, or
	(at your option) any later version.

	This program is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU General Public License for more details.

	You should have received a copy of the GNU General Public License
	along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""

import argparse
import sys
import os
import subprocess
import resource
import mmap
import signal
import math
from enum import Enum

def mem_str(mem):
	return str(mem * 1024 // resource.getpagesize()) + "p/" + str(round(mem / 1024,2)) + "M"

def cpu_mem_str(cpu,mem):
	return mem_str(mem) + " " + str(round(cpu,4)) + "s"

class Verdict(Enum):
	AC = 0
	RTE = 1 << 0
	MLE = 1 << 1
	OLE = 1 << 2
	TLE = 1 << 3
	RE = 1 << 4
	WA = 1 << 5

def iswhite(ch):
	return chr(ch) in [' ', '\t', '\v', '\f', '\n', '\r']

def isline(ch):
	return chr(ch) in ['\n', '\r']

def identical(judge, process):
	return judge == process

def standard(judge, process):
	try:
		#Use a C subprogram instead python for speed increase
		from _checker import standard as c_standard
		return c_standard(judge, process)
	except ImportError:
		j = 0
		p = 0
		jlen = len(judge)
		plen = len(process)
		while j < jlen and iswhite(judge[j]):
			j += 1
		while p < plen and iswhite(process[p]):
			p += 1
		while True:
			nj = False
			np = False
			while j < jlen and iswhite(judge[j]):
				nj |= isline(judge[j])
				j += 1
			while p < plen and iswhite(process[p]):
				np |= isline(process[p])
				p += 1
			if j == jlen or p == plen:
				return j == jlen and p == plen
			if nj != np:
				return False
			while j < jlen and not iswhite(judge[j]):
				if p >= plen:
					return False
				if judge[j] != process[p]:
					return False
				j += 1
				p += 1

def check(judge_output, process_output, checker=standard):
	return checker(judge_output, process_output)

def judge_case(file_i,file_o,cpu_limit,mem_limit,executable,checker, out_limit=None):
	r,w = os.pipe()
	os.set_inheritable(w,True)
	os.set_inheritable(2,True)
	pid = os.fork()
	if pid == 0:
		stdin = os.open(file_i,os.O_RDONLY)
		if stdin != 0:
			os.dup2(stdin,0)
			os.close(stdin)
		os.dup2(w,1)
		cpu_limit_sec = math.ceil(cpu_limit)
		resource.setrlimit(resource.RLIMIT_AS,(mem_limit,mem_limit))
		resource.setrlimit(resource.RLIMIT_CPU,(cpu_limit_sec,cpu_limit_sec))
		os.execv(executable,[executable])
	submission = pid
	os.close(w)
	if out_limit == None:
		out_limit = os.stat(file_o).st_size
	result = os.read(r,mmap.PAGESIZE)
	while len(result) < out_limit:
		buf = os.read(r,mmap.PAGESIZE)
		if len(buf) == 0:
			break
		result += buf
	os.kill(submission,signal.SIGKILL)
	_,return_status,rusage = os.wait4(submission,0)
	# TODO: Other processes randomly sending signals can break this
	if not (len(result) <= out_limit):
		return Verdict.OLE,rusage.ru_utime,rusage.ru_maxrss,return_status
	if os.WIFSIGNALED(return_status) and os.WTERMSIG(return_status) == signal.SIGKILL or rusage.ru_utime > cpu_limit:
		return Verdict.TLE,rusage.ru_utime,rusage.ru_maxrss,return_status
	if os.WIFSIGNALED(return_status):
		return Verdict.RTE,rusage.ru_utime,rusage.ru_maxrss,return_status
	if os.WEXITSTATUS(return_status) != 0:
		return Verdict.RE,rusage.ru_utime,rusage.ru_maxrss,return_status
	answer = open(file_o,"rb").read()
	if not check(answer, result, checker):
		return Verdict.WA,rusage.ru_utime,rusage.ru_maxrss,return_status
	return Verdict.AC,rusage.ru_utime,rusage.ru_maxrss,return_status

def main():
	checkers = {
		'standard' : standard,
		'identical': identical,
	}
	parser = argparse.ArgumentParser(description="An offline judging tool.")
	parser.add_argument("test_cases", help="Directory that contains the test cases to judge with.")
	parser.add_argument("cpu_limit", type=float, help="Time limit.")
	parser.add_argument("mem_limit", help="Memory limit.")
	parser.add_argument("executable", help="File to run.")
	parser.add_argument("checker", default="standard", nargs="?", choices=checkers.keys(), help="Checker to be used to compare outputs.")
	args = parser.parse_args()
	MEMORY_UNIT = {
		"B" : 2**0,
		"K" : 2**10,
		"M" : 2**20,
		"G" : 2**30,
	}
	RTE_MSG = {
		4 : "Illegal instruction",
		6 : "Aborted",
		8 : "Floating point exception",
		11 : "Segmentation fault"
	}
	case_path = args.test_cases
	cpu_limit = args.cpu_limit
	mem_limit = (int(args.mem_limit[:-1]) * MEMORY_UNIT[args.mem_limit[-1:]]) & -mmap.PAGESIZE
	executable = args.executable
	checker = checkers[args.checker]
	num_ac = 0
	status_mask = 0
	cpu_total = 0.0
	mem_total = 0
	num_cases = 0
	def recursive_judge(case,depth=-1):
		nonlocal cpu_total,mem_total,status_mask,num_cases,num_ac
		if os.path.isdir(case):
			if depth != -1:
				print(" " * 4 * depth + " - " + os.path.basename(case))
			for filename in sorted(os.listdir(case)):
				recursive_judge(os.path.join(case,filename),depth+1)
		elif case.endswith(".in") and os.path.isfile(case) and os.path.isfile(case[:-2] + "out"):
			file_out = case[:-2] + "out"
			verdict,cpu_usage,mem_usage,return_status = judge_case(case,file_out,cpu_limit,mem_limit,executable,checker,out_limit=os.path.getsize(file_out) * 4)
			cpu_total += cpu_usage
			mem_total = max(mem_usage,mem_total)
			status_mask |= verdict.value
			if verdict == Verdict.AC:
				num_ac += 1
			num_cases += 1
			case_name_verdict = " " * 4 * depth + os.path.basename(case[:-3]) + ": " + verdict.name
			usage_str = cpu_mem_str(cpu_usage,mem_usage)
			if verdict == Verdict.WA or verdict == Verdict.AC or verdict == Verdict.OLE:
				print(case_name_verdict + " " + usage_str)
			elif verdict == Verdict.TLE:
				print(case_name_verdict + "(>" + str(cpu_limit) + ") " + mem_str(mem_usage))
			elif verdict == Verdict.RTE:
				signal = os.WTERMSIG(return_status)
				print(case_name_verdict + "(" + (RTE_MSG[signal] if signal in RTE_MSG else "signal " + str(status)) + ") " + usage_str)
			elif verdict == Verdict.RE:
				print(case_name_verdict + "(" + str(os.WEXITSTATUS(return_status)) + ") " + usage_str)
	recursive_judge(case_path)
	print(str(num_ac) + "/" + str(num_cases) + " " + Verdict(status_mask & -status_mask).name + " " + cpu_mem_str(cpu_total,mem_total))

if __name__ == "__main__":
	try:
		sys.exit(main())
	except KeyboardInterrupt:
		sys.exit(1)