{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE RankNTypes #-}

module GHC.Data.List.Infinite
  ( Infinite (..)
  , head, tail
  , filter
  , (++)
  , unfoldr
  , (!!)
  , groupBy
  , dropList
  , iterate
  , concatMap
  , allListsOf
  , toList
  , repeat
  ) where

import Prelude ((-), Applicative (..), Bool (..), Foldable, Functor (..), Int, Maybe (..), Traversable (..), flip, otherwise)
import Control.Category (Category (..))
import Control.Monad (guard)
import qualified Data.Foldable as F
import Data.List.NonEmpty (NonEmpty (..))
import qualified GHC.Base as List (build)

data Infinite a = Inf a (Infinite a)
  deriving (forall a. Eq a => a -> Infinite a -> Bool
forall a. Num a => Infinite a -> a
forall a. Ord a => Infinite a -> a
forall m. Monoid m => Infinite m -> m
forall a. Infinite a -> Bool
forall a. Infinite a -> Int
forall a. Infinite a -> [a]
forall a. (a -> a -> a) -> Infinite a -> a
forall m a. Monoid m => (a -> m) -> Infinite a -> m
forall b a. (b -> a -> b) -> b -> Infinite a -> b
forall a b. (a -> b -> b) -> b -> Infinite a -> b
forall (t :: * -> *).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
product :: forall a. Num a => Infinite a -> a
$cproduct :: forall a. Num a => Infinite a -> a
sum :: forall a. Num a => Infinite a -> a
$csum :: forall a. Num a => Infinite a -> a
minimum :: forall a. Ord a => Infinite a -> a
$cminimum :: forall a. Ord a => Infinite a -> a
maximum :: forall a. Ord a => Infinite a -> a
$cmaximum :: forall a. Ord a => Infinite a -> a
elem :: forall a. Eq a => a -> Infinite a -> Bool
$celem :: forall a. Eq a => a -> Infinite a -> Bool
length :: forall a. Infinite a -> Int
$clength :: forall a. Infinite a -> Int
null :: forall a. Infinite a -> Bool
$cnull :: forall a. Infinite a -> Bool
toList :: forall a. Infinite a -> [a]
$ctoList :: forall a. Infinite a -> [a]
foldl1 :: forall a. (a -> a -> a) -> Infinite a -> a
$cfoldl1 :: forall a. (a -> a -> a) -> Infinite a -> a
foldr1 :: forall a. (a -> a -> a) -> Infinite a -> a
$cfoldr1 :: forall a. (a -> a -> a) -> Infinite a -> a
foldl' :: forall b a. (b -> a -> b) -> b -> Infinite a -> b
$cfoldl' :: forall b a. (b -> a -> b) -> b -> Infinite a -> b
foldl :: forall b a. (b -> a -> b) -> b -> Infinite a -> b
$cfoldl :: forall b a. (b -> a -> b) -> b -> Infinite a -> b
foldr' :: forall a b. (a -> b -> b) -> b -> Infinite a -> b
$cfoldr' :: forall a b. (a -> b -> b) -> b -> Infinite a -> b
foldr :: forall a b. (a -> b -> b) -> b -> Infinite a -> b
$cfoldr :: forall a b. (a -> b -> b) -> b -> Infinite a -> b
foldMap' :: forall m a. Monoid m => (a -> m) -> Infinite a -> m
$cfoldMap' :: forall m a. Monoid m => (a -> m) -> Infinite a -> m
foldMap :: forall m a. Monoid m => (a -> m) -> Infinite a -> m
$cfoldMap :: forall m a. Monoid m => (a -> m) -> Infinite a -> m
fold :: forall m. Monoid m => Infinite m -> m
$cfold :: forall m. Monoid m => Infinite m -> m
Foldable, forall a b. a -> Infinite b -> Infinite a
forall a b. (a -> b) -> Infinite a -> Infinite b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> Infinite b -> Infinite a
$c<$ :: forall a b. a -> Infinite b -> Infinite a
fmap :: forall a b. (a -> b) -> Infinite a -> Infinite b
$cfmap :: forall a b. (a -> b) -> Infinite a -> Infinite b
Functor, Functor Infinite
Foldable Infinite
forall (t :: * -> *).
Functor t
-> Foldable t
-> (forall (f :: * -> *) a b.
    Applicative f =>
    (a -> f b) -> t a -> f (t b))
-> (forall (f :: * -> *) a. Applicative f => t (f a) -> f (t a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> t a -> m (t b))
-> (forall (m :: * -> *) a. Monad m => t (m a) -> m (t a))
-> Traversable t
forall (m :: * -> *) a. Monad m => Infinite (m a) -> m (Infinite a)
forall (f :: * -> *) a.
Applicative f =>
Infinite (f a) -> f (Infinite a)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Infinite a -> m (Infinite b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Infinite a -> f (Infinite b)
sequence :: forall (m :: * -> *) a. Monad m => Infinite (m a) -> m (Infinite a)
$csequence :: forall (m :: * -> *) a. Monad m => Infinite (m a) -> m (Infinite a)
mapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Infinite a -> m (Infinite b)
$cmapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Infinite a -> m (Infinite b)
sequenceA :: forall (f :: * -> *) a.
Applicative f =>
Infinite (f a) -> f (Infinite a)
$csequenceA :: forall (f :: * -> *) a.
Applicative f =>
Infinite (f a) -> f (Infinite a)
traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Infinite a -> f (Infinite b)
$ctraverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Infinite a -> f (Infinite b)
Traversable)

head :: Infinite a -> a
head :: forall a. Infinite a -> a
head (Inf a
a Infinite a
_) = a
a
{-# NOINLINE [1] head #-}

tail :: Infinite a -> Infinite a
tail :: forall a. Infinite a -> Infinite a
tail (Inf a
_ Infinite a
as) = Infinite a
as
{-# NOINLINE [1] tail #-}

{-# RULES
"head/build" forall (g :: forall b . (a -> b -> b) -> b) . head (build g) = g \ x _ -> x
  #-}

instance Applicative Infinite where
    pure :: forall a. a -> Infinite a
pure = forall a. a -> Infinite a
repeat
    Inf a -> b
f Infinite (a -> b)
fs <*> :: forall a b. Infinite (a -> b) -> Infinite a -> Infinite b
<*> Inf a
a Infinite a
as = forall a. a -> Infinite a -> Infinite a
Inf (a -> b
f a
a) (Infinite (a -> b)
fs forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Infinite a
as)

mapMaybe :: (a -> Maybe b) -> Infinite a -> Infinite b
mapMaybe :: forall a b. (a -> Maybe b) -> Infinite a -> Infinite b
mapMaybe a -> Maybe b
f = Infinite a -> Infinite b
go
  where
    go :: Infinite a -> Infinite b
go (Inf a
a Infinite a
as) = let bs :: Infinite b
bs = Infinite a -> Infinite b
go Infinite a
as in case a -> Maybe b
f a
a of
        Maybe b
Nothing -> Infinite b
bs
        Just b
b -> forall a. a -> Infinite a -> Infinite a
Inf b
b Infinite b
bs
{-# NOINLINE [1] mapMaybe #-}

{-# RULES
"mapMaybe" [~1] forall f as . mapMaybe f as = build \ c -> foldr (mapMaybeFB c f) as
"mapMaybeList" [1] forall f . foldr (mapMaybeFB Inf f) = mapMaybe f
  #-}

{-# INLINE [0] mapMaybeFB #-}
mapMaybeFB :: (b -> r -> r) -> (a -> Maybe b) -> a -> r -> r
mapMaybeFB :: forall b r a. (b -> r -> r) -> (a -> Maybe b) -> a -> r -> r
mapMaybeFB b -> r -> r
cons a -> Maybe b
f a
a r
bs = case a -> Maybe b
f a
a of
    Maybe b
Nothing -> r
bs
    Just b
r -> b -> r -> r
cons b
r r
bs

filter :: (a -> Bool) -> Infinite a -> Infinite a
filter :: forall a. (a -> Bool) -> Infinite a -> Infinite a
filter a -> Bool
f = forall a b. (a -> Maybe b) -> Infinite a -> Infinite b
mapMaybe (\ a
a -> a
a forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ forall (f :: * -> *). Alternative f => Bool -> f ()
guard (a -> Bool
f a
a))
{-# INLINE filter #-}

infixr 5 ++
(++) :: Foldable f => f a -> Infinite a -> Infinite a
++ :: forall (f :: * -> *) a.
Foldable f =>
f a -> Infinite a -> Infinite a
(++) = forall a b c. (a -> b -> c) -> b -> a -> c
flip (forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
F.foldr forall a. a -> Infinite a -> Infinite a
Inf)

unfoldr :: (b -> (a, b)) -> b -> Infinite a
unfoldr :: forall b a. (b -> (a, b)) -> b -> Infinite a
unfoldr b -> (a, b)
f b
b = forall a. (forall b. (a -> b -> b) -> b) -> Infinite a
build \ a -> b -> b
c -> let go :: b -> b
go b
b = case b -> (a, b)
f b
b of (a
a, b
b') -> a
a a -> b -> b
`c` b -> b
go b
b' in b -> b
go b
b
{-# INLINE unfoldr #-}

(!!) :: Infinite a -> Int -> a
Inf a
a Infinite a
_ !! :: forall a. Infinite a -> Int -> a
!! Int
0 = a
a
Inf a
_ Infinite a
as !! Int
n = Infinite a
as forall a. Infinite a -> Int -> a
!! (Int
nforall a. Num a => a -> a -> a
-Int
1)

groupBy :: (a -> a -> Bool) -> Infinite a -> Infinite (NonEmpty a)
groupBy :: forall a. (a -> a -> Bool) -> Infinite a -> Infinite (NonEmpty a)
groupBy a -> a -> Bool
eq = Infinite a -> Infinite (NonEmpty a)
go
  where
    go :: Infinite a -> Infinite (NonEmpty a)
go (Inf a
a Infinite a
as) = forall a. a -> Infinite a -> Infinite a
Inf (a
aforall a. a -> [a] -> NonEmpty a
:|[a]
bs) (Infinite a -> Infinite (NonEmpty a)
go Infinite a
cs)
      where ([a]
bs, Infinite a
cs) = forall a. (a -> Bool) -> Infinite a -> ([a], Infinite a)
span (a -> a -> Bool
eq a
a) Infinite a
as

span :: (a -> Bool) -> Infinite a -> ([a], Infinite a)
span :: forall a. (a -> Bool) -> Infinite a -> ([a], Infinite a)
span a -> Bool
p = forall a b. (a -> Maybe b) -> Infinite a -> ([b], Infinite a)
spanJust (\ a
a -> a
a forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ forall (f :: * -> *). Alternative f => Bool -> f ()
guard (a -> Bool
p a
a))
{-# INLINE span #-}

spanJust :: (a -> Maybe b) -> Infinite a -> ([b], Infinite a)
spanJust :: forall a b. (a -> Maybe b) -> Infinite a -> ([b], Infinite a)
spanJust a -> Maybe b
p = Infinite a -> ([b], Infinite a)
go
  where
    go :: Infinite a -> ([b], Infinite a)
go as :: Infinite a
as@(Inf a
a Infinite a
as')
      | Just b
b <- a -> Maybe b
p a
a = let ([b]
bs, Infinite a
cs) = Infinite a -> ([b], Infinite a)
go Infinite a
as' in (b
bforall a. a -> [a] -> [a]
:[b]
bs, Infinite a
cs)
      | Bool
otherwise = ([], Infinite a
as)

iterate :: (a -> a) -> a -> Infinite a
iterate :: forall a. (a -> a) -> a -> Infinite a
iterate a -> a
f = a -> Infinite a
go where go :: a -> Infinite a
go a
a = forall a. a -> Infinite a -> Infinite a
Inf a
a (a -> Infinite a
go (a -> a
f a
a))
{-# NOINLINE [1] iterate #-}

{-# RULES
"iterate" [~1] forall f a . iterate f a = build (\ c -> iterateFB c f a)
"iterateFB" [1] iterateFB Inf = iterate
  #-}

iterateFB :: (a -> b -> b) -> (a -> a) -> a -> b
iterateFB :: forall a b. (a -> b -> b) -> (a -> a) -> a -> b
iterateFB a -> b -> b
c a -> a
f a
a = a -> b
go a
a
  where go :: a -> b
go a
a = a
a a -> b -> b
`c` a -> b
go (a -> a
f a
a)
{-# INLINE [0] iterateFB #-}

concatMap :: Foldable f => (a -> f b) -> Infinite a -> Infinite b
concatMap :: forall (f :: * -> *) a b.
Foldable f =>
(a -> f b) -> Infinite a -> Infinite b
concatMap a -> f b
f = Infinite a -> Infinite b
go where go :: Infinite a -> Infinite b
go (Inf a
a Infinite a
as) = a -> f b
f a
a forall (f :: * -> *) a.
Foldable f =>
f a -> Infinite a -> Infinite a
++ Infinite a -> Infinite b
go Infinite a
as
{-# NOINLINE [1] concatMap #-}

{-# RULES "concatMap" forall f as . concatMap f as = build \ c -> foldr (\ x b -> F.foldr c b (f x)) as #-}

{-# SPECIALIZE concatMap :: (a -> [b]) -> Infinite a -> Infinite b #-}

foldr :: (a -> b -> b) -> Infinite a -> b
foldr :: forall a b. (a -> b -> b) -> Infinite a -> b
foldr a -> b -> b
f = Infinite a -> b
go where go :: Infinite a -> b
go (Inf a
a Infinite a
as) = a -> b -> b
f a
a (Infinite a -> b
go Infinite a
as)
{-# INLINE [0] foldr #-}

build :: (forall b . (a -> b -> b) -> b) -> Infinite a
build :: forall a. (forall b. (a -> b -> b) -> b) -> Infinite a
build forall b. (a -> b -> b) -> b
g = forall b. (a -> b -> b) -> b
g forall a. a -> Infinite a -> Infinite a
Inf
{-# INLINE [1] build #-}

-- Analogous to 'foldr'/'build' fusion for '[]'
{-# RULES
"foldr/build" forall f (g :: forall b . (a -> b -> b) -> b) . foldr f (build g) = g f
"foldr/id" foldr Inf = id

"foldr/cons/build" forall f a (g :: forall b . (a -> b -> b) -> b) . foldr f (Inf a (build g)) = f a (g f)
  #-}

{-# RULES
"map" [~1] forall f (as :: Infinite a) . fmap f as = build \ c -> foldr (mapFB c f) as
"mapFB" forall c f g . mapFB (mapFB c f) g = mapFB c (f . g)
"mapFB/id" forall c . mapFB c (\ x -> x) = c
  #-}

mapFB :: (b -> c -> c) -> (a -> b) -> a -> c -> c
mapFB :: forall b c a. (b -> c -> c) -> (a -> b) -> a -> c -> c
mapFB b -> c -> c
c a -> b
f = \ a
x c
ys -> b -> c -> c
c (a -> b
f a
x) c
ys
{-# INLINE [0] mapFB #-}

dropList :: [a] -> Infinite b -> Infinite b
dropList :: forall a b. [a] -> Infinite b -> Infinite b
dropList [] Infinite b
bs = Infinite b
bs
dropList (a
_:[a]
as) (Inf b
_ Infinite b
bs) = forall a b. [a] -> Infinite b -> Infinite b
dropList [a]
as Infinite b
bs

-- | Compute all lists of the given alphabet.
-- For example: @'allListsOf' "ab" = ["a", "b", "aa", "ba", "ab", "bb", "aaa", "baa", "aba", ...]@
allListsOf :: [a] -> Infinite [a]
allListsOf :: forall a. [a] -> Infinite [a]
allListsOf [a]
as = forall (f :: * -> *) a b.
Foldable f =>
(a -> f b) -> Infinite a -> Infinite b
concatMap (\ [a]
bs -> [a
aforall a. a -> [a] -> [a]
:[a]
bs | a
a <- [a]
as]) ([] forall a. a -> Infinite a -> Infinite a
`Inf` forall a. [a] -> Infinite [a]
allListsOf [a]
as)

-- See Note [Fusion for `Infinite` lists].
toList :: Infinite a -> [a]
toList :: forall a. Infinite a -> [a]
toList = \ Infinite a
as -> forall a. (forall b. (a -> b -> b) -> b -> b) -> [a]
List.build (\ a -> b -> b
c b
_ -> forall a b. (a -> b -> b) -> Infinite a -> b
foldr a -> b -> b
c Infinite a
as)
{-# INLINE toList #-}

repeat :: a -> Infinite a
repeat :: forall a. a -> Infinite a
repeat a
a = Infinite a
as where as :: Infinite a
as = forall a. a -> Infinite a -> Infinite a
Inf a
a Infinite a
as
{-# INLINE [0] repeat #-}

repeatFB :: (a -> b -> b) -> a -> b
repeatFB :: forall a b. (a -> b -> b) -> a -> b
repeatFB a -> b -> b
c a
x = b
xs where xs :: b
xs = a -> b -> b
c a
x b
xs
{-# INLINE [0] repeatFB #-}

{-# RULES
"repeat" [~1] forall a . repeat a = build \ c -> repeatFB c a
"repeatFB" [1] repeatFB Inf = repeat
  #-}

{-
Note [Fusion for `Infinite` lists]
~~~~~~~~~~~~~~~~~~~~
We use RULES to support foldr/build fusion for Infinite lists, analogously to the RULES in
GHC.Base to support fusion for regular lists. In particular, we define the following:
• `build :: (forall b . (a -> b -> b) -> b) -> Infinite a`
• `foldr :: (a -> b -> b) -> Infinite a -> b`
• A RULE `foldr f (build g) = g f`
• `Infinite`-producing functions in terms of `build`, and `Infinite`-consuming functions in
  terms of `foldr`

This can work across data types. For example, consider `toList :: Infinite a -> [a]`.
We want 'toList' to be both a good consumer (of 'Infinite' lists) and a good producer (of '[]').
Ergo, we define it in terms of 'Infinite.foldr' and `List.build`.

For a bigger example, consider `List.map f (toList (Infinite.map g as))`

We want to fuse away the intermediate `Infinite` structure between `Infnite.map` and `toList`,
and the list structure between `toList` and `List.map`. And indeed we do: see test
"InfiniteListFusion".
-}