change again_end=Reroll to reroll if the total number of dice would…

… exceed `again_count`

the behavior in the depth mode is changed back so that the reroll happens on the last level rather than one level extra

src/icepool/__init__.py

@@ -78,11 +78,11 @@
 
 For finer control over rolling processes, use e.g. `Die.map()` instead.
 
-#### `again_count` mode
+#### Count mode
 
-Effectively, we start with one roll queued and execute one 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.
+When `again_count` is provided, we start with one roll queued and execute one 
+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 we execute `again_count` rolls without running out, the result is the sum of 
 the rolls plus the number of leftover rolls @ `again_end`.
@@ -93,16 +93,26 @@
 * Binary `+`
 * `n @ AgainExpression`, where `n` is a non-negative `int` or `Population`.
 
-#### `again_depth` mode
+Furthermore, the `+` operator is assumed to be associative and commutative.
+
+#### Depth mode
 
 When `again_depth=0`, `again_end` is directly substituted
-for each occurence of `Again`. Otherwise, the result for `again_depth-1` is
-substituted for each occurence of `Again`.
+for each occurence of `Again`. For other values of `again_depth`, the result for
+`again_depth-1` is substituted for each occurence of `Again`.
+
+If `again_end=icepool.Reroll`, then any `AgainExpression`s in the final depth
+are rerolled.
+
+#### Rerolls
 
-### `Reroll`
+`Reroll` only rerolls that particular die, not the entire process. Any such
+rerolls do not count against the `again_count` or `again_depth` limit.
 
-If `again_end=icepool.Reroll`, then `again_end` rolls one final time for each
-instance of `Again`, rerolling until a non-`AgainExpression` is reached.
+If `again_end=icepool.Reroll`:
+* Count mode: Any result that would cause the number of rolls to exceed
+    `again_count` is rerolled.
+* Depth mode: Any `AgainExpression`s in the final depth level are rerolled.
 """
 
 from icepool.population.die_with_truth import DieWithTruth

src/icepool/population/again.py

@@ -33,11 +33,7 @@ def __init__(self,
                 from rolling again. Only applicable if `function` is provided.
             truth_value: The truth value of the resulting object, if applicable.
                 You probably don't need to use this externally.
-            is_additive: Whether the only operators applied to
-                sub-`AgainExpression`s containing `Again` are:
-                * Binary `+`
-                * `n @ AgainExpression`, where `n` is a non-negative `int` or
-                    `Population`.
+            is_additive: Whether the expression is compatible with `again_count`.
         """
         self._func = function
         self._args = args
@@ -261,7 +257,7 @@ def compute_terminal(outcomes: Sequence, times: Sequence[int],
         return not_again_die, not_again_die.zero_outcome(
         ), not_again_die.zero_outcome()
     elif again_end is icepool.Reroll:
-        return not_again_die, not_again_die, not_again_die.zero_outcome()
+        return not_again_die, icepool.Reroll, not_again_die.zero_outcome()
     else:
         return not_again_die, again_end, again_end * 0
 
@@ -280,10 +276,10 @@ def make_again_count_outcome(outcome, zero):
                 raise ValueError(
                     'again_count mode cannot be used with a non-additive AgainExpression.'
                 )
-            return icepool.vectorize(outcome._evaluate(zero), 0,
-                                     outcome._again_count())
+            return icepool.tupleize(outcome._evaluate(zero), False,
+                                    outcome._again_count())
         else:
-            return icepool.vectorize(zero, 1, 0)
+            return icepool.tupleize(zero, True, 0)
 
     # Flat total, added again count.
     again_count_die: icepool.Die[tuple[Any, int]] = icepool.Die(
@@ -292,25 +288,34 @@ def make_again_count_outcome(outcome, zero):
 
     # State: flat total, number of terminal dice, remaining number of agains
     initial_state: icepool.Die[tuple[Any, int,
-                                     int]] = icepool.Die([(zero, 0, 1)])
+                                     int]] = icepool.Die([(0, zero, 0, 1)])
 
-    def next_again_count_state(flat, terminal: int, remaining_again: int,
+    def next_again_count_state(step, flat, terminal: int, again: int,
                                roll_result: tuple[Any, int, int]):
-        if remaining_again == 0:
-            return flat, terminal, remaining_again
-        add_flat, add_terminal, add_again = roll_result
-        return (flat + add_flat, terminal + add_terminal,
-                remaining_again - 1 + add_again)
-
-    if again_end is icepool.Reroll:
-        again_end = not_again_die
-    else:
-        again_end = icepool.implicit_convert_to_die(again_end)
+        if step == -1:
+            return step, flat, terminal, again
+        add_flat, is_terminal, add_again = roll_result
+        step += 1
+        flat += add_flat
+        terminal += is_terminal
+        again += add_again - 1
+        if step + again > again_count + 1:
+            return icepool.Reroll
+        if again == 0:
+            return -1, flat, terminal, again
+        return (step, flat, terminal, again)
 
     final_state: icepool.Die[tuple[Any, int, int]] = initial_state.map(
-        next_again_count_state, again_count_die, star=True, repeat=again_count)
+        next_again_count_state,
+        again_count_die,
+        star=True,
+        repeat=again_count + 1)
 
-    def finalize(flat, terminal: int, remaining_again: int):
+    if again_end is icepool.Reroll:
+        again_end = zero
+    again_end = icepool.implicit_convert_to_die(again_end)
+
+    def finalize(step, flat, terminal: int, remaining_again: int):
         return flat + terminal @ not_again_die + remaining_again @ again_end
 
     return final_state.map(finalize, star=True)
@@ -331,9 +336,12 @@ def evaluate_agains_using_depth(outcomes: Sequence, times: Sequence[int],
     return tail
 
 
-def replace_again(outcome, die: 'icepool.Die'):
+def replace_again(outcome, repl):
     if isinstance(outcome, AgainExpression):
-        return outcome._evaluate(die)
+        if repl is icepool.Reroll:
+            return icepool.Reroll
+        else:
+            return outcome._evaluate(repl)
     else:
         # tuple or simple arg that is not Again.
         return outcome

tests/again_test.py

@@ -65,13 +65,12 @@ def test_again_plus_again_depth_1():
 
 def test_again_reroll_depth_0():
     die = Die([1, 2, 3, 4, 5, 6 + Again], again_depth=0, again_end=Reroll)
-    assert die == icepool.d6.map({6: 6 + d(5)})
+    assert die == d(5)
 
 
 def test_again_reroll_depth_1():
     die = Die([1, 2, 3, 4, 5, 6 + Again], again_depth=1, again_end=Reroll)
-    second = icepool.d6.map({6: 6 + d(5)})
-    assert die == icepool.d6.map({6: 6 + second})
+    assert die == d6.map({6: 6 + d(5)})
 
 
 def test_again_infinity():
@@ -97,3 +96,13 @@ def test_is_additive():
     assert not (2 * Again).is_additive
     assert (d6 @ Again).is_additive
     assert not ((d6 - 2) @ Again).is_additive
+
+
+def test_again_count_0():
+    die = Die([1, 2, 3, 4, 5, 6 + Again], again_count=0, again_end=Reroll)
+    assert die == d(5)
+
+
+def test_again_count_1():
+    die = Die([1, 2, 3, 4, 5, 6 + Again], again_count=1, again_end=Reroll)
+    assert die == d6.map({6: 6 + d(5)})