From 1647ab75d9a2578cc952eb9fd1011bfb8092df27 Mon Sep 17 00:00:00 2001
From: Albert Julius Liu <ajul1987+highdiceroller@gmail.com>
Date: Tue, 4 Jun 2024 01:28:25 -0700
Subject: [PATCH] `again_count` rerolls the entire process when exceeded

---
 src/icepool/__init__.py         |  5 ++++-
 src/icepool/population/again.py | 25 ++++++++++++++-----------
 tests/again_test.py             |  8 +++++---
 3 files changed, 23 insertions(+), 15 deletions(-)

diff --git a/src/icepool/__init__.py b/src/icepool/__init__.py
index 1ef80142..915acde3 100644
--- a/src/icepool/__init__.py
+++ b/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
 
diff --git a/src/icepool/population/again.py b/src/icepool/population/again.py
index 516ec000..0ce20582 100644
--- a/src/icepool/population/again.py
+++ b/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)
diff --git a/tests/again_test.py b/tests/again_test.py
index d4236015..375481e7 100644
--- a/tests/again_test.py
+++ b/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