Repository: HypothesisWorks/hypothesis

Issue:
Better internal implementation for 'filtered choice'
As [discussed in the shrinking guide](https://github.com/HypothesisWorks/hypothesis/blob/master/guides/strategies-that-shrink.rst#make-your-own-luck), we have an internal trick that can make sampling from a filtered list much more efficient.  It would be great to translate `sampled_from(...).filter(...)` into this form internally!  

---------------

On the implementation side, the most elegant way to do this is probably to extend `FilteredStrategy.do_draw` with special cases for strategies that are choosing from a fixed number of elements, i.e.:

1. Choose a random element, and if it satisfies the filter return it.  Try this three times, because rejection sampling has low overhead.
2. [new bit] If the wrapped strategy is a `sampled_from` of reasonable size (e.g. <1k).  This could extend to `booleans()` and `just()` in future
    - We're going to calculate a list of allowed values, pick of of them, and write that element's index in the *full* list to the buffer (see shrinking guide for details)
    - Start by choosing an index `k` into the filtered list.  We don't know how many there are, so just assume that it's unfiltered-1 (as we've failed a random draw).
    - Create the filtered list until we get to the `k`th element.  This is up to twice as fast *on average* as creating the full list - less if very few elements are allowed, but probably more when we're shrinking.
    - If there are less than `k` allowed elements, pick a new `k` which is a valid index.  If there are no valid elements, fall out of this condition - it'll be noted and marked invalid as if rejection sampling failed and we didn't try this special case.
    - Write the index of the chosen element in the unfiltered list to the buffer, and return the element.
3. Give up, because if there is a valid element we can't find it.

Where this gets a bit hairy is in the interaction with `LaxyStrategy` and delayed validation, so we'll probably need to fiddle with the delegation of the `do_draw` or `filter` methods a bit to have it work.  That's why we have this write-up of the discussion from #1862.


Discussion:
Per discussion on #1887, we should ensure that this is also implemented for unique collections via `lists(..., unique=True)`.

Resolving patch:
diff --git a/hypothesis-python/src/hypothesis/searchstrategy/lazy.py b/hypothesis-python/src/hypothesis/searchstrategy/lazy.py
--- a/hypothesis-python/src/hypothesis/searchstrategy/lazy.py
+++ b/hypothesis-python/src/hypothesis/searchstrategy/lazy.py
@@ -155,6 +155,11 @@ def __repr__(self):
     def do_draw(self, data):
         return data.draw(self.wrapped_strategy)
 
+    def do_filtered_draw(self, data, filter_strategy):
+        return self.wrapped_strategy.do_filtered_draw(
+            data=data, filter_strategy=filter_strategy
+        )
+
     @property
     def label(self):
         return self.wrapped_strategy.label
diff --git a/hypothesis-python/src/hypothesis/searchstrategy/misc.py b/hypothesis-python/src/hypothesis/searchstrategy/misc.py
--- a/hypothesis-python/src/hypothesis/searchstrategy/misc.py
+++ b/hypothesis-python/src/hypothesis/searchstrategy/misc.py
@@ -18,7 +18,8 @@
 from __future__ import absolute_import, division, print_function
 
 import hypothesis.internal.conjecture.utils as d
-from hypothesis.searchstrategy.strategies import SearchStrategy
+from hypothesis.internal.compat import bit_length, hrange
+from hypothesis.searchstrategy.strategies import SearchStrategy, filter_not_satisfied
 
 
 class BoolStrategy(SearchStrategy):
@@ -85,3 +86,66 @@ def calc_is_cacheable(self, recur):
 
     def do_draw(self, data):
         return d.choice(data, self.elements)
+
+    def do_filtered_draw(self, data, filter_strategy):
+        # Set of indices that have been tried so far, so that we never test
+        # the same element twice during a draw.
+        known_bad_indices = set()
+
+        def check_index(i):
+            """Return ``True`` if the element at ``i`` satisfies the filter
+            condition.
+            """
+            if i in known_bad_indices:
+                return False
+            ok = filter_strategy.condition(self.elements[i])
+            if not ok:
+                if not known_bad_indices:
+                    filter_strategy.note_retried(data)
+                known_bad_indices.add(i)
+            return ok
+
+        # Start with ordinary rejection sampling. It's fast if it works, and
+        # if it doesn't work then it was only a small amount of overhead.
+        for _ in hrange(3):
+            i = d.integer_range(data, 0, len(self.elements) - 1)
+            if check_index(i):
+                return self.elements[i]
+
+        # Figure out the bit-length of the index that we will write back after
+        # choosing an allowed element.
+        write_length = bit_length(len(self.elements))
+
+        # Impose an arbitrary cutoff to prevent us from wasting too much time
+        # on very large element lists.
+        cutoff = 10000
+
+        # Before building the list of allowed indices, speculatively choose
+        # one of them. We don't yet know how many allowed indices there will be,
+        # so this choice might be out-of-bounds, but that's OK.
+        max_good_indices = min(len(self.elements) - len(known_bad_indices), cutoff)
+        speculative_index = d.integer_range(data, 0, max_good_indices - 1)
+
+        # Calculate the indices of allowed values, so that we can choose one
+        # of them at random. But if we encounter the speculatively-chosen one,
+        # just use that and return immediately.
+        allowed_indices = []
+        for i in hrange(min(len(self.elements), cutoff)):
+            if check_index(i):
+                allowed_indices.append(i)
+                if len(allowed_indices) > speculative_index:
+                    # Early-exit case: We reached the speculative index, so
+                    # we just return the corresponding element.
+                    data.draw_bits(write_length, forced=i)
+                    return self.elements[i]
+
+        # The speculative index didn't work out, but at this point we've built
+        # the complete list of allowed indices, so we can just choose one of
+        # them.
+        if allowed_indices:
+            i = d.choice(data, allowed_indices)
+            data.draw_bits(write_length, forced=i)
+            return self.elements[i]
+        # If there are no allowed indices, the filter couldn't be satisfied.
+
+        return filter_not_satisfied
diff --git a/hypothesis-python/src/hypothesis/searchstrategy/strategies.py b/hypothesis-python/src/hypothesis/searchstrategy/strategies.py
--- a/hypothesis-python/src/hypothesis/searchstrategy/strategies.py
+++ b/hypothesis-python/src/hypothesis/searchstrategy/strategies.py
@@ -349,6 +349,12 @@ def filter(self, condition):
         """
         return FilteredStrategy(condition=condition, strategy=self)
 
+    def do_filtered_draw(self, data, filter_strategy):
+        # Hook for strategies that want to override the behaviour of
+        # FilteredStrategy. Most strategies don't, so by default we delegate
+        # straight back to the default filtered-draw implementation.
+        return filter_strategy.default_do_filtered_draw(data)
+
     @property
     def branches(self):
         # type: () -> List[SearchStrategy[Ex]]
@@ -585,6 +591,9 @@ def branches(self):
         ]
 
 
+filter_not_satisfied = UniqueIdentifier("filter not satisfied")
+
+
 class FilteredStrategy(SearchStrategy):
     def __init__(self, strategy, condition):
         super(FilteredStrategy, self).__init__()
@@ -610,6 +619,20 @@ def do_validate(self):
 
     def do_draw(self, data):
         # type: (ConjectureData) -> Ex
+        result = self.filtered_strategy.do_filtered_draw(
+            data=data, filter_strategy=self
+        )
+        if result is not filter_not_satisfied:
+            return result
+
+        data.note_event("Aborted test because unable to satisfy %r" % (self,))
+        data.mark_invalid()
+        raise AssertionError("Unreachable, for Mypy")  # pragma: no cover
+
+    def note_retried(self, data):
+        data.note_event(lazyformat("Retried draw from %r to satisfy filter", self))
+
+    def default_do_filtered_draw(self, data):
         for i in hrange(3):
             start_index = data.index
             value = data.draw(self.filtered_strategy)
@@ -617,16 +640,13 @@ def do_draw(self, data):
                 return value
             else:
                 if i == 0:
-                    data.note_event(
-                        lazyformat("Retried draw from %r to satisfy filter", self)
-                    )
+                    self.note_retried(data)
                 # This is to guard against the case where we consume no data.
                 # As long as we consume data, we'll eventually pass or raise.
                 # But if we don't this could be an infinite loop.
                 assume(data.index > start_index)
-        data.note_event("Aborted test because unable to satisfy %r" % (self,))
-        data.mark_invalid()
-        raise AssertionError("Unreachable, for Mypy")  # pragma: no cover
+
+        return filter_not_satisfied
 
     @property
     def branches(self):


Test patch:
diff --git a/hypothesis-python/tests/cover/test_sampled_from.py b/hypothesis-python/tests/cover/test_sampled_from.py
--- a/hypothesis-python/tests/cover/test_sampled_from.py
+++ b/hypothesis-python/tests/cover/test_sampled_from.py
@@ -21,7 +21,8 @@
 import enum
 
 from hypothesis import given
-from hypothesis.errors import InvalidArgument
+from hypothesis.errors import FailedHealthCheck, InvalidArgument
+from hypothesis.internal.compat import hrange
 from hypothesis.strategies import sampled_from
 from tests.common.utils import checks_deprecated_behaviour, fails_with
 
@@ -51,3 +52,19 @@ def test_can_sample_enums(member):
 @checks_deprecated_behaviour
 def test_sampling_empty_is_deprecated():
     assert sampled_from([]).is_empty
+
+
+@fails_with(FailedHealthCheck)
+@given(sampled_from(hrange(10)).filter(lambda x: x < 0))
+def test_unsat_filtered_sampling(x):
+    assert False
+
+
+def test_easy_filtered_sampling():
+    x = sampled_from(hrange(100)).filter(lambda x: x == 0).example()
+    assert x == 0
+
+
+@given(sampled_from(hrange(100)).filter(lambda x: x == 99))
+def test_filtered_sampling_finds_rare_value(x):
+    assert x == 99
diff --git a/hypothesis-python/tests/nocover/test_sampled_from.py b/hypothesis-python/tests/nocover/test_sampled_from.py
new file mode 100644
--- /dev/null
+++ b/hypothesis-python/tests/nocover/test_sampled_from.py
@@ -0,0 +1,45 @@
+# coding=utf-8
+#
+# This file is part of Hypothesis, which may be found at
+# https://github.com/HypothesisWorks/hypothesis/
+#
+# Most of this work is copyright (C) 2013-2019 David R. MacIver
+# (david@drmaciver.com), but it contains contributions by others. See
+# CONTRIBUTING.rst for a full list of people who may hold copyright, and
+# consult the git log if you need to determine who owns an individual
+# contribution.
+#
+# This Source Code Form is subject to the terms of the Mozilla Public License,
+# v. 2.0. If a copy of the MPL was not distributed with this file, You can
+# obtain one at https://mozilla.org/MPL/2.0/.
+#
+# END HEADER
+
+from __future__ import absolute_import, division, print_function
+
+import pytest
+
+import hypothesis.strategies as st
+from hypothesis import given
+from hypothesis.internal.compat import hrange
+from tests.common.utils import counts_calls
+
+
+@pytest.mark.parametrize("n", [100, 10 ** 5, 10 ** 6, 2 ** 25])
+def test_filter_large_lists(n):
+    filter_limit = 100 * 10000
+
+    @counts_calls
+    def cond(x):
+        assert cond.calls < filter_limit
+        return x % 2 != 0
+
+    s = st.sampled_from(hrange(n)).filter(cond)
+
+    @given(s)
+    def run(x):
+        assert x % 2 != 0
+
+    run()
+
+    assert cond.calls < filter_limit