[ fix #2034 ] Productive cantor for Colist1

libs/base/Data/Colist1.idr

@@ -4,6 +4,7 @@ import Data.Colist
 import Data.List
 import Data.List1
 import Data.Nat
+import Data.Stream
 import public Data.Zippable
 
 %default total
@@ -114,6 +115,11 @@ public export
 appendl : Colist1 a -> List a -> Colist1 a
 appendl (x ::: xs) ys = x ::: appendl xs ys
 
+||| Take a `Colist1` apart
+public export
+uncons : Colist1 a -> (a, Colist a)
+uncons (h ::: tl) = (h, tl)
+
 ||| Extract the first element from a `Colist1`
 public export
 head : Colist1 a -> a
@@ -212,61 +218,52 @@ Zippable Colist1 where
 -- Interleavings
 --------------------------------------------------------------------------------
 
--- zig, zag, and cantor are generalised version of the Stream functions
--- defined in the paper
+-- zig, zag, and cantor are taken from the paper
 -- Applications of Applicative Proof Search
 -- by Liam O'Connor
 
--- Unfortunately cannot be put in `Data.Colist` because it's using `Colist1`
--- internally.
+public export
+zig : List1 (Colist1 a) -> Colist (Colist1 a) -> Colist a
+
+public export
+zag : List1 a -> List (Colist1 a) -> Colist (Colist1 a) -> Colist a
+
+zig xs = zag (head <$> xs) (mapMaybe (fromColist . tail) $ forget xs)
+
+zag (x ::: []) [] [] = x :: []
+zag (x ::: []) (z :: zs) [] = x :: zig (z ::: zs) []
+zag (x ::: []) zs (l :: ls) = x :: zig (l ::: zs) ls
+zag (x ::: (y :: xs)) zs ls = x :: zag (y ::: xs) zs ls
+
+public export
+cantor : Colist1 (Colist1 a) -> Colist1 a
+cantor (xxs ::: []) = xxs
+cantor ((x ::: xs) ::: (yys :: zzss))
+  = x ::: zig (yys ::: mapMaybe fromColist [xs]) zzss
+
 namespace Colist
 
-  public export
-  zig : List1 (Colist1 a) -> List (Colist1 a) -> Colist a
-
-  public export
-  zag : List1 a -> List1 (Colist a) -> List (Colist1 a) -> Colist a
-
-  zig xs =
-    let (hds, tls) = unzipWith (\ (x ::: xs) => (x, xs)) xs in
-    zag hds tls
-
-  zag (x ::: []) zs [] = x ::
-    let Just zs = List.toList1' $ mapMaybe fromColist (forget zs)
-          | Nothing => []
-    in zig zs []
-  zag (x ::: []) zs (l :: ls) = x ::
-    let zs = mapMaybe fromColist (forget zs)
-    in zig (l ::: zs) ls
-  zag (x ::: (y :: xs)) zs ls = x :: zag (y ::: xs) zs ls
-
   public export
   cantor : List (Colist a) -> Colist a
   cantor xs =
     let Just (l ::: ls) = List.toList1' $ mapMaybe fromColist xs
           | Nothing => []
-    in zig (l ::: []) ls
+    in zig (l ::: []) (fromList ls)
 
-public export
-zig : List1 (Colist1 a) -> Colist (Colist1 a) -> Colist1 a
-
-public export
-zag : List1 a -> List1 (Colist a) -> Colist (Colist1 a) -> Colist1 a
-
-zig xs = zag (head <$> xs) (tail <$> xs)
-
-zag (x ::: []) zs [] = x :::
-  let Just zs = List.toList1' (mapMaybe fromColist (forget zs))
-        | Nothing => []
-  in Colist.zig zs []
-zag (x ::: []) zs (l :: ls) =  x :::
-  let zs = mapMaybe fromColist (forget zs)
-  in forgetInf (zig (l ::: zs) ls)
-zag (x ::: (y :: xs)) zs ls = x ::: forgetInf (zag (y ::: xs) zs ls)
-
-public export
-cantor : Colist1 (Colist1 a) -> Colist1 a
-cantor (l ::: ls) = zig (l ::: []) ls
+  -- Exploring the (truncated) Nat*Nat top right quadrant of the plane
+  -- using Cantor's zig-zag traversal:
+  example :
+     let nats : Colist Nat; nats = fromStream Stream.nats in
+     take 10 (Colist.cantor [ map (0,) nats
+                            , map (1,) nats
+                            , map (2,) nats
+                            , map (3,) nats])
+     === [ (0, 0)
+         , (1, 0), (0, 1)
+         , (2, 0), (1, 1), (0, 2)
+         , (3, 0), (2, 1), (1, 2), (0, 3)
+         ]
+  example = Refl
 
 namespace DPair
 
@@ -302,3 +299,19 @@ namespace Pair
   public export
   plane : Colist1 a -> (a -> Colist1 b) -> Colist1 (a, b)
   plane = Pair.planeWith (,)
+
+--------------------------------------------------------------------------------
+-- Example
+--------------------------------------------------------------------------------
+
+-- Exploring the Nat*Nat top right quadrant of the plane
+-- using Cantor's zig-zag traversal:
+example :
+     let nats1 = fromStream Stream.nats in
+     Colist1.take 10 (Pair.plane nats1 (const nats1))
+     === (0, 0) :::
+         [ (1, 0), (0, 1)
+         , (2, 0), (1, 1), (0, 2)
+         , (3, 0), (2, 1), (1, 2), (0, 3)
+         ]
+example = Refl

libs/base/Data/List.idr

@@ -501,7 +501,7 @@ toList1 [] impossible
 toList1 (x :: xs) = x ::: xs
 
 ||| Convert to a non-empty list, returning Nothing if the list is empty.
-export
+public export
 toList1' : (l : List a) -> Maybe (List1 a)
 toList1' [] = Nothing
 toList1' (x :: xs) = Just (x ::: xs)

libs/base/Data/Stream.idr

@@ -114,20 +114,6 @@ public export
 cantor : Stream (Stream a) -> Stream a
 cantor (l :: ls) = zig (l ::: []) ls
 
--- Exploring the Nat*Nat top right quadrant of the plane
--- using Cantor's zig-zag traversal:
-example :
-  let quadrant : Stream (Stream (Nat, Nat))
-      quadrant = map (\ i => map (i,) Stream.nats) Stream.nats
-  in
-     take 10 (cantor quadrant)
-     === [ (0, 0)
-         , (1, 0), (0, 1)
-         , (2, 0), (1, 1), (0, 2)
-         , (3, 0), (2, 1), (1, 2), (0, 3)
-         ]
-example = Refl
-
 namespace DPair
 
   ||| Explore the plane corresponding to all possible pairings
@@ -163,6 +149,21 @@ namespace Pair
   plane : Stream a -> (a -> Stream b) -> Stream (a, b)
   plane = Pair.planeWith (,)
 
+--------------------------------------------------------------------------------
+-- Example
+--------------------------------------------------------------------------------
+
+-- Exploring the Nat*Nat top right quadrant of the plane
+-- using Cantor's zig-zag traversal:
+example :
+     take 10 (plane Stream.nats (const Stream.nats))
+     === [ (0, 0)
+         , (1, 0), (0, 1)
+         , (2, 0), (1, 1), (0, 2)
+         , (3, 0), (2, 1), (1, 2), (0, 3)
+         ]
+example = Refl
+
 --------------------------------------------------------------------------------
 -- Implementations
 --------------------------------------------------------------------------------