153 lines
4.6 KiB
Python
153 lines
4.6 KiB
Python
from enum import Enum, auto
|
||
from functools import partial
|
||
from itertools import product
|
||
from multiprocessing import Pool
|
||
from subprocess import DEVNULL, PIPE, Popen, TimeoutExpired
|
||
from time import time
|
||
from typing import List
|
||
|
||
chars = '0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!"#$%&\'()*+,-./:;<=>?@[\\]^_`{|}~ '
|
||
|
||
|
||
class Status(Enum):
|
||
"""Possible correctness statuses for a program."""
|
||
Correct = auto() # returns the expected output and a 0 exit code
|
||
Invalid = auto() # has a syntax error or uses an undefined command
|
||
WrongAnswer = auto() # returns an unknown exit code or invalid output
|
||
Timeout = auto() # ran for more time than allowed by the timeout value
|
||
|
||
|
||
def check_pair(script, instr, outstr, timeout) -> Status:
|
||
"""
|
||
Check that a Bash script outputs a given string when given a input string.
|
||
|
||
:param script: script to execute
|
||
:param instr: string given to the script’s stdin
|
||
:param outstr: expected stdout value
|
||
:param timeout: maximum allowed time in seconds
|
||
:returns: status indicating how the script behaved
|
||
"""
|
||
process = Popen(
|
||
[
|
||
"/bin/bash", "--restricted", "-c", "--",
|
||
"trap 'kill -9 $(jobs -p) && wait' SIGINT SIGTERM EXIT;\n"
|
||
+ script,
|
||
],
|
||
stdin=PIPE,
|
||
stdout=PIPE,
|
||
stderr=DEVNULL,
|
||
)
|
||
|
||
try:
|
||
stdout, stderr = process.communicate(instr.encode(), timeout)
|
||
|
||
if process.returncode in (2, 126, 127, 128):
|
||
return Status.Invalid
|
||
|
||
if process.returncode != 0 or stdout != outstr.encode():
|
||
return Status.WrongAnswer
|
||
|
||
return Status.Correct
|
||
except TimeoutExpired:
|
||
try:
|
||
process.terminate()
|
||
stdout, stderr = process.communicate()
|
||
except ProcessLookupError:
|
||
pass
|
||
|
||
return Status.Timeout
|
||
|
||
|
||
def check_script(pairs, timeout, script) -> Status:
|
||
"""
|
||
Check that a Bash script satisfies a set of test cases.
|
||
|
||
:param pairs: input/expected output pairs
|
||
:param timeout: maximum allowed time in seconds
|
||
:param script: script to test
|
||
:returns: status indicating how the script behaved
|
||
"""
|
||
for pair in pairs:
|
||
status = check_pair(script, *pair, timeout)
|
||
if status != Status.Correct:
|
||
return script, status
|
||
return script, status
|
||
|
||
|
||
def generate_scripts(max_length):
|
||
"""
|
||
Generate all scripts up to a given length.
|
||
|
||
:param max_length: maximum length to generate
|
||
:yields: generated scripts
|
||
"""
|
||
for length in range(max_length + 1):
|
||
for letters in product(chars, repeat=length):
|
||
yield "".join(letters)
|
||
|
||
|
||
def find_script(
|
||
pairs,
|
||
max_length,
|
||
processes,
|
||
timeout,
|
||
out_valid_prefix,
|
||
out_log,
|
||
) -> List[str]:
|
||
"""
|
||
Find scripts that satisfy the given set of test cases.
|
||
|
||
:param pairs: input/expected output pairs
|
||
:param max_length: maximum script length to test
|
||
:param processes: number of parallel processes to spawn
|
||
:param timeout: maximum allowed time in seconds for each script run
|
||
:param out_valid_prefix: store valid scripts (includes correct scripts
|
||
and incorrect scripts that do not contain syntax or runtime errors)
|
||
into files starting with this prefix
|
||
:param out_log: stream to which progress logs are written
|
||
:returns: list of matching scripts
|
||
"""
|
||
candidates = []
|
||
out_valid = []
|
||
bound_check_script = partial(check_script, pairs, timeout)
|
||
|
||
for i in range(max_length + 1):
|
||
out_valid.append(open(out_valid_prefix + str(i), "w"))
|
||
|
||
chars_count = len(chars)
|
||
num_tasks = int((chars_count ** (max_length + 1) - 1) / (chars_count - 1))
|
||
done_tasks = 0
|
||
start_time = time()
|
||
|
||
with Pool(processes) as pool:
|
||
for script, status in pool.imap_unordered(
|
||
bound_check_script,
|
||
generate_scripts(max_length),
|
||
chunksize=10,
|
||
):
|
||
done_tasks += 1
|
||
|
||
if done_tasks % 10000 == 0:
|
||
print(
|
||
f"Progress: {done_tasks}/{num_tasks} \
|
||
{done_tasks / num_tasks * 100:.1f}% \
|
||
(running for {time() - start_time:.1f}s)",
|
||
file=out_log, flush=True
|
||
)
|
||
|
||
if status == Status.Correct:
|
||
print(
|
||
f"> Found candidate: '{script}'",
|
||
file=out_log, flush=True
|
||
)
|
||
candidates.append(script)
|
||
|
||
if status != Status.Invalid:
|
||
print(script, file=out_valid[len(script)], flush=True)
|
||
|
||
for file in out_valid:
|
||
file.close()
|
||
|
||
print(f"Finished in {time() - start_time:.1f}s", file=out_log, flush=True)
|
||
return candidates
|