<commit_before>import collections
import claripy

class SimVariable(object):
    def __init__(self):
        pass

class SimTemporaryVariable(SimVariable):
    def __init__(self, tmp_id):
        SimVariable.__init__(self)

        self.tmp_id = tmp_id

    def __repr__(self):
        s = "<tmp %d>" % (self.tmp_id)

        return s

    def __hash__(self):
        return hash('tmp_%d' % (self.tmp_id))

    def __eq__(self, other):
        if isinstance(other, SimTemporaryVariable):
            return hash(self) == hash(other)

        else:
            return False

class SimRegisterVariable(SimVariable):
    def __init__(self, reg_offset, size):
        SimVariable.__init__(self)

        self.reg = reg_offset
        self.size = size

    def __repr__(self):
        s = "<%d %d>" % (self.reg, self.size)

        return s

    def __hash__(self):
        return hash('reg_%d_%d' % (self.reg, self.size))

    def __eq__(self, other):
        if isinstance(other, SimRegisterVariable):
            return hash(self) == hash(other)

        else:
            return False

class SimMemoryVariable(SimVariable):
    def __init__(self, addr, size):
        SimVariable.__init__(self)

        self.addr = addr

        if isinstance(size, claripy.ast.BV) and not size.symbolic:
            # Convert it to a concrete number
            size = size._model_concrete.value

        self.size = size

    def __repr__(self):
        if type(self.size) in (int, long):
            size = '%d' % self.size
        else:
            size = '%s' % self.size

        if type(self.addr) in (int, long):
            s = "<0x%x %s>" % (self.addr, size)
        else:
            s = "<%s %s>" % (self.addr, size)

        return s

    def __hash__(self):
        if isinstance(self.addr, claripy.ast.BV):
            addr_hash = hash(self.addr._model_concrete)
        else:
            addr_hash = hash(self.addr)
        return hash((addr_hash, hash(self.size)))

    def __eq__(self, other):
        if isinstance(other, SimMemoryVariable):
            return hash(self) == hash(other)

        else:
            return False

class SimVariableSet(collections.MutableSet):
    def __init__(self):
        self.register_variables = set()
        self.memory_variables = set()

    def add(self, item):
        if item not in self:
            if isinstance(item, SimRegisterVariable):
                self.register_variables.add(item)
            elif isinstance(item, SimMemoryVariable):
                self.memory_variables.add(item)
            else:
                # TODO:
                raise Exception('')

    def discard(self, item):
        if item in self:
            if isinstance(item, SimRegisterVariable):
                self.register_variables.discard(item)
            elif isinstance(item, SimMemoryVariable):
                self.memory_variables.discard(item)
            else:
                # TODO:
                raise Exception('')

    def __len__(self):
        return len(self.register_variables) + len(self.memory_variables)

    def __iter__(self):
        for i in self.register_variables: yield i
        for i in self.memory_variables: yield i

    def add_memory_variables(self, addrs, size):
        for a in addrs:
            var = SimMemoryVariable(a, size)
            self.add(var)

    def copy(self):
        s = SimVariableSet()
        s.register_variables |= self.register_variables
        s.memory_variables |= self.memory_variables

        return s

    def complement(self, other):
        """
        Calculate the complement of `self` and `other`
        :param other: Another SimVariableSet instance
        :return: The complement result
        """

        s = SimVariableSet()
        s.register_variables = self.register_variables - other.register_variables
        s.memory_variables = self.memory_variables - other.memory_variables

        return s

    def __contains__(self, item):
        if isinstance(item, SimRegisterVariable):
            for v in self.register_variables:
                # TODO: Make it better!
                if v.reg == item.reg:
                    return True
            return False
        elif isinstance(item, SimMemoryVariable):
            # TODO: Make it better!
            a = item.addr
            if isinstance(a, (tuple, list)): a = a[-1]

            for v in self.memory_variables:
                b = v.addr
                if isinstance(b, (tuple, list)): b = b[-1]

                if (isinstance(a, claripy.ast.Base) or isinstance(b, claripy.ast.Base)) and (a == b).is_true():
                    return True
                elif a == b:
                    return True<commit_msg>Properly using claripy backend model to calculate hash of variable addresses.<commit_after>import collections
import claripy

class SimVariable(object):
    def __init__(self):
        pass

class SimTemporaryVariable(SimVariable):
    def __init__(self, tmp_id):
        SimVariable.__init__(self)

        self.tmp_id = tmp_id

    def __repr__(self):
        s = "<tmp %d>" % (self.tmp_id)

        return s

    def __hash__(self):
        return hash('tmp_%d' % (self.tmp_id))

    def __eq__(self, other):
        if isinstance(other, SimTemporaryVariable):
            return hash(self) == hash(other)

        else:
            return False

class SimRegisterVariable(SimVariable):
    def __init__(self, reg_offset, size):
        SimVariable.__init__(self)

        self.reg = reg_offset
        self.size = size

    def __repr__(self):
        s = "<%d %d>" % (self.reg, self.size)

        return s

    def __hash__(self):
        return hash('reg_%d_%d' % (self.reg, self.size))

    def __eq__(self, other):
        if isinstance(other, SimRegisterVariable):
            return hash(self) == hash(other)

        else:
            return False

class SimMemoryVariable(SimVariable):
    def __init__(self, addr, size):
        SimVariable.__init__(self)

        self.addr = addr

        if isinstance(size, claripy.ast.BV) and not size.symbolic:
            # Convert it to a concrete number
            size = size._model_concrete.value

        self.size = size

    def __repr__(self):
        if type(self.size) in (int, long):
            size = '%d' % self.size
        else:
            size = '%s' % self.size

        if type(self.addr) in (int, long):
            s = "<0x%x %s>" % (self.addr, size)
        else:
            s = "<%s %s>" % (self.addr, size)

        return s

    def __hash__(self):
        if isinstance(self.addr, AddressWrapper):
            addr_hash = hash(self.addr)
        elif type(self.addr) in (int, long):
            addr_hash = self.addr
        elif self.addr._model_concrete is not self.addr:
            addr_hash = hash(self.addr._model_concrete)
        elif self.addr._model_vsa is not self.addr:
            addr_hash = hash(self.addr._model_vsa)
        elif self.addr._model_z3 is not self.addr:
            addr_hash = hash(self.addr._model_z3)
        else:
            addr_hash = hash(self.addr)
        return hash((addr_hash, hash(self.size)))

    def __eq__(self, other):
        if isinstance(other, SimMemoryVariable):
            return hash(self) == hash(other)

        else:
            return False

class SimVariableSet(collections.MutableSet):
    def __init__(self):
        self.register_variables = set()
        self.memory_variables = set()

    def add(self, item):
        if item not in self:
            if isinstance(item, SimRegisterVariable):
                self.register_variables.add(item)
            elif isinstance(item, SimMemoryVariable):
                self.memory_variables.add(item)
            else:
                # TODO:
                raise Exception('')

    def discard(self, item):
        if item in self:
            if isinstance(item, SimRegisterVariable):
                self.register_variables.discard(item)
            elif isinstance(item, SimMemoryVariable):
                self.memory_variables.discard(item)
            else:
                # TODO:
                raise Exception('')

    def __len__(self):
        return len(self.register_variables) + len(self.memory_variables)

    def __iter__(self):
        for i in self.register_variables: yield i
        for i in self.memory_variables: yield i

    def add_memory_variables(self, addrs, size):
        for a in addrs:
            var = SimMemoryVariable(a, size)
            self.add(var)

    def copy(self):
        s = SimVariableSet()
        s.register_variables |= self.register_variables
        s.memory_variables |= self.memory_variables

        return s

    def complement(self, other):
        """
        Calculate the complement of `self` and `other`
        :param other: Another SimVariableSet instance
        :return: The complement result
        """

        s = SimVariableSet()
        s.register_variables = self.register_variables - other.register_variables
        s.memory_variables = self.memory_variables - other.memory_variables

        return s

    def __contains__(self, item):
        if isinstance(item, SimRegisterVariable):
            for v in self.register_variables:
                # TODO: Make it better!
                if v.reg == item.reg:
                    return True
            return False
        elif isinstance(item, SimMemoryVariable):
            # TODO: Make it better!
            a = item.addr
            if isinstance(a, (tuple, list)): a = a[-1]

            for v in self.memory_variables:
                b = v.addr
                if isinstance(b, (tuple, list)): b = b[-1]

                if (isinstance(a, claripy.ast.Base) or isinstance(b, claripy.ast.Base)) and (a == b).is_true():
                    return True
                elif a == b:
                    return True

from .storage.memory import AddressWrapper
{'fraction_non_alphanumeric': 0.09834075273168758, 'fraction_numerical': 0.0008093889113719142, 'mean_word_length': 1.8911962562152675, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7519059', 'n_tokens_mistral': 2820, 'n_tokens_neox': 2710, 'n_words': 727}