From 58185a06ceb48e4abf28b6f8e52812e1b0db1d58 Mon Sep 17 00:00:00 2001
From: John Benediktsson <mrjbq7@gmail.com>
Date: Thu, 31 May 2012 08:17:08 -0700
Subject: [PATCH] sequences.extras: simpler round-robin.

---
 extra/sequences/extras/extras-tests.factor |  3 ++-
 extra/sequences/extras/extras.factor       | 19 ++++++++-----------
 2 files changed, 10 insertions(+), 12 deletions(-)

diff --git a/extra/sequences/extras/extras-tests.factor b/extra/sequences/extras/extras-tests.factor
index 2ef0bfe15f..1c9a376ca3 100644
--- a/extra/sequences/extras/extras-tests.factor
+++ b/extra/sequences/extras/extras-tests.factor
@@ -72,4 +72,5 @@ IN: sequences.extras.tests
 { 8 } [ 3 iota dup [ 1 + * ] 2map-sum ] unit-test
 { 4 } [ "hello" "jello" [ = ] 2count ] unit-test
 
-{ "ADEBFC" } [ { "ABC" "D" "EF" } round-robin ] unit-test
+{ { } } [ { } round-robin ] unit-test
+{ "ADEBFC" } [ { "ABC" "D" "EF" } round-robin >string ] unit-test
diff --git a/extra/sequences/extras/extras.factor b/extra/sequences/extras/extras.factor
index 0759d8172f..b84955d843 100644
--- a/extra/sequences/extras/extras.factor
+++ b/extra/sequences/extras/extras.factor
@@ -184,15 +184,12 @@ PRIVATE>
 : 2count ( ... seq1 seq2 quot: ( ... elt1 elt2 -- ... ? ) -- ... n )
     [ 1 0 ? ] compose 2map-sum ; inline
 
-:: round-robin-as ( seqs exemplar -- newseq )
-    seqs length :> len
-    0 0 seqs sum-lengths [
-        f [
-            drop dup len >= [ drop 1 + 0 ] when
-            2dup seqs nth-unsafe ?nth
-            [ 1 + ] [ dup not ] bi*
-        ] loop
-    ] exemplar replicate-as 2nip ;
+: max-lengths ( seq -- n )
+    [ length ] [ max ] map-reduce ;
 
-: round-robin ( seqs -- newseq )
-    [ { } ] [ dup first round-robin-as ] if-empty ;
+: round-robin ( seq -- newseq )
+    [ { } ] [
+        dup [ max-lengths ] [ length ] bi [ iota ] bi@
+        [ [ 2array ] with map ] curry map concat swap
+        [ [ first2 ] dip nth-unsafe ?nth ] curry map sift
+    ] if-empty ;