From 463c2718105ca142905e6e1b0ba524329dad3972 Mon Sep 17 00:00:00 2001
From: Albert Julius Liu <ajul1987+highdiceroller@gmail.com>
Date: Fri, 30 Aug 2024 11:31:01 -0700
Subject: [PATCH] add `sequence()` method to `Die` and `Deck`

---
 src/icepool/population/deck.py | 27 ++++++++++++++++++++++++++-
 src/icepool/population/die.py  |  9 +++++++++
 tests/vector_test.py           | 23 ++++++++++++++++++++++-
 3 files changed, 57 insertions(+), 2 deletions(-)

diff --git a/src/icepool/population/deck.py b/src/icepool/population/deck.py
index 0adf19a8..b327781c 100644
--- a/src/icepool/population/deck.py
+++ b/src/icepool/population/deck.py
@@ -13,7 +13,7 @@
 
 from collections import Counter, defaultdict
 from functools import cached_property
-from typing import Any, Callable, Iterable, Iterator, Mapping, Sequence, Type, overload
+from typing import Any, Callable, Iterable, Iterator, Mapping, MutableSequence, Sequence, Type, overload
 
 
 class Deck(Population[T_co]):
@@ -292,6 +292,31 @@ def transition_function(outcome):
             [transition_function(outcome) for outcome in self.outcomes()],
             times=self.quantities())
 
+    @cached_property
+    def _sequence_cache(
+            self) -> 'MutableSequence[icepool.Die[tuple[T_co, ...]]]':
+        return [icepool.Die([()])]
+
+    def sequence(self, deals: int, /) -> 'icepool.Die[tuple[T_co, ...]]':
+        """Possible sequences produced by dealing from this deck a number of times.
+        
+        This is extremely expensive computationally. If you don't care about
+        order, use `deal()` instead.
+        """
+        if deals < 0:
+            raise ValueError('The number of deals cannot be negative.')
+        for i in range(len(self._sequence_cache), deals + 1):
+
+            def transition(curr):
+                remaining = icepool.Die(self - curr)
+                return icepool.map(lambda curr, next: curr + (next, ), curr,
+                                   remaining)
+
+            result: 'icepool.Die[tuple[T_co, ...]]' = self._sequence_cache[
+                i - 1].map(transition)
+            self._sequence_cache.append(result)
+        return result
+
     @cached_property
     def _hash_key(self) -> tuple:
         return Deck, tuple(self.items())
diff --git a/src/icepool/population/die.py b/src/icepool/population/die.py
index c42760cc..ab7372ed 100644
--- a/src/icepool/population/die.py
+++ b/src/icepool/population/die.py
@@ -794,6 +794,15 @@ def __rmatmul__(self, other: 'int | Die[int]') -> 'Die':
         other = implicit_convert_to_die(other)
         return other.__matmul__(self)
 
+    def sequence(self, rolls: int) -> 'icepool.Die[tuple[T_co, ...]]':
+        """Possible sequences produced by rolling this die a number of times.
+        
+        This is extremely expensive computationally. If possible, use `reduce()`
+        instead; if you don't care about order, `Die.pool()` is better.
+        """
+        return icepool.cartesian_product(*(self for _ in range(rolls)),
+                                         outcome_type=tuple)  # type: ignore
+
     def pool(self, rolls: int | Sequence[int] = 1, /) -> 'icepool.Pool[T_co]':
         """Creates a `Pool` from this `Die`.
 
diff --git a/tests/vector_test.py b/tests/vector_test.py
index 0f656f88..1726c7df 100644
--- a/tests/vector_test.py
+++ b/tests/vector_test.py
@@ -1,7 +1,7 @@
 import icepool
 import pytest
 
-from icepool import d, d4, d6, d8, vectorize, Die, Vector
+from icepool import d, d4, d6, d8, vectorize, Die, Vector, map, Deck
 from collections import namedtuple
 
 
@@ -150,3 +150,24 @@ def test_named_tuple():
 
 def test_auto_tupleize():
     assert Die([(d6, d6, 6)]).map(sum) == 2 @ d6 + 6
+
+
+def test_die_sequence():
+    assert d6.sequence(3) == map(lambda a, b, c: (a, b, c), d6, d6, d6)
+
+
+def test_deck_sequence():
+    result = Deck(range(5)).sequence(3)
+    assert len(result) == 5 * 4 * 3
+    assert result.map(
+        lambda a, b, c: a != b and a != c and b != c).probability(True) == 1
+
+
+def test_deck_sequence_with_dups():
+    result = Deck([0, 0, 1, 1]).sequence(2)
+    assert result == Die({
+        (0, 0): 2,
+        (0, 1): 4,
+        (1, 0): 4,
+        (1, 1): 2,
+    })