<commit_before>from random import Random
from contextlib import contextmanager
from .operations import (
    ChooseFrom,
    ReadAndWrite,
    Check,
    PushRandom,
    BinaryOperator,
    UnaryOperator,
    Dup,
    Drop,
    Swap,
    Rot,
)
from collections import namedtuple, Counter
import operator
import traceback
import math


ProgramStep = namedtuple(
    "ProgramStep",
    ("definitions", "arguments", "operation")
)


class TestMachineError(Exception):
    pass


class FrozenVarStack(TestMachineError):
    def __init__(self):
        super(FrozenVarStack, self).__init__("Cannot modify frozen varstack")


class NoFailingProgram(TestMachineError):
    pass


class VarStack(object):
    def __init__(self, name, context):
        self.name = name
        self.context = context
        self.data = []
        self.names = []
        self.frozen = False

    def _integrity_check(self):
        assert len(self.data) == len(self.names)

    @contextmanager
    def freeze(self):
        self.frozen = True
        try:
            yield
        finally:
            self.frozen = False

    def modification(self):
        if self.frozen:
            raise FrozenVarStack()

    def pop(self):
        self.modification()
        self._integrity_check()
        result = self.data.pop()
        self.context.on_read(self.names.pop())
        return result

    def push(self, head):
        self.modification()
        self._integrity_check()
        self.data.append(head)
        v = self.context.newvar()
        self.names.append(v)
        self.context.on_write(v)

    def dup(self):
        self.modification()
        self._integrity_check()
        self.names.append(self.names[-1])
        self.data.append(self.data[-1])

    def peek(self, index=0):
        self._integrity_check()
        i = -1 - index
        self.context.on_read(self.names[i])
        return self.data[i]

    def has(self, count):
        self._integrity_check()
        return len(self.data) >= count


class RunContext(object):
    def __init__(self, random=None):
        self.random = random or Random()
        self.varstacks = {}
        self.var_index = 0
        self.reset_tracking()
        self.log = []

    def reset_tracking(self):
        self.values_read = []
        self.values_written = []

    def run_program(self, program):
        for operation in program:
            self.execute(operation)

    def execute(self, operation):
        self.reset_tracking()
        try:
            operation.invoke(self)
            self.log.append(ProgramStep(
                operation=operation,
                definitions=tuple(self.values_written),
                arguments=tuple(self.values_read)
            ))
        except Exception:
            self.log.append(ProgramStep(
                operation=operation,
                definitions=(),
                arguments=tuple(self.values_read)
            ))
            raise

    def __repr__(self):
        return "RunContext(%s)" % (
            ', '.join(
                "%s=%r" % (v.name, len(v.data))
                for v in self.varstacks.values()
            )
        )

    def newvar(self):
        self.var_index += 1
        return "t%d" % (self.var_index,)

    def on_read(self, var):
        self.values_read.append(var)

    def on_write(self, var):
        self.values_written.append(var)

    def varstack(self, name):
        try:
            return self.varstacks[name]
        except KeyError:
            varstack = VarStack(name, self)
            self.varstacks[name] = varstack
            return varstack

    def read(self, argspec):
        result = []
        seen = Counter()
        for a in argspec:
            result.append(self.varstack(a).peek(seen[a]))
            seen[a] += 1
        return tuple(result)


class TestMachine(object):
    def __init__(
        self,
        n_iters=500,
        prog_length=200,
        good_enough=10,
        print_output=True,
    ):
        self.languages = []
        self.n_iters = n_iters
        self.prog_length = prog_length
        self.good_enough = good_enough
        self.print_output = print_output

    def inform(self, message, *args, **kwargs):
        assert not (args and kwargs)
        if self.print_output:
            print (message % (args or kwargs))

    def operation(self, *args, **kwargs):
        """
        Add an operation which pops arguments from each of the varstacks named
        in args, passes the result in that order to function and pushes the
        result of the invocation onto target. If target is None the result is
        ignored.
        """
        self.add_language(ReadAndWrite(
            *args, **kwargs
        ))

    def check(self, *args, **kwargs):
        """
        Add an operation which reads from the varstacks in args in order,
        without popping their result and passes them in order to test. If test
        returns something truthy this operation passes, else it will fail.
        """
        self.add_language(Check(*args, **kwargs))

    def generate(self, produce, target, name=None):
        """
        Add a generator for operations which produces values by calling
        produce with a Random instance and pushes them onto target.
        """
        self.add_language(
            PushRandom(produce=produce, target=target, name=name)
        )

    def basic_operations(self, varstack):
        self.add_language(Dup(varstack))
        self.add_language(Drop(varstack))
        self.add_language(Swap(varstack))
        self.add_language(Rot(varstack))

    def ints(self, target):
        self.basic_operations(target)
        self.arithmetic_operations(target)
        self.generate(lambda r: r.randint(0, 10 ** 6), target)
        self.generate(lambda r: r.randint(-10, 10), target)

    def lists(self, source, target):
        self.basic_operations(target)
        self.generate(lambda r: [], target)
        self.operation(
            function=lambda x, y: x.append(y),
            argspec=(target, source),
            target=None,
            name="append",
            pattern="%s.append(%s)"
        )
        self.operation(
            function=lambda x: [x],
            argspec=(source,),
            target=target,
            name="singleton",
            pattern="[%s]",
        )
        self.operation(
            function=lambda x, y: [x, y],
            argspec=(source, source),
            target=target,
            name="pair",
            pattern="[%s, %s]",
        )
        self.operation(
            function=list,
            argspec=(target,),
            target=target
        )
        self.binary_operation(operator.add, target, "+")

    def arithmetic_operations(self, varstack):
        self.binary_operation(operator.add, varstack, "+")
        self.binary_operation(operator.sub, varstack, "-")
        self.binary_operation(operator.mul, varstack, "*")
        self.binary_operation(
            operator.div, varstack, "/",
            precondition=lambda x, y: y != 0
        )
        self.unary_operation(operator.neg, varstack, "-")

    def power_operation(self, varstack):
        self.binary_operation(
            operator.pow, varstack, "**",
            precondition=lambda x, y: (x > 0) or (math.floor(y) == y)
        )

    def add_language(self, language):
        self.languages.append(language)

    @property
    def language(self):
        return ChooseFrom(self.languages)

    def find_failing_program(
        self,
    ):<commit_msg>use a better test for non-zeroness<commit_after>from random import Random
from contextlib import contextmanager
from .operations import (
    ChooseFrom,
    ReadAndWrite,
    Check,
    PushRandom,
    BinaryOperator,
    UnaryOperator,
    Dup,
    Drop,
    Swap,
    Rot,
)
from collections import namedtuple, Counter
import operator
import traceback
import math


ProgramStep = namedtuple(
    "ProgramStep",
    ("definitions", "arguments", "operation")
)


class TestMachineError(Exception):
    pass


class FrozenVarStack(TestMachineError):
    def __init__(self):
        super(FrozenVarStack, self).__init__("Cannot modify frozen varstack")


class NoFailingProgram(TestMachineError):
    pass


class VarStack(object):
    def __init__(self, name, context):
        self.name = name
        self.context = context
        self.data = []
        self.names = []
        self.frozen = False

    def _integrity_check(self):
        assert len(self.data) == len(self.names)

    @contextmanager
    def freeze(self):
        self.frozen = True
        try:
            yield
        finally:
            self.frozen = False

    def modification(self):
        if self.frozen:
            raise FrozenVarStack()

    def pop(self):
        self.modification()
        self._integrity_check()
        result = self.data.pop()
        self.context.on_read(self.names.pop())
        return result

    def push(self, head):
        self.modification()
        self._integrity_check()
        self.data.append(head)
        v = self.context.newvar()
        self.names.append(v)
        self.context.on_write(v)

    def dup(self):
        self.modification()
        self._integrity_check()
        self.names.append(self.names[-1])
        self.data.append(self.data[-1])

    def peek(self, index=0):
        self._integrity_check()
        i = -1 - index
        self.context.on_read(self.names[i])
        return self.data[i]

    def has(self, count):
        self._integrity_check()
        return len(self.data) >= count


class RunContext(object):
    def __init__(self, random=None):
        self.random = random or Random()
        self.varstacks = {}
        self.var_index = 0
        self.reset_tracking()
        self.log = []

    def reset_tracking(self):
        self.values_read = []
        self.values_written = []

    def run_program(self, program):
        for operation in program:
            self.execute(operation)

    def execute(self, operation):
        self.reset_tracking()
        try:
            operation.invoke(self)
            self.log.append(ProgramStep(
                operation=operation,
                definitions=tuple(self.values_written),
                arguments=tuple(self.values_read)
            ))
        except Exception:
            self.log.append(ProgramStep(
                operation=operation,
                definitions=(),
                arguments=tuple(self.values_read)
            ))
            raise

    def __repr__(self):
        return "RunContext(%s)" % (
            ', '.join(
                "%s=%r" % (v.name, len(v.data))
                for v in self.varstacks.values()
            )
        )

    def newvar(self):
        self.var_index += 1
        return "t%d" % (self.var_index,)

    def on_read(self, var):
        self.values_read.append(var)

    def on_write(self, var):
        self.values_written.append(var)

    def varstack(self, name):
        try:
            return self.varstacks[name]
        except KeyError:
            varstack = VarStack(name, self)
            self.varstacks[name] = varstack
            return varstack

    def read(self, argspec):
        result = []
        seen = Counter()
        for a in argspec:
            result.append(self.varstack(a).peek(seen[a]))
            seen[a] += 1
        return tuple(result)


class TestMachine(object):
    def __init__(
        self,
        n_iters=500,
        prog_length=200,
        good_enough=10,
        print_output=True,
    ):
        self.languages = []
        self.n_iters = n_iters
        self.prog_length = prog_length
        self.good_enough = good_enough
        self.print_output = print_output

    def inform(self, message, *args, **kwargs):
        assert not (args and kwargs)
        if self.print_output:
            print (message % (args or kwargs))

    def operation(self, *args, **kwargs):
        """
        Add an operation which pops arguments from each of the varstacks named
        in args, passes the result in that order to function and pushes the
        result of the invocation onto target. If target is None the result is
        ignored.
        """
        self.add_language(ReadAndWrite(
            *args, **kwargs
        ))

    def check(self, *args, **kwargs):
        """
        Add an operation which reads from the varstacks in args in order,
        without popping their result and passes them in order to test. If test
        returns something truthy this operation passes, else it will fail.
        """
        self.add_language(Check(*args, **kwargs))

    def generate(self, produce, target, name=None):
        """
        Add a generator for operations which produces values by calling
        produce with a Random instance and pushes them onto target.
        """
        self.add_language(
            PushRandom(produce=produce, target=target, name=name)
        )

    def basic_operations(self, varstack):
        self.add_language(Dup(varstack))
        self.add_language(Drop(varstack))
        self.add_language(Swap(varstack))
        self.add_language(Rot(varstack))

    def ints(self, target):
        self.basic_operations(target)
        self.arithmetic_operations(target)
        self.generate(lambda r: r.randint(0, 10 ** 6), target)
        self.generate(lambda r: r.randint(-10, 10), target)

    def lists(self, source, target):
        self.basic_operations(target)
        self.generate(lambda r: [], target)
        self.operation(
            function=lambda x, y: x.append(y),
            argspec=(target, source),
            target=None,
            name="append",
            pattern="%s.append(%s)"
        )
        self.operation(
            function=lambda x: [x],
            argspec=(source,),
            target=target,
            name="singleton",
            pattern="[%s]",
        )
        self.operation(
            function=lambda x, y: [x, y],
            argspec=(source, source),
            target=target,
            name="pair",
            pattern="[%s, %s]",
        )
        self.operation(
            function=list,
            argspec=(target,),
            target=target
        )
        self.binary_operation(operator.add, target, "+")

    def arithmetic_operations(self, varstack):
        self.binary_operation(operator.add, varstack, "+")
        self.binary_operation(operator.sub, varstack, "-")
        self.binary_operation(operator.mul, varstack, "*")
        self.binary_operation(
            operator.div, varstack, "/",
            precondition=lambda x, y: bool(y)
        )
        self.unary_operation(operator.neg, varstack, "-")

    def power_operation(self, varstack):
        self.binary_operation(
            operator.pow, varstack, "**",
            precondition=lambda x, y: (x > 0) or (math.floor(y) == y)
        )

    def add_language(self, language):
        self.languages.append(language)

    @property
    def language(self):
        return ChooseFrom(self.languages)

    def find_failing_program(
        self,
    ):
{'fraction_non_alphanumeric': 0.1069961574135418, 'fraction_numerical': 0.003246323042268451, 'mean_word_length': 1.974384236453202, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3071733', 'n_tokens_mistral': 4400, 'n_tokens_neox': 4244, 'n_words': 990}