<commit_before>import weakref
import itertools

class ModelCache(object):
    _defaults = { 0, 0.0, True }

    def __init__(self, model):
        self.model = model
        self.replacements = weakref.WeakKeyDictionary()

    def __hash__(self):
        if not hasattr(self, '_hash'):
            self._hash = hash(frozenset(self.model.items())) #pylint:disable=attribute-defined-outside-init
        return self._hash

    def __eq__(self, other):
        return self.model == other.model

    def __getstate__(self):
        return (self.model,)

    def __setstate__(self, s):
        self.model = s[0]
        self.replacements = weakref.WeakKeyDictionary()

    #
    # Splitting support
    #

    def filter(self, variables):
        return ModelCache({ k:v for k,v in self.model.iteritems() if k in variables })

    @staticmethod
    def combine(*models):
        return ModelCache(dict(itertools.chain.from_iterable(m.model.iteritems() for m in models)))

    #
    # Model-driven evaluation
    #

    def _leaf_op(self, a):
        return (
            all_operations.BVV(self.model.get(a.args[0], 0), a.length) if a.op == 'BVS' else
            all_operations.BoolV(self.model.get(a.args[0], True)) if a.op == 'BoolS' else
            all_operations.FPV(self.model.get(a.args[0], 0.0), a.args[1]) if a.op == 'FPS' else
            a
        )

    def eval_ast(self, ast):
        new_ast = ast._replace(self.replacements, leaf_operation=self._leaf_op)
        return backends.concrete.eval(new_ast, 1)[0]

    def eval_constraints(self, constraints):
        return all(self.eval_ast(c) for c in constraints)

    def eval_list(self, asts):
        return tuple(self.eval_ast(c) for c in asts )

class ModelCacheMixin(object):
    def __init__(self, *args, **kwargs):
        super(ModelCacheMixin, self).__init__(*args, **kwargs)
        self._models = set()
        self._exhausted = False
        self._eval_exhausted = weakref.WeakKeyDictionary()
        self._max_exhausted = weakref.WeakKeyDictionary()
        self._min_exhausted = weakref.WeakKeyDictionary()

    def _blank_copy(self, c):
        super(ModelCacheMixin, self)._blank_copy(c)
        c._models = set()
        c._exhausted = False
        c._eval_exhausted = weakref.WeakKeyDictionary()
        c._max_exhausted = weakref.WeakKeyDictionary()
        c._min_exhausted = weakref.WeakKeyDictionary()

    def _copy(self, c):
        super(ModelCacheMixin, self)._copy(c)
        c._models = set(self._models)
        c._exhausted = self._exhausted
        c._eval_exhausted = weakref.WeakKeyDictionary(self._eval_exhausted)
        c._max_exhausted = weakref.WeakKeyDictionary(self._max_exhausted)
        c._min_exhausted = weakref.WeakKeyDictionary(self._min_exhausted)

    def _ana_getstate(self):
        return (
            self._models,
            self._exhausted,
            self._eval_exhausted.keys(),
            self._max_exhausted.keys(),
            self._min_exhausted.keys(),
            super(ModelCacheMixin, self)._ana_getstate()
        )

    def _ana_setstate(self, s):
        (
            self._models,
            self._exhausted,
            _eval_exhausted,
            _max_exhausted,
            _min_exhausted,
            base_state
        ) = s
        super(ModelCacheMixin, self)._ana_setstate(base_state)
        self._eval_exhausted = weakref.WeakKeyDictionary((k,True) for k in _eval_exhausted)
        self._max_exhausted = weakref.WeakKeyDictionary((k,True) for k in _max_exhausted)
        self._min_exhausted = weakref.WeakKeyDictionary((k,True) for k in _min_exhausted)

    #
    # Model cleaning
    #

    def add(self, constraints, invalidate_cache=True, **kwargs):
        if len(constraints) == 0:
            return constraints

        old_vars = frozenset(self.variables)
        added = super(ModelCacheMixin, self).add(constraints, **kwargs)
        if len(added) == 0:
            return added

        new_vars = any(a.variables - old_vars for a in added)
        if new_vars or invalidate_cache:
            still_valid = set(self._get_models(extra_constraints=added))
            if len(still_valid) != len(self._models):
                self._exhausted = False
                self._eval_exhausted.clear()
                self._max_exhausted.clear()
                self._min_exhausted.clear()
                self._models = still_valid

        return added

    def split(self):
        results = super(ModelCacheMixin, self).split()
        for r in results:
            r._models = { m.filter(r.variables) for m in self._models }
        return results

    def combine(self, others):
        combined = super(ModelCacheMixin, self).combine(others)

        if any(len(o._models) == 0 for o in others) or len(self._models) == 0:
            # this would need a solve anyways, so screw it
            return combined

        vars_count = len(self.variables) + sum(len(s.variables) for s in others)
        all_vars = self.variables.union(*[s.variables for s in others])
        if vars_count != len(all_vars):
            # this is the case where there are variables missing from the models.
            # We'll need more intelligence here to handle it
            return combined

        models = [ self._models ]
        models.extend(o._models for o in others)
        combined._models.update(ModelCache.combine(*product) for product in itertools.product(*models))
        return combined

    def update(self, other):
        """
        Updates this cache mixin with results discovered by the other split off one.
        """

        self._models.update(other._models)
        self._eval_exhausted.update(other._eval_exhausted)
        self._max_exhausted.update(other._max_exhausted)
        self._min_exhausted.update(other._min_exhausted)

    #
    # Cache retrieval
    #

    def _model_hook(self, m):
        self._models.add(ModelCache(m))

    def _get_models(self, extra_constraints=()):
        for m in self._models:
            if m.eval_constraints(extra_constraints):
                yield m

    def _get_batch_solutions(self, asts, n=None, extra_constraints=()):
        results = set()

        for m in self._get_models(extra_constraints):
            results.add(m.eval_list(asts))
            if len(results) == n:
                break

        return results

    def _get_solutions(self, e, n=None, extra_constraints=()):
        return tuple(v[0] for v in self._get_batch_solutions(
            [e], n=n, extra_constraints=extra_constraints
        ))


    #
    # Cached functions
    #


    def satisfiable(self, extra_constraints=(), **kwargs):
        for _ in self._get_models(extra_constraints=extra_constraints):
            return True<commit_msg>put a limit on the number of models that are kept around, to avoid product explosions for composite solvers<commit_after>import weakref
import itertools

class ModelCache(object):
    _defaults = { 0, 0.0, True }

    def __init__(self, model):
        self.model = model
        self.replacements = weakref.WeakKeyDictionary()

    def __hash__(self):
        if not hasattr(self, '_hash'):
            self._hash = hash(frozenset(self.model.items())) #pylint:disable=attribute-defined-outside-init
        return self._hash

    def __eq__(self, other):
        return self.model == other.model

    def __getstate__(self):
        return (self.model,)

    def __setstate__(self, s):
        self.model = s[0]
        self.replacements = weakref.WeakKeyDictionary()

    #
    # Splitting support
    #

    def filter(self, variables):
        return ModelCache({ k:v for k,v in self.model.iteritems() if k in variables })

    @staticmethod
    def combine(*models):
        return ModelCache(dict(itertools.chain.from_iterable(m.model.iteritems() for m in models)))

    #
    # Model-driven evaluation
    #

    def _leaf_op(self, a):
        return (
            all_operations.BVV(self.model.get(a.args[0], 0), a.length) if a.op == 'BVS' else
            all_operations.BoolV(self.model.get(a.args[0], True)) if a.op == 'BoolS' else
            all_operations.FPV(self.model.get(a.args[0], 0.0), a.args[1]) if a.op == 'FPS' else
            a
        )

    def eval_ast(self, ast):
        new_ast = ast._replace(self.replacements, leaf_operation=self._leaf_op)
        return backends.concrete.eval(new_ast, 1)[0]

    def eval_constraints(self, constraints):
        return all(self.eval_ast(c) for c in constraints)

    def eval_list(self, asts):
        return tuple(self.eval_ast(c) for c in asts )

class ModelCacheMixin(object):
    _MODEL_LIMIT = 257

    def __init__(self, *args, **kwargs):
        super(ModelCacheMixin, self).__init__(*args, **kwargs)
        self._models = set()
        self._exhausted = False
        self._eval_exhausted = weakref.WeakKeyDictionary()
        self._max_exhausted = weakref.WeakKeyDictionary()
        self._min_exhausted = weakref.WeakKeyDictionary()

    def _blank_copy(self, c):
        super(ModelCacheMixin, self)._blank_copy(c)
        c._models = set()
        c._exhausted = False
        c._eval_exhausted = weakref.WeakKeyDictionary()
        c._max_exhausted = weakref.WeakKeyDictionary()
        c._min_exhausted = weakref.WeakKeyDictionary()

    def _copy(self, c):
        super(ModelCacheMixin, self)._copy(c)
        c._models = set(self._models)
        c._exhausted = self._exhausted
        c._eval_exhausted = weakref.WeakKeyDictionary(self._eval_exhausted)
        c._max_exhausted = weakref.WeakKeyDictionary(self._max_exhausted)
        c._min_exhausted = weakref.WeakKeyDictionary(self._min_exhausted)

    def _ana_getstate(self):
        return (
            self._models,
            self._exhausted,
            self._eval_exhausted.keys(),
            self._max_exhausted.keys(),
            self._min_exhausted.keys(),
            super(ModelCacheMixin, self)._ana_getstate()
        )

    def _ana_setstate(self, s):
        (
            self._models,
            self._exhausted,
            _eval_exhausted,
            _max_exhausted,
            _min_exhausted,
            base_state
        ) = s
        super(ModelCacheMixin, self)._ana_setstate(base_state)
        self._eval_exhausted = weakref.WeakKeyDictionary((k,True) for k in _eval_exhausted)
        self._max_exhausted = weakref.WeakKeyDictionary((k,True) for k in _max_exhausted)
        self._min_exhausted = weakref.WeakKeyDictionary((k,True) for k in _min_exhausted)

    #
    # Model cleaning
    #

    def add(self, constraints, invalidate_cache=True, **kwargs):
        if len(constraints) == 0:
            return constraints

        old_vars = frozenset(self.variables)
        added = super(ModelCacheMixin, self).add(constraints, **kwargs)
        if len(added) == 0:
            return added

        new_vars = any(a.variables - old_vars for a in added)
        if new_vars or invalidate_cache:
            still_valid = set(self._get_models(extra_constraints=added))
            if len(still_valid) != len(self._models):
                self._exhausted = False
                self._eval_exhausted.clear()
                self._max_exhausted.clear()
                self._min_exhausted.clear()
                self._models = still_valid

        return added

    def split(self):
        results = super(ModelCacheMixin, self).split()
        for r in results:
            r._models = { m.filter(r.variables) for m in self._models }
        return results

    def combine(self, others):
        combined = super(ModelCacheMixin, self).combine(others)

        if any(len(o._models) == 0 for o in others) or len(self._models) == 0:
            # this would need a solve anyways, so screw it
            return combined

        vars_count = len(self.variables) + sum(len(s.variables) for s in others)
        all_vars = self.variables.union(*[s.variables for s in others])
        if vars_count != len(all_vars):
            # this is the case where there are variables missing from the models.
            # We'll need more intelligence here to handle it
            return combined

        model_lists = [ self._models ]
        model_lists.extend(o._models for o in others)
        combined._models.update(
            ModelCache.combine(*product) for product in
            itertools.islice(itertools.product(*model_lists), self._MODEL_LIMIT)
        )
        return combined

    def update(self, other):
        """
        Updates this cache mixin with results discovered by the other split off one.
        """

        self._models.update(other._models)
        self._eval_exhausted.update(other._eval_exhausted)
        self._max_exhausted.update(other._max_exhausted)
        self._min_exhausted.update(other._min_exhausted)

    #
    # Cache retrieval
    #

    def _model_hook(self, m):
        if len(self._models) < self._MODEL_LIMIT:
            self._models.add(ModelCache(m))

    def _get_models(self, extra_constraints=()):
        for m in self._models:
            if m.eval_constraints(extra_constraints):
                yield m

    def _get_batch_solutions(self, asts, n=None, extra_constraints=()):
        results = set()

        for m in self._get_models(extra_constraints):
            results.add(m.eval_list(asts))
            if len(results) == n:
                break

        return results

    def _get_solutions(self, e, n=None, extra_constraints=()):
        return tuple(v[0] for v in self._get_batch_solutions(
            [e], n=n, extra_constraints=extra_constraints
        ))


    #
    # Cached functions
    #


    def satisfiable(self, extra_constraints=(), **kwargs):
        for _ in self._get_models(extra_constraints=extra_constraints):
            return True
{'fraction_non_alphanumeric': 0.10122252947169262, 'fraction_numerical': 0.0028380148450007276, 'mean_word_length': 2.3766584766584766, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 3, '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': '4376987', 'n_tokens_mistral': 4300, 'n_tokens_neox': 4154, 'n_words': 898}