angr.utils.cowdict 源代码

from __future__ import annotations
from collections import ChainMap




class ChainMapCOW(ChainMap):
    """
    Implements a copy-on-write version of ChainMap that supports auto-collapsing.
    """



    def __init__(self, *args, collapse_threshold=None):
        super().__init__(*args)
        self.dirty = False
        self.collapse_threshold = collapse_threshold




    def copy(self):
        self.dirty = True
        return self




    def clean(self):
        if self.dirty:
            # collapse?
            if self.collapse_threshold is not None and len(self.maps) >= self.collapse_threshold:
                collapsed = {}
                for m in reversed(self.maps):
                    collapsed.update(m)
                return ChainMapCOW(collapsed, collapse_threshold=self.collapse_threshold)
            ch = self.new_child()
            ch.collapse_threshold = self.collapse_threshold
            return ch
        return self






class DefaultChainMapCOW(ChainMapCOW):
    """
    Implements a copy-on-write version of ChainMap with default values that supports auto-collapsing.
    """



    def __init__(self, *args, default_factory=None, collapse_threshold=None):
        super().__init__(*args, collapse_threshold=collapse_threshold)
        self.default_factory = default_factory


    def __getitem__(self, key):
        try:
            return super().__getitem__(key)
        except KeyError:
            self.__setitem__(key, self.default_factory())
            return super().__getitem__(key)



    def clean(self):
        if self.dirty:
            # collapse?
            if self.collapse_threshold is not None and len(self.maps) >= self.collapse_threshold:
                collapsed = {}
                for m in reversed(self.maps):
                    collapsed.update(m)
                return DefaultChainMapCOW(
                    collapsed, default_factory=self.default_factory, collapse_threshold=self.collapse_threshold
                )
            r = self.new_child()
            r.default_factory = self.default_factory
            r.collapse_threshold = self.collapse_threshold
            return r
        return self