From cefcea84f072c64641a4a095ae2d20e13573bcc2 Mon Sep 17 00:00:00 2001
From: Reagan <reaganjlee@gmail.com>
Date: Fri, 11 Oct 2024 00:32:14 -0700
Subject: [PATCH 01/13] Merge `Bundle` as subclass of `SampledFromStrategy`

---
 hypothesis-python/RELEASE.rst                 |  7 ++
 hypothesis-python/src/hypothesis/stateful.py  | 89 +++++++++++++----
 .../strategies/_internal/strategies.py        | 29 ++++--
 .../tests/cover/test_stateful.py              | 99 +++++++++++++++++++
 4 files changed, 201 insertions(+), 23 deletions(-)
 create mode 100644 hypothesis-python/RELEASE.rst

diff --git a/hypothesis-python/RELEASE.rst b/hypothesis-python/RELEASE.rst
new file mode 100644
index 0000000000..c8739170f4
--- /dev/null
+++ b/hypothesis-python/RELEASE.rst
@@ -0,0 +1,7 @@
+RELEASE_TYPE: minor
+
+This release changes :class:`hypothesis.stateful.Bundle` to use the internals of
+:func:`~hypothesis.strategies.sampled_from`, improving the `filter` and `map` methods.
+In addition to performance improvements, you can now ``consumes(some_bundle).filter(...)``!
+
+Thanks to Reagan Lee for this feature (:issue:`3944`).
\ No newline at end of file
diff --git a/hypothesis-python/src/hypothesis/stateful.py b/hypothesis-python/src/hypothesis/stateful.py
index 7c60d2752f..aac7119a53 100644
--- a/hypothesis-python/src/hypothesis/stateful.py
+++ b/hypothesis-python/src/hypothesis/stateful.py
@@ -17,12 +17,25 @@
 """
 import collections
 import inspect
+import sys
 from collections.abc import Iterable, Sequence
 from copy import copy
-from functools import lru_cache
+from functools import lru_cache, partial
 from io import StringIO
 from time import perf_counter
-from typing import Any, Callable, ClassVar, Optional, Union, overload
+from typing import (
+    Any,
+    Callable,
+    ClassVar,
+    Dict,
+    Iterable,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+    overload,
+)
 from unittest import TestCase
 
 import attr
@@ -56,6 +69,7 @@
     Ex,
     Ex_Inv,
     OneOfStrategy,
+    SampledFromStrategy,
     SearchStrategy,
     check_strategy,
 )
@@ -469,7 +483,7 @@ def __repr__(self) -> str:
 self_strategy = st.runner()
 
 
-class Bundle(SearchStrategy[Ex]):
+class Bundle(SampledFromStrategy[Ex]):
     """A collection of values for use in stateful testing.
 
     Bundles are a kind of strategy where values can be added by rules,
@@ -492,32 +506,72 @@ class MyStateMachine(RuleBasedStateMachine):
     """
 
     def __init__(
-        self, name: str, *, consume: bool = False, draw_references: bool = True
+        self,
+        name: str,
+        *,
+        consume: bool = False,
+        draw_references: bool = True,
+        **kwargs,
     ) -> None:
+        super().__init__(
+            [...], **kwargs
+        )  # Some random items that'll get replaced in do_draw
         self.name = name
         self.consume = consume
         self.draw_references = draw_references
 
+        self.bundle = None
+
+        # Shrink towards the right rather than the left. This makes it easier
+        # to delete data generated earlier, as when the error is towards the
+        # end there can be a lot of hard to remove padding.
+        self._SHRINK_TOWARDS = sys.maxsize
+
+    def reference_to_val_func(self, dic, item):
+        assert isinstance(item, int)
+
+        element = self.bundle[item]
+
+        assert isinstance(element, VarReference)
+        return dic.get(element.name)
+
     def do_draw(self, data):
         machine = data.draw(self_strategy)
-
         bundle = machine.bundle(self.name)
         if not bundle:
             data.mark_invalid(f"Cannot draw from empty bundle {self.name!r}")
-        # Shrink towards the right rather than the left. This makes it easier
-        # to delete data generated earlier, as when the error is towards the
-        # end there can be a lot of hard to remove padding.
-        position = data.draw_integer(0, len(bundle) - 1, shrink_towards=len(bundle))
+
+        # We use both self.bundle and self.elements to make sure an index is used to safely pop
+        self.bundle = bundle
+        self.elements = range(len(bundle))
+
+        self.reference_to_value = partial(
+            self.reference_to_val_func, machine.names_to_values
+        )
+
+        idx = super().do_draw(data)
+        reference = bundle[idx]
         if self.consume:
-            reference = bundle.pop(
-                position
-            )  # pragma: no cover  # coverage is flaky here
-        else:
-            reference = bundle[position]
+            bundle.pop(idx)  # pragma: no cover # coverage is flaky here
+        return reference
+
+    def filter(self, condition):
+        return type(self)(
+            self.name,
+            consume=self.consume,
+            draw_references=self.draw_references,
+            transformations=(*self._transformations, ("filter", condition)),
+            repr=self.repr_,
+        )
 
-        if self.draw_references:
-            return reference
-        return machine.names_to_values[reference.name]
+    def map(self, pack):
+        return type(self)(
+            self.name,
+            consume=self.consume,
+            draw_references=self.draw_references,
+            transformations=(*self._transformations, ("map", pack)),
+            repr=self.repr_,
+        )
 
     def __repr__(self):
         consume = self.consume
@@ -562,6 +616,7 @@ def consumes(bundle: Bundle[Ex]) -> SearchStrategy[Ex]:
     return type(bundle)(
         name=bundle.name,
         consume=True,
+        transformations=bundle._transformations,
     )
 
 
diff --git a/hypothesis-python/src/hypothesis/strategies/_internal/strategies.py b/hypothesis-python/src/hypothesis/strategies/_internal/strategies.py
index c040f748a5..ac98008f7a 100644
--- a/hypothesis-python/src/hypothesis/strategies/_internal/strategies.py
+++ b/hypothesis-python/src/hypothesis/strategies/_internal/strategies.py
@@ -473,13 +473,14 @@ def is_simple_data(value):
         return False
 
 
-class SampledFromStrategy(SearchStrategy):
+class SampledFromStrategy(SearchStrategy[Ex]):
     """A strategy which samples from a set of elements. This is essentially
     equivalent to using a OneOfStrategy over Just strategies but may be more
     efficient and convenient.
     """
 
     _MAX_FILTER_CALLS = 10_000
+    _SHRINK_TOWARDS = 0
 
     def __init__(self, elements, repr_=None, transformations=()):
         super().__init__()
@@ -487,6 +488,7 @@ def __init__(self, elements, repr_=None, transformations=()):
         assert self.elements
         self.repr_ = repr_
         self._transformations = transformations
+        self.reference_to_value = lambda x: x
 
     def map(self, pack):
         return type(self)(
@@ -552,7 +554,11 @@ def do_draw(self, data):
         return result
 
     def get_element(self, i):
-        return self._transform(self.elements[i])
+        element = self.elements[i]
+        value = self._transform(self.reference_to_value(element))
+        if is_identity_function(self.reference_to_value):
+            return value
+        return element if value is not filter_not_satisfied else filter_not_satisfied
 
     def do_filtered_draw(self, data):
         # Set of indices that have been tried so far, so that we never test
@@ -562,7 +568,9 @@ def do_filtered_draw(self, data):
         # 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 range(3):
-            i = data.draw_integer(0, len(self.elements) - 1)
+            i = data.draw_integer(
+                0, len(self.elements) - 1, shrink_towards=self._SHRINK_TOWARDS
+            )
             if i not in known_bad_indices:
                 element = self.get_element(i)
                 if element is not filter_not_satisfied:
@@ -583,7 +591,9 @@ def do_filtered_draw(self, data):
         # 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.
-        speculative_index = data.draw_integer(0, max_good_indices - 1)
+        speculative_index = data.draw_integer(
+            0, max_good_indices - 1, shrink_towards=self._SHRINK_TOWARDS
+        )
 
         # Calculate the indices of allowed values, so that we can choose one
         # of them at random. But if we encounter the speculatively-chosen one,
@@ -598,14 +608,21 @@ def do_filtered_draw(self, data):
                     if len(allowed) > speculative_index:
                         # Early-exit case: We reached the speculative index, so
                         # we just return the corresponding element.
-                        data.draw_integer(0, len(self.elements) - 1, forced=i)
+                        data.draw_integer(
+                            0,
+                            len(self.elements) - 1,
+                            forced=i,
+                            shrink_towards=self._SHRINK_TOWARDS,
+                        )
                         return element
 
         # The speculative index didn't work out, but at this point we've built
         # and can choose from the complete list of allowed indices and elements.
         if allowed:
             i, element = data.choice(allowed)
-            data.draw_integer(0, len(self.elements) - 1, forced=i)
+            data.draw_integer(
+                0, len(self.elements) - 1, forced=i, shrink_towards=self._SHRINK_TOWARDS
+            )
             return element
         # If there are no allowed indices, the filter couldn't be satisfied.
         return filter_not_satisfied
diff --git a/hypothesis-python/tests/cover/test_stateful.py b/hypothesis-python/tests/cover/test_stateful.py
index 2dd51081cc..4afbc8c36f 100644
--- a/hypothesis-python/tests/cover/test_stateful.py
+++ b/hypothesis-python/tests/cover/test_stateful.py
@@ -1320,6 +1320,105 @@ def rule1(self, data):
 TestLotsOfEntropyPerStepMachine = LotsOfEntropyPerStepMachine.TestCase
 
 
+def test_filter():
+    class Machine(RuleBasedStateMachine):
+        a = Bundle("a")
+
+        @initialize(target=a)
+        def initialize(self):
+            return multiple(1, 2, 3)
+
+        @rule(
+            a1=a.filter(lambda x: x < 2),
+            a2=a.filter(lambda x: x > 2),
+            a3=a,
+        )
+        def fail_fast(self, a1, a2, a3):
+            raise AssertionError
+
+    Machine.TestCase.settings = NO_BLOB_SETTINGS
+    with pytest.raises(AssertionError) as err:
+        run_state_machine_as_test(Machine)
+
+    result = "\n".join(err.value.__notes__)
+    assert (
+        result
+        == """
+Falsifying example:
+state = Machine()
+a_0, a_1, a_2 = state.initialize()
+state.fail_fast(a1=a_0, a2=a_2, a3=a_2)
+state.teardown()
+""".strip()
+    )
+
+
+def test_consumes_filter():
+    class Machine(RuleBasedStateMachine):
+        a = Bundle("a")
+
+        @initialize(target=a)
+        def initialize(self):
+            return multiple(1, 2, 3)
+
+        @rule(
+            a1=consumes(a).filter(lambda x: x < 2),
+            a2=consumes(a).filter(lambda x: x > 2),
+            a3=consumes(a),
+        )
+        def fail_fast(self, a1, a2, a3):
+            raise AssertionError
+
+    Machine.TestCase.settings = NO_BLOB_SETTINGS
+    with pytest.raises(AssertionError) as err:
+        run_state_machine_as_test(Machine)
+
+    result = "\n".join(err.value.__notes__)
+    assert (
+        result
+        == """
+Falsifying example:
+state = Machine()
+a_0, a_1, a_2 = state.initialize()
+state.fail_fast(a1=a_0, a2=a_2, a3=a_1)
+state.teardown()
+""".strip()
+    )
+
+
+def test_map_with_filter():
+    class Machine(RuleBasedStateMachine):
+        a = Bundle("a")
+
+        @initialize(target=a)
+        def initialize(self):
+            return multiple(2, 4)
+
+        @rule(
+            a1=a.map(lambda x: x**2).filter(lambda x: x < 3**2),
+            a2=consumes(a).map(lambda x: x**2).filter(lambda x: x > 3**2),
+            a3=consumes(a),
+        )
+        def fail_fast(self, a1, a2, a3):
+            raise AssertionError
+
+    Machine.TestCase.settings = NO_BLOB_SETTINGS
+    with pytest.raises(AssertionError) as err:
+        run_state_machine_as_test(Machine)
+
+    result = "\n".join(err.value.__notes__)
+    assert (
+        result
+        == """
+Falsifying example:
+state = Machine()
+a_0, a_1 = state.initialize()
+state.fail_fast(a1=a_0, a2=a_1, a3=a_0)
+state.teardown()
+""".strip()
+    )
+
+
 def test_flatmap():
     class Machine(RuleBasedStateMachine):
         buns = Bundle("buns")

From 29f16074179607f03c205a118fb9a3e716a061b8 Mon Sep 17 00:00:00 2001
From: Reagan <reaganjlee@gmail.com>
Date: Fri, 11 Oct 2024 01:02:02 -0700
Subject: [PATCH 02/13] fix + add test

---
 hypothesis-python/src/hypothesis/stateful.py   | 12 +++++++++---
 hypothesis-python/tests/cover/test_stateful.py | 11 ++++++++++-
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/hypothesis-python/src/hypothesis/stateful.py b/hypothesis-python/src/hypothesis/stateful.py
index aac7119a53..53a9f26d32 100644
--- a/hypothesis-python/src/hypothesis/stateful.py
+++ b/hypothesis-python/src/hypothesis/stateful.py
@@ -553,6 +553,8 @@ def do_draw(self, data):
         reference = bundle[idx]
         if self.consume:
             bundle.pop(idx)  # pragma: no cover # coverage is flaky here
+        if not self.draw_references:
+            return machine.names_to_values[reference.name]
         return reference
 
     def filter(self, condition):
@@ -561,7 +563,7 @@ def filter(self, condition):
             consume=self.consume,
             draw_references=self.draw_references,
             transformations=(*self._transformations, ("filter", condition)),
-            repr=self.repr_,
+            repr_=self.repr_,
         )
 
     def map(self, pack):
@@ -570,7 +572,7 @@ def map(self, pack):
             consume=self.consume,
             draw_references=self.draw_references,
             transformations=(*self._transformations, ("map", pack)),
-            repr=self.repr_,
+            repr_=self.repr_,
         )
 
     def __repr__(self):
@@ -593,7 +595,11 @@ def available(self, data):
     def flatmap(self, expand):
         if self.draw_references:
             return type(self)(
-                self.name, consume=self.consume, draw_references=False
+                self.name,
+                consume=self.consume,
+                draw_references=False,
+                transformations=self._transformations,
+                repr_=self.repr_,
             ).flatmap(expand)
         return super().flatmap(expand)
 
diff --git a/hypothesis-python/tests/cover/test_stateful.py b/hypothesis-python/tests/cover/test_stateful.py
index 4afbc8c36f..b522b37dc5 100644
--- a/hypothesis-python/tests/cover/test_stateful.py
+++ b/hypothesis-python/tests/cover/test_stateful.py
@@ -1419,7 +1419,7 @@ def fail_fast(self, a1, a2, a3):
     )
 
 
-def test_flatmap():
+def test_flatmap_with_filter():
     class Machine(RuleBasedStateMachine):
         buns = Bundle("buns")
 
@@ -1432,6 +1432,15 @@ def use_flatmap(self, bun):
             assert isinstance(bun, int)
             return bun
 
+        @rule(
+            target=buns,
+            bun=buns.flatmap(lambda x: just(x + 1)).filter(lambda x: x > 0),
+        )
+        def use_flatmap(self, bun):
+            assert isinstance(bun, int)
+            assert bun > 0
+            return bun
+
         @rule(bun=buns)
         def use_directly(self, bun):
             assert isinstance(bun, int)

From f0d8f71175f977300ffc15c3ecfce37d9c5a431e Mon Sep 17 00:00:00 2001
From: Reagan <reaganjlee@gmail.com>
Date: Fri, 11 Oct 2024 01:04:09 -0700
Subject: [PATCH 03/13] format

---
 hypothesis-python/src/hypothesis/stateful.py | 14 +-------------
 1 file changed, 1 insertion(+), 13 deletions(-)

diff --git a/hypothesis-python/src/hypothesis/stateful.py b/hypothesis-python/src/hypothesis/stateful.py
index 53a9f26d32..a577da6206 100644
--- a/hypothesis-python/src/hypothesis/stateful.py
+++ b/hypothesis-python/src/hypothesis/stateful.py
@@ -23,19 +23,7 @@
 from functools import lru_cache, partial
 from io import StringIO
 from time import perf_counter
-from typing import (
-    Any,
-    Callable,
-    ClassVar,
-    Dict,
-    Iterable,
-    List,
-    Optional,
-    Sequence,
-    Tuple,
-    Union,
-    overload,
-)
+from typing import Any, Callable, ClassVar, Optional, Union, overload
 from unittest import TestCase
 
 import attr

From 7cca6e5bbae70877d2be754a677b35242b461401 Mon Sep 17 00:00:00 2001
From: Reagan <reaganjlee@gmail.com>
Date: Fri, 11 Oct 2024 12:45:52 -0700
Subject: [PATCH 04/13] mypy + linter

---
 hypothesis-python/src/hypothesis/stateful.py   | 9 ++++++---
 hypothesis-python/tests/cover/test_stateful.py | 2 +-
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/hypothesis-python/src/hypothesis/stateful.py b/hypothesis-python/src/hypothesis/stateful.py
index a577da6206..798aa52427 100644
--- a/hypothesis-python/src/hypothesis/stateful.py
+++ b/hypothesis-python/src/hypothesis/stateful.py
@@ -23,7 +23,7 @@
 from functools import lru_cache, partial
 from io import StringIO
 from time import perf_counter
-from typing import Any, Callable, ClassVar, Optional, Union, overload
+from typing import Any, Callable, ClassVar, Optional, Tuple, Union, overload
 from unittest import TestCase
 
 import attr
@@ -499,10 +499,13 @@ def __init__(
         *,
         consume: bool = False,
         draw_references: bool = True,
-        **kwargs,
+        repr_: Optional[str] = None,
+        transformations: Iterable[Tuple[str, Callable]] = (),
     ) -> None:
         super().__init__(
-            [...], **kwargs
+            [...],
+            repr_=repr_,
+            transformations=transformations,
         )  # Some random items that'll get replaced in do_draw
         self.name = name
         self.consume = consume
diff --git a/hypothesis-python/tests/cover/test_stateful.py b/hypothesis-python/tests/cover/test_stateful.py
index b522b37dc5..f6172ddedd 100644
--- a/hypothesis-python/tests/cover/test_stateful.py
+++ b/hypothesis-python/tests/cover/test_stateful.py
@@ -1436,7 +1436,7 @@ def use_flatmap(self, bun):
             target=buns,
             bun=buns.flatmap(lambda x: just(x + 1)).filter(lambda x: x > 0),
         )
-        def use_flatmap(self, bun):
+        def use_flatmap_filtered(self, bun):
             assert isinstance(bun, int)
             assert bun > 0
             return bun

From 52e0fc769748f5eda79674f8ec1e8e4b4fa217cb Mon Sep 17 00:00:00 2001
From: Reagan <reaganjlee@gmail.com>
Date: Fri, 11 Oct 2024 17:59:08 -0700
Subject: [PATCH 05/13] ruff

---
 hypothesis-python/src/hypothesis/stateful.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/hypothesis-python/src/hypothesis/stateful.py b/hypothesis-python/src/hypothesis/stateful.py
index 798aa52427..56e025b3d1 100644
--- a/hypothesis-python/src/hypothesis/stateful.py
+++ b/hypothesis-python/src/hypothesis/stateful.py
@@ -23,7 +23,7 @@
 from functools import lru_cache, partial
 from io import StringIO
 from time import perf_counter
-from typing import Any, Callable, ClassVar, Optional, Tuple, Union, overload
+from typing import Any, Callable, ClassVar, Optional, Union, overload
 from unittest import TestCase
 
 import attr
@@ -500,7 +500,7 @@ def __init__(
         consume: bool = False,
         draw_references: bool = True,
         repr_: Optional[str] = None,
-        transformations: Iterable[Tuple[str, Callable]] = (),
+        transformations: Iterable[tuple[str, Callable]] = (),
     ) -> None:
         super().__init__(
             [...],

From 52e965f9f892c0e8d91cc293a87c7d924fa7163c Mon Sep 17 00:00:00 2001
From: Reagan <reaganjlee@gmail.com>
Date: Fri, 11 Oct 2024 19:21:14 -0700
Subject: [PATCH 06/13] improve speed 20%

---
 hypothesis-python/src/hypothesis/stateful.py | 30 +++++++++-----------
 1 file changed, 14 insertions(+), 16 deletions(-)

diff --git a/hypothesis-python/src/hypothesis/stateful.py b/hypothesis-python/src/hypothesis/stateful.py
index 56e025b3d1..f57f1fe2cc 100644
--- a/hypothesis-python/src/hypothesis/stateful.py
+++ b/hypothesis-python/src/hypothesis/stateful.py
@@ -512,40 +512,38 @@ def __init__(
         self.draw_references = draw_references
 
         self.bundle = None
+        self.machine = None
+        self.reference_to_value = self.reference_to_val_func
 
         # Shrink towards the right rather than the left. This makes it easier
         # to delete data generated earlier, as when the error is towards the
         # end there can be a lot of hard to remove padding.
         self._SHRINK_TOWARDS = sys.maxsize
 
-    def reference_to_val_func(self, dic, item):
+    def reference_to_val_func(self, item):
         assert isinstance(item, int)
-
         element = self.bundle[item]
-
         assert isinstance(element, VarReference)
-        return dic.get(element.name)
+        return self.machine.names_to_values.get(element.name)
 
     def do_draw(self, data):
-        machine = data.draw(self_strategy)
-        bundle = machine.bundle(self.name)
-        if not bundle:
+        self.machine = data.draw(self_strategy)
+        self.bundle = self.machine.bundle(self.name)
+        if not self.bundle:
             data.mark_invalid(f"Cannot draw from empty bundle {self.name!r}")
 
         # We use both self.bundle and self.elements to make sure an index is used to safely pop
-        self.bundle = bundle
-        self.elements = range(len(bundle))
-
-        self.reference_to_value = partial(
-            self.reference_to_val_func, machine.names_to_values
-        )
+        self.elements = range(len(self.bundle))
 
         idx = super().do_draw(data)
-        reference = bundle[idx]
+        reference = self.bundle[idx]
+
         if self.consume:
-            bundle.pop(idx)  # pragma: no cover # coverage is flaky here
+            self.bundle.pop(idx)  # pragma: no cover # coverage is flaky here
+
         if not self.draw_references:
-            return machine.names_to_values[reference.name]
+            return self.machine.names_to_values[reference.name]
+
         return reference
 
     def filter(self, condition):

From 24c544dbdb6372bfa0d20658d3efa2db937f830a Mon Sep 17 00:00:00 2001
From: Reagan <reaganjlee@gmail.com>
Date: Fri, 11 Oct 2024 19:25:33 -0700
Subject: [PATCH 07/13] ruff

---
 hypothesis-python/src/hypothesis/stateful.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/hypothesis-python/src/hypothesis/stateful.py b/hypothesis-python/src/hypothesis/stateful.py
index f57f1fe2cc..73f8da0856 100644
--- a/hypothesis-python/src/hypothesis/stateful.py
+++ b/hypothesis-python/src/hypothesis/stateful.py
@@ -20,7 +20,7 @@
 import sys
 from collections.abc import Iterable, Sequence
 from copy import copy
-from functools import lru_cache, partial
+from functools import lru_cache
 from io import StringIO
 from time import perf_counter
 from typing import Any, Callable, ClassVar, Optional, Union, overload

From 7ec4c6962efb484d9c4214ae99a8638f224a56f5 Mon Sep 17 00:00:00 2001
From: Reagan <reaganjlee@gmail.com>
Date: Thu, 17 Oct 2024 00:54:00 -0700
Subject: [PATCH 08/13] test changes

---
 .../tests/cover/test_stateful.py              | 84 ++++++++++++++++++-
 1 file changed, 82 insertions(+), 2 deletions(-)

diff --git a/hypothesis-python/tests/cover/test_stateful.py b/hypothesis-python/tests/cover/test_stateful.py
index f6172ddedd..0e40b62a9e 100644
--- a/hypothesis-python/tests/cover/test_stateful.py
+++ b/hypothesis-python/tests/cover/test_stateful.py
@@ -1386,7 +1386,7 @@ def fail_fast(self, a1, a2, a3):
     )
 
 
-def test_map_with_filter():
+def test_consumes_map_with_filter():
     class Machine(RuleBasedStateMachine):
         a = Bundle("a")
 
@@ -1419,7 +1419,7 @@ def fail_fast(self, a1, a2, a3):
     )
 
 
-def test_flatmap_with_filter():
+def test_flatmap_with_combinations():
     class Machine(RuleBasedStateMachine):
         buns = Bundle("buns")
 
@@ -1441,6 +1441,86 @@ def use_flatmap_filtered(self, bun):
             assert bun > 0
             return bun
 
+        @rule(
+            target=buns,
+            bun=buns.flatmap(lambda x: just(x + 1)).map(lambda x: -(x + 1)),
+        )
+        def use_flatmap_mapped(self, bun):
+            assert isinstance(bun, int)
+            assert bun < 0
+            return bun
+
+        @rule(bun=buns)
+        def use_directly(self, bun):
+            assert isinstance(bun, int)
+
+    Machine.TestCase.settings = Settings(stateful_step_count=5, max_examples=10)
+    run_state_machine_as_test(Machine)
+
+
+def test_map_with_combinations():
+    class Machine(RuleBasedStateMachine):
+        buns = Bundle("buns")
+
+        @initialize(target=buns)
+        def create_bun(self):
+            return 1
+
+        @rule(bun=buns.map(lambda x: -x))
+        def use_map_base(self, bun):
+            assert isinstance(bun, int)
+            assert bun < 0
+
+        @rule(
+            bun=buns.map(lambda x: -x).filter(lambda x: x < -1),
+        )
+        def use_flatmap_filtered(self, bun):
+            assert isinstance(bun, int)
+            assert bun < -1
+
+        @rule(
+            bun=buns.map(lambda x: -x).flatmap(lambda x: just(abs(x))),
+        )
+        def use_flatmap_mapped(self, bun):
+            assert isinstance(bun, int)
+            assert bun > 0
+
+        @rule(bun=buns)
+        def use_directly(self, bun):
+            assert isinstance(bun, int)
+            assert bun > 0
+
+    Machine.TestCase.settings = Settings(stateful_step_count=5, max_examples=10)
+    run_state_machine_as_test(Machine)
+
+
+def test_filter_with_combinations():
+    class Machine(RuleBasedStateMachine):
+        buns = Bundle("buns")
+
+        @initialize(target=buns)
+        def create_bun(self):
+            return multiple(0, 1, 2)
+
+        @rule(bun=buns.filter(lambda x: x > 0))
+        def use_filter_base(self, bun):
+            assert isinstance(bun, int)
+            assert bun > 0
+
+        @rule(
+            bun=buns.filter(lambda x: x > 0).flatmap(lambda x: just(-x)),
+        )
+        def use_filter_flatmapped(self, bun):
+            assert isinstance(bun, int)
+            assert bun < 0
+
+        @rule(
+            bun=buns.filter(lambda x: x > 0).map(lambda x: -x),
+        )
+        def use_flatmap_mapped(self, bun):
+            assert isinstance(bun, int)
+            assert bun < 0
+
         @rule(bun=buns)
         def use_directly(self, bun):
             assert isinstance(bun, int)

From 5f461ff09a3575e078d3e7eca95a8d1391217644 Mon Sep 17 00:00:00 2001
From: Reagan <reaganjlee@gmail.com>
Date: Thu, 17 Oct 2024 00:54:58 -0700
Subject: [PATCH 09/13] move more complicated `get_element()` to stateful.py

---
 hypothesis-python/src/hypothesis/stateful.py           | 10 +++++++++-
 .../src/hypothesis/strategies/_internal/strategies.py  |  7 +------
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/hypothesis-python/src/hypothesis/stateful.py b/hypothesis-python/src/hypothesis/stateful.py
index 73f8da0856..5acb9ebedb 100644
--- a/hypothesis-python/src/hypothesis/stateful.py
+++ b/hypothesis-python/src/hypothesis/stateful.py
@@ -60,6 +60,7 @@
     SampledFromStrategy,
     SearchStrategy,
     check_strategy,
+    filter_not_satisfied,
 )
 from hypothesis.vendor.pretty import RepresentationPrinter
 
@@ -526,6 +527,11 @@ def reference_to_val_func(self, item):
         assert isinstance(element, VarReference)
         return self.machine.names_to_values.get(element.name)
 
+    def get_element(self, i):
+        element = self.elements[i]
+        value = self._transform(self.reference_to_value(element))
+        return element if value is not filter_not_satisfied else filter_not_satisfied
+
     def do_draw(self, data):
         self.machine = data.draw(self_strategy)
         self.bundle = self.machine.bundle(self.name)
@@ -542,7 +548,9 @@ def do_draw(self, data):
             self.bundle.pop(idx)  # pragma: no cover # coverage is flaky here
 
         if not self.draw_references:
-            return self.machine.names_to_values[reference.name]
+            return self.machine.names_to_values[
+                reference.name
+            ]  # currently missing the mapped case
 
         return reference
 
diff --git a/hypothesis-python/src/hypothesis/strategies/_internal/strategies.py b/hypothesis-python/src/hypothesis/strategies/_internal/strategies.py
index ac98008f7a..a6fc51cefc 100644
--- a/hypothesis-python/src/hypothesis/strategies/_internal/strategies.py
+++ b/hypothesis-python/src/hypothesis/strategies/_internal/strategies.py
@@ -488,7 +488,6 @@ def __init__(self, elements, repr_=None, transformations=()):
         assert self.elements
         self.repr_ = repr_
         self._transformations = transformations
-        self.reference_to_value = lambda x: x
 
     def map(self, pack):
         return type(self)(
@@ -554,11 +553,7 @@ def do_draw(self, data):
         return result
 
     def get_element(self, i):
-        element = self.elements[i]
-        value = self._transform(self.reference_to_value(element))
-        if is_identity_function(self.reference_to_value):
-            return value
-        return element if value is not filter_not_satisfied else filter_not_satisfied
+        return self._transform(self.elements[i])
 
     def do_filtered_draw(self, data):
         # Set of indices that have been tried so far, so that we never test

From 90f436076dfaa57b050f7e857d2e32e2c7419112 Mon Sep 17 00:00:00 2001
From: Reagan <reaganjlee@gmail.com>
Date: Thu, 17 Oct 2024 16:26:18 -0700
Subject: [PATCH 10/13] `VarReferenceMapping` for deterministic printing +
 maintaining transformed values

---
 hypothesis-python/src/hypothesis/stateful.py | 46 +++++++++++---------
 1 file changed, 25 insertions(+), 21 deletions(-)

diff --git a/hypothesis-python/src/hypothesis/stateful.py b/hypothesis-python/src/hypothesis/stateful.py
index 5acb9ebedb..77bd2ed668 100644
--- a/hypothesis-python/src/hypothesis/stateful.py
+++ b/hypothesis-python/src/hypothesis/stateful.py
@@ -187,12 +187,12 @@ def output(s):
                 try:
                     data = dict(data)
                     for k, v in list(data.items()):
-                        if isinstance(v, VarReference):
-                            data[k] = machine.names_to_values[v.name]
+                        if isinstance(v, VarReferenceMapping):
+                            data[k] = v.value
                         elif isinstance(v, list) and all(
-                            isinstance(item, VarReference) for item in v
+                            isinstance(item, VarReferenceMapping) for item in v
                         ):
-                            data[k] = [machine.names_to_values[item.name] for item in v]
+                            data[k] = [item.value for item in v]
 
                     label = f"execute:rule:{rule.function.__name__}"
                     start = perf_counter()
@@ -295,12 +295,12 @@ def __init__(self) -> None:
             )
 
     def _pretty_print(self, value):
-        if isinstance(value, VarReference):
-            return value.name
+        if isinstance(value, VarReferenceMapping):
+            return value.reference.name
         elif isinstance(value, list) and all(
-            isinstance(item, VarReference) for item in value
+            isinstance(item, VarReferenceMapping) for item in value
         ):
-            return "[" + ", ".join([item.name for item in value]) + "]"
+            return "[" + ", ".join([item.reference.name for item in value]) + "]"
         self.__stream.seek(0)
         self.__stream.truncate(0)
         self.__printer.output_width = 0
@@ -514,23 +514,22 @@ def __init__(
 
         self.bundle = None
         self.machine = None
-        self.reference_to_value = self.reference_to_val_func
 
         # Shrink towards the right rather than the left. This makes it easier
         # to delete data generated earlier, as when the error is towards the
         # end there can be a lot of hard to remove padding.
         self._SHRINK_TOWARDS = sys.maxsize
 
-    def reference_to_val_func(self, item):
-        assert isinstance(item, int)
-        element = self.bundle[item]
-        assert isinstance(element, VarReference)
-        return self.machine.names_to_values.get(element.name)
+    def get_transformed_value(self, reference):
+        assert isinstance(reference, VarReference)
+        return self._transform(self.machine.names_to_values.get(reference.name))
 
     def get_element(self, i):
-        element = self.elements[i]
-        value = self._transform(self.reference_to_value(element))
-        return element if value is not filter_not_satisfied else filter_not_satisfied
+        idx = self.elements[i]
+        assert isinstance(idx, int)
+        reference = self.bundle[idx]
+        value = self.get_transformed_value(reference)
+        return idx if value is not filter_not_satisfied else filter_not_satisfied
 
     def do_draw(self, data):
         self.machine = data.draw(self_strategy)
@@ -548,11 +547,10 @@ def do_draw(self, data):
             self.bundle.pop(idx)  # pragma: no cover # coverage is flaky here
 
         if not self.draw_references:
-            return self.machine.names_to_values[
-                reference.name
-            ]  # currently missing the mapped case
+            return self.get_transformed_value(reference)
 
-        return reference
+        # we need both reference and the value itself to pretty-print deterministically and maintain any transformations that is bundle-specific
+        return VarReferenceMapping(reference, self.get_transformed_value(reference))
 
     def filter(self, condition):
         return type(self)(
@@ -886,6 +884,12 @@ class VarReference:
     name = attr.ib()
 
 
+@attr.s()
+class VarReferenceMapping:
+    reference = attr.ib(type=VarReference)
+    value = attr.ib()
+
+
 # There are multiple alternatives for annotating the `precond` type, all of them
 # have drawbacks. See https://github.com/HypothesisWorks/hypothesis/pull/3068#issuecomment-906642371
 def precondition(precond: Callable[[Any], bool]) -> Callable[[TestFunc], TestFunc]:

From 0194f90e9ef24886d2b2e9145d9d8ae3ba7bf424 Mon Sep 17 00:00:00 2001
From: Reagan <reaganjlee@gmail.com>
Date: Thu, 17 Oct 2024 16:26:57 -0700
Subject: [PATCH 11/13] tests

---
 .../tests/cover/test_stateful.py              | 53 +++++++++++++++----
 1 file changed, 44 insertions(+), 9 deletions(-)

diff --git a/hypothesis-python/tests/cover/test_stateful.py b/hypothesis-python/tests/cover/test_stateful.py
index 0e40b62a9e..23b5c9a82e 100644
--- a/hypothesis-python/tests/cover/test_stateful.py
+++ b/hypothesis-python/tests/cover/test_stateful.py
@@ -1434,25 +1434,26 @@ def use_flatmap(self, bun):
 
         @rule(
             target=buns,
-            bun=buns.flatmap(lambda x: just(x + 1)).filter(lambda x: x > 0),
+            bun=buns.flatmap(lambda x: just(-x)).filter(lambda x: x < -1),
         )
         def use_flatmap_filtered(self, bun):
             assert isinstance(bun, int)
-            assert bun > 0
-            return bun
+            assert bun < -1
+            return -bun
 
         @rule(
             target=buns,
-            bun=buns.flatmap(lambda x: just(x + 1)).map(lambda x: -(x + 1)),
+            bun=buns.flatmap(lambda x: just(x + 1)).map(lambda x: -x),
         )
         def use_flatmap_mapped(self, bun):
             assert isinstance(bun, int)
             assert bun < 0
-            return bun
+            return -bun
 
         @rule(bun=buns)
         def use_directly(self, bun):
             assert isinstance(bun, int)
+            assert bun >= 0
 
     Machine.TestCase.settings = Settings(stateful_step_count=5, max_examples=10)
     run_state_machine_as_test(Machine)
@@ -1479,7 +1480,7 @@ def use_flatmap_filtered(self, bun):
             assert bun < -1
 
         @rule(
-            bun=buns.map(lambda x: -x).flatmap(lambda x: just(abs(x))),
+            bun=buns.map(lambda x: -x).flatmap(lambda x: just(abs(x) + 1)),
         )
         def use_flatmap_mapped(self, bun):
             assert isinstance(bun, int)
@@ -1500,7 +1501,7 @@ class Machine(RuleBasedStateMachine):
 
         @initialize(target=buns)
         def create_bun(self):
-            return multiple(0, 1, 2)
+            return multiple(0, -1, -2)
 
         @rule(bun=buns.filter(lambda x: x > 0))
         def use_filter_base(self, bun):
@@ -1515,11 +1516,11 @@ def use_filter_flatmapped(self, bun):
             assert bun < 0
 
         @rule(
-            bun=buns.filter(lambda x: x > 0).map(lambda x: -x),
+            bun=buns.filter(lambda x: x < 0).map(lambda x: -x),
         )
         def use_flatmap_mapped(self, bun):
             assert isinstance(bun, int)
-            assert bun < 0
+            assert bun > 0
 
         @rule(bun=buns)
         def use_directly(self, bun):
@@ -1527,3 +1528,37 @@ def use_directly(self, bun):
 
     Machine.TestCase.settings = Settings(stateful_step_count=5, max_examples=10)
     run_state_machine_as_test(Machine)
+
+
+def test_mapped_values_assigned_properly():
+    class Machine(RuleBasedStateMachine):
+        a = Bundle("a")
+
+        @initialize(target=a)
+        def initialize(self):
+            return multiple("ret1", "ret2")
+
+        @rule(
+            a1=a,
+            a2=a.map(lambda x: x + x),
+            a3=consumes(a).map(lambda x: x + x),
+            a4=a,
+        )
+        def fail_fast(self, a1, a2, a3, a4):
+            raise AssertionError
+
+    Machine.TestCase.settings = NO_BLOB_SETTINGS
+    with pytest.raises(AssertionError) as err:
+        run_state_machine_as_test(Machine)
+
+    result = "\n".join(err.value.__notes__)
+    assert (
+        result
+        == """
+Falsifying example:
+state = Machine()
+a_0, a_1 = state.initialize()
+state.fail_fast(a1=a_1, a2=a_1, a3=a_1, a4=a_0)
+state.teardown()
+""".strip()
+    )

From e74c1e2d121d5b52b26471d2c159dcc15ccab660 Mon Sep 17 00:00:00 2001
From: Reagan <reaganjlee@gmail.com>
Date: Thu, 17 Oct 2024 16:27:58 -0700
Subject: [PATCH 12/13] ruff format

---
 hypothesis-python/src/hypothesis/stateful.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/hypothesis-python/src/hypothesis/stateful.py b/hypothesis-python/src/hypothesis/stateful.py
index 77bd2ed668..f548cbb9c1 100644
--- a/hypothesis-python/src/hypothesis/stateful.py
+++ b/hypothesis-python/src/hypothesis/stateful.py
@@ -549,7 +549,8 @@ def do_draw(self, data):
         if not self.draw_references:
             return self.get_transformed_value(reference)
 
-        # we need both reference and the value itself to pretty-print deterministically and maintain any transformations that is bundle-specific
+        # we need both reference and the value itself to pretty-print deterministically
+        # and maintain any transformations that is bundle-specific
         return VarReferenceMapping(reference, self.get_transformed_value(reference))
 
     def filter(self, condition):

From a3c0c31fd352bcae2bc1aea555475986dda87f63 Mon Sep 17 00:00:00 2001
From: Reagan <reaganjlee@gmail.com>
Date: Thu, 17 Oct 2024 16:57:25 -0700
Subject: [PATCH 13/13] mypy

---
 hypothesis-python/src/hypothesis/stateful.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/hypothesis-python/src/hypothesis/stateful.py b/hypothesis-python/src/hypothesis/stateful.py
index f548cbb9c1..512117e0b5 100644
--- a/hypothesis-python/src/hypothesis/stateful.py
+++ b/hypothesis-python/src/hypothesis/stateful.py
@@ -887,8 +887,8 @@ class VarReference:
 
 @attr.s()
 class VarReferenceMapping:
-    reference = attr.ib(type=VarReference)
-    value = attr.ib()
+    reference: VarReference = attr.ib()
+    value: Any = attr.ib()
 
 
 # There are multiple alternatives for annotating the `precond` type, all of them