<commit_before>import collections
import cooldict
import claripy
import cffi
import cle

from ..s_errors import SimMemoryError
from .. import s_options as options
from .memory_object import SimMemoryObject

_ffi = cffi.FFI()

import logging
l = logging.getLogger('simuvex.storage.paged_memory')

#pylint:disable=unidiomatic-typecheck

class SimPagedMemory(collections.MutableMapping):
    def __init__(self, backer=None, pages=None, name_mapping=None, hash_mapping=None, page_size=None):
        self._backer = { } if backer is None else backer
        self._pages = { } if pages is None else pages
        self._page_size = 0x1000 if page_size is None else page_size
        self.state = None

        # reverse mapping
        self._name_mapping = cooldict.BranchingDict() if name_mapping is None else name_mapping
        self._hash_mapping = cooldict.BranchingDict() if hash_mapping is None else hash_mapping
        self._updated_mappings = set()

    def __getstate__(self):
        return {
            '_backer': self._backer,
            '_pages': self._pages,
            '_page_size': self._page_size,
            'state': self.state,
            '_name_mapping': self._name_mapping,
            '_hash_mapping': self._hash_mapping,
        }

    def __setstate__(self, s):
        self.__dict__.update(s)

    def branch(self):
        new_pages = { k:v.branch() for k,v in self._pages.iteritems() }
        new_name_mapping = self._name_mapping.branch() if options.REVERSE_MEMORY_NAME_MAP in self.state.options else self._name_mapping
        new_hash_mapping = self._hash_mapping.branch() if options.REVERSE_MEMORY_HASH_MAP in self.state.options else self._hash_mapping

        m = SimPagedMemory(backer=self._backer,
                           pages=new_pages,
                           page_size=self._page_size,
                           name_mapping=new_name_mapping,
                           hash_mapping=new_hash_mapping)
        return m

    def __getitem__(self, addr):
        page_num = addr / self._page_size
        page_idx = addr % self._page_size
        #print "GET", addr, page_num, page_idx

        if page_num not in self._pages:
            self._initialize_page(page_num)
        return self._pages[page_num][page_idx]

    def __setitem__(self, addr, v):
        page_num = addr / self._page_size
        page_idx = addr % self._page_size
        #print "SET", addr, page_num, page_idx

        self._update_mappings(addr, v.object)
        if page_num not in self._pages:
            self._initialize_page(page_num)
        self._pages[page_num][page_idx] = v
        #print "...",id(self._pages[page_num])

    def __delitem__(self, addr):
        raise Exception("For performance reasons, deletion is not supported. Contact Yan if this needs to change.")
        # Specifically, the above is for two reasons:
        #
        #    1. deleting stuff out of memory doesn't make sense
        #    2. if the page throws a key error, the backer dict is accessed. Thus, deleting things would simply
        #       change them back to what they were in the backer dict

        #page_num = addr / self._page_size
        #page_idx = addr % self._page_size
        ##print "DEL", addr, page_num, page_idx

        #if page_num not in self._pages:
        #    self._pages[page_num] = cooldict.BranchingDict(d=self._backer)
        #del self._pages[page_num][page_idx]

    def store_memory_object(self, mo, overwrite=True):
        '''
        This function optimizes a large store by storing a single reference to
        the SimMemoryObject instead of one for each byte.

        @param memory_object: the memory object to store
        '''

        self._update_range_mappings(mo.base, mo.object, mo.length)

        mo_start = mo.base
        mo_end = mo.base + mo.length
        page_start = mo_start - mo_start%self._page_size
        page_end = mo_end + (self._page_size - mo_end%self._page_size) if mo_end % self._page_size else mo_end
        pages = [ b for b in range(page_start, page_end, self._page_size) ]

        for p in pages:
            page_num = p / self._page_size
            if page_num not in self._pages:
                self._initialize_page(page_num)
            page = self._pages[page_num]

            #print (mo.base, mo.length), (p, p + self._page_size)

            if mo.base <= p and mo.base + mo.length >= p + self._page_size:
                # takes up the whole page
                page.sinkhole(mo, wipe=overwrite)
            else:
                for a in range(max(mo.base, p), min(mo.base+mo.length, p+self._page_size)):
                    if overwrite or a%self._page_size not in page:
                        page[a%self._page_size] = mo

    def _initialize_page(self, n):
        new_page = cooldict.SinkholeCOWDict()
        self._pages[n] = new_page

        new_page_addr = n*self._page_size
        if self._backer is None:
            pass
        elif isinstance(self._backer, cle.Clemory):
            # first, find the right clemory backer
            for addr, backer in self._backer.cbackers:
                start_backer = new_page_addr - addr

                if start_backer < 0 and abs(start_backer) > self._page_size:
                    continue
                if start_backer > len(backer):
                    continue

                snip_start = max(0, start_backer)
                write_start = max(new_page_addr, addr + snip_start)
                write_size = self._page_size - write_start%self._page_size

                #print hex(addr), hex(snip_start), write_size, hex(write_start)
                #import ipdb; ipdb.set_trace()

                snip = _ffi.buffer(backer)[snip_start:snip_start+write_size]
                mo = SimMemoryObject(claripy.BVV(snip), write_start)
                self.store_memory_object(mo)
        elif len(self._backer) < self._page_size:
            for i in self._backer:
                if new_page_addr <= i and i <= new_page_addr + self._page_size:
                    self.store_memory_object(SimMemoryObject(claripy.BVV(self._backer[i]), i))
        elif len(self._backer) > self._page_size:
            for i in range(self._page_size):
                try:
                    self.store_memory_object(SimMemoryObject(claripy.BVV(self._backer[i]), new_page_addr+i))
                except KeyError:
                    pass

    def __contains__(self, addr):
        try:
            self.__getitem__(addr)
            return True
        except KeyError:
            return False

    def __iter__(self):
        for k in self._backer:
            yield k
        for p in self._pages:
            for a in self._pages[p]:
                yield p*self._page_size + a

    def __len__(self):
        return len(self._backer) + sum(len(v) for v in self._pages.itervalues())

    def changed_bytes(self, other):
        '''
        Gets the set of changed bytes between self and other.

        @param other: the other SimPagedMemory
        @returns a set of differing bytes
        '''
        if self._page_size != other._page_size:
            raise SimMemoryError("SimPagedMemory page sizes differ. This is asking for disaster.")
<commit_msg>cleaned up SimPagedMemory's iterators and len() in relation to sinkholing<commit_after>import collections
import cooldict
import claripy
import cffi
import cle

from ..s_errors import SimMemoryError
from .. import s_options as options
from .memory_object import SimMemoryObject

_ffi = cffi.FFI()

import logging
l = logging.getLogger('simuvex.storage.paged_memory')

#pylint:disable=unidiomatic-typecheck

class SimPagedMemory(collections.MutableMapping):
    def __init__(self, backer=None, pages=None, name_mapping=None, hash_mapping=None, page_size=None):
        self._backer = { } if backer is None else backer
        self._pages = { } if pages is None else pages
        self._page_size = 0x1000 if page_size is None else page_size
        self.state = None

        # reverse mapping
        self._name_mapping = cooldict.BranchingDict() if name_mapping is None else name_mapping
        self._hash_mapping = cooldict.BranchingDict() if hash_mapping is None else hash_mapping
        self._updated_mappings = set()

    def __getstate__(self):
        return {
            '_backer': self._backer,
            '_pages': self._pages,
            '_page_size': self._page_size,
            'state': self.state,
            '_name_mapping': self._name_mapping,
            '_hash_mapping': self._hash_mapping,
        }

    def __setstate__(self, s):
        self.__dict__.update(s)

    def branch(self):
        new_pages = { k:v.branch() for k,v in self._pages.iteritems() }
        new_name_mapping = self._name_mapping.branch() if options.REVERSE_MEMORY_NAME_MAP in self.state.options else self._name_mapping
        new_hash_mapping = self._hash_mapping.branch() if options.REVERSE_MEMORY_HASH_MAP in self.state.options else self._hash_mapping

        m = SimPagedMemory(backer=self._backer,
                           pages=new_pages,
                           page_size=self._page_size,
                           name_mapping=new_name_mapping,
                           hash_mapping=new_hash_mapping)
        return m

    def __getitem__(self, addr):
        page_num = addr / self._page_size
        page_idx = addr % self._page_size
        #print "GET", addr, page_num, page_idx

        if page_num not in self._pages:
            self._initialize_page(page_num)
        return self._pages[page_num][page_idx]

    def __setitem__(self, addr, v):
        page_num = addr / self._page_size
        page_idx = addr % self._page_size
        #print "SET", addr, page_num, page_idx

        self._update_mappings(addr, v.object)
        if page_num not in self._pages:
            self._initialize_page(page_num)
        self._pages[page_num][page_idx] = v
        #print "...",id(self._pages[page_num])

    def __delitem__(self, addr):
        raise Exception("For performance reasons, deletion is not supported. Contact Yan if this needs to change.")
        # Specifically, the above is for two reasons:
        #
        #    1. deleting stuff out of memory doesn't make sense
        #    2. if the page throws a key error, the backer dict is accessed. Thus, deleting things would simply
        #       change them back to what they were in the backer dict

        #page_num = addr / self._page_size
        #page_idx = addr % self._page_size
        ##print "DEL", addr, page_num, page_idx

        #if page_num not in self._pages:
        #    self._pages[page_num] = cooldict.BranchingDict(d=self._backer)
        #del self._pages[page_num][page_idx]

    def store_memory_object(self, mo, overwrite=True):
        '''
        This function optimizes a large store by storing a single reference to
        the SimMemoryObject instead of one for each byte.

        @param memory_object: the memory object to store
        '''

        self._update_range_mappings(mo.base, mo.object, mo.length)

        mo_start = mo.base
        mo_end = mo.base + mo.length
        page_start = mo_start - mo_start%self._page_size
        page_end = mo_end + (self._page_size - mo_end%self._page_size) if mo_end % self._page_size else mo_end
        pages = [ b for b in range(page_start, page_end, self._page_size) ]

        for p in pages:
            page_num = p / self._page_size
            if page_num not in self._pages:
                self._initialize_page(page_num)
            page = self._pages[page_num]

            #print (mo.base, mo.length), (p, p + self._page_size)

            if mo.base <= p and mo.base + mo.length >= p + self._page_size:
                # takes up the whole page
                page.sinkhole(mo, wipe=overwrite)
            else:
                for a in range(max(mo.base, p), min(mo.base+mo.length, p+self._page_size)):
                    if overwrite or a%self._page_size not in page:
                        page[a%self._page_size] = mo

    def _initialize_page(self, n):
        new_page = cooldict.SinkholeCOWDict()
        self._pages[n] = new_page

        new_page_addr = n*self._page_size
        if self._backer is None:
            pass
        elif isinstance(self._backer, cle.Clemory):
            # first, find the right clemory backer
            for addr, backer in self._backer.cbackers:
                start_backer = new_page_addr - addr

                if start_backer < 0 and abs(start_backer) > self._page_size:
                    continue
                if start_backer > len(backer):
                    continue

                snip_start = max(0, start_backer)
                write_start = max(new_page_addr, addr + snip_start)
                write_size = self._page_size - write_start%self._page_size

                #print hex(addr), hex(snip_start), write_size, hex(write_start)
                #import ipdb; ipdb.set_trace()

                snip = _ffi.buffer(backer)[snip_start:snip_start+write_size]
                mo = SimMemoryObject(claripy.BVV(snip), write_start)
                self.store_memory_object(mo)
        elif len(self._backer) < self._page_size:
            for i in self._backer:
                if new_page_addr <= i and i <= new_page_addr + self._page_size:
                    self.store_memory_object(SimMemoryObject(claripy.BVV(self._backer[i]), i))
        elif len(self._backer) > self._page_size:
            for i in range(self._page_size):
                try:
                    self.store_memory_object(SimMemoryObject(claripy.BVV(self._backer[i]), new_page_addr+i))
                except KeyError:
                    pass

    def __contains__(self, addr):
        try:
            self.__getitem__(addr)
            return True
        except KeyError:
            return False

    def __iter__(self):
        for k in self._backer:
            yield k
        for p in self._pages:
            if not self._pages[p]._sinkholed:
                for a in self._pages[p]:
                    yield p*self._page_size + a
            else:
                for a in range(self._page_size):
                    yield p*self._page_size + a

    def __len__(self):
        return sum((len(v) if not self._pages._sinkholed else self._page_size) for v in self._pages.itervalues())

    def changed_bytes(self, other):
        '''
        Gets the set of changed bytes between self and other.

        @param other: the other SimPagedMemory
        @returns a set of differing bytes
        '''
        if self._page_size != other._page_size:
            raise SimMemoryError("SimPagedMemory page sizes differ. This is asking for disaster.")

{'fraction_non_alphanumeric': 0.08399182561307902, 'fraction_numerical': 0.001226158038147139, 'mean_word_length': 2.107089947089947, 'pattern_counts': {'":': 0, '<': 13, '<?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': '340269', 'n_tokens_mistral': 4431, 'n_tokens_neox': 4286, 'n_words': 1122}