again_count rerolls the entire process when exceeded

src/icepool/__init__.py

@@ -84,7 +84,8 @@
 roll at a time. For every `Again` we roll, we queue another roll.
 If we run out of rolls, we sum the rolls to find the result. If the total number
 of rolls (not including the initial roll) would exceed `again_count`, we reroll
-the last die.
+the entire process, effectively conditioning the process on not rolling more
+than `again_count` extra dice.
 
 This mode only allows "additive" expressions to be used with `Again`, which
 means that only the following operators are allowed:
@@ -93,6 +94,8 @@
 * `n @ AgainExpression`, where `n` is a non-negative `int` or `Population`.
 
 Furthermore, the `+` operator is assumed to be associative and commutative.
+For example, `str` or `tuple` outcomes will not produce elements with a definite
+order.
 
 #### Depth mode
 

src/icepool/population/again.py

@@ -270,30 +270,33 @@ def make_step_outcome(outcome, zero):
                     'again_count mode cannot be used with a non-additive AgainExpression.'
                 )
             return icepool.vectorize(outcome._evaluate(zero), 0,
-                                     outcome._again_count() - 1, -1)
+                                     outcome._again_count() - 1)
         else:
-            return icepool.vectorize(zero, 1, -1, -1)
+            return icepool.vectorize(zero, 1, -1)
 
     # Flat total, added again count.
     step_die: icepool.Die = icepool.Die(
         [make_step_outcome(outcome, zero) for outcome in outcomes], times)
 
-    @cache
-    def evaluate(state):
-        flat, terminal, again, count = state
+    def step(state, roll):
+        flat, terminal, again = state
         if again == 0:
             return state
-        if count < 0:
-            return icepool.Reroll
         if terminal + again > again_count + 1:
             return icepool.Reroll
-        return (step_die + state).map(evaluate, star=False)
+        state += roll
+        return state
 
-    initial_state = icepool.Vector([zero, 0, 1, again_count])
+    initial_state: icepool.Die = icepool.Die([icepool.Vector([zero, 0, 1])])
 
-    final_state = evaluate(initial_state)
+    final_state: icepool.Die = initial_state.map(step,
+                                                 step_die,
+                                                 star=False,
+                                                 repeat=again_count + 1)
 
-    def finalize(flat, terminal, again, count):
+    def finalize(flat, terminal, again):
+        if again > 0:
+            return icepool.Reroll
         return flat + terminal @ not_again_die
 
     return final_state.map(finalize, star=True)

tests/again_test.py

@@ -101,17 +101,19 @@ def test_is_additive():
 @pytest.mark.parametrize('n', [0, 1, 2])
 def test_again_count(n):
     count = Die([1, 2, 3, 4, 5, 6 + Again], again_count=n)
-    depth = Die([1, 2, 3, 4, 5, 6 + Again], again_depth=n, again_end=Reroll)
+    depth = d6.explode(depth=n)
+    depth = depth.reroll([depth.max_outcome()], depth=None)
     assert count == depth
 
 
 def test_again_count_double_blocked():
     result = Die([1, 2, 3, 4, 5, 6 + Again + Again], again_count=1)
     expected = d(5)
-    assert result == expected
+    assert result.simplify() == expected
 
 
 def test_again_count_double():
     result = Die([1, 2, 3, 4, 5, 6 + Again + Again], again_count=2)
-    expected = d6.map({6: 6 + 2 @ d(5)})
+    bonus = 2 @ d6.map({6: 100})
+    expected = d6.map({6: 6 + bonus}).reroll(lambda x: x > 100, depth=None)
     assert result == expected