From 50fb3dab59fe4b44faed83024bcfd3a4df6c661c Mon Sep 17 00:00:00 2001
From: Sergiu Ivanov <sivanov@colimite.fr>
Date: Wed, 9 Feb 2022 23:55:20 +0100
Subject: [PATCH] utils: Type collect-by-key/sets.

---
 scribblings/utils.scrbl | 11 +++++++++++
 utils.rkt               | 31 ++++++++++++++++---------------
 2 files changed, 27 insertions(+), 15 deletions(-)

diff --git a/scribblings/utils.scrbl b/scribblings/utils.scrbl
index db99b99..d0fefbc 100644
--- a/scribblings/utils.scrbl
+++ b/scribblings/utils.scrbl
@@ -341,6 +341,17 @@ produced by this function are suitable for graph constructors.
 (collect-by-key '(a b a) '(1 2 3))
 ]}
 
+@defproc[(collect-by-key/sets [keys (Listof Any)] [vals (Listof Any)])
+			      (Values (Listof Any) (Listof (Setof Any)))]{
+
+Like @racket[collect-by-key], but produce a list of sets instead of a list
+of lists.
+
+@examples[#:eval utils-evaluator
+(collect-by-key/sets '(a b a) '(1 2 3))
+]}
+
+
 @section{Functions and procedures}
 
 @section{Randomness}
diff --git a/utils.rkt b/utils.rkt
index 008f987..11dcac6 100644
--- a/utils.rkt
+++ b/utils.rkt
@@ -22,7 +22,8 @@
    handle-org-booleans map-sexp read-org-sexp unorg unstringify-pairs
    read-org-variable-mapping unorgv read-symbol-list drop-first-last
    list-sets->list-strings pretty-print-set pretty-print-set-sets
-   update-vertices/unweighted update-graph dotit collect-by-key)
+   update-vertices/unweighted update-graph dotit collect-by-key
+   collect-by-key/sets)
 
   (define-type Variable Symbol)
   (define-type (VariableMapping A) (Immutable-HashTable Variable A))
@@ -396,6 +397,17 @@
       (define-values (e2 l2) (collect-by-key '((1 2) (1 2)) '(a b)))
       (check-equal? e1 '((1 2) (1 3))) (check-equal? l1 '((a) (b)))
       (check-equal? e2 '((1 2))) (check-equal? l2 '((b a)))))
+
+  (: collect-by-key/sets (-> (Listof Any) (Listof Any)
+                             (Values (Listof Any) (Listof (Setof Any)))))
+  (define (collect-by-key/sets edges labels)
+    (define-values (es ls) (collect-by-key edges labels))
+    (values es ((inst map (Setof Any) (Listof Any)) list->set ls)))
+
+  (module+ test
+    (test-case "collect-by-key/sets"
+      (define-values (e3 l3) (collect-by-key/sets '(a b a) '(1 2 1)))
+      (check-equal? e3 '(b a)) (check-equal? l3 (list (set 2) (set 1)))))
   )
 
 (require 'typed)
@@ -404,15 +416,14 @@
          map-sexp read-org-sexp unorg unstringify-pairs
          read-org-variable-mapping unorgv read-symbol-list drop-first-last
          list-sets->list-strings pretty-print-set pretty-print-set-sets
-         update-vertices/unweighted update-graph dotit collect-by-key)
+         update-vertices/unweighted update-graph dotit collect-by-key
+         collect-by-key/sets)
 
 ;;; Untyped section.
 
 (provide
  ;; Functions
- (contract-out [collect-by-key/sets (-> (listof any/c) (listof any/c) (values (listof any/c) (listof (set/c any/c))))]
-
-               [ht-values/list->set (-> (hash/c any/c (listof any/c)) (hash/c any/c (set/c any/c)))]
+ (contract-out [ht-values/list->set (-> (hash/c any/c (listof any/c)) (hash/c any/c (set/c any/c)))]
                [hash->list/ordered (-> hash? (listof (cons/c any/c any/c)))]
                [multi-split-at (-> (listof (listof any/c)) number?
                                    (values (listof (listof any/c)) (listof (listof any/c))))]
@@ -451,16 +462,6 @@
 ;;; Additional list and hash map utilities
 ;;; ======================================
 
-;;; Like collect-by-key, but returns a list of sets of weights.
-(define (collect-by-key/sets edges labels)
-  (let-values ([(es ls) (collect-by-key edges labels)])
-    (values es (map list->set ls))))
-
-(module+ test
-  (test-case "collect-by-key/sets"
-    (define-values (e3 l3) (collect-by-key/sets '(a b a) '(1 2 1)))
-    (check-equal? e3 '(b a)) (check-equal? l3 (list (set 2) (set 1)))))
-
 ;;; Converts the values of a hash table from lists to sets.
 (define (ht-values/list->set ht)
   (for/hash ([(k v) (in-hash ht)])