{-# LANGUAGE DeriveFunctor       #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- | Possibly infinite streams of @'Maybe' a@s.
module System.FS.Sim.Stream (
    -- * Streams
    Stream
    -- * Running
  , runStream
    -- * Construction
  , always
  , empty
  , mkInfinite
  , repeating
  , unsafeMkFinite
    -- * Query
  , null
    -- * Generation and shrinking
  , genFinite
  , genInfinite
  , genMaybe
  , genMaybe'
  , shrinkStream
  ) where

import           Control.Monad (replicateM)
import           Prelude hiding (null)
import qualified Test.QuickCheck as QC
import           Test.QuickCheck (Gen)

{-------------------------------------------------------------------------------
  Streams
-------------------------------------------------------------------------------}

-- | A 'Stream' is a stream of @'Maybe' a@s, which is /possibly/ infinite or
-- /definitely/ finite.
--
-- Finiteness is tracked internally and used for 'QC.shrink'ing and the 'Show'
-- instance.
data Stream a = Stream {
      -- | Info about the size of the stream.
      forall a. Stream a -> InternalInfo
_streamInternalInfo :: InternalInfo
    , forall a. Stream a -> [Maybe a]
_getStream          :: [Maybe a]
    }
  deriving (forall a b. (a -> b) -> Stream a -> Stream b)
-> (forall a b. a -> Stream b -> Stream a) -> Functor Stream
forall a b. a -> Stream b -> Stream a
forall a b. (a -> b) -> Stream a -> Stream b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> Stream a -> Stream b
fmap :: forall a b. (a -> b) -> Stream a -> Stream b
$c<$ :: forall a b. a -> Stream b -> Stream a
<$ :: forall a b. a -> Stream b -> Stream a
Functor

-- | Tag for 'Stream's that describes whether it is either /definitely/ a finite
-- stream, or /possibly/ an infinite stream.
--
-- Useful for the 'Show' instance of 'Stream': when a 'Stream' is /definitely/
-- finite, we can safely print the full stream.
data InternalInfo = Infinite | Finite

-- | Fully shows a 'Stream' if it is /definitely/ finite, or prints a
-- placeholder string if it is /possibly/ infinite.
instance Show a => Show (Stream a) where
  showsPrec :: Int -> Stream a -> ShowS
showsPrec Int
n (Stream InternalInfo
info [Maybe a]
xs) = case InternalInfo
info of
      InternalInfo
Infinite -> (String
"<infinite stream>" String -> ShowS
forall a. [a] -> [a] -> [a]
++)
      InternalInfo
Finite   -> (if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
10 then (Char
'('Char -> ShowS
forall a. a -> [a] -> [a]
:) else ShowS
forall a. a -> a
id)
                ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Maybe a] -> ShowS
forall a. Show a => a -> ShowS
shows [Maybe a]
xs
                ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
" ++ ..." String -> ShowS
forall a. [a] -> [a] -> [a]
++)
                ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
10 then (Char
')'Char -> ShowS
forall a. a -> [a] -> [a]
:) else ShowS
forall a. a -> a
id)

{-------------------------------------------------------------------------------
  Running
-------------------------------------------------------------------------------}

-- | Advance the 'Stream'. Return the @'Maybe' a@ and the remaining 'Stream'.
--
-- Returns 'Nothing' by default if the 'Stream' is empty.
runStream :: Stream a -> (Maybe a, Stream a)
runStream :: forall a. Stream a -> (Maybe a, Stream a)
runStream s :: Stream a
s@(Stream InternalInfo
_    []    ) = (Maybe a
forall a. Maybe a
Nothing, Stream a
s)
runStream   (Stream InternalInfo
info (Maybe a
a:[Maybe a]
as)) = (Maybe a
a, InternalInfo -> [Maybe a] -> Stream a
forall a. InternalInfo -> [Maybe a] -> Stream a
Stream InternalInfo
info [Maybe a]
as)

{-------------------------------------------------------------------------------
  Construction
-------------------------------------------------------------------------------}

-- | Make an empty 'Stream'.
empty :: Stream a
empty :: forall a. Stream a
empty = InternalInfo -> [Maybe a] -> Stream a
forall a. InternalInfo -> [Maybe a] -> Stream a
Stream InternalInfo
Finite []

-- | Make a 'Stream' that always generates the given @a@.
always :: a -> Stream a
always :: forall a. a -> Stream a
always a
x = InternalInfo -> [Maybe a] -> Stream a
forall a. InternalInfo -> [Maybe a] -> Stream a
Stream InternalInfo
Infinite (Maybe a -> [Maybe a]
forall a. a -> [a]
repeat (a -> Maybe a
forall a. a -> Maybe a
Just a
x))

-- | Make a 'Stream' that infinitely repeats the given list.
repeating :: [Maybe a] -> Stream a
repeating :: forall a. [Maybe a] -> Stream a
repeating [Maybe a]
xs = InternalInfo -> [Maybe a] -> Stream a
forall a. InternalInfo -> [Maybe a] -> Stream a
Stream InternalInfo
Infinite ([Maybe a] -> Stream a) -> [Maybe a] -> Stream a
forall a b. (a -> b) -> a -> b
$ [[Maybe a]] -> [Maybe a]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([Maybe a] -> [[Maybe a]]
forall a. a -> [a]
repeat [Maybe a]
xs)

-- | UNSAFE: Make a 'Stream' that is marked as definitely finite.
--
-- This is unsafe since a user can pass in any list, and evaluating
-- 'Test.QuickCheck.shrink' or 'show' on the resulting 'Stream' will diverge. It
-- is the user's responsibility to only pass in a finite list.
unsafeMkFinite :: [Maybe a] -> Stream a
unsafeMkFinite :: forall a. [Maybe a] -> Stream a
unsafeMkFinite = InternalInfo -> [Maybe a] -> Stream a
forall a. InternalInfo -> [Maybe a] -> Stream a
Stream InternalInfo
Finite

-- | Make a 'Stream' that is marked as possibly infinite.
mkInfinite :: [Maybe a] -> Stream a
mkInfinite :: forall a. [Maybe a] -> Stream a
mkInfinite = InternalInfo -> [Maybe a] -> Stream a
forall a. InternalInfo -> [Maybe a] -> Stream a
Stream InternalInfo
Infinite

{-------------------------------------------------------------------------------
  Query
-------------------------------------------------------------------------------}

-- | Return 'True' if the stream is empty.
--
-- A stream consisting of only 'Nothing's (even if it is only one) is not
-- considered to be empty.
null :: Stream a -> Bool
null :: forall a. Stream a -> Bool
null (Stream InternalInfo
_ []) = Bool
True
null Stream a
_             = Bool
False

{-------------------------------------------------------------------------------
  Generation and shrinking
-------------------------------------------------------------------------------}

-- | Shrink a stream like it is an 'Test.QuickCheck.InfiniteList'.
--
-- Possibly infinite streams are shrunk differently than lists that are
-- definitely finite, which is to ensure that shrinking terminates.
-- * Possibly infinite streams are shrunk by taking finite prefixes of the
--  argument stream. As such, shrinking a possibly infinite stream creates
--  definitely finite streams.
-- * Definitely finite streams are shrunk like lists are shrunk normally,
--   preserving that the created streams are still definitely finite.
shrinkStream :: Stream a -> [Stream a]
shrinkStream :: forall a. Stream a -> [Stream a]
shrinkStream (Stream InternalInfo
info [Maybe a]
xs0) = case InternalInfo
info of
    InternalInfo
Infinite -> InternalInfo -> [Maybe a] -> Stream a
forall a. InternalInfo -> [Maybe a] -> Stream a
Stream InternalInfo
Finite ([Maybe a] -> Stream a) -> [[Maybe a]] -> [Stream a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int -> [Maybe a] -> [Maybe a]
forall a. Int -> [a] -> [a]
take Int
n [Maybe a]
xs0 | Int
n <- (Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int
2Int -> Int -> Int
forall a b. (Num a, Integral b) => a -> b -> a
^) [Int
0 :: Int ..]]
    InternalInfo
Finite   -> InternalInfo -> [Maybe a] -> Stream a
forall a. InternalInfo -> [Maybe a] -> Stream a
Stream InternalInfo
Finite ([Maybe a] -> Stream a) -> [[Maybe a]] -> [Stream a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Maybe a -> [Maybe a]) -> [Maybe a] -> [[Maybe a]]
forall a. (a -> [a]) -> [a] -> [[a]]
QC.shrinkList ([Maybe a] -> Maybe a -> [Maybe a]
forall a b. a -> b -> a
const []) [Maybe a]
xs0

-- | Make a @'Maybe' a@ generator based on an @a@ generator.
--
-- Each element has a chance of being either 'Nothing' or an element generated
-- with the given @a@ generator (wrapped in a 'Just').
--
-- The first argument is the likelihood (as used by 'QC.frequency') of a
-- 'Just' where 'Nothing' has likelihood 2.
genMaybe ::
     Int   -- ^ Likelihood of 'Nothing'
  -> Int   -- ^ Likelihood of @'Just' a@
  -> Gen a
  -> Gen (Maybe a)
genMaybe :: forall a. Int -> Int -> Gen a -> Gen (Maybe a)
genMaybe Int
nLi Int
jLi Gen a
genA = [(Int, Gen (Maybe a))] -> Gen (Maybe a)
forall a. HasCallStack => [(Int, Gen a)] -> Gen a
QC.frequency
    [ (Int
nLi, Maybe a -> Gen (Maybe a)
forall a. a -> Gen a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
forall a. Maybe a
Nothing)
    , (Int
jLi, a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> Gen a -> Gen (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen a
genA)
    ]

-- | Like 'genMaybe', but with the likelihood of 'Nothing' fixed to @2@. 'QC.frequency'
genMaybe' ::
     Int   -- ^ Likelihood of @'Just' a@
  -> Gen a
  -> Gen (Maybe a)
genMaybe' :: forall a. Int -> Gen a -> Gen (Maybe a)
genMaybe' = Int -> Int -> Gen a -> Gen (Maybe a)
forall a. Int -> Int -> Gen a -> Gen (Maybe a)
genMaybe Int
2

-- | Generate a finite 'Stream' of length @n@.
genFinite ::
     Int           -- ^ Requested size of finite stream. Tip: use 'genMaybe'.
  -> Gen (Maybe a)
  -> Gen (Stream a)
genFinite :: forall a. Int -> Gen (Maybe a) -> Gen (Stream a)
genFinite Int
n Gen (Maybe a)
gen = InternalInfo -> [Maybe a] -> Stream a
forall a. InternalInfo -> [Maybe a] -> Stream a
Stream InternalInfo
Finite ([Maybe a] -> Stream a) -> Gen [Maybe a] -> Gen (Stream a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Gen (Maybe a) -> Gen [Maybe a]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
n Gen (Maybe a)
gen

-- | Generate an infinite 'Stream'.
genInfinite ::
     Gen (Maybe a)  -- ^ Tip: use 'genMaybe'.
  -> Gen (Stream a)
genInfinite :: forall a. Gen (Maybe a) -> Gen (Stream a)
genInfinite Gen (Maybe a)
gen = InternalInfo -> [Maybe a] -> Stream a
forall a. InternalInfo -> [Maybe a] -> Stream a
Stream InternalInfo
Infinite ([Maybe a] -> Stream a) -> Gen [Maybe a] -> Gen (Stream a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen (Maybe a) -> Gen [Maybe a]
forall a. Gen a -> Gen [a]
QC.listOf Gen (Maybe a)
gen