From 8cd75b0fa3ce05b6ab83a9c0eb6a80ef83ad335b Mon Sep 17 00:00:00 2001
From: Sergiu Ivanov <sivanov@colimite.fr>
Date: Fri, 4 Mar 2022 17:24:18 +0100
Subject: [PATCH] utils: Add cartesian-product-2/stream.

---
 scribblings/utils.scrbl | 20 ++++++++++++++++++++
 utils.rkt               | 29 +++++++++++++++++++++++++++--
 2 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/scribblings/utils.scrbl b/scribblings/utils.scrbl
index 2a5fba7..858b173 100644
--- a/scribblings/utils.scrbl
+++ b/scribblings/utils.scrbl
@@ -427,6 +427,26 @@ range @racket[min] to @racket[max]-1.}
 (stream->list (stream-take (in-random 5 10) 5))
 ]}
 
+@defproc[(cartesian-product-2/stream [s1 (Sequenceof a)]
+				     [s2 (Sequenceof b)])
+				     (Sequenceof (Pair a b))]{
+
+Generates a stream containing all the pairs of the elements from @racket[s1]
+and @racket[s2].  The elements of @racket[s2] are enumerated in order for every
+element of the @racket[s1], taken in order as well.
+
+@examples[#:eval utils-evaluator
+(require typed/racket/stream)
+(stream->list (cartesian-product-2/stream (in-range 1 5) '(a b)))
+]
+
+The streams can be infinite.  If the second stream is infinite, only the first
+element of @racket[s1] will be enumerated.
+
+@examples[#:eval utils-evaluator
+(stream->list (stream-take (cartesian-product-2/stream '(a b) (in-naturals)) 10))
+]}
+
 @section{Functions and procedures}
 
 @section{Randomness}
diff --git a/utils.rkt b/utils.rkt
index f515e22..c4994ee 100644
--- a/utils.rkt
+++ b/utils.rkt
@@ -24,7 +24,7 @@
    list-sets->list-strings pretty-print-set pretty-print-set-sets
    update-vertices/unweighted update-graph dotit collect-by-key
    collect-by-key/sets ht-values/list->set hash->list/ordered
-   multi-split-at lists-transpose in-random)
+   multi-split-at lists-transpose in-random cartesian-product-2/stream)
 
   (define-type Variable Symbol)
   (define-type (VariableMapping A) (Immutable-HashTable Variable A))
@@ -475,6 +475,31 @@
                       0.22552976312818723
                       0.21646500425988832
                       0.15188352823997242))))
+
+  (: cartesian-product-2/stream (All (a b) (-> (Sequenceof a) (Sequenceof b) (Sequenceof (Pair a b)))))
+  (define (cartesian-product-2/stream s1 s2)
+    (: cp2-store (All (a b) (-> (Sequenceof a) (Sequenceof b) (Sequenceof b)
+                                (Sequenceof (Pair a b)))))
+    ;; The recursive implementation using s2-store as an accumulator.
+    ;; Main idea: combine the elements of s1 with the element of s2
+    ;; until they are exhausted, then restart with the next element of
+    ;; s1 and the original content of s2.
+    (define (cp2-store s1 s2 s2-store)
+      (cond
+        [(stream-empty? s1) (stream)]
+        [(stream-empty? s2) (cp2-store (stream-rest s1) s2-store s2-store)]
+        [else
+         (stream-cons (cons (stream-first s1) (stream-first s2))
+                      (cp2-store s1 (stream-rest s2) s2-store))]))
+    (cp2-store s1 s2 s2))
+
+  (module+ test
+    (test-case "cartesian-product-2/stream"
+      (check-equal? (stream->list (cartesian-product-2/stream (in-range 1 5) '(a b)))
+                    '((1 . a) (1 . b) (2 . a) (2 . b) (3 . a) (3 . b) (4 . a) (4 . b)))
+      (check-equal?
+       (stream->list (stream-take (cartesian-product-2/stream '(a b) (in-naturals)) 10))
+       '((a . 0) (a . 1) (a . 2) (a . 3) (a . 4) (a . 5) (a . 6) (a . 7) (a . 8) (a . 9)))))
   )
 
 (require 'typed)
@@ -485,7 +510,7 @@
          list-sets->list-strings pretty-print-set pretty-print-set-sets
          update-vertices/unweighted update-graph dotit collect-by-key
          collect-by-key/sets ht-values/list->set hash->list/ordered
-         multi-split-at lists-transpose in-random)
+         multi-split-at lists-transpose in-random cartesian-product-2/stream)
 
 ;;; Untyped section.