{-# 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 (Foldable, Functor, Traversable)

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

tail :: Infinite a -> Infinite a
tail (Inf _ as) = 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 = repeat
    Inf f fs <*> Inf a as = Inf (f a) (fs <*> as)

mapMaybe :: (a -> Maybe b) -> Infinite a -> Infinite b
mapMaybe f = go
  where
    go (Inf a as) = let bs = go as in case f a of
        Nothing -> bs
        Just b -> Inf 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 cons f a bs = case f a of
    Nothing -> bs
    Just r -> cons r bs

filter :: (a -> Bool) -> Infinite a -> Infinite a
filter f = mapMaybe (\ a -> a <$ guard (f a))
{-# INLINE filter #-}

infixr 5 ++
(++) :: Foldable f => f a -> Infinite a -> Infinite a
(++) = flip (F.foldr Inf)

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

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

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

span :: (a -> Bool) -> Infinite a -> ([a], Infinite a)
span p = spanJust (\ a -> a <$ guard (p a))
{-# INLINE span #-}

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

iterate :: (a -> a) -> a -> Infinite a
iterate f = go where go a = Inf a (go (f 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 c f a = go a
  where go a = a `c` go (f a)
{-# INLINE [0] iterateFB #-}

concatMap :: Foldable f => (a -> f b) -> Infinite a -> Infinite b
concatMap f = go where go (Inf a as) = f a ++ go 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 f = go where go (Inf a as) = f a (go as)
{-# INLINE [0] foldr #-}

build :: (forall b . (a -> b -> b) -> b) -> Infinite a
build g = g 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 c f = \ x ys -> c (f x) ys
{-# INLINE [0] mapFB #-}

dropList :: [a] -> Infinite b -> Infinite b
dropList [] bs = bs
dropList (_:as) (Inf _ bs) = dropList as 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 as = concatMap (\ bs -> [a:bs | a <- as]) ([] `Inf` allListsOf as)

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

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

repeatFB :: (a -> b -> b) -> a -> b
repeatFB c x = xs where xs = c x 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".
-}