<commit_before>import cffi
from typing import Tuple, Type, Dict, Optional, Iterable, List, Any, Set
import logging

from angr.storage.memory_mixins import MemoryMixin
from angr.storage.memory_mixins.paged_memory.pages import PageType, ListPage, UltraPage
from ....errors import SimMemoryError

# yeet
ffi = cffi.FFI()

l = logging.getLogger(__name__)


class PagedMemoryMixin(MemoryMixin):
    """
    A bottom-level storage mechanism. Dispatches reads to individual pages, the type of which is the PAGE_TYPE class
    variable.
    """
    PAGE_TYPE: Type[PageType] = None  # must be provided in subclass

    def __init__(self,  page_size=0x1000, default_permissions=3, permissions_map=None, page_kwargs=None, **kwargs):
        super().__init__(**kwargs)
        self.page_size = page_size
        self._extra_page_kwargs = page_kwargs if page_kwargs is not None else {}

        self._permissions_map = permissions_map if permissions_map is not None else {}
        self._default_permissions = default_permissions
        self._pages: Dict[int, Optional[PageType]] = {}

    def __del__(self):
        # a thought: we could support mapping pages in multiple places in memory if here we just kept a set of released
        # page ids and never released any page more than once
        for page in self._pages.values():
            page.release_shared()

    @MemoryMixin.memo
    def copy(self, memo):
        o = super().copy(memo)

        o.page_size = self.page_size
        o._pages = dict(self._pages)
        o._permissions_map = self._permissions_map
        o._default_permissions = self._default_permissions

        for page in o._pages.values():
            if page is not None:
                page.acquire_shared()

        return o

    def _get_page(self, pageno: int, writing: bool, **kwargs) -> PageType:
        force_default = True
        # force_default means don't consult any "backers"
        # if None is stored explicitly in _pages, it means it was unmapped explicitly, so don't consult backers
        try:
            page = self._pages[pageno]
        except KeyError:
            page = None
            force_default = False

        if page is None:
            page = self._initialize_page(pageno, force_default=force_default, **kwargs)
            self._pages[pageno] = page

        if writing:
            page = page.acquire_unique()
            self._pages[pageno] = page
        return page

    def _initialize_default_page(self, pageno: int, permissions=None, **kwargs) -> PageType:
        # the difference between _initialize_default_page and _initialize_page with force_default=True
        # is that the latter may segfault. this strictly gives you a new page.
        kwargs['allow_default'] = True
        return PagedMemoryMixin._initialize_page(self, pageno, permissions=permissions, **kwargs)

    def _initialize_page(self, pageno: int, permissions=None, allow_default=True, **kwargs) -> PageType:
        if not allow_default:
            raise SimMemoryError("I have been instructed not to create a default page")

        return self.PAGE_TYPE(**self._page_kwargs(pageno, permissions))

    def _page_kwargs(self, pageno, permissions):
        # permissions lookup: let permissions arg override everything else
        # then try the permissions map
        # then fall back to the default permissions
        if permissions is None:
            permissions = self._default_permissions
            addr = pageno * self.page_size

            for (start, end), perms in self._permissions_map.items():
                if start <= addr <= end:
                    permissions = perms
                    break

        return dict(
            memory=self,
            memory_id='%s_%d' % (self.id, pageno),
            permissions=permissions,
            **self._extra_page_kwargs
        )

    def _divide_addr(self, addr: int) -> Tuple[int, int]:
        return divmod(addr, self.page_size)

    def load(self, addr: int, size: int=None, endness=None, **kwargs):
        if endness is None:
            endness = self.endness

        if type(size) is not int:
            raise TypeError("Need size to be resolved to an int by this point")

        if type(addr) is not int:
            raise TypeError("Need addr to be resolved to an int by this point")

        pageno, pageoff = self._divide_addr(addr)
        vals = []

        # fasttrack basic case
        if pageoff + size <= self.page_size:
            page = self._get_page(pageno, False, **kwargs)
            vals.append(page.load(pageoff, size=size, endness=endness, page_addr=pageno*self.page_size, memory=self, cooperate=True, **kwargs))

        else:
            max_pageno = (1 << self.state.arch.bits) // self.page_size
            bytes_done = 0
            while bytes_done < size:
                page = self._get_page(pageno, False, **kwargs)
                sub_size = min(self.page_size-pageoff, size-bytes_done)
                vals.append(page.load(pageoff, size=sub_size, endness=endness, page_addr=pageno*self.page_size, memory=self, cooperate=True, **kwargs))

                bytes_done += sub_size
                pageno = (pageno + 1) % max_pageno
                pageoff = 0

        out = self.PAGE_TYPE._compose_objects(vals, size, endness, memory=self, **kwargs)
        l.debug("%s.load(%#x, %d, %s) = %s", self.id, addr, size, endness, out)
        return out

    def store(self, addr: int, data, size: int=None, endness=None, **kwargs):
        if endness is None:
            endness = self.endness

        if type(size) is not int:
            raise TypeError("Need size to be resolved to an int by this point")

        if type(addr) is not int:
            raise TypeError("Need addr to be resolved to an int by this point")

        l.debug("%s.store(%#x, %s, %s)", self.id, addr, data, endness)

        pageno, pageoff = self._divide_addr(addr)
        sub_gen = self.PAGE_TYPE._decompose_objects(addr, data, endness, memory=self, **kwargs)
        next(sub_gen)

        # fasttrack basic case
        if pageoff + size <= self.page_size:
            sub_data = sub_gen.send(size)
            page = self._get_page(pageno, True, **kwargs)
            page.store(pageoff, sub_data, size=size, endness=endness, page_addr=pageno*self.page_size, memory=self, cooperate=True, **kwargs)
            sub_gen.close()
            return

        max_pageno = (1 << self.state.arch.bits) // self.page_size
        bytes_done = 0
        while bytes_done < size:
            # if we really want we could add an optimization where writing an entire page creates a new page object
            # instead of overwriting the old one. would have to be careful about maintaining the contract for _get_page
            page = self._get_page(pageno, True, **kwargs)
            sub_size = min(self.page_size-pageoff, size-bytes_done)

            sub_data = sub_gen.send(sub_size)

            page.store(pageoff, sub_data, size=sub_size, endness=endness, page_addr=pageno*self.page_size, memory=self, cooperate=True, **kwargs)

            bytes_done += sub_size
            pageno = (pageno + 1) % max_pageno
            pageoff = 0

        sub_gen.close()

    def merge(self, others: Iterable['PagedMemoryMixin'], merge_conditions, common_ancestor=None):
        changed_pages = set()
        for o in others:
            changed_pages |= self.changed_pages(o)

        self._merge(others, changed_pages, merge_conditions)

    def _merge(self, others: Iterable['PagedMemoryMixin'], changed_pages: Iterable[int], merge_conditions):

        if merge_conditions is None:<commit_msg>Address Audrey's comments.<commit_after>import cffi
from typing import Tuple, Type, Dict, Optional, Iterable, List, Any, Set
import logging

from angr.storage.memory_mixins import MemoryMixin
from angr.storage.memory_mixins.paged_memory.pages import PageType, ListPage, UltraPage
from ....errors import SimMemoryError

# yeet
ffi = cffi.FFI()

l = logging.getLogger(__name__)


class PagedMemoryMixin(MemoryMixin):
    """
    A bottom-level storage mechanism. Dispatches reads to individual pages, the type of which is the PAGE_TYPE class
    variable.
    """
    PAGE_TYPE: Type[PageType] = None  # must be provided in subclass

    def __init__(self,  page_size=0x1000, default_permissions=3, permissions_map=None, page_kwargs=None, **kwargs):
        super().__init__(**kwargs)
        self.page_size = page_size
        self._extra_page_kwargs = page_kwargs if page_kwargs is not None else {}

        self._permissions_map = permissions_map if permissions_map is not None else {}
        self._default_permissions = default_permissions
        self._pages: Dict[int, Optional[PageType]] = {}

    def __del__(self):
        # a thought: we could support mapping pages in multiple places in memory if here we just kept a set of released
        # page ids and never released any page more than once
        for page in self._pages.values():
            page.release_shared()

    @MemoryMixin.memo
    def copy(self, memo):
        o = super().copy(memo)

        o.page_size = self.page_size
        o._pages = dict(self._pages)
        o._permissions_map = self._permissions_map
        o._default_permissions = self._default_permissions

        for page in o._pages.values():
            if page is not None:
                page.acquire_shared()

        return o

    def _get_page(self, pageno: int, writing: bool, **kwargs) -> PageType:
        force_default = True
        # force_default means don't consult any "backers"
        # if None is stored explicitly in _pages, it means it was unmapped explicitly, so don't consult backers
        try:
            page = self._pages[pageno]
        except KeyError:
            page = None
            force_default = False

        if page is None:
            page = self._initialize_page(pageno, force_default=force_default, **kwargs)
            self._pages[pageno] = page

        if writing:
            page = page.acquire_unique()
            self._pages[pageno] = page
        return page

    def _initialize_default_page(self, pageno: int, permissions=None, **kwargs) -> PageType:
        # the difference between _initialize_default_page and _initialize_page with force_default=True
        # is that the latter may segfault. this strictly gives you a new page.
        kwargs['allow_default'] = True
        return PagedMemoryMixin._initialize_page(self, pageno, permissions=permissions, **kwargs)

    def _initialize_page(self, pageno: int, permissions=None, allow_default=True, **kwargs) -> PageType:
        if not allow_default:
            raise SimMemoryError("I have been instructed not to create a default page")

        return self.PAGE_TYPE(**self._page_kwargs(pageno, permissions))

    def _page_kwargs(self, pageno, permissions):
        # permissions lookup: let permissions arg override everything else
        # then try the permissions map
        # then fall back to the default permissions
        if permissions is None:
            permissions = self._default_permissions
            addr = pageno * self.page_size

            for (start, end), perms in self._permissions_map.items():
                if start <= addr <= end:
                    permissions = perms
                    break

        return dict(
            memory=self,
            memory_id='%s_%d' % (self.id, pageno),
            permissions=permissions,
            **self._extra_page_kwargs
        )

    def _divide_addr(self, addr: int) -> Tuple[int, int]:
        return divmod(addr, self.page_size)

    def load(self, addr: int, size: int=None, endness=None, **kwargs):
        if endness is None:
            endness = self.endness

        if type(size) is not int:
            raise TypeError("Need size to be resolved to an int by this point")

        if type(addr) is not int:
            raise TypeError("Need addr to be resolved to an int by this point")

        pageno, pageoff = self._divide_addr(addr)
        vals = []

        # fasttrack basic case
        if pageoff + size <= self.page_size:
            page = self._get_page(pageno, False, **kwargs)
            vals.append(page.load(pageoff, size=size, endness=endness, page_addr=pageno*self.page_size, memory=self, cooperate=True, **kwargs))

        else:
            max_pageno = (1 << self.state.arch.bits) // self.page_size
            bytes_done = 0
            while bytes_done < size:
                page = self._get_page(pageno, False, **kwargs)
                sub_size = min(self.page_size-pageoff, size-bytes_done)
                vals.append(page.load(pageoff, size=sub_size, endness=endness, page_addr=pageno*self.page_size, memory=self, cooperate=True, **kwargs))

                bytes_done += sub_size
                pageno = (pageno + 1) % max_pageno
                pageoff = 0

        out = self.PAGE_TYPE._compose_objects(vals, size, endness, memory=self, **kwargs)
        l.debug("%s.load(%#x, %d, %s) = %s", self.id, addr, size, endness, out)
        return out

    def store(self, addr: int, data, size: int=None, endness=None, **kwargs):
        if endness is None:
            endness = self.endness

        if type(size) is not int:
            raise TypeError("Need size to be resolved to an int by this point")

        if type(addr) is not int:
            raise TypeError("Need addr to be resolved to an int by this point")

        l.debug("%s.store(%#x, %s, %s)", self.id, addr, data, endness)

        pageno, pageoff = self._divide_addr(addr)
        sub_gen = self.PAGE_TYPE._decompose_objects(addr, data, endness, memory=self, **kwargs)
        next(sub_gen)

        # fasttrack basic case
        if pageoff + size <= self.page_size:
            sub_data = sub_gen.send(size)
            page = self._get_page(pageno, True, **kwargs)
            page.store(pageoff, sub_data, size=size, endness=endness, page_addr=pageno*self.page_size, memory=self, cooperate=True, **kwargs)
            sub_gen.close()
            return

        max_pageno = (1 << self.state.arch.bits) // self.page_size
        bytes_done = 0
        while bytes_done < size:
            # if we really want we could add an optimization where writing an entire page creates a new page object
            # instead of overwriting the old one. would have to be careful about maintaining the contract for _get_page
            page = self._get_page(pageno, True, **kwargs)
            sub_size = min(self.page_size-pageoff, size-bytes_done)

            sub_data = sub_gen.send(sub_size)

            page.store(pageoff, sub_data, size=sub_size, endness=endness, page_addr=pageno*self.page_size, memory=self, cooperate=True, **kwargs)

            bytes_done += sub_size
            pageno = (pageno + 1) % max_pageno
            pageoff = 0

        sub_gen.close()

    def merge(self, others: Iterable['PagedMemoryMixin'], merge_conditions, common_ancestor=None):
        changed_pages = set()
        for o in others:
            changed_pages |= self.changed_pages(o)

        if merge_conditions is None:
{'fraction_non_alphanumeric': 0.09160153662736786, 'fraction_numerical': 0.0018545502715591469, 'mean_word_length': 2.481438782568596, 'pattern_counts': {'":': 0, '<': 23, '<?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': '7658597', 'n_tokens_mistral': 4476, 'n_tokens_neox': 4334, 'n_words': 1285}