arxpy/differential/characteristic.py
class RelatedKeyCh(object):
    """Represent related-key characteristics of block ciphers.

    A related-key characteristic of a `Cipher` is a pair `BvCharacteristic`,
    one over the `KeySchedule` of the cipher, and another one over the
    the `Encryption` function of the cipher, where the output differences
    of the key schedule characteristic are used as round key differences
    in the encryption characteristic.

    The master key differences start with the prefix ``"dmk"``,
    the round key differences start with the prefix ``"dk"``,
    the plaintext differences start with the prefix ``"dp"``
    and the non-linear differences start with the prefix ``"dx"``.

        >>> from arxpy.bitvector.core import Variable
        >>> from arxpy.differential.difference import XorDiff
        >>> from arxpy.differential.characteristic import RelatedKeyCh
        >>> from arxpy.primitives.primitives import Cipher
        >>> from arxpy.primitives.lea import LeaCipher
        >>> issubclass(LeaCipher, Cipher)
        True
        >>> LeaCipher.set_rounds(1)
        >>> rkch = RelatedKeyCh(LeaCipher, XorDiff)
        >>> rkch .key_schedule_ch.ssa  # doctest: +NORMALIZE_WHITESPACE
        {'input_vars': (dmk0, dmk1, dmk2, dmk3),
        'output_vars': (dk1, dk3, dk5, dk3, dk7, dk3),
        'assignments': ((dk0, 0xc3efe9db + dmk0), (dk1, dk0 <<< 1), (dk2, 0x87dfd3b7 + dmk1), (dk3, dk2 <<< 3),
        (dk4, 0x0fbfa76f + dmk2), (dk5, dk4 <<< 6), (dk6, 0x1f7f4ede + dmk3), (dk7, dk6 <<< 11))}
        >>> rkch.key_schedule_ch.input_diff
        (XorDiff(dmk0), XorDiff(dmk1), XorDiff(dmk2), XorDiff(dmk3))
        >>> rkch.key_schedule_ch.output_diff  # doctest: +NORMALIZE_WHITESPACE
        [[XorDiff(dk1), XorDiff(dk0 <<< 1)], [XorDiff(dk3), XorDiff(dk2 <<< 3)],
        [XorDiff(dk5), XorDiff(dk4 <<< 6)], [XorDiff(dk3), XorDiff(dk2 <<< 3)],
        [XorDiff(dk7), XorDiff(dk6 <<< 11)], [XorDiff(dk3), XorDiff(dk2 <<< 3)]]
        >>> rkch.key_schedule_ch.nonlinear_diffs  # doctest: +NORMALIZE_WHITESPACE
        OrderedDict([(XorDiff(dk0), XDCA_0xc3efe9db(XorDiff(dmk0))),
        (XorDiff(dk2), XDCA_0x87dfd3b7(XorDiff(dmk1))),
        (XorDiff(dk4), XDCA_0x0fbfa76f(XorDiff(dmk2))),
        (XorDiff(dk6), XDCA_0x1f7f4ede(XorDiff(dmk3)))])
        >>> rkch.encryption_ch.ssa  # doctest: +NORMALIZE_WHITESPACE
        {'input_vars': (dp0, dp1, dp2, dp3),
        'output_vars': (dx3, dx7, dx11, dp0),
        'assignments': ((dx0, dk1 ^ dp0), (dx1, dk3 ^ dp1), (dx2, dx0 + dx1), (dx3, dx2 <<< 9),
        (dx4, dk5 ^ dp1), (dx5, dk3 ^ dp2), (dx6, dx4 + dx5), (dx7, dx6 >>> 5), (dx8, dk7 ^ dp2),
        (dx9, dk3 ^ dp3), (dx10, dx8 + dx9), (dx11, dx10 >>> 3))}
        >>> rkch.encryption_ch.input_diff
        (XorDiff(dp0), XorDiff(dp1), XorDiff(dp2), XorDiff(dp3))
        >>> rkch.encryption_ch.output_diff # doctest: +NORMALIZE_WHITESPACE
        [[XorDiff(dx3), XorDiff(dx2 <<< 9)], [XorDiff(dx7), XorDiff(dx6 >>> 5)],
        [XorDiff(dx11), XorDiff(dx10 >>> 3)], [XorDiff(dp0), XorDiff(dp0)]]
        >>> rkch.encryption_ch.nonlinear_diffs  # doctest: +NORMALIZE_WHITESPACE
        OrderedDict([(XorDiff(dx2), XDA(XorDiff(dp0 ^ (dk0 <<< 1)), XorDiff(dp1 ^ (dk2 <<< 3)))),
        (XorDiff(dx6), XDA(XorDiff(dp1 ^ (dk4 <<< 6)), XorDiff(dp2 ^ (dk2 <<< 3)))),
        (XorDiff(dx10), XDA(XorDiff(dp2 ^ (dk6 <<< 11)), XorDiff(dp3 ^ (dk2 <<< 3))))])

    Attributes:
        key_schedule_ch: the `BvCharacteristic` over the key schedule
        encryption_ch: the `BvCharacteristic` over the encryption function
    """

    def __init__(self, bv_cipher, diff_type):
        assert issubclass(bv_cipher, primitives.Cipher)
        assert issubclass(diff_type, difference.Difference)

        func = bv_cipher.key_schedule
        prefix = "dk"
        input_diff_names = tuple(["dmk" + str(i) for i in range(len(func.input_widths))])
        ks_ch = BvCharacteristic(func, diff_type, input_diff_names, prefix)

        class Encryption(bv_cipher.encryption):
            round_keys = ks_ch.ssa["output_vars"]

        func = Encryption
        prefix = "dx"
        input_diff_names = ["dp" + str(i) for i in range(len(func.input_widths))]
        round_key_diff = {}
        for var, diff in ks_ch.output_diff:
            round_key_diff[var.val] = diff
        encryption_ch = BvCharacteristic(func, diff_type, input_diff_names,
                                         prefix, round_key_diff)

        self.diff_type = diff_type
        self.key_schedule_ch = ks_ch
        self.encryption_ch = encryption_ch
        self._cipher = bv_cipher

    def empirical_weight(self, key_input_diff, key_output_diff, key_samples,
                         enc_input_diff, enc_output_diff, enc_samples, precision=1):
        """Return the empirical weight of a given differential for multiple keys.

        This method returns the differential probability weight for the
        key schedule characteristic (see `BvCharacteristic.empirical_weight`)
        and the `collections.Counter` storing the distribution of weights for the encryption
        characteristic (see `SingleKeyCh.empirical_weight`).

            >>> from arxpy.bitvector.core import Variable, Constant
            >>> from arxpy.differential.difference import XorDiff
            >>> from arxpy.differential.characteristic import RelatedKeyCh
            >>> from arxpy.primitives.lea import LeaCipher
            >>> LeaCipher.set_rounds(1)
            >>> rkch = RelatedKeyCh(LeaCipher, XorDiff)
            >>> zero, one = XorDiff(Constant(0, 32)), XorDiff(Constant(1, 32))
            >>> kid, kod = [zero]*4, [zero]*6
            >>> eid, eod = [zero]*4, [zero]*4
            >>> rkch.empirical_weight(kid, kod, 10, eid, eod, 100)
            (0.0, Counter({0.0: 10}))
            >>> kid, kod = [zero]*4, [one]*6
            >>> eid, eod = [zero]*4, [one]*4
            >>> rkch.empirical_weight(kid, kod, 10, eid, eod, 100)
            (inf, Counter({inf: 10}))

        """
        key_weight = self.key_schedule_ch.empirical_weight(key_input_diff, key_output_diff, key_samples)
        # noinspection PyProtectedMember
        enc_counter = self.encryption_ch._empirical_weight_distribution(self._cipher,
                                                                        enc_input_diff, enc_output_diff, enc_samples,
                                                                        key_samples, precision, key_output_diff)
        return key_weight, enc_counter

    def signature(self, ch_signature_type):
        """Return the signature of the related-key characteristic.

        The signature of a related-key characteristic is the
        concatenation of the key schedule and encryption signatures.

        See also `BvCharacteristic.signature`.

            >>> from arxpy.bitvector.core import Variable
            >>> from arxpy.differential.difference import XorDiff
            >>> from arxpy.differential.characteristic import RelatedKeyCh, ChSignatureType
            >>> from arxpy.primitives.primitives import Cipher
            >>> from arxpy.primitives.lea import LeaCipher
            >>> LeaCipher.set_rounds(1)
            >>> rkch = RelatedKeyCh(LeaCipher, XorDiff)
            >>> rkch.signature(ChSignatureType.Full)  # doctest:+NORMALIZE_WHITESPACE
            [dmk0, dmk1, dmk2, dmk3, dk0, dk2, dk4, dk6, dp0, dp1, dp2, dp3, dx2, dx6, dx10]
            >>> rkch.signature(ChSignatureType.InputOutput)  # doctest:+NORMALIZE_WHITESPACE
            [dmk0, dmk1, dmk2, dmk3, dk1, dk3, dk5, dk3, dk7, dk3, dp0, dp1, dp2, dp3, dx3, dx7, dx11]

        """
        return self.key_schedule_ch.signature(ch_signature_type) + self.encryption_ch.signature(ch_signature_type)

    def _to_dict(self):
        dict_ch = {
            "key_schedule_ch": self.key_schedule_ch._to_dict(),
            "encryption_ch": self.encryption_ch._to_dict(),
        }
        return dict_ch

    def __str__(self):
        return pprint.pformat(self._to_dict(), width=100, compact=True)
arxpy/differential/tests/test_characteristic.py
class Cipher1(Cipher):  # noqa: D101
    key_schedule = KeySchedule1
    encryption = Encryption1
    rounds = 1

    @classmethod
    def set_rounds(cls, new_rounds):
        assert new_rounds == 1
arxpy/differential/characteristic.py
class BvCharacteristic(object):
    """Represent characteristics of bit-vector functions.

    Given a bit-vector function `BvFunction` :math:`f`,
    a characteristic is a trail of differences obtained by
    propagating an input difference over :math:`f`.

    In particular, a characteristic is composed of the
    input difference and the output difference of each
    non-linear operation.

    This class manages symbolic characteristics,
    where the input difference is given symbolically
    and the intermediate differences are `Term`
    that depend on the input difference.

        >>> from arxpy.bitvector.core import Variable
        >>> from arxpy.differential.difference import XorDiff, RXDiff
        >>> from arxpy.differential.characteristic import BvCharacteristic
        >>> from arxpy.primitives.primitives import BvFunction
        >>> from arxpy.primitives.chaskey import ChaskeyPi
        >>> issubclass(ChaskeyPi, BvFunction)
        True
        >>> ChaskeyPi.set_rounds(1)
        >>> ch = BvCharacteristic(ChaskeyPi, XorDiff, ["dv0", "dv1", "dv2", "dv3"])
        >>> ch.ssa  # doctest: +NORMALIZE_WHITESPACE
        {'input_vars': (dv0, dv1, dv2, dv3),
        'output_vars': (d7, d12, d13, d9),
        'assignments': ((d0, dv0 + dv1), (d1, dv1 <<< 5), (d2, d0 ^ d1), (d3, d0 <<< 16), (d4, dv2 + dv3),
        (d5, dv3 <<< 8), (d6, d4 ^ d5), (d7, d3 + d6), (d8, d6 <<< 13), (d9, d7 ^ d8), (d10, d2 + d4),
        (d11, d2 <<< 7), (d12, d10 ^ d11), (d13, d10 <<< 16))}
        >>> ch.input_diff
        (XorDiff(dv0), XorDiff(dv1), XorDiff(dv2), XorDiff(dv3))
        >>> ch.nonlinear_diffs # doctest: +NORMALIZE_WHITESPACE
        OrderedDict([(XorDiff(d0), XDA(XorDiff(dv0), XorDiff(dv1))),
        (XorDiff(d4), XDA(XorDiff(dv2), XorDiff(dv3))),
        (XorDiff(d7), XDA(XorDiff(d0 <<< 16), XorDiff(d4 ^ (dv3 <<< 8)))),
        (XorDiff(d10), XDA(XorDiff(d0 ^ (dv1 <<< 5)), XorDiff(d4)))])
        >>> ch.output_diff # doctest: +NORMALIZE_WHITESPACE
        [[XorDiff(d7), XorDiff(d7)],
        [XorDiff(d12), XorDiff(d10 ^ ((d0 ^ (dv1 <<< 5)) <<< 7))],
        [XorDiff(d13), XorDiff(d10 <<< 16)],
        [XorDiff(d9), XorDiff(d7 ^ ((d4 ^ (dv3 <<< 8)) <<< 13))]]
        >>> ch = BvCharacteristic(ChaskeyPi, RXDiff, ["dv0", "dv1", "dv2", "dv3"])
        >>> ch.input_diff
        (RXDiff(dv0), RXDiff(dv1), RXDiff(dv2), RXDiff(dv3))
        >>> ch.nonlinear_diffs # doctest: +NORMALIZE_WHITESPACE
        OrderedDict([(RXDiff(d0), RXDA(RXDiff(dv0), RXDiff(dv1))),
        (RXDiff(d4), RXDA(RXDiff(dv2), RXDiff(dv3))),
        (RXDiff(d7), RXDA(RXDiff(d0 <<< 16), RXDiff(d4 ^ (dv3 <<< 8)))),
        (RXDiff(d10), RXDA(RXDiff(d0 ^ (dv1 <<< 5)), RXDiff(d4)))])
        >>> ch.output_diff # doctest: +NORMALIZE_WHITESPACE
        [[RXDiff(d7), RXDiff(d7)],
        [RXDiff(d12), RXDiff(d10 ^ ((d0 ^ (dv1 <<< 5)) <<< 7))],
        [RXDiff(d13), RXDiff(d10 <<< 16)],
        [RXDiff(d9), RXDiff(d7 ^ ((d4 ^ (dv3 <<< 8)) <<< 13))]]

    Attributes:
        func: the `BvFunction`
        diff_type: the `Difference` of the characteristic
        input_diff: a list containing the input symbolic differences
        nonlinear_diffs: an `collections.OrderedDict` mapping non-linear symbolic differences
            to their corresponding `Derivative`
        output_diff: a list, where the i-th element is a pair containing
            the i-th output symbolic difference and its value
    """

    def __init__(self, func, diff_type, input_diff_names, prefix="d", initial_var2diff=None):
        assert issubclass(func, primitives.BvFunction)
        assert issubclass(diff_type, difference.Difference)

        assert len(input_diff_names) == len(func.input_widths)
        input_diff = []
        for name, width in zip(input_diff_names, func.input_widths):
            input_diff.append(diff_type(core.Variable(name, width)))
        input_diff = tuple(input_diff)

        self.func = func
        self.diff_type = diff_type
        self.input_diff = input_diff

        # Propagate the input difference through the function

        names = [d.val.name for d in self.input_diff]
        ssa = self.func.ssa(names, id_prefix=prefix)
        self.ssa = ssa
        self._prefix = prefix
        self._input_diff_names = input_diff_names

        for var in ssa["output_vars"]:
            if isinstance(var, core.Constant):
                raise ValueError("constant outputs (independent of the inputs) are not supported")

        var2diff = {}  # Variable to Difference
        for var, diff in zip(ssa["input_vars"], self.input_diff):
            var2diff[var] = diff

        if initial_var2diff is not None:
            for var in initial_var2diff:
                if str(var) in names:
                    raise ValueError("the input differences cannot be replaced by initial_var2diff")
            var2diff.update(initial_var2diff)

        self.nonlinear_diffs = collections.OrderedDict()
        for var, expr in ssa["assignments"]:
            expr_args = []
            for arg in expr.args:
                if isinstance(arg, int):
                    expr_args.append(arg)  # 'int' object has no attribute 'xreplace'
                else:
                    expr_args.append(arg.xreplace(var2diff))

            if all(not isinstance(arg, diff_type) for arg in expr_args):
                # symbolic computations with the key
                var2diff[var] = expr
                continue

            if all(isinstance(arg, diff_type) for arg in expr_args):
                der = self.diff_type.derivative(type(expr), expr_args)
            else:
                def contains_key_var(term):
                    from sympy import basic
                    for sub in basic.preorder_traversal(term):
                        if sub in func.round_keys:
                            return True
                    else:
                        return False

                if type(expr) == operation.BvAdd and hasattr(func, 'round_keys') and \
                        all(isinstance(r, core.Variable) for r in func.round_keys) and \
                        any(contains_key_var(a) for a in expr_args):
                    # temporary solution to Derivative(BvAddCte_k(x)) != Derivative(x + k)
                    # with x a Diff and k a key variable
                    keyed_indices = []
                    for i, a in enumerate(expr_args):
                        if contains_key_var(a):
                            keyed_indices.append(i)
                    if len(keyed_indices) != 1 or expr_args[keyed_indices[0]] not in func.round_keys:
                        raise NotImplementedError("invalid expression: op={}, args={}".format(
                            type(expr).__name__, expr_args))
                    # expr_args[keyed_indices[0]] replaced to the zero diff
                    zero_diff = diff_type(core.Constant(0, expr_args[keyed_indices[0]].width))
                    der = self.diff_type.derivative(type(expr), [expr_args[(keyed_indices[0] + 1) % 2], zero_diff])
                elif hasattr(expr, "xor_derivative"):
                    # temporary solution to operations containing a custom derivative
                    input_diff_expr = []
                    for i, arg in enumerate(expr_args):
                        if isinstance(arg, diff_type):
                            input_diff_expr.append(arg)
                        else:
                            assert isinstance(arg, core.Term)  # int arguments currently not supported
                            input_diff_expr.append(diff_type.from_pair(arg, arg))
                    der = self.diff_type.derivative(type(expr), input_diff_expr)
                else:
                    fixed_args = []
                    for i, arg in enumerate(expr_args):
                        if not isinstance(arg, diff_type):
                            fixed_args.append(arg)
                        else:
                            fixed_args.append(None)
                    new_op = extraop.make_partial_operation(type(expr), tuple(fixed_args))
                    der = self.diff_type.derivative(new_op, [arg for arg in expr_args if isinstance(arg, diff_type)])

            if isinstance(der, derivative.Derivative):
                diff = self.diff_type(var)
                var2diff[var] = diff
                self.nonlinear_diffs[diff] = der
            else:
                var2diff[var] = der

        self._var2diff = var2diff

        self.output_diff = []
        for var in ssa["output_vars"]:
            self.output_diff.append([self.diff_type(var), var2diff[var]])

    def empirical_weight(self, input_diff, output_diff, pair_samples):
        """Return the empirical weight of a given differential.

        Given a differential (a pair of input and output differences),
        the differential probability is the fraction of input pairs
        with the given input difference leading to output pairs
        with the given output difference.

        This method returns an approximation of the weight of the
        differential probability by sampling a given number
        of input pairs.

        If no correct output pairs are found, `math.inf` is returned.

            >>> from arxpy.bitvector.core import Constant
            >>> from arxpy.differential.difference import XorDiff, RXDiff
            >>> from arxpy.differential.characteristic import BvCharacteristic
            >>> from arxpy.primitives.chaskey import ChaskeyPi
            >>> ChaskeyPi.set_rounds(1)
            >>> ch = BvCharacteristic(ChaskeyPi, XorDiff, ["dv" + str(i) for i in range(4)])
            >>> zero, one = XorDiff(Constant(0, 32)), XorDiff(Constant(1, 32))
            >>> ch.empirical_weight([zero, zero, zero, zero], [zero, zero, zero, zero], 100)
            0.0
            >>> ch.empirical_weight([zero, zero, zero, zero], [one, one, one, one], 100)
            inf
            >>> ch = BvCharacteristic(ChaskeyPi, RXDiff, ["dv" + str(i) for i in range(4)])
            >>> zero, one = RXDiff(Constant(0, 32)), RXDiff(Constant(1, 32))
            >>> 4 - 1 <= ch.empirical_weight([zero]*4, [zero]*4, 3 * 2**6) <= 8
            True
            >>> ch.empirical_weight([zero]*4, [one]*4, 3 * 2**6)
            inf

        """
        assert isinstance(input_diff, collections.abc.Sequence)
        assert isinstance(output_diff, collections.abc.Sequence)
        assert all(isinstance(d, difference.Difference) for d in input_diff)
        assert all(isinstance(d, difference.Difference) for d in output_diff)
        assert all(isinstance(d.val, core.Constant) for d in input_diff)
        assert all(isinstance(d.val, core.Constant) for d in output_diff)

        assert len(input_diff) == len(self.input_diff)
        assert len(output_diff) == len(self.output_diff)
        assert len(self.ssa["input_vars"]) == len(input_diff)
        assert len(self.ssa["output_vars"]) == len(output_diff)

        with context.Simplification(False):
            input_widths = [d.val.width for d in self.input_diff]
            if pair_samples >= 2**sum(input_widths):
                iterators = [range(2 ** w) for w in input_widths]
                list_pairs = []
                for x in itertools.product(*iterators):
                    pt = [core.Constant(x_i, w) for x_i, w in zip(x, input_widths)]
                    other_pt = [diff.get_pair_element(pt[i]) for i, diff in enumerate(input_diff)]
                    list_pairs.append([pt, other_pt])
                pair_samples = len(list_pairs)
                assert pair_samples == 2**sum(input_widths)
            else:
                list_pairs = []
                for _ in range(pair_samples):
                    pt = []
                    other_pt = []
                    for diff in input_diff:
                        random_int = random.randrange(2 ** diff.val.width)
                        random_bv = core.Constant(random_int, diff.val.width)
                        pt.append(random_bv)
                        other_pt.append(diff.get_pair_element(random_bv))
                    list_pairs.append([pt, other_pt])

            correct_pairs = 0

            for index_input in range(pair_samples):
                pt, other_pt = list_pairs[index_input]
                ct = self.func(*pt)
                other_ct = self.func(*other_pt)

                assert all(isinstance(x, core.Constant) for x in ct), str(ct)
                assert all(isinstance(x, core.Constant) for x in other_ct), str(other_ct)

                for i, diff in enumerate(output_diff):
                    # noinspection PyUnresolvedReferences
                    if self.diff_type.from_pair(ct[i], other_ct[i]) != diff:
                        break
                else:
                    correct_pairs += 1

            if correct_pairs == 0:
                weight = math.inf
            else:
                weight = abs(-math.log(correct_pairs * 1.0 / pair_samples, 2))

        return weight

    def _empirical_weight_distribution(self, cipher, input_diff, output_diff, pair_samples, key_samples,
                                       precision=1, rk_output_diff=None):
        # this function is not part of SingleKeyCh since it must be accessible
        # for the encryption characteristic of RelatedKeyCh (which is a
        # plain BvCharacteristic)
        assert isinstance(input_diff, collections.abc.Sequence)
        assert isinstance(output_diff, collections.abc.Sequence)
        assert all(isinstance(d, difference.Difference) for d in input_diff)
        assert all(isinstance(d, difference.Difference) for d in output_diff)
        assert all(isinstance(d.val, core.Constant) for d in input_diff)
        assert all(isinstance(d.val, core.Constant) for d in output_diff)

        assert len(input_diff) == len(self.input_diff)
        assert len(output_diff) == len(self.output_diff)
        assert len(self.ssa["input_vars"]) == len(input_diff)
        assert len(self.ssa["output_vars"]) == len(output_diff)

        old_round_keys = self.func.round_keys

        empirical_weights = collections.Counter()

        if rk_output_diff is not None:
            class RelatedFunc(self.func):
                pass
        else:
            RelatedFunc = self.func

        with context.Simplification(False):
            input_widths = [d.val.width for d in self.input_diff]
            if pair_samples >= 2**sum(input_widths):
                iterators = [range(2 ** w) for w in input_widths]
                list_pairs = []
                for x in itertools.product(*iterators):
                    pt = [core.Constant(x_i, w) for x_i, w in zip(x, input_widths)]
                    other_pt = [diff.get_pair_element(pt[i]) for i, diff in enumerate(input_diff)]
                    list_pairs.append([pt, other_pt])
                pair_samples = len(list_pairs)
                assert pair_samples == 2**sum(input_widths)
            else:
                list_pairs = []
                for _ in range(pair_samples):
                    pt = []
                    other_pt = []
                    for diff in input_diff:
                        random_int = random.randrange(2 ** diff.val.width)
                        random_bv = core.Constant(random_int, diff.val.width)
                        pt.append(random_bv)
                        other_pt.append(diff.get_pair_element(random_bv))
                    list_pairs.append([pt, other_pt])

            for _ in range(key_samples):
                master_key = []
                for width in cipher.key_schedule.input_widths:
                    master_key.append(core.Constant(random.randrange(2 ** width), width))
                self.func.round_keys = cipher.key_schedule(*master_key)
                assert all(isinstance(rk, core.Constant) for rk in self.func.round_keys), str(self.func.round_keys)

                if rk_output_diff is not None:
                    RelatedFunc.round_keys = [d.get_pair_element(r) for r, d in zip(self.func.round_keys, rk_output_diff)]
                    assert all(isinstance(rk, core.Constant) for rk in RelatedFunc.round_keys), str(RelatedFunc.round_keys)

                correct_pairs = 0

                for index_input in range(pair_samples):
                    pt, other_pt = list_pairs[index_input]
                    ct = self.func(*pt)
                    other_ct = RelatedFunc(*other_pt)

                    assert all(isinstance(x, core.Constant) for x in ct), str(ct)
                    assert all(isinstance(x, core.Constant) for x in other_ct), str(other_ct)

                    for i, diff in enumerate(output_diff):
                        # noinspection PyUnresolvedReferences
                        if self.diff_type.from_pair(ct[i], other_ct[i]) != diff:
                            break
                    else:
                        correct_pairs += 1

                if correct_pairs == 0:
                    weight = math.inf
                else:
                    weight = abs(-math.log(correct_pairs * 1.0 / pair_samples, 2))
                # weight = float(("{0:."+str(precision)+"f}").format(weight))
                weight = round(weight, precision)
                empirical_weights[weight] += 1

        self.func.round_keys = old_round_keys

        return empirical_weights

    def signature(self, ch_signature_type):
        """Return the signature of the characteristic.

        The signature is a "hash" of the characteristic used for comparing.

        For the type of the signature, see `ChSignatureType`.

            >>> from arxpy.bitvector.core import Variable
            >>> from arxpy.differential.difference import XorDiff
            >>> from arxpy.differential.characteristic import BvCharacteristic, ChSignatureType
            >>> from arxpy.primitives.primitives import BvFunction
            >>> from arxpy.primitives.chaskey import ChaskeyPi
            >>> issubclass(ChaskeyPi, BvFunction)
            True
            >>> ChaskeyPi.set_rounds(1)
            >>> ch = BvCharacteristic(ChaskeyPi, XorDiff, ["dv0", "dv1", "dv2", "dv3"])
            >>> ch.signature(ChSignatureType.Full)
            [dv0, dv1, dv2, dv3, d0, d4, d7, d10]
            >>> ch.signature(ChSignatureType.InputOutput)
            [dv0, dv1, dv2, dv3, d7, d12, d13, d9]

        """
        if ch_signature_type == ChSignatureType.Full:
            return [d.val for d in self.input_diff] + [d.val for d in self.nonlinear_diffs]
        elif ch_signature_type == ChSignatureType.InputOutput:
            # sig = [d for d in self.input_diff]
            sig_var = [d.val for d in self.input_diff]

            for out_diff, _ in self.output_diff:
                for aux_var in self._var2diff[out_diff.val].val.atoms(core.Variable):
                    if aux_var not in sig_var:
                        # sig.append(out_diff)
                        sig_var.append(out_diff.val)
                        break

            return sig_var
        else:
            raise ValueError("invalid ch_signature_type: {}".format(ch_signature_type))

    def _to_dict(self):
        dict_ch = {
            "ssa": self.ssa,
            "input_diff": self.input_diff,
            "output_diff": self.output_diff,
            "nonlinear_diffs": self.nonlinear_diffs,
        }
        return dict_ch

    def __str__(self):
        return pprint.pformat(self._to_dict(), width=100, compact=True)
arxpy/differential/difference.py
class XorDiff(Difference):
    """Represent XOR differences.

    The XOR difference of two `Term` is given by the XOR
    of the terms. In other words, the *difference operation*
    of `XorDiff` is the `BvXor` (see `Difference`).

        >>> from arxpy.bitvector.core import Constant, Variable
        >>> from arxpy.differential.difference import XorDiff
        >>> x, y = Constant(0b000, 3), Constant(0b000, 3)
        >>> alpha = XorDiff.from_pair(x, y)
        >>> alpha
        XorDiff(0b000)
        >>> alpha.get_pair_element(x)
        0b000
        >>> x, y = Constant(0b010, 3), Constant(0b101, 3)
        >>> alpha = XorDiff.from_pair(x, y)
        >>> alpha
        XorDiff(0b111)
        >>> alpha.get_pair_element(x)
        0b101
        >>> k = Variable("k", 8)
        >>> alpha = XorDiff.from_pair(k, k)
        >>> alpha
        XorDiff(0x00)
        >>> alpha.get_pair_element(k)
        k
    """

    diff_op = operation.BvXor
    inv_diff_op = operation.BvXor

    @classmethod
    def derivative(cls, op, input_diff):
        """Return the derivative of ``op`` at the point ``input_diff``.

        See `Difference.derivative` for more information.

            >>> from arxpy.bitvector.core import Variable, Constant
            >>> from arxpy.bitvector.operation import BvAdd, BvXor, RotateLeft, BvSub
            >>> from arxpy.bitvector.extraop import make_partial_operation
            >>> from arxpy.differential.difference import XorDiff
            >>> d1, d2 = XorDiff(Variable("d1", 8)), XorDiff(Variable("d2", 8))
            >>> XorDiff.derivative(BvXor, [d1, d2])
            XorDiff(d1 ^ d2)
            >>> Xor1 = make_partial_operation(BvXor, tuple([None, Constant(1, 8)]))
            >>> XorDiff.derivative(Xor1, d1)
            XorDiff(d1)
            >>> Rotate1 = make_partial_operation(RotateLeft, tuple([None, 1]))
            >>> XorDiff.derivative(Rotate1, d1)
            XorDiff(d1 <<< 1)
            >>> XorDiff.derivative(BvAdd, [d1, d2])
            XDA(XorDiff(d1), XorDiff(d2))
            >>> XorDiff.derivative(BvSub, [d1, d2])
            XDS(XorDiff(d1), XorDiff(d2))
            >>> CteAdd1 = make_partial_operation(BvAdd, tuple([None, Constant(1, 8)]))
            >>> XorDiff.derivative(CteAdd1, d1)
            XDCA_0x01(XorDiff(d1))

        """
        input_diff = _tuplify(input_diff)
        assert len(input_diff) == sum(op.arity)

        msg = "invalid arguments: op={}, input_diff={}".format(
            op.__name__,
            [d.vrepr() if isinstance(d, core.Term) else d for d in input_diff])

        if not all(isinstance(diff, cls) for diff in input_diff):
            raise ValueError(msg)

        if op == operation.BvNot:
            return input_diff[0]

        if op == operation.BvXor:
            return cls(op(*[d.val for d in input_diff]))

        if op == operation.Concat:
            return cls(op(*[d.val for d in input_diff]))

        if op == operation.BvAdd:
            from arxpy.differential import derivative
            return derivative.XDA(input_diff)

        if op == operation.BvSub:
            from arxpy.differential import derivative
            return derivative.XDS(input_diff)

        if issubclass(op, extraop.PartialOperation):
            if op.base_op == operation.BvXor:
                assert len(input_diff) == 1
                d1 = input_diff[0]
                val = op.fixed_args[0] if op.fixed_args[0] is not None else op.fixed_args[1]
                d2 = cls.from_pair(val, val)
                input_diff = [d1, d2]
                return cls(op.base_op(*[d.val for d in input_diff]))

            if op.base_op == operation.BvAnd:
                assert len(input_diff) == 1
                d1 = input_diff[0]
                val = op.fixed_args[0] if op.fixed_args[0] is not None else op.fixed_args[1]
                if isinstance(val, core.Constant):
                    return cls(op.base_op(d1.val, val))

            if op.base_op in [operation.RotateLeft, operation.RotateRight]:
                if op.fixed_args[0] is None and op.fixed_args[1] is not None:
                    assert len(input_diff) == 1
                    d = input_diff[0]
                    return cls(op.base_op(d.val, op.fixed_args[1]))
                else:
                    raise ValueError(msg)

            if op.base_op in [operation.BvShl, operation.BvLshr]:
                if op.fixed_args[0] is None and op.fixed_args[1] is not None:
                    assert len(input_diff) == 1
                    d = input_diff[0]
                    return cls(op.base_op(d.val, op.fixed_args[1]))
                else:
                    raise ValueError(msg)

            if op.base_op == operation.Extract:
                if op.fixed_args[0] is None and op.fixed_args[1] is not None and op.fixed_args[2] is not None:
                    assert len(input_diff) == 1
                    d = input_diff[0]
                    return cls(op.base_op(d.val, op.fixed_args[1], op.fixed_args[2]))
                else:
                    raise ValueError(msg)

            if op.base_op == operation.Concat:
                assert len(input_diff) == 1
                d1 = input_diff[0]
                if op.fixed_args[0] is not None:
                    val = op.fixed_args[0]
                    input_diff = [cls.from_pair(val, val), d1]
                else:
                    val = op.fixed_args[1]
                    input_diff = [d1, cls.from_pair(val, val)]
                return cls(op.base_op(*[d.val for d in input_diff]))

            if op.base_op == operation.BvAdd:
                assert len(input_diff) == 1
                d = input_diff[0]
                cte = op.fixed_args[0] if op.fixed_args[0] is not None else op.fixed_args[1]
                from arxpy.differential import derivative
                return derivative.XDCA(d, cte)
            else:
                raise ValueError(msg)

        if hasattr(op, "xor_derivative"):
            return op.xor_derivative(input_diff)

        raise ValueError(msg)
arxpy/differential/difference.py
class RXDiff(Difference):
    """Represent rotational-XOR (RX) differences.

    The pair ``(x, (x <<< 1) ^ d)`` has RX difference ``d``.
    In other words,  the RX difference of two `Term` ``x`` and ``y``
    is defined as ``(x <<< 1) ^ y``.

    See `Difference` for more information.

        >>> from arxpy.bitvector.core import Constant, Variable
        >>> from arxpy.differential.difference import RXDiff
        >>> x, y = Constant(0b000, 3), Constant(0b000, 3)
        >>> alpha = RXDiff.from_pair(x, y)
        >>> alpha
        RXDiff(0b000)
        >>> alpha.get_pair_element(x)
        0b000
        >>> x, y = Constant(0b000, 3), Constant(0b001, 3)
        >>> alpha = RXDiff.from_pair(x, y)
        >>> alpha
        RXDiff(0b001)
        >>> alpha.get_pair_element(x)
        0b001
        >>> k = Variable("k", 8)
        >>> alpha = RXDiff.from_pair(k, k)
        >>> alpha
        RXDiff(k ^ (k <<< 1))
        >>> alpha.get_pair_element(k)
        k
    """

    diff_op = RXOp
    inv_diff_op = RXInvOp

    @classmethod
    def derivative(cls, op, input_diff):
        """Return the derivative of ``op`` at the point ``input_diff``.

        See `Difference.derivative` for more information.

            >>> from arxpy.bitvector.core import Variable, Constant
            >>> from arxpy.bitvector.operation import BvAdd, BvXor, RotateLeft
            >>> from arxpy.bitvector.extraop import make_partial_operation
            >>> from arxpy.differential.difference import RXDiff
            >>> d1, d2 = RXDiff(Variable("d1", 8)), RXDiff(Variable("d2", 8))
            >>> RXDiff.derivative(BvXor, [d1, d2])
            RXDiff(d1 ^ d2)
            >>> Xor1 = make_partial_operation(BvXor, tuple([None, Constant(1, 8)]))
            >>> RXDiff.derivative(Xor1, d1)
            RXDiff(0x03 ^ d1)
            >>> Rotate1 = make_partial_operation(RotateLeft, tuple([None, 1]))
            >>> RXDiff.derivative(Rotate1, d1)
            RXDiff(d1 <<< 1)
            >>> RXDiff.derivative(BvAdd, [d1, d2])
            RXDA(RXDiff(d1), RXDiff(d2))

        """
        input_diff = _tuplify(input_diff)
        assert len(input_diff) == sum(op.arity)

        msg = "invalid arguments: op={}, input_diff={}".format(
            op.__name__,
            [d.vrepr() if isinstance(d, core.Term) else d for d in input_diff])

        if not all(isinstance(diff, cls) for diff in input_diff):
            raise ValueError(msg)

        if op == operation.BvNot:
            return input_diff[0]

        if op == operation.BvXor:
            return cls(op(*[d.val for d in input_diff]))

        if op == operation.BvAdd:
            from arxpy.differential import derivative
            return derivative.RXDA(input_diff)

        # Concact, BvSub

        if issubclass(op, extraop.PartialOperation):
            if op.base_op == operation.BvXor:
                assert len(input_diff) == 1
                d1 = input_diff[0]
                val = op.fixed_args[0] if op.fixed_args[0] is not None else op.fixed_args[1]
                d2 = cls.from_pair(val, val)
                input_diff = [d1, d2]
                return cls(op.base_op(*[d.val for d in input_diff]))

            if op.base_op in [operation.RotateLeft, operation.RotateRight]:
                if op.fixed_args[0] is None and op.fixed_args[1] is not None:
                    assert len(input_diff) == 1
                    d = input_diff[0]
                    return cls(op.base_op(d.val, op.fixed_args[1]))
                else:
                    raise ValueError(msg)

            # BvShl, Extract

        if hasattr(op, "rx_derivative"):
            return op.rx_derivative(input_diff)

        raise ValueError(msg)
arxpy/primitives/primitives.py
class BvFunction(object):
    """Represent (iterated) fixed-width bit-vector functions.

    A `BvFunction` takes fixed-width `Constant` operands and return a
    tuple of fixed-width `Constant`. An iterated bit-vector function
    contains a subroutine that is iterated a certain number of *rounds*,
    which can be changed using `set_rounds`.

    Similar to `Operation`, `BvFunction` is evaluated
    using the operator ``()`` and provides *Automatic Constant Conversion*.
    Note that `BvFunction` only accepts `Constant` operands and
    always return a tuple, as opposed to `Operation` that accepts
    `Term` and scalar operands and returns a single `Term`.

        >>> from arxpy.primitives.primitives import BvFunction
        >>> from arxpy.primitives.chaskey import ChaskeyPi
        >>> issubclass(ChaskeyPi, BvFunction)
        True
        >>> ChaskeyPi(0, 0, 0, 0)  # automatic conversion from int to Constant
        (0x00000000, 0x00000000, 0x00000000, 0x00000000)

    Attributes:
        input_widths: a list containing the widths of the inputs
        output_widths: a list containing the widths of the outputs
        rounds: the number of iterations

    """
    input_widths = None
    output_widths = None
    rounds = None

    def __new__(cls, *args, **options):
        if len(cls.input_widths) != len(args):
            raise ValueError("{} requires {} inputs but {} were given: {}".format(
                cls.__name__, len(cls.input_widths), len(args), args))
        newargs = []
        for arg, width in zip(args, cls.input_widths):
            newargs.append(core.bitvectify(arg, width))
        args = newargs

        if all(isinstance(arg, core.Constant) for arg in args) or \
                options.pop("symbolic_inputs", False):
            result = cls.eval(*args)
        else:
            raise TypeError("expected bit-vector constant arguments")

        assert isinstance(result, collections.abc.Sequence)
        assert len(cls.output_widths) == len(result)

        output = []
        for r, width in zip(result, cls.output_widths):
            output.append(core.bitvectify(r, width))

        return tuple(output)

    @classmethod
    def eval(cls, *args):
        """Evaluate the function (internal method)."""
        raise NotImplementedError("subclasses need to override this method")

    @classmethod
    def set_rounds(cls, new_rounds):
        """Change the number of rounds and adjust the input/output widths."""
        raise NotImplementedError("subclasses need to override this method")

    # noinspection PyArgumentList
    @classmethod
    def ssa(cls, input_names, id_prefix):
        """Return a static single assignment program representing the function.

        Args:
            input_names: the names  for the input variables
            id_prefix: the prefix to denote the intermediate variables

        Return:
            : a dictionary with three keys

            - *input_vars*: a list of `Variable` representing the inputs
            - *output_vars*: a list of `Variable` representing the outputs
            - *assignments*: an ordered sequence of pairs
              (`Variable`, `Operation`) representing each assignment
              of the SSA program.

        ::

                >>> from arxpy.primitives.chaskey import ChaskeyPi
                >>> ChaskeyPi.set_rounds(1)
                >>> ChaskeyPi.ssa(["v0", "v1", "v2", "v3"], "x")  # doctest: +NORMALIZE_WHITESPACE
                {'input_vars': (v0, v1, v2, v3),
                'output_vars': (x7, x12, x13, x9),
                'assignments': ((x0, v0 + v1), (x1, v1 <<< 5), (x2, x0 ^ x1), (x3, x0 <<< 16), (x4, v2 + v3),
                (x5, v3 <<< 8), (x6, x4 ^ x5), (x7, x3 + x6), (x8, x6 <<< 13), (x9, x7 ^ x8), (x10, x2 + x4),
                (x11, x2 <<< 7), (x12, x10 ^ x11), (x13, x10 <<< 16))}

        """
        input_vars = []
        for name, width in zip(input_names, cls.input_widths):
            input_vars.append(core.Variable(name, width))
        input_vars = tuple(input_vars)

        table = context.MemoizationTable(id_prefix=id_prefix)

        with context.Memoization(table):
            # noinspection PyArgumentList
            output_vars = cls(*input_vars, symbolic_inputs=True)

        ssa_dict = {
            "input_vars": input_vars,
            "output_vars": output_vars,
            "assignments": tuple(table.items())
        }

        for var, expr in ssa_dict["assignments"]:
            for arg in expr.args:
                if isinstance(arg, operation.Operation):
                    raise ValueError("assignment {} <- {} was not decomposed".format(var, expr))

        to_delete = []
        vars_needed = set()
        for var in output_vars:
            vars_needed.add(var)
        for var, expr in reversed(ssa_dict["assignments"]):
            if var in vars_needed:
                for arg in expr.atoms(core.Variable):
                    vars_needed.add(arg)
            else:
                to_delete.append((var, expr))
                # raise ValueError("assignment {} <- {} is redundant in \n{}".format(var, expr, ssa_dict))

        if len(to_delete) > 0:
            import warnings
            warnings.warn("removing redundant assignments {} in \n{}".format(to_delete, ssa_dict))
            ssa_dict["assignments"] = list(ssa_dict["assignments"])
            for assignment in to_delete:
                ssa_dict["assignments"].remove(assignment)
            ssa_dict["assignments"] = tuple(ssa_dict["assignments"])

        return ssa_dict
arxpy/smt/search.py
class SearchSkCh(SearchCh):
    """Represent the problem of finding a `SingleKeyCh` with conditions on the weight.

    Args:
        skch (SingleKeyCh): a symbolic single-key characteristic of a `Cipher`
        der_mode (DerMode): one of the modes available for the derivative weights
        allow_zero_input_diff(bool): if ``True``, allow the input difference to be zero

    See also `SearchCh`.

        >>> from arxpy.differential.difference import XorDiff
        >>> from arxpy.differential.characteristic import SingleKeyCh
        >>> from arxpy.smt.search import SearchSkCh
        >>> from arxpy.primitives import speck
        >>> Speck32 = speck.get_Speck_instance(speck.SpeckInstance.speck_32_64)
        >>> Speck32.set_rounds(1)
        >>> ch = SingleKeyCh(Speck32, XorDiff)
        >>> search_problem = SearchSkCh(ch)
        >>> search_problem.formula_size()
        349
        >>> search_problem.error()
        0
        >>> print(search_problem.hrepr(False))
        assert ~(0x00000000 == (dp0 :: dp1))
        assert 0x0000 == ((~(... << ...) ^ (dx1 << 0x0001)) & (~(... << ...) ^ ((... >>> ...) << 0x0001)) & ((dp1 << 0x0001) ^ dx1 ^ dp1 ^ (... >>> ...)))
        assert w0 == PopCount(~((dx1 ^ ~...) & (~... ^ (... >>> ...)))[14:])
        assert w == w0
        minimize w
        >>> search_problem.solve(0).ch_weight
        0x0

    """

    def __init__(self, skch, der_mode=DerMode.Default, allow_zero_input_diff=False,
                 initial_constraints=None):
        assert isinstance(skch, characteristic.SingleKeyCh)
        super().__init__(skch, der_mode=der_mode, allow_zero_input_diff=allow_zero_input_diff,
                         initial_constraints=initial_constraints)


    def solve(self, initial_weight, solver_name="btor", search_mode=SkChSearchMode.Optimal, check=False,
              return_generator=False, verbose_level=0, filename=None):
        """Solve the SMT problem associated  to the search of a valid `SingleKeyCh`.

        See also `SearchCh.solve`.
        """
        assert initial_weight < sum(self.ch.func.input_widths) + math.floor(self.error())
        assert isinstance(search_mode, SkChSearchMode)

        return self._incremental_solve(initial_weight, solver_name, search_mode, check, return_generator,
                                       verbose_level, filename)
arxpy/smt/search.py
class DerMode(enum.Enum):
    """Represent the different constraints available for the derivative weights of a characteristic.

    Attributes:
        Default: no additional constraint is added
        ProbabilityOne: fix all derivative weights to zero
        Valid: only the validity constraint is added (weight is ignored)
        XDCA_Approx: for the weights of `XDCA`, the ``precision`` is set to ``0``

    """
    Default = enum.auto()
    ProbabilityOne = enum.auto()
    Valid = enum.auto()
    XDCA_Approx = enum.auto()
arxpy/smt/search.py
class RkChSearchMode(enum.Enum):
    """Represent the different options for searching `RelatedKeyCh`.

    A related-key characteristic is said to be *valid* if its encryption weight
    is lower than the bit-size of the plaintext difference and
    its key schedule weight is lower than the bit-size of the master-key difference.

    Attributes:
        FirstMinSum: returns the 1st related-key characteristic found
            by iteratively searching with increasing *weight*, where the *weight*
            is taken as the sum of the key schedule weight and the encryption weight.
            The characteristic found is returned as a `RkChFound` object,
            but it may not have optimal theoretical weight if it it contains
            a non-exact derivative weight or it may be empirically invalid.
            If no characteristic is found, ``None`` is returned.
        FirstMinSumValid: similar to `FirstMinSum`, but if the characteristic found is
            empirically invalid (after computing its empirical weight), the search continues.
        OptimalMinSum: similar to `FirstMinSum`, but several characteristics are found to ensure
            that the one returned has the optimal theoretical weight.
            (only differs from `FirstMinSum` if there is a non-exact derivative weight).
        OptimalMinSumDifferential: similar to `OptimalMinSum`, but once a characteristic is found,
            only characteristics with different input/output differences are searched.
            (to ensure optimality).
        FirstValidKeyMinEnc: similar to `FirstMinSum`,
            but the *weight* is taken as the encryption weight and the
            key schedule weight is restricted to be smaller than the key size.
        FirstValidKeyMinEncValid: similar to `FirstMinSumValid`,
            but the *weight* is defined as `FirstValidKeyMinEnc`.
        OptimalValidKeyMinEnc: similar to `OptimalMinSum`,
            but the *weight* is defined as `FirstValidKeyMinEnc`.
        OptimalValidKeyMinEncDifferential: similar to `OptimalMinSumDifferential`,
            but the *weight* is  defined as `FirstValidKeyMinEnc`.
        FirstFixEncMinKey: similar to `FirstMinSum`,
            but the *weight* is taken as the key schedule weight and
            the encryption weight is fixed to a given value.
        FirstFixEncMinKeyValid: similar to `FirstMinSumValid`,
            but the *weight* is defined as `FirstFixEncMinKey`.
        OptimalFixEncMinKey: similar to `OptimalMinSum`,
            but the *weight* is defined as `FirstFixEncMinKey`.
        OptimalFixEncMinKeyDifferential: similar to `OptimalMinSumDifferential`,
            but the *weight* is defined as `FirstFixEncMinKey`.
        AllOptimalMinSum: similar to `OptimalMinSum`, but the search never stops.
        AllValid: a non-stop search of valid characteristics with no
            additional constraints.

    """
    FirstMinSum = enum.auto()
    FirstMinSumValid = enum.auto()
    OptimalMinSum = enum.auto()
    OptimalMinSumDifferential = enum.auto()
    AllOptimalMinSum = enum.auto()
    FirstValidKeyMinEnc = enum.auto()
    FirstValidKeyMinEncValid = enum.auto()
    OptimalValidKeyMinEnc = enum.auto()
    OptimalValidKeyMinEncDifferential = enum.auto()
    FirstFixEncMinKey = enum.auto()
    FirstFixEncMinKeyValid = enum.auto()
    OptimalFixEncMinKey = enum.auto()
    OptimalFixEncMinKeyDifferential = enum.auto()
    AllValid = enum.auto()
arxpy/bitvector/operation.py
class BvComp(Operation):
    """Equality operator.

    Provides Automatic Constant Conversion. See `Operation` for more
    information.

        >>> from arxpy.bitvector.core import Constant, Variable
        >>> from arxpy.bitvector.operation import BvComp
        >>> BvComp(Constant(1, 8), Constant(2, 8))
        0b0
        >>> BvComp(Constant(1, 8), 2)
        0b0
        >>> BvComp(Constant(1, 8), Variable("y", 8))
        0x01 == y

    The operator == is used for exact structural equality testing and
    it returns either True or False. On the other hand, BvComp
    performs symbolic equality testing and it leaves the relation unevaluated
    if it cannot prove the objects are equal (or unequal).

        >>> Variable("x", 8) == Variable("y", 8)
        False
        >>> BvComp(Variable("x", 8), Variable("y", 8))  # symbolic equality
        x == y

    """

    arity = [2, 0]
    is_symmetric = True
    is_simple = True
    infix_symbol = "=="

    @classmethod
    def condition(cls, x, y):
        return x.width == y.width

    @classmethod
    def output_width(cls, x, y):
        return 1

    @classmethod
    def eval(cls, x, y):
        zero = core.Constant(0, 1)
        one = core.Constant(1, 1)

        if x is y:
            return one
        elif isinstance(x, core.Constant) and isinstance(y, core.Constant):
            return one if x.val == y.val else zero
arxpy/differential/tests/test_characteristic.py
class Cipher2(Cipher):  # noqa: D101
    key_schedule = KeySchedule2
    encryption = Encryption2
    rounds = 1

    @classmethod
    def set_rounds(cls, new_rounds):
        assert new_rounds == 1
arxpy/smt/search.py
class SearchRkCh(object):
    """Represent the problem of finding an `RelatedKeyCh` with conditions on the weight.

    Args:
        rkch (RelatedKeyCh): a symbolic related-key characteristic of a `Cipher`
        key_der_mode (DerMode): the derivative mode for the key schedule characteristic
        enc_der_mode (DerMode): the derivative mode for the encryption characteristic
        allow_zero_key_input_diff (bool): if ``True``, allow the input difference
            of the key schedule to be zero
        allow_zero_enc_input_diff (bool): if ``True``, allow the input difference
            of the encryption to be zero

    See also `SearchCh`.

        >>> from arxpy.differential.difference import XorDiff
        >>> from arxpy.differential.characteristic import RelatedKeyCh
        >>> from arxpy.smt.search import SearchRkCh
        >>> from arxpy.primitives.lea import LeaCipher
        >>> LeaCipher.set_rounds(1)
        >>> rkch = RelatedKeyCh(LeaCipher, XorDiff)
        >>> search_problem = SearchRkCh(rkch)
        >>> search_problem.formula_size()
        50495
        >>> search_problem.key_schedule_problem.error(), search_problem.encryption_problem.error()
        (4.199250014423123, 0)
        >>> print(search_problem.hrepr(False))
        assert ~(0x00000000000000000000000000000000 == (dmk0 :: dmk1 :: dmk2 :: dmk3))
        assert 0x00000000 == (((... << ...) & (... << ...) & ~(... ^ ...) & (~(... << ...) | (~... ^ ... ^ ...))) | (~(... << ...) & ~(... << ...) & (dk0 ^ dmk0)))
        assert wk0_tmp_0lz == LeadingZeros(~((~... & ~...) << 0x00000001))
        assert wk0_tmp_1rev == Reverse((~dk0 & ~dmk0) << 0x00000001)
        assert wk0_tmp_2rev == Reverse(~(... >> ...) & ((... & ...) << 0x00000001) & (~(... ^ ...) >> 0x00000001))
        assert wk0_tmp_3rev == Reverse(wk0_tmp_2rev ^ wk0_tmp_1rev ^ (wk0_tmp_1rev + wk0_tmp_2rev))
        assert wk0_tmp_4rev == Reverse((((... & ...) + (... & ...)) & ~(... | ...)) | ((... & ... & (... >> ...)) - ((... + ...) & (... | ...))))
        assert wk0 == (((PopCountDiff((... & ...) | (... << ...), (... & ...) ^ ... ^ ...)) :: 0b000) - (0b00 :: ((0b0 :: PopCount(...)) + (PopCount(...) :: 0b0))))
        assert 0x00000000 == (((... << ...) & (... << ...) & ~(... ^ ...) & (~(... << ...) | (~... ^ ... ^ ...))) | (~(... << ...) & ~(... << ...) & (dk2 ^ dmk1)))
        assert wk1_tmp_0lz == LeadingZeros(~((~... & ~...) << 0x00000001))
        assert wk1_tmp_1rev == Reverse((~dk2 & ~dmk1) << 0x00000001)
        assert wk1_tmp_2rev == Reverse(~(... >> ...) & ((... & ...) << 0x00000001) & (~(... ^ ...) >> 0x00000001))
        assert wk1_tmp_3rev == Reverse(wk1_tmp_2rev ^ wk1_tmp_1rev ^ (wk1_tmp_1rev + wk1_tmp_2rev))
        assert wk1_tmp_4rev == Reverse((((... & ...) + (... & ...)) & ~(... | ...)) | ((... & ... & (... >> ...)) - ((... + ...) & (... | ...))))
        assert wk1 == (((PopCountDiff((... & ...) | (... << ...), (... & ...) ^ ... ^ ...)) :: 0b000) - (0b00 :: ((0b0 :: PopCount(...)) + (PopCount(...) :: 0b0))))
        assert 0x00000000 == (((... << ...) & (... << ...) & ~(... ^ ...) & (~(... << ...) | (~... ^ ... ^ ...))) | (~(... << ...) & ~(... << ...) & (dk4 ^ dmk2)))
        assert wk2_tmp_0lz == LeadingZeros(~((~... & ~...) << 0x00000001))
        assert wk2_tmp_1rev == Reverse((~dk4 & ~dmk2) << 0x00000001)
        assert wk2_tmp_2rev == Reverse(~(... >> ...) & ((... & ...) << 0x00000001) & (~(... ^ ...) >> 0x00000001))
        assert wk2_tmp_3rev == Reverse(wk2_tmp_2rev ^ wk2_tmp_1rev ^ (wk2_tmp_1rev + wk2_tmp_2rev))
        assert wk2_tmp_4rev == Reverse((((... & ...) + (... & ...)) & ~(... | ...)) | ((... & ... & (... >> ...)) - ((... + ...) & (... | ...))))
        assert wk2 == (((PopCountDiff((... & ...) | (... << ...), (... & ...) ^ ... ^ ...)) :: 0b000) - (0b00 :: ((0b0 :: PopCount(...)) + (PopCount(...) :: 0b0))))
        assert (0b0000000000000000000000000000000 == ((... & ... & ~... & (~... | (... ^ ...))) | (~... & ~... & ((dk6[:1]) ^ (dmk3[:1]))))) & ((dk6[0]) == (dmk3[0]))
        assert wk3_tmp_0lz == LeadingZeros(~((~... & ~...) << 0b0000000000000000000000000000001))
        assert wk3_tmp_1rev == Reverse((~(dk6[:1]) & ~(dmk3[:1])) << 0b0000000000000000000000000000001)
        assert wk3_tmp_2rev == Reverse(~(... >> ...) & ((... & ...) << 0b0000000000000000000000000000001) & (~(... ^ ...) >> 0b0000000000000000000000000000001))
        assert wk3_tmp_3rev == Reverse(wk3_tmp_2rev ^ wk3_tmp_1rev ^ (wk3_tmp_1rev + wk3_tmp_2rev))
        assert wk3_tmp_4rev == Reverse((((... & ...) + (... & ...)) & ~(... | ...)) | ((... & ... & (... >> ...)) - ((... + ...) & (... | ...))))
        assert wk3 == (((PopCountDiff((... & ...) | (... << ...), (... & ...) ^ ... ^ ...)) :: 0b000) - (0b00 :: ((0b0 :: PopCount(...)) + (PopCount(...) :: 0b0))))
        assert wk == (((0b0 :: wk0) + (0b0 :: wk1) + (0b0 :: wk2) + (0b00 :: wk3))[:3])
        minimize wk
        assert 0x00000000 == ((~(... << ...) ^ (dx2 << 0x00000001)) & (~(... << ...) ^ ((... ^ ...) << 0x00000001)) & (((dp0 ^ (... <<< ...)) << 0x00000001) ^ dx2 ^ dp0 ^ (... <<< ...) ^ dp1 ^ (... <<< ...)))
        assert 0x00000000 == ((~(... << ...) ^ (dx6 << 0x00000001)) & (~(... << ...) ^ ((... ^ ...) << 0x00000001)) & (((dp1 ^ (... <<< ...)) << 0x00000001) ^ dx6 ^ dp1 ^ (... <<< ...) ^ dp2 ^ (... <<< ...)))
        assert 0x00000000 == ((~(... << ...) ^ (dx10 << 0x00000001)) & (~(... << ...) ^ ((... ^ ...) << 0x00000001)) & (((dp2 ^ (... <<< ...)) << 0x00000001) ^ dx10 ^ dp2 ^ (... <<< ...) ^ dp3 ^ (... <<< ...)))
        assert w0w1w2 == (PopCountSum3(~((dx10 ^ ~...) & (~... ^ ... ^ ...))[30:], ~((dx2 ^ ~...) & (~... ^ ... ^ ...))[30:], ~((dx6 ^ ~...) & (~... ^ ... ^ ...))[30:]))
        assert w == w0w1w2
        minimize w
        >>> rkch_found = search_problem.solve(0, 0, solver_name="btor")
        >>> rkch_found.key_ch_found.ch_weight, rkch_found.enc_ch_found.ch_weight
        (0b0000000, 0b0000000)

    """

    def __init__(self, rkch, key_der_mode=DerMode.Default, enc_der_mode=DerMode.Default,
                 allow_zero_key_input_diff=False, allow_zero_enc_input_diff=True,
                 initial_key_constraints=None, initial_enc_constraints=None):
        assert isinstance(rkch, characteristic.RelatedKeyCh)

        env = environment.reset_env()
        self._env = env

        self.key_schedule_problem = SearchCh(
            ch=rkch.key_schedule_ch, allow_zero_input_diff=allow_zero_key_input_diff,
            der_mode=key_der_mode, weight_prefix="wk", env=env,
            initial_constraints=initial_key_constraints)

        self.encryption_problem = SearchCh(
            ch=rkch.encryption_ch, allow_zero_input_diff=allow_zero_enc_input_diff,
            der_mode=enc_der_mode, weight_prefix="w", env=env,
            initial_constraints=initial_enc_constraints)

        self.rkch = rkch

    def formula_size(self, measure=None):
        """Return the size of the underlying SMT problem.

        See  also `SearchCh.formula_size`.
        """
        return self.key_schedule_problem.formula_size(measure) + self.encryption_problem.formula_size(measure)

    def hrepr(self, full_repr=False):
        """Return a human readable representing of the SMT problem.

        See  also `SearchCh.hrepr`.
        """
        key_hrepr = self.key_schedule_problem.hrepr(full_repr)
        enc_hrepr = self.encryption_problem.hrepr(full_repr)
        return key_hrepr + "\n" + enc_hrepr

    def solve(self, initial_ew, initial_kw, solver_name="btor", search_mode=RkChSearchMode.OptimalMinSum,
              check=False, verbose_level=0, filename=None):
        """Solve the SMT problem associated to the search a valid `RelatedKeyCh`.

        Args:
             initial_ew(int): the initial encryption weight for starting the iterative search
             initial_kw(int): the initial key schedule weight for starting the iterative search
             solver_name(str): the name of the solver (according to pySMT) to be used
             search_mode(RkChSearchMode): one of the search modes available
             check(bool): if ``True``, `RkChFound.check_empirical_weight` will be called
                after a characteristic is found. If it is not valid, the search will continue.
             verbose_level(int): an integer between ``0`` (no verbose) and ``3`` (full verbose).
             filename(str): if not ``None``, the output will be  printed to the given file
                rather than the to stdout.

        """
        kp = self.key_schedule_problem
        ep = self.encryption_problem
        assert initial_ew < sum(ep.ch.func.input_widths) + math.floor(ep.error())
        assert initial_kw < sum(kp.ch.func.input_widths) + math.floor(kp.error())
        assert isinstance(search_mode, RkChSearchMode)

        return self._incremental_solve(initial_ew, initial_kw, solver_name, search_mode, check,
                                       verbose_level, filename)

    # noinspection PyPep8
    def _incremental_solve(self, initial_ew, initial_kw, solver_name, search_mode, check,
                           verbose_level, filename):
        strict_shift = True if solver_name == "btor" else False
        bv2pysmt = functools.partial(types.bv2pysmt, env=self._env, strict_shift=strict_shift)
        smart_print = _get_smart_print(filename)

        kp = self.key_schedule_problem
        ep = self.encryption_problem

        def sum_weights(key_weight, enc_weight):
            ch_max_weight = kp.ch_max_weight + ep.ch_max_weight
            ch_width = max(ch_max_weight.bit_length(), key_weight.width, enc_weight.width)
            return operation.ZeroExtend(key_weight, ch_width - key_weight.width) + \
                   operation.ZeroExtend(enc_weight, ch_width - enc_weight.width)

        differences_in_model = []
        for d in itertools.chain(kp.ch.input_diff, kp.ch.nonlinear_diffs.keys(),
                                 ep.ch.input_diff, ep.ch.nonlinear_diffs.keys()):
            differences_in_model.append(d)

        with self._env.factory.Solver(name=solver_name, logic=logics.QF_BV) as solver:
            for c in kp.constraints:
                solver.add_assertion(bv2pysmt(c, boolean=True))

            for c in ep.constraints:
                solver.add_assertion(bv2pysmt(c, boolean=True))

            key_max_error = kp.error()
            key_orig_upper_bound = sum(kp.ch.func.input_widths) + math.ceil(key_max_error)
            key_upper_bound = min(key_orig_upper_bound, kp.ch_max_weight + 1)

            enc_max_error = ep.error()
            enc_orig_upper_bound = sum(ep.ch.func.input_widths) + math.ceil(enc_max_error)
            enc_upper_bound = min(enc_orig_upper_bound, ep.ch_max_weight + 1)

            if verbose_level >= 1:
                smart_print("Key schedule ch. upper bound: {}, max error: {}".format(
                    key_upper_bound, key_max_error
                ))
                smart_print("Encryption ch. upper bound: {}, max error: {}\n".format(
                    enc_upper_bound, enc_max_error
                ))

            if search_mode in [RkChSearchMode.FirstMinSum,
                               RkChSearchMode.FirstValidKeyMinEnc,
                               RkChSearchMode.FirstFixEncMinKey]:

                if search_mode == RkChSearchMode.FirstMinSum:
                    target_weight = int(initial_ew) + int(initial_kw)
                    upper_bound = min(
                        enc_upper_bound + key_upper_bound,
                        kp.ch_max_weight + ep.ch_max_weight + 1  # to not count +1 twice
                    )
                    get_assertion = lambda tw: operation.BvComp(sum_weights(kp.ch_weight, ep.ch_weight), tw)
                elif search_mode == RkChSearchMode.FirstValidKeyMinEnc:
                    target_weight = initial_ew
                    upper_bound = enc_upper_bound
                    if key_orig_upper_bound <= kp.ch_max_weight:
                        solver.add_assertion(bv2pysmt(kp.ch_weight < key_orig_upper_bound))
                        if verbose_level >= 1:
                            smart_print("Added assertion:", kp.ch_weight < key_orig_upper_bound)
                    get_assertion = lambda tw: operation.BvComp(ep.ch_weight, tw)
                elif search_mode == RkChSearchMode.FirstFixEncMinKey:
                    target_weight = initial_kw
                    upper_bound = key_upper_bound
                    solver.add_assertion(bv2pysmt(
                        operation.BvComp(ep.ch_weight, initial_ew), boolean=True))
                    if verbose_level >= 1:
                        smart_print("Added assertion:", operation.BvComp(ep.ch_weight, initial_ew))
                    get_assertion = lambda tw: operation.BvComp(kp.ch_weight, tw)

                while target_weight < upper_bound:
                    assertion = get_assertion(target_weight)
                    if verbose_level >= 1:
                        smart_print(_get_time(), "| Solving", assertion)
                    satisfiable = solver.solve([bv2pysmt(assertion, boolean=True)])

                    if satisfiable:
                        model = types.pysmt_model2bv_model(solver.get_model(), differences_in_model)
                        if search_mode == RkChSearchMode.FirstMinSum:
                            assert sum_weights(model[kp.ch_weight], model[ep.ch_weight]) == target_weight
                        elif search_mode == RkChSearchMode.FirstValidKeyMinEnc:
                            assert model[ep.ch_weight] == target_weight
                        elif search_mode == RkChSearchMode.FirstFixEncMinKey:
                            assert model[kp.ch_weight] == target_weight

                        rkch_found = RkChFound(self.rkch, self.key_schedule_problem, self.encryption_problem, model)

                        if check:
                            try:
                                rkch_found.check_empirical_weight(verbose_level, filename)
                            except ExactWeightError:
                                rkch_found.key_ch_found.emp_weight = math.inf
                                rkch_found.enc_ch_found.emp_weight = math.inf

                        return rkch_found
                    else:
                        target_weight += 1
                else:
                    if verbose_level >= 1:
                        smart_print(_get_time(), "| Unsatisfiable")
                    return None

            if search_mode == RkChSearchMode.AllValid:
                last_rkch_found = None
                rkch_sig = self.rkch.signature(characteristic.ChSignatureType.Full)

                if key_orig_upper_bound <= kp.ch_max_weight:
                    solver.add_assertion(bv2pysmt(kp.ch_weight < key_orig_upper_bound))
                    if verbose_level >= 1:
                        smart_print("Added assertion:", kp.ch_weight < key_orig_upper_bound)
                if enc_orig_upper_bound <= ep.ch_max_weight:
                    solver.add_assertion(bv2pysmt(ep.ch_weight < enc_orig_upper_bound))
                    if verbose_level >= 1:
                        smart_print("Added assertion:", ep.ch_weight < enc_orig_upper_bound)

                for _ in MAX_CH_FOUND:
                    if last_rkch_found is not None:
                        candidate_sig = last_rkch_found.signature(characteristic.ChSignatureType.Full)
                        # disable simplification due to recursion error
                        with context.Simplification(False):
                            c = ~operation.BvComp(rkch_sig[0], candidate_sig[0])
                            for i in range(1, len(rkch_sig)):
                                c |= ~operation.BvComp(rkch_sig[i], candidate_sig[i])
                        solver.add_assertion(bv2pysmt(c, boolean=True))
                        if verbose_level >= 3:
                            smart_print("Added assertion:", c)
                    if verbose_level >= 1:
                        smart_print(_get_time(), "| Solving")
                    satisfiable = solver.solve()

                    if satisfiable:
                        model = types.pysmt_model2bv_model(solver.get_model(), differences_in_model)
                        last_rkch_found = RkChFound(self.rkch, self.key_schedule_problem, self.encryption_problem, model)

                        valid_ch = True
                        if check:
                            try:
                                last_rkch_found.check_empirical_weight(verbose_level, filename)
                            except ExactWeightError:
                                valid_ch = False
                                if verbose_level >= 1:
                                    smart_print(_get_time(), "| Found invalid characteristic")
                                    smart_print(last_rkch_found)
                                    if verbose_level >= 2:
                                        smart_print(last_rkch_found.vrepr())

                        if valid_ch:
                            smart_print(_get_time(), "| Found characteristic")
                            smart_print(last_rkch_found)
                            if verbose_level >= 2:
                                smart_print(last_rkch_found.vrepr())
                    else:
                        if verbose_level >= 1:
                            if last_rkch_found is not None:
                                smart_print(_get_time(), "| No more characteristics found")
                            else:
                                smart_print(_get_time(), "| No characteristic found")
                        smart_print()
                        return

            if search_mode in [RkChSearchMode.OptimalMinSum,
                               RkChSearchMode.OptimalMinSumDifferential,
                               RkChSearchMode.OptimalValidKeyMinEnc,
                               RkChSearchMode.OptimalValidKeyMinEncDifferential,
                               RkChSearchMode.OptimalFixEncMinKey,
                               RkChSearchMode.OptimalFixEncMinKeyDifferential,
                               RkChSearchMode.FirstMinSumValid,
                               RkChSearchMode.FirstFixEncMinKeyValid,
                               RkChSearchMode.FirstValidKeyMinEncValid,
                               RkChSearchMode.AllOptimalMinSum]:
                min_exact_weight = math.inf
                best_rkch_found = None
                last_rkch_found = None
                signature_type = None

                if search_mode in [RkChSearchMode.OptimalMinSum,
                                   RkChSearchMode.OptimalMinSumDifferential,
                                   RkChSearchMode.FirstMinSumValid,
                                   RkChSearchMode.AllOptimalMinSum]:
                    target_weight = int(initial_ew) + int(initial_kw)
                    upper_bound = min(
                        enc_upper_bound + key_upper_bound,
                        kp.ch_max_weight + ep.ch_max_weight + 1  # + 1, not + 2
                    )
                    if search_mode in [RkChSearchMode.OptimalMinSum, RkChSearchMode.AllOptimalMinSum]:
                        signature_type = characteristic.ChSignatureType.Full
                    elif search_mode in [RkChSearchMode.OptimalMinSumDifferential, RkChSearchMode.FirstMinSumValid]:
                        signature_type = characteristic.ChSignatureType.InputOutput
                    else:
                        raise ValueError("invalid search_mode")
                    rkch_sig = self.rkch.signature(signature_type)

                    get_assertion = lambda tw: operation.BvComp(sum_weights(kp.ch_weight, ep.ch_weight), tw)
                    get_sig = lambda rkf: rkf.signature(signature_type)
                    get_sig_str = lambda rkf: rkf.signature(signature_type, return_str=True)
                    get_exact_weight = lambda rkf: rkf.key_ch_found.get_exact_weight() + rkf.enc_ch_found.get_exact_weight()
                    if search_mode == RkChSearchMode.AllOptimalMinSum:
                        get_upper_bound = lambda rkf, pub: pub
                    else:
                        get_upper_bound = lambda rkf, pub: min(
                            kp._new_upper_bound(rkf.key_ch_found, key_upper_bound) +
                            ep._new_upper_bound(rkf.enc_ch_found, enc_upper_bound),
                            pub
                        )

                elif search_mode in [RkChSearchMode.OptimalValidKeyMinEnc,
                                     RkChSearchMode.OptimalValidKeyMinEncDifferential,
                                     RkChSearchMode.FirstValidKeyMinEncValid]:
                    target_weight = initial_ew
                    upper_bound = enc_upper_bound
                    if key_orig_upper_bound <= kp.ch_max_weight:
                        solver.add_assertion(bv2pysmt(kp.ch_weight < key_orig_upper_bound))
                        if verbose_level >= 1:
                            smart_print("Added assertion:", kp.ch_weight < key_orig_upper_bound)
                    if search_mode == RkChSearchMode.OptimalValidKeyMinEnc:
                        signature_type = characteristic.ChSignatureType.Full
                    else:
                        signature_type = characteristic.ChSignatureType.InputOutput
                    rkch_sig = self.rkch.encryption_ch.signature(signature_type)

                    get_assertion = lambda tw: operation.BvComp(ep.ch_weight, tw)
                    get_sig = lambda rkf: rkf.enc_ch_found.signature(signature_type)
                    get_sig_str = lambda rkf: rkf.enc_ch_found.signature(signature_type, return_str=True)
                    get_exact_weight = lambda rkf: rkf.enc_ch_found.get_exact_weight()
                    get_upper_bound = lambda rkf, pub: ep._new_upper_bound(rkf.enc_ch_found, pub)

                elif search_mode in [RkChSearchMode.OptimalFixEncMinKey,
                                     RkChSearchMode.OptimalFixEncMinKeyDifferential,
                                     RkChSearchMode.FirstFixEncMinKeyValid]:
                    target_weight = initial_kw
                    upper_bound = key_upper_bound
                    solver.add_assertion(bv2pysmt(
                        operation.BvComp(ep.ch_weight, initial_ew), boolean=True))
                    if verbose_level >= 1:
                        smart_print("Added assertion:", operation.BvComp(ep.ch_weight, initial_ew))

                    if search_mode == RkChSearchMode.OptimalFixEncMinKey:
                        signature_type = characteristic.ChSignatureType.Full
                    else:
                        signature_type = characteristic.ChSignatureType.InputOutput
                    rkch_sig = self.rkch.key_schedule_ch.signature(signature_type)

                    get_assertion = lambda tw: operation.BvComp(kp.ch_weight, tw)
                    get_sig = lambda rkf: rkf.key_ch_found.signature(signature_type)
                    get_sig_str = lambda rkf: rkf.key_ch_found.signature(signature_type, return_str=True)
                    get_exact_weight = lambda rkf: rkf.key_ch_found.get_exact_weight()
                    get_upper_bound = lambda rkf, pub: kp._new_upper_bound(rkf.key_ch_found, pub)
                else:
                    raise ValueError("invalid search mode: {}".format(search_mode))

                solver.push()
                solver.add_assertion(bv2pysmt(get_assertion(target_weight), boolean=True))

                while target_weight < upper_bound:
                    if last_rkch_found is not None:
                        candidate_sig = get_sig(last_rkch_found)
                        assert len(candidate_sig) > 0
                        c = ~operation.BvComp(rkch_sig[0], candidate_sig[0])
                        for i in range(1, len(rkch_sig)):
                            c |= ~operation.BvComp(rkch_sig[i], candidate_sig[i])
                        solver.add_assertion(bv2pysmt(c, boolean=True))
                        if verbose_level >= 3:
                            smart_print("Added assertion:", c)
                    if verbose_level >= 1:
                        smart_print(_get_time(), "| Solving", get_assertion(target_weight))
                    satisfiable = solver.solve()

                    if satisfiable:
                        model = types.pysmt_model2bv_model(solver.get_model(), differences_in_model)

                        last_rkch_found = RkChFound(self.rkch, self.key_schedule_problem, self.encryption_problem, model)

                        exact_weight = get_exact_weight(last_rkch_found)

                        valid_ch = True

                        if exact_weight < min_exact_weight:
                            if check:
                                try:
                                    last_rkch_found.check_empirical_weight(verbose_level, filename)
                                except ExactWeightError:
                                    valid_ch = False
                                    if verbose_level >= 1:
                                        smart_print(_get_time(), "| Found invalid characteristic")
                                        if verbose_level >= 2:
                                            smart_print(last_rkch_found.vrepr())

                            if valid_ch:
                                best_rkch_found = last_rkch_found
                                min_exact_weight = exact_weight
                                if verbose_level >= 1:
                                    smart_print(_get_time(), "| Found better characteristic")
                                    smart_print(last_rkch_found)
                                    if verbose_level >= 2:
                                        smart_print(last_rkch_found.vrepr())
                        else:
                            if verbose_level >= 1:
                                smart_print(_get_time(), "| Found worse characteristic (not checked)")
                                if verbose_level >= 2:
                                    smart_print(last_rkch_found.vrepr())

                        if valid_ch:
                            if search_mode in [RkChSearchMode.FirstMinSumValid,
                                               RkChSearchMode.FirstValidKeyMinEncValid,
                                               RkChSearchMode.FirstFixEncMinKeyValid]:
                                return last_rkch_found
                            else:
                                new_upper_bound = get_upper_bound(last_rkch_found, upper_bound)
                                if int(exact_weight) == 0 or target_weight >= upper_bound:
                                    break
                                if new_upper_bound != upper_bound and verbose_level >= 1:
                                    smart_print("New upper bound:", new_upper_bound)
                                upper_bound = new_upper_bound
                    else:
                        if verbose_level >= 1:
                            if last_rkch_found is not None:
                                smart_print(_get_time(), "| No more characteristics found")
                            else:
                                smart_print(_get_time(), "| No characteristic found")
                            if verbose_level >= 3 and last_rkch_found is not None:
                                smart_print("Removed assertions")
                            smart_print()

                        solver.pop()

                        target_weight += 1
                        last_rkch_found = None
                        solver.push()
                        solver.add_assertion(bv2pysmt(get_assertion(target_weight), boolean=True))

                if best_rkch_found is not None:
                    return best_rkch_found
                else:
                    if verbose_level >= 1:
                        smart_print(_get_time(), "| No valid characteristic found")
                    return None
arxpy/smt/search.py
class SearchCh(object):
    """Represent the SMT problem of finding a `BvCharacteristic` with conditions on the weight.

    Args:
        ch (BvCharacteristic): a symbolic characteristic of a `BvFunction`
        der_mode (DerMode): one of the modes available for the derivative weights
        weight_prefix (str): the prefix to label weight variables
        allow_zero_input_diff(bool): if ``True``, allow the input difference to be zero
        initial_constraints (list): a list of constraints (given as `Term`) to add to the SMT problem
        env (pysmt.environment.Environment): if ``None``, a new pySMT environment is created

    ::

        >>> from arxpy.bitvector.core import Variable
        >>> from arxpy.differential.difference import XorDiff
        >>> from arxpy.differential.characteristic import BvCharacteristic
        >>> from arxpy.primitives.chaskey import ChaskeyPi
        >>> ChaskeyPi.set_rounds(1)
        >>> ch = BvCharacteristic(ChaskeyPi, XorDiff, ["dv0", "dv1", "dv2", "dv3"])
        >>> search_problem = SearchCh(ch)
        >>> search_problem.formula_size()
        2804
        >>> search_problem.error()
        0
        >>> print(search_problem.hrepr(False))
        assert ~(0x00000000000000000000000000000000 == (dv0 :: dv1 :: dv2 :: dv3))
        assert 0x00000000 == ((~(... << ...) ^ (d0 << 0x00000001)) & (~(... << ...) ^ (dv1 << 0x00000001)) & ((dv0 << 0x00000001) ^ d0 ^ dv0 ^ dv1))
        assert 0x00000000 == ((~(... << ...) ^ (d4 << 0x00000001)) & (~(... << ...) ^ (dv3 << 0x00000001)) & ((dv2 << 0x00000001) ^ d4 ^ dv2 ^ dv3))
        assert 0x00000000 == ((~(... << ...) ^ (d7 << 0x00000001)) & (~(... << ...) ^ ((... ^ ...) << 0x00000001)) & (((d0 <<< 16) << 0x00000001) ^ d7 ^ ... ^ ... ^ (... <<< ...)))
        assert 0x00000000 == ((~(... << ...) ^ (d10 << 0x00000001)) & (~(... << ...) ^ (d4 << 0x00000001)) & (((d0 ^ (... <<< ...)) << 0x00000001) ^ d10 ^ d4 ^ ... ^ ...))
        assert w0w1w2 == (PopCountSum3(~((d0 ^ ~...) & (dv1 ^ ~...))[30:], ~((d4 ^ ~...) & (dv3 ^ ~...))[30:], ~((d7 ^ ~...) & (~... ^ ... ^ ...))[30:]))
        assert w3 == PopCount(~((d10 ^ ~...) & (d4 ^ ~...))[30:])
        assert w == (w0w1w2 + (0b00 :: w3))
        minimize w
        >>> search_problem.solve(0).ch_weight
        0b0000000

    """

    def __init__(self, ch, der_mode=DerMode.Default, weight_prefix="w",
                 allow_zero_input_diff=False, initial_constraints=None, env=None,):
        assert isinstance(ch, characteristic.BvCharacteristic)
        assert isinstance(der_mode, DerMode)

        if der_mode == DerMode.XDCA_Approx:
            class XDCA_Approx(derivative.XDCA):
                precision = 0

            for diff in ch.nonlinear_diffs:
                der = ch.nonlinear_diffs[diff]
                if isinstance(der, derivative.XDCA):
                    der = XDCA_Approx(der.input_diff, der.op.constant)
                    ch.nonlinear_diffs[diff] = der
                    assert der.precision == 0

        if initial_constraints is None:
            constraints = []
        else:
            constraints = initial_constraints[:]

        if not allow_zero_input_diff:
            input_diff = functools.reduce(operation.Concat, [d.val for d in ch.input_diff])
            constraints.append(operation.BvNot(operation.BvComp(input_diff, core.Constant(0, input_diff.width))))

        self.ch = ch
        self.constraints = constraints
        self.weight_prefix = weight_prefix
        self.der_mode = der_mode

        with context.NotEvaluation([extraop.Reverse, extraop.LeadingZeros,
                                    extraop.PopCount, extraop.PopCountDiff, extraop.PopCountSum2, extraop.PopCountSum3]):
            self._generate_constraints()

        if env is None:
            self._env = environment.reset_env()

    def _generate_constraints(self):
        """Generate the bit-vector constraints of the SMT problem."""
        if not self.ch.nonlinear_diffs:
            constraints = []
            ch_weight = core.Variable(self.weight_prefix, 1)
            constraints.append(operation.BvComp(ch_weight, 0))
            self.constraints.extend(constraints)
            self.ch_weight = ch_weight
            self.der_weights = []
            self.ch_max_weight = 0
            return

        if self.der_mode == DerMode.ProbabilityOne:
            constraints = []
            for (diff, der) in self.ch.nonlinear_diffs.items():
                constraints.append(der.has_probability_one(diff))

            ch_weight = core.Variable(self.weight_prefix, 1)
            constraints.append(operation.BvComp(ch_weight, 0))

            self.constraints.extend(constraints)
            self.ch_weight = ch_weight
            self.der_weights = []
            self.ch_max_weight = 0
            return

        if self.der_mode == DerMode.Valid:
            constraints = []
            for (diff, der) in self.ch.nonlinear_diffs.items():
                constraints.append(der.is_possible(diff))

            ch_weight = core.Variable(self.weight_prefix, 1)
            constraints.append(operation.BvComp(ch_weight, 0))

            self.constraints.extend(constraints)
            self.ch_weight = ch_weight
            self.der_weights = []
            self.ch_max_weight = 0
            return

        max_num_frac_bits = max(der.num_frac_bits() for der in self.ch.nonlinear_diffs.values())
        ch_max_weight = 0
        for der in self.ch.nonlinear_diffs.values():
            ch_max_weight += der.max_weight() << (max_num_frac_bits - der.num_frac_bits())
        ch_max_weight = ch_max_weight >> max_num_frac_bits   # ch_weight ignore fraction bits
        ch_width = max(ch_max_weight.bit_length(), 1)
        for diff, der in self.ch.nonlinear_diffs.items():
            if isinstance(der, derivative.XDCA):
                # the width of XDAC.weight is "greater" than its max_weight
                # (the MSB bits of XDAC.weight are not used)
                int_width = der.weight(diff)[0].width - der.num_frac_bits()
            else:
                int_width = der.weight(diff).width - der.num_frac_bits()
            ch_width = max(ch_width, int_width)

        constraints = []
        # fb2width2hw is a dictionary mapping frac_bits -> input_width -> (weight_var, PopCount)
        # note only PopCounts with the same number of frac_bits and similar args are grouped together
        fb2width2hw = {}
        original_der_weights = []
        der_weights = []

        def zero_ext_right(var, num_zeros):
            """Expand with zeros to the right."""
            if num_zeros == 0:
                return var
            else:
                return operation.Concat(var, core.Constant(0, num_zeros))

        for i, (diff, der) in enumerate(self.ch.nonlinear_diffs.items()):
            constraints.append(der.is_possible(diff))

            if isinstance(der, derivative.XDCA):
                prefix = self.weight_prefix + str(i) + "_tmp"
                # with context.NotEvaluation(extraop.PopCount):
                assert extraop.PopCount in context.NotEvaluation.current_context
                weight_value = der.weight(diff, prefix=prefix)
                if isinstance(weight_value, collections.abc.Sequence):
                    assert len(weight_value) == 2
                    weight_value, assertions = weight_value
                    weight_var = core.Variable(self.weight_prefix + str(i), weight_value.width)
                    constraints.extend(assertions)
            else:
                # with context.NotEvaluation(extraop.PopCount):
                assert extraop.PopCount in context.NotEvaluation.current_context
                weight_value = der.weight(diff)
                weight_var = core.Variable(self.weight_prefix + str(i), weight_value.width)

            if weight_value == 0:
                continue

            num_fb = der.num_frac_bits()
            if isinstance(weight_value, extraop.PopCount):
                if num_fb not in fb2width2hw:
                    fb2width2hw[num_fb] = {}
                input_hw_width = weight_value.args[0].width
                if input_hw_width not in fb2width2hw[num_fb]:
                    fb2width2hw[num_fb][input_hw_width] = []
                fb2width2hw[num_fb][input_hw_width] += [[weight_var, weight_value]]
            else:
                constraints.append(operation.BvComp(weight_var, weight_value))
                original_der_weights.append(weight_var)
                assert ch_width >= weight_var.width - der.num_frac_bits()  # der_int_width
                weight_var = operation.ZeroExtend(weight_var, ch_width - (weight_var.width - der.num_frac_bits()))
                weight_var = zero_ext_right(weight_var, max_num_frac_bits - der.num_frac_bits())
                der_weights.append(weight_var)

        def grouped(iterable, n):
            """Group iterable in packs of n elements.

            If len(iterable) % n != 0, the last elements are not included.
            Example: grouped([s0,s1,s2,s3,s4], 2) -> ((s0, s1), (s2, s3))
            """
            return zip(*[iter(iterable)] * n)

        for num_fb in fb2width2hw:
            for weightvar_hwvalue_pairs in fb2width2hw[num_fb].values():
                for ((w1, p1), (w2, p2), (w3, p3)) in grouped(weightvar_hwvalue_pairs, 3):
                    assert p1.width == p2.width == p3.width
                    weight_value = extraop.PopCountSum3(p1.args[0], p2.args[0], p3.args[0])
                    weight_var = core.Variable(w1.name + w2.name + w3.name, weight_value.width)
                    constraints.append(operation.BvComp(weight_var, weight_value))
                    original_der_weights.append(weight_var)
                    assert ch_width >= weight_var.width - num_fb
                    weight_var = operation.ZeroExtend(weight_var, ch_width - (weight_var.width - num_fb))
                    weight_var = zero_ext_right(weight_var, max_num_frac_bits - num_fb)
                    der_weights.append(weight_var)
                if len(weightvar_hwvalue_pairs) % 3 == 2:
                    w1, p1 = weightvar_hwvalue_pairs[-2]
                    w2, p2 = weightvar_hwvalue_pairs[-1]
                    assert p1.width == p2.width
                    weight_value = extraop.PopCountSum2(p1.args[0], p2.args[0])
                    weight_var = core.Variable(w1.name + w2.name, weight_value.width)
                    constraints.append(operation.BvComp(weight_var, weight_value))
                    original_der_weights.append(weight_var)
                    assert ch_width >= weight_var.width - num_fb
                    weight_var = operation.ZeroExtend(weight_var, ch_width - (weight_var.width - num_fb))
                    weight_var = zero_ext_right(weight_var, max_num_frac_bits - num_fb)
                    der_weights.append(weight_var)
                elif len(weightvar_hwvalue_pairs) % 3 == 1:
                    w1, p1 = weightvar_hwvalue_pairs[-1]
                    weight_var = w1
                    weight_value = p1
                    constraints.append(operation.BvComp(weight_var, weight_value))
                    original_der_weights.append(weight_var)
                    assert ch_width >= weight_var.width - num_fb
                    weight_var = operation.ZeroExtend(weight_var, ch_width - (weight_var.width - num_fb))
                    weight_var = zero_ext_right(weight_var, max_num_frac_bits - num_fb)
                    der_weights.append(weight_var)

        # value of the characteristic weight
        sum_w = sum(der_weights)[:max_num_frac_bits]
        ch_weight = core.Variable(self.weight_prefix, sum_w.width)
        constraints.append(operation.BvComp(ch_weight, sum_w))

        self.constraints.extend(constraints)
        self.ch_weight = ch_weight
        self.ch_max_weight = ch_max_weight
        self.der_weights = original_der_weights

    def solve(self, initial_weight, solver_name="btor", search_mode=ChSearchMode.Optimal, check=False,
              return_generator=False, verbose_level=0, filename=None):
        """Solve the SMT problem associated to the search of a valid `BvCharacteristic`.

        Args:
             initial_weight(int): the initial weight for starting the iterative search
             solver_name(str): the name of the solver (according to pySMT) to be used
             search_mode(ChSearchMode): one of the search modes available
             check(bool): if ``True``, `ChFound.check_empirical_weight` will be called
                after a characteristic is found. If it is not valid, the search will continue.
             return_generator(bool): if ``True``, return a Python generator with the
                characteristics found (only valid in AllOptimal)
             verbose_level(int): an integer between ``0`` (no verbose) and ``3`` (full verbose).
             filename(str): if not ``None``, the output will be  printed to the given file
                rather than the to stdout.

        """
        assert initial_weight < sum(self.ch.func.input_widths) + math.floor(self.error())
        assert type(self) == SearchCh
        assert isinstance(search_mode, ChSearchMode)

        # don't merge incremental_solve() in solve() to allow different
        # ways of solving in the future
        return self._incremental_solve(initial_weight, solver_name, search_mode, check, return_generator,
                                       verbose_level, filename)

    def _incremental_solve(self, initial_weight, solver_name, search_mode, check, return_generator,
                           verbose_level, filename):
        strict_shift = True if solver_name == "btor" else False  # e.g., btor_rol: width must be a power of 2
        bv2pysmt = functools.partial(types.bv2pysmt, env=self._env, strict_shift=strict_shift)

        smart_print = _get_smart_print(filename)

        differences_in_model = []
        for d in itertools.chain(self.ch.input_diff, self.ch.nonlinear_diffs.keys()):
            differences_in_model.append(d)

        if type(self) == SearchCh:
            search_mode_class = ChSearchMode
            ch_found_class = ChFound
        elif type(self) == SearchSkCh:
            search_mode_class = SkChSearchMode
            ch_found_class = SkChFound
        else:
            raise ValueError("invalid subclass of SearchCh")

        if return_generator and not search_mode_class.AllOptimal:
            raise ValueError("return_generator can only be enabled with the search mode AllOptimal")

        # with self._env.factory.Solver(name=solver_name, logic=logics.QF_BV) as solver:
        solver = self._env.factory.Solver(name=solver_name, logic=logics.QF_BV)
        for c in self.constraints:
            solver.add_assertion(bv2pysmt(c, boolean=True))

        max_error = self.error()
        orig_upper_bound = sum(self.ch.func.input_widths) + math.ceil(max_error)  # strict upper bound
        upper_bound = min(orig_upper_bound, self.ch_max_weight + 1)

        if verbose_level >= 1:
            smart_print("Upper bound: {}, max error: {}\n".format(upper_bound, max_error))

        if search_mode == search_mode_class.FirstCh:
            target_weight = initial_weight

            while target_weight < upper_bound:
                eq_weight = operation.BvComp(self.ch_weight, target_weight)
                if verbose_level >= 1:
                    smart_print(_get_time(), "| Solving", eq_weight)
                satisfiable = solver.solve([bv2pysmt(eq_weight, boolean=True)])

                if satisfiable:
                    model = types.pysmt_model2bv_model(solver.get_model(), differences_in_model)
                    assert model[self.ch_weight] == target_weight

                    ch_found = ch_found_class(self.ch, self.ch_weight, self.der_weights, model)

                    if check:
                        try:
                            ch_found.check_empirical_weight(verbose_level, filename)
                        except ExactWeightError:
                            ch_found.emp_weight = math.inf

                    solver.exit()
                    return ch_found
                else:
                    target_weight += 1
            else:
                if verbose_level >= 1:
                    smart_print(_get_time(), "| Unsatisfiable")

                solver.exit()
                return None

        if search_mode == search_mode_class.AllValid:
            last_ch_found = None
            ch_sig = self.ch.signature(characteristic.ChSignatureType.Full)

            if orig_upper_bound <= self.ch_max_weight:
                solver.add_assertion(bv2pysmt(self.ch_weight < orig_upper_bound))
                if verbose_level >= 1:
                    smart_print("Added assertion:", self.ch_weight < orig_upper_bound)

            for _ in MAX_CH_FOUND:
                if last_ch_found is not None:
                    candidate_sig = last_ch_found.signature(characteristic.ChSignatureType.Full)
                    # disable simplification due to recursion error
                    with context.Simplification(False):
                        c = ~operation.BvComp(ch_sig[0], candidate_sig[0])
                        for i in range(1, len(ch_sig)):
                            c |= ~operation.BvComp(ch_sig[i], candidate_sig[i])
                    solver.add_assertion(bv2pysmt(c, boolean=True))
                    if verbose_level >= 3:
                        smart_print("Added assertion:", c)
                if verbose_level >= 1:
                    smart_print(_get_time(), "| Solving")
                satisfiable = solver.solve()

                if satisfiable:
                    model = types.pysmt_model2bv_model(solver.get_model(), differences_in_model)
                    last_ch_found = ch_found_class(self.ch, self.ch_weight, self.der_weights, model)

                    valid_ch = True
                    if check:
                        try:
                            last_ch_found.check_empirical_weight(verbose_level, filename)
                        except ExactWeightError:
                            valid_ch = False
                            if verbose_level >= 1:
                                smart_print(_get_time(), "| Found invalid characteristic")
                                smart_print(last_ch_found)
                                if verbose_level >= 2:
                                    smart_print(last_ch_found.vrepr())

                    if valid_ch:
                        smart_print(_get_time(), "| Found characteristic")
                        smart_print(last_ch_found)
                        if verbose_level >= 2:
                            smart_print(last_ch_found.vrepr())
                else:
                    if verbose_level >= 1:
                        if last_ch_found is not None:
                            smart_print(_get_time(), "| No more characteristics found")
                        else:
                            smart_print(_get_time(), "| No characteristic found")
                        smart_print()
                    solver.exit()
                    return

        if search_mode in [search_mode_class.Optimal, search_mode_class.OptimalDifferential,
                           search_mode_class.FirstChValid, search_mode_class.AllOptimal]:
            def iterate_ch(upper_bound):
                solver.push()

                min_exact_weight = math.inf
                last_ch_found = None
                best_ch_found = None
                target_weight = initial_weight

                if search_mode in [search_mode_class.Optimal, search_mode_class.AllOptimal]:
                    signature_type = characteristic.ChSignatureType.Full
                elif search_mode in [search_mode_class.FirstChValid, search_mode_class.OptimalDifferential]:
                    signature_type = characteristic.ChSignatureType.InputOutput

                ch_sig = self.ch.signature(signature_type)

                solver.add_assertion(bv2pysmt(
                    operation.BvComp(self.ch_weight, target_weight), boolean=True))

                while target_weight < upper_bound:
                    if last_ch_found is not None:
                        candidate_sig = last_ch_found.signature(signature_type)
                        assert len(candidate_sig) > 0
                        c = ~operation.BvComp(ch_sig[0], candidate_sig[0])
                        for i in range(1, len(ch_sig)):
                            c |= ~operation.BvComp(ch_sig[i], candidate_sig[i])
                        solver.add_assertion(bv2pysmt(c, boolean=True))
                        if verbose_level >= 3:
                            smart_print("Added assertion:", c)
                    if verbose_level >= 1:
                        smart_print(_get_time(), "| Solving", operation.BvComp(self.ch_weight, target_weight))

                    satisfiable = solver.solve()

                    if satisfiable:
                        model = types.pysmt_model2bv_model(solver.get_model(), differences_in_model)
                        assert model[self.ch_weight] == target_weight
                        last_ch_found = ch_found_class(self.ch, self.ch_weight, self.der_weights, model)

                        exact_weight = last_ch_found.get_exact_weight()

                        valid_ch = True

                        if exact_weight < min_exact_weight:
                            if check:
                                try:
                                    last_ch_found.check_empirical_weight(verbose_level, filename)
                                except ExactWeightError:
                                    valid_ch = False
                                    if verbose_level >= 1:
                                        smart_print(_get_time(), "| Found invalid characteristic")
                                        if verbose_level >= 2:
                                            smart_print(last_ch_found.vrepr())

                            if valid_ch:
                                best_ch_found = last_ch_found
                                min_exact_weight = exact_weight
                                if verbose_level >= 1:
                                    smart_print(_get_time(), "| Found better characteristic")
                                    smart_print(last_ch_found)
                                    if verbose_level >= 2:
                                        smart_print(last_ch_found.vrepr())
                        else:
                            if verbose_level >= 1:
                                smart_print(_get_time(), "| Found worse characteristic (not checked)")
                                if verbose_level >= 2:
                                    smart_print(last_ch_found.vrepr())

                        if valid_ch:
                            if return_generator:
                                yield last_ch_found
                            elif search_mode == search_mode_class.FirstChValid:
                                solver.exit()
                                return last_ch_found
                            elif search_mode in [search_mode_class.Optimal, search_mode_class.OptimalDifferential]:
                                new_upper_bound = self._new_upper_bound(last_ch_found, upper_bound)
                                if int(exact_weight) == 0 or target_weight >= new_upper_bound:
                                    break
                                if new_upper_bound != upper_bound and verbose_level >= 1:
                                    smart_print("New upper bound:", new_upper_bound)
                                upper_bound = new_upper_bound
                    else:
                        if verbose_level >= 1:
                            if last_ch_found is not None:
                                smart_print(_get_time(), "| No more characteristics found")
                            else:
                                smart_print(_get_time(), "| No characteristic found")

                            if verbose_level >= 3 and last_ch_found is not None:
                                smart_print("Removed assertions")

                            smart_print()

                        solver.pop()

                        target_weight += 1
                        last_ch_found = None
                        solver.push()
                        solver.add_assertion(bv2pysmt(
                            operation.BvComp(self.ch_weight, target_weight), boolean=True))

                solver.exit()

                if best_ch_found is not None:
                    return best_ch_found
                else:
                    if verbose_level >= 1:
                        smart_print(_get_time(), "| No valid characteristic found")
                    return

            if return_generator:
                return iterate_ch(upper_bound)
            else:
                try:
                    next(iterate_ch(upper_bound))
                except StopIteration as result:
                    solver.exit()
                    return result.value

        if search_mode == search_mode_class.TopDifferentials:
            index_ch_found = 0
            diff2weight = dict()
            topdiff = tuple(sorted(diff2weight.items(), key=lambda x: x[1])[:10])

            def sum_weights(w0, w1):
                if w0 == w1:
                    if w0 < 1:
                        return 0
                    else:
                        return w0 - 1

                # return - math.log2(2**(-w0) + 2**(-w1))

                # a < b
                # 2^(-a) + 2^(-b) = 1/2^a + 1/2^b = 2^(b-a)/2^(b) + 1/2^b
                #                 = (2^(b-a) + 1) / 2^b
                # -log2(2^(-a) + 2^(-b)) = (-1) * (log2(2^(b-a) + 1) - b)
                if w0 < w1:
                    a, b = w0, w1
                else:
                    a, b = w1, w0
                try:
                    return -(math.log2(2**(b - a) + 1) - b)
                except OverflowError:
                    from decimal import Decimal
                    a, b = Decimal(a), Decimal(b)
                    return -(Decimal(2**(b-a) + 1).ln()/Decimal(2).ln() - b)

            solver.push()

            min_exact_weight = math.inf
            last_ch_found = None
            best_ch_found = None
            target_weight = initial_weight

            signature_type = characteristic.ChSignatureType.InputOutput

            ch_sig = self.ch.signature(signature_type)

            solver.add_assertion(bv2pysmt(
                operation.BvComp(self.ch_weight, target_weight), boolean=True))

            while target_weight < upper_bound:
                if last_ch_found is not None:
                    candidate_sig = last_ch_found.signature(signature_type)
                    assert len(candidate_sig) > 0
                    c = ~operation.BvComp(ch_sig[0], candidate_sig[0])
                    for i in range(1, len(ch_sig)):
                        c |= ~operation.BvComp(ch_sig[i], candidate_sig[i])
                    solver.add_assertion(bv2pysmt(c, boolean=True))
                    if verbose_level >= 3:
                        smart_print("Added assertion:", c)
                if verbose_level >= 1 and last_ch_found is None:
                    # only print the target_weight for the 1st characteristic
                    smart_print(_get_time(), "| Solving", operation.BvComp(self.ch_weight, target_weight))

                satisfiable = solver.solve()

                if satisfiable:
                    model = types.pysmt_model2bv_model(solver.get_model(), differences_in_model)
                    assert model[self.ch_weight] == target_weight
                    last_ch_found = ch_found_class(self.ch, self.ch_weight, self.der_weights, model)

                    index_ch_found += 1

                    exact_weight = last_ch_found.get_exact_weight()

                    valid_ch = True
                    if check:
                        try:
                            last_ch_found.check_empirical_weight(verbose_level, filename)
                        except ExactWeightError:
                            valid_ch = False
                            if verbose_level >= 1:
                                smart_print(_get_time(), "| Found invalid characteristic")
                                if verbose_level >= 2:
                                    smart_print(last_ch_found.vrepr())
                    if not valid_ch:
                        continue

                    if exact_weight < min_exact_weight:
                        best_ch_found = last_ch_found
                        min_exact_weight = exact_weight
                        if verbose_level >= 1:
                            smart_print(_get_time(), "| Found better characteristic")
                            smart_print(last_ch_found)
                            if verbose_level >= 2:
                                smart_print(last_ch_found.vrepr())

                    input_diff = [value.val for var, value in last_ch_found.input_diff]
                    output_diff = [value.val for var, value in last_ch_found.output_diff]
                    in_out_diff = tuple(input_diff + output_diff)

                    if in_out_diff not in diff2weight:
                        diff2weight[in_out_diff] = exact_weight
                    else:
                        diff2weight[in_out_diff] = sum_weights(diff2weight[in_out_diff], exact_weight)
                        if verbose_level >= 2:
                            smart_print("diff {} | sum_weights({}, {}) = {}".format(
                                in_out_diff,
                                diff2weight[in_out_diff],
                                exact_weight,
                                sum_weights(diff2weight[in_out_diff], exact_weight)
                            ))

                    # TODO: (v2) store dictionary in folder

                    newtopdiff = tuple(sorted(diff2weight.items(), key=lambda x: x[1])[:10])

                    if verbose_level >= 1 and (index_ch_found & (index_ch_found - 1)) == 0:
                        smart_print(_get_time(), "| Found 2^{} characteristics".format(int(math.log2(index_ch_found))))

                    if len(topdiff) == 9 and len(newtopdiff) == 10:
                        if verbose_level >= 1:
                            smart_print(_get_time(), "| First top 10 differentials")
                            smart_print(newtopdiff)
                    elif len(newtopdiff) == 10 and newtopdiff != topdiff:
                        if verbose_level >= 1:
                            differing = [i_d for i_d, (old_d, new_d) in enumerate(zip(topdiff, newtopdiff)) if old_d != new_d]
                            smart_print(_get_time(), "| Found new top 10 differentials differing in", differing)
                            smart_print(newtopdiff)
                            # smart_print(diff2weight)
                    topdiff = newtopdiff

                    if index_ch_found > MAX_CH_FOUND:
                        return best_ch_found, diff2weight

                else: # if not satisfiable
                    if verbose_level >= 1:
                        if last_ch_found is not None:
                            smart_print(_get_time(), "| No more characteristics found")
                        else:
                            smart_print(_get_time(), "| No characteristic found")

                        if verbose_level >= 3 and last_ch_found is not None:
                            smart_print("Removed assertions")

                        smart_print()

                    solver.pop()

                    target_weight += 1
                    last_ch_found = None
                    solver.push()
                    solver.add_assertion(bv2pysmt(
                        operation.BvComp(self.ch_weight, target_weight), boolean=True))

            solver.exit()

            if best_ch_found is not None:
                return best_ch_found, diff2weight
            else:
                if verbose_level >= 1:
                    smart_print(_get_time(), "| No valid characteristic found")
                return

    def formula_size(self, measure=None):
        """Return the size of the underlying SMT problem.

        See `pysmt.oracles.SizeOracle` for choosing the ``measure``.
        """
        env = environment.reset_env()
        bv2pysmt = functools.partial(types.bv2pysmt, env=env)
        assertions = [bv2pysmt(c, boolean=True) for c in self.constraints]
        formula = env.formula_manager.And(*assertions)
        return env.sizeo.get_size(formula, measure)

    def hrepr(self, full_repr=False):
        """Return a human readable representing of the SMT problem.

        If ``full_repr`` is False, the short string representation `srepr` is used.
        """
        representation = []
        for c in self.constraints:
            if full_repr:
                s = str(c)
            else:
                s = c.srepr()
            representation.append("assert " + s)
        representation.append("minimize {}".format(self.ch_weight))
        return "\n".join(representation)

    def error(self):
        """Return an upper bound of the weight error of the SMT problem."""
        if self.der_mode == DerMode.ProbabilityOne:
            return 0
        elif self.der_mode == DerMode.Valid:
            return 0
        else:
            e = 0
            for der in self.ch.nonlinear_diffs.values():
                e += der.error()
            return e

    def _new_upper_bound(self, ch_found, previous_bound=math.inf):
        """Computes the new upper bound given the previous solution.

        Let me the maximum error of the characteristic (after applying the ceil function).

        For the previous solution, let:
         - eiw the integer part of the exact weight
         - ew the exact weight

        Assuming there exists a better new solution, let denote by eiw' and ew'
        the integer part and the full exact weight for this new solution.

        The worst case (when the new upper bound reduces the least)
        would be when the better solution has exact weight close
        to the previous solution, that is, eiw - 1 < ew' < eiw.

        The weight of this better solution must belong to the following
        interval with integer endpoints: [eiw - 1, eiw - 1 + me]
        (assuming the exact weight is smaller than the weight).

        Therefore the SMT should find characteristics with integer
        part of the ch_weight up to (but not included) the value eiw + me.
        """
        exact_weight = ch_found.get_exact_weight()
        new_bound = math.floor(exact_weight) + math.ceil(self.error())
        return min(previous_bound, new_bound, self.ch_max_weight + 1)
arxpy/smt/search.py
class SkChSearchMode(enum.Enum):
    """Represent the different options for searching `SingleKeyCh`.

    The options are the same as `ChSearchMode`, but `SkChFound` objects are returned instead of `ChFound`.
    """
    FirstCh = enum.auto()
    FirstChValid = enum.auto()
    Optimal = enum.auto()
    OptimalDifferential = enum.auto()
    AllOptimal = enum.auto()
    AllValid = enum.auto()
    TopDifferentials = enum.auto()
arxpy/differential/tests/test_characteristic.py
class MyFunction(BvFunction):  # noqa: D101
    input_widths = [8, 8]
    output_widths = [8]
    rounds = 1

    @classmethod
    def eval(cls, x, y):
        x = RotateLeft(x ^ Constant(1, 8), 1)
        y = ~y
        return tuple([x + y])
arxpy/smt/search.py
def _get_smart_print(filename=None):
    def smart_print(*msg, **kwargs):
        if filename is not None:
            with open(filename, "a") as fh:
                print(*msg, file=fh, flush=True, **kwargs)
        else:
            print(*msg, flush=True, **kwargs)
    return smart_print
arxpy/smt/search.py
class ChSearchMode(enum.Enum):
    """Represent the different options for searching `BvCharacteristic`.

    A characteristic is said to be *valid* if its weight
    :math:`w = - \log_2(p)`, where *p* is the probability of the characteristic,
    is lower than the bit-size of the input difference.

    A derivative weight is said to be *non-exact* if its weight
    does not match its exact weight.

    Attributes:
        FirstCh: returns the 1st characteristic found by iteratively searching with increasing weight.
            The characteristic found is returned as a `ChFound` object, but it may not have
            optimal theoretical weight if it contains a non-exact derivative weight
            or it may be empirically invalid. If no characteristic is found,
            ``None`` is returned.
        FirstChValid: similar to `FirstCh`, but if the characteristic found is
            empirically invalid (after computing its empirical weight), the search continues.
        Optimal: similar to `FirstCh`, but several characteristics are found to ensure
            that the one returned has the optimal theoretical weight.
            (only differs from `FirstCh` if there is a non-exact derivative weight).
        OptimalDifferential: similar to `Optimal`, but once a characteristic is found,
            only characteristics with different input/output differences are searched.
        AllOptimal: similar to `Optimal`, but the search never stops.
        TopDifferentials: searches for `MAX_CH_FOUND` characteristics
            (as in `FirstCh`) and return the differentials obtained
            by gathering those characteristics.
        AllValid: a non-stop search of valid characteristics with no
            additional constraints.

    """
    FirstCh = enum.auto()
    FirstChValid = enum.auto()
    Optimal = enum.auto()
    OptimalDifferential = enum.auto()
    AllOptimal = enum.auto()
    AllValid = enum.auto()
    TopDifferentials = enum.auto()
arxpy/differential/characteristic.py
class SingleKeyCh(BvCharacteristic):
    """Represent single-key characteristics of block ciphers.

    A single-key characteristic of a `Cipher` is a `BvCharacteristic`
    over the `Encryption` function of the cipher.

    The plaintext differences start with the prefix ``"dp"``
    and the non-linear differences start with the prefix ``"dx"``.

        >>> from arxpy.bitvector.core import Variable
        >>> from arxpy.differential.difference import XorDiff
        >>> from arxpy.differential.characteristic import SingleKeyCh
        >>> from arxpy.primitives.primitives import Cipher
        >>> from arxpy.primitives import speck
        >>> Speck32 = speck.get_Speck_instance(speck.SpeckInstance.speck_32_64)
        >>> issubclass(Speck32, Cipher)
        True
        >>> Speck32.set_rounds(1)
        >>> ch = SingleKeyCh(Speck32, XorDiff)
        >>> ch .ssa  # doctest: +NORMALIZE_WHITESPACE
        {'input_vars': (dp0, dp1), 'output_vars': (dx2, dx4),
        'assignments': ((dx0, dp0 >>> 7), (dx1, dp1 + dx0), (dx2, dx1 ^ k0), (dx3, dp1 <<< 2), (dx4, dx2 ^ dx3))}
        >>> ch.input_diff
        (XorDiff(dp0), XorDiff(dp1))
        >>> ch.nonlinear_diffs
        OrderedDict([(XorDiff(dx1), XDA(XorDiff(dp1), XorDiff(dp0 >>> 7)))])
        >>> ch.output_diff
        [[XorDiff(dx2), XorDiff(dx1)], [XorDiff(dx4), XorDiff(dx1 ^ (dp1 <<< 2))]]

    """

    def __init__(self, bv_cipher, diff_type):
        assert issubclass(bv_cipher, primitives.Cipher)
        assert issubclass(diff_type, difference.Difference)

        rk = []
        for i, width in enumerate(bv_cipher.key_schedule.output_widths):
            rk.append(core.Variable("k" + str(i), width))

        class Encryption(bv_cipher.encryption):
            round_keys = tuple(rk)

        func = Encryption
        num_inputs = len(func.input_widths)
        input_diff_names = ["dp" + str(i) for i in range(num_inputs)]
        prefix = "dx"
        super().__init__(func, diff_type, input_diff_names, prefix)
        self._cipher = bv_cipher

    def empirical_weight(self, input_diff, output_diff, pair_samples, key_samples,
                         precision=1, rk_diffs=None):
        """Return the empirical weight distribution of a given differential.

        This method returns a `collections.Counter` storing the distribution of
        differential probability weights over the given number of keys.

        The weights are rounded to the given number of precision
        digits after the decimal point.

        See also `BvCharacteristic.empirical_weight`.

            >>> from arxpy.bitvector.core import Constant
            >>> from arxpy.differential.difference import XorDiff
            >>> from arxpy.differential.characteristic import SingleKeyCh
            >>> from arxpy.primitives import speck
            >>> Speck32 = speck.get_Speck_instance(speck.SpeckInstance.speck_32_64)
            >>> Speck32.set_rounds(1)
            >>> ch = SingleKeyCh(Speck32, XorDiff)
            >>> zero, one = XorDiff(Constant(0, 16)), XorDiff(Constant(1, 16))
            >>> ch.empirical_weight([zero, zero], [zero, zero], 100, 10)
            Counter({0.0: 10})
            >>> ch.empirical_weight([zero, zero], [one, one], 100, 10)
            Counter({inf: 10})

        """
        return self._empirical_weight_distribution(self._cipher, input_diff, output_diff, pair_samples, key_samples,
                                                   precision, rk_diffs)
arxpy/bitvector/operation.py
class Concat(Operation):
    """Concatenation operation.

    Given the bit-vectors :math:`(x_{n-1}, \dots, x_0)` and
    :math:`(y_{m-1}, \dots, y_0)`, ``Concat(x, y)`` returns the bit-vector
    :math:`(x_{n-1}, \dots, x_0, y_{m-1}, \dots, y_0)`.

        >>> from arxpy.bitvector.core import Constant, Variable
        >>> from arxpy.bitvector.operation import Concat
        >>> Concat(Constant(0x12, 8), Constant(0x345, 12))
        0x12345
        >>> Concat(Variable("x", 8), Variable("y", 8))
        x :: y

    """

    arity = [2, 0]
    is_symmetric = False
    infix_symbol = "::"

    @classmethod
    def output_width(cls, x, y):
        return x.width + y.width

    @classmethod
    def eval(cls, x, y):
        def doit(x, y):
            """Concatenation when both operands are int."""
            return int(x.bin() + y.bin()[2:], 2)

        if isinstance(x, core.Constant) and isinstance(y, core.Constant):
            return core.Constant(doit(x, y), cls.output_width(x, y))
        elif isinstance(x, core.Constant) and isinstance(y, Concat) and \
                isinstance(y.args[0], core.Constant):
            return Concat(Concat(x, y.args[0]), y.args[1])
        elif isinstance(y, core.Constant) and isinstance(x, Concat) and \
                isinstance(x.args[1], core.Constant):
            return Concat(x.args[0], Concat(x.args[1], y))
        elif isinstance(x, Extract) and isinstance(y, Extract):
            # x[5:4] concat x[3:2] = x[5:2]
            if x.args[0] == y.args[0] and x.args[2] == y.args[1] + 1:
                return Extract(x.args[0], x.args[1], y.args[2])
arxpy/primitives/primitives.py
class Cipher(object):
    """Represent (iterated) block ciphers.

    A (iterated) block cipher consists of `KeySchedule` function
    that computes round keys from a master key and an `Encryption`
    function that computes a ciphertext from a given plaintext
    and the round keys.

    Given a ``cipher``, it can be evaluated with the operator ``()``
    by passing it as arguments the plaintext and the master key,
    that is, ``cipher(plaintext, masterkey)`` returns the ciphertext.

        >>> from arxpy.primitives.primitives import Cipher
        >>> from arxpy.primitives import speck
        >>> Speck32 = speck.get_Speck_instance(speck.SpeckInstance.speck_32_64)
        >>> issubclass(Speck32, Cipher)
        True
        >>> plaintext = [0, 0]
        >>> masterkey = [0, 0, 0, 0]
        >>> Speck32(plaintext, masterkey)
        (0x2bb9, 0xc642)

    Attributes:
        key_schedule: the `KeySchedule` function of the cipher
        encryption: the `Encryption` function of the cipher

    """
    key_schedule = None
    encryption = None
    rounds = None

    _minimum_rounds = 1  # for testing

    def __new__(cls, plaintext, masterkey, **options):
        assert isinstance(plaintext, collections.abc.Sequence)
        assert isinstance(masterkey, collections.abc.Sequence)
        assert cls.rounds >= cls._minimum_rounds

        previous_round_keys = cls.encryption.round_keys

        round_keys = cls.key_schedule(*masterkey, **options)
        cls.encryption.round_keys = round_keys
        result = cls.encryption(*plaintext, **options)

        cls.encryption.round_keys = previous_round_keys

        return result

    @classmethod
    def set_rounds(cls, new_rounds):
        """Change the number of rounds and adjust the input/output widths."""
        raise NotImplementedError("subclasses need to override this method")


Question: import doctest
import collections
import functools
import math
import unittest
import arxpy.smt.search
from arxpy.bitvector.operation import Concat, BvComp
from arxpy.primitives.primitives import BvFunction, Cipher
from arxpy.differential.difference import XorDiff, RXDiff
from arxpy.differential.characteristic import (
    BvCharacteristic, SingleKeyCh, RelatedKeyCh
)
from arxpy.smt.search import (
    SearchCh, SearchSkCh, SearchRkCh, DerMode, ChSearchMode, SkChSearchMode, RkChSearchMode, _get_smart_print
)
from arxpy.differential.tests.test_characteristic import (
    MyFunction, Cipher1, Cipher2
)
"""Tests for the Optimization module."""







VERBOSE_LEVEL = 0
# 0: quiet
# 1: basic info
# 2: + ssa, model
# 3: + hrepr
# 4: + full hrepr


ChOption = collections.namedtuple('ChOptions', ['der_mode', 'search_mode'])
SkChOption = collections.namedtuple('SingleKeyOptions', ['der_mode', 'search_mode'])
RkChOption = collections.namedtuple('RelatedKeyOptions', ['key_der_mode', 'enc_der_mode',
                                                          'search_mode', 'initial_ew', 'initial_kw'])


CH_OPTIONS = [
    ChOption(DerMode.ProbabilityOne, ChSearchMode.Optimal),
    ChOption(DerMode.XDCA_Approx, ChSearchMode.FirstCh),
    ChOption(DerMode.Default, ChSearchMode.Optimal),
    ChOption(DerMode.Default, ChSearchMode.OptimalDifferential),
]

SKCH_OPTIONS = [
    SkChOption(DerMode.ProbabilityOne, SkChSearchMode.Optimal),
    SkChOption(DerMode.XDCA_Approx, SkChSearchMode.FirstCh),
    SkChOption(DerMode.Default, SkChSearchMode.Optimal),
    SkChOption(DerMode.Default, SkChSearchMode.OptimalDifferential),
]

RKCH_OPTIONS = [
    RkChOption(DerMode.ProbabilityOne, DerMode.ProbabilityOne, RkChSearchMode.OptimalMinSum, 0, 0),
    RkChOption(DerMode.ProbabilityOne, DerMode.ProbabilityOne, RkChSearchMode.OptimalMinSumDifferential, 0, 0),
    RkChOption(DerMode.Default, DerMode.Default, RkChSearchMode.FirstMinSum, 0, 0),
    RkChOption(DerMode.ProbabilityOne, DerMode.Default, RkChSearchMode.FirstValidKeyMinEnc, 0, 0),
    RkChOption(DerMode.Default, DerMode.Default, RkChSearchMode.OptimalFixEncMinKey, 1, 0),
]


def test_search_ch_skch(bvf_cipher, diff_type, initial_weight, solver_name, rounds, der_mode, search_mode, check,
                        verbose_level, filename):
    smart_print = _get_smart_print(filename)

    if rounds is not None:
        bvf_cipher.set_rounds(rounds)

    if issubclass(bvf_cipher, BvFunction):
        num_inputs = len(bvf_cipher.input_widths)
        input_diff_names = ["dp" + str(i) for i in range(num_inputs)]
        ch = BvCharacteristic(bvf_cipher, diff_type, input_diff_names)
    else:
        assert issubclass(bvf_cipher, Cipher)
        ch = SingleKeyCh(bvf_cipher, diff_type)

    if verbose_level >= 1:
        str_rounds = "" if rounds is None else "{} rounds".format(rounds)
        smart_print(str_rounds, bvf_cipher.__name__, diff_type.__name__, type(ch).__name__)
        if verbose_level >= 2:
            smart_print("Characteristic:")
            smart_print(ch)

    if issubclass(bvf_cipher, BvFunction):
        problem = SearchCh(ch, der_mode=der_mode)
    else:

Answer:         problem = SearchSkCh(ch, der_mode=der_mode)