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

-- | Finite and infinite streams of @'Maybe' a@s.
module System.FS.Sim.Stream (
    -- * Streams
    Stream (..)
  , InternalInfo (..)
    -- * Running
  , runStream
  , runStreamN
  , runStreamIndefinitely
    -- * Construction
  , always
  , empty
  , repeating
  , unsafeMkInfinite
  , unsafeMkFinite
    -- * Modify
  , filter
    -- * Query
  , null
  , isFinite
  , isInfinite
    -- * Generation and shrinking
  , genFinite
  , genFiniteN
  , genInfinite
  , genMaybe
  , shrinkStream
  , liftShrinkStream
  ) where

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

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

-- | A stream of @'Maybe' a@s that can be infinite.
data Stream a =
  -- | UNSAFE: when constructing, modifying, or accessing the internals of a
  -- 'Stream', it is the responsibility of the user to preserve the following
  -- invariant:
  --
  -- INVARIANT: if the stream is marked as 'Infinite', then the internal list
  -- should be infinite. If the stream is marked as 'Finite', then the internal
  -- list should finite.
  --
  -- * If the internal list is infinite but marked as 'Finite', then 'QC.shrink'
  --   or 'show' on the corresponding stream will diverge.
  --
  -- * If the internal list is finite but marked as 'Infinite', then 'QC.shrink'
  --   on the corresponding stream will degrade to an infinite list of empty
  --   streams.
  UnsafeStream {
      -- | UNSAFE: see 'UnsafeStream' for more information.
      --
      -- Info about the finiteness of the stream. It is used for 'QC.shrink'ing
      -- and the 'Show' instance.
      forall a. Stream a -> InternalInfo
unsafeStreamInternalInfo :: InternalInfo
      -- | UNSAFE: see 'UnsafeStream' for more information.
      --
      -- The internal list underlying the stream.
    , forall a. Stream a -> [Maybe a]
unsafeStreamList         :: [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 finite or infinite.
--
-- Useful for the 'Show' instance of 'Stream': when a 'Stream' is finite, we can
-- safely print the full stream.
data InternalInfo = Infinite | Finite

-- | Fully shows a 'Stream' if it is finite, or prints a placeholder string if
-- it is infinite.
instance Show a => Show (Stream a) where
  showsPrec :: Int -> Stream a -> ShowS
showsPrec Int
n (UnsafeStream 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
-------------------------------------------------------------------------------}

-- | \( O(1) \): 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@(UnsafeStream InternalInfo
_    []    ) = (Maybe a
forall a. Maybe a
Nothing, Stream a
s)
runStream   (UnsafeStream InternalInfo
info (Maybe a
a:[Maybe a]
as)) = (Maybe a
a, InternalInfo -> [Maybe a] -> Stream a
forall a. InternalInfo -> [Maybe a] -> Stream a
UnsafeStream InternalInfo
info [Maybe a]
as)

-- | \( O(n) \): like 'runStream', but advancing the stream @n@ times.
--
-- If @n<=0@, then the stream is advanced @0@ times.
runStreamN :: Int -> Stream a -> ([Maybe a], Stream a)
runStreamN :: forall a. Int -> Stream a -> ([Maybe a], Stream a)
runStreamN Int
n Stream a
s
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = ([], Stream a
s)
  | Bool
otherwise =
      let (Maybe a
x, Stream a
s') = Stream a -> (Maybe a, Stream a)
forall a. Stream a -> (Maybe a, Stream a)
runStream Stream a
s
          ([Maybe a]
xs, Stream a
s'') = Int -> Stream a -> ([Maybe a], Stream a)
forall a. Int -> Stream a -> ([Maybe a], Stream a)
runStreamN (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Stream a
s'
      in  (Maybe a
xMaybe a -> [Maybe a] -> [Maybe a]
forall a. a -> [a] -> [a]
:[Maybe a]
xs, Stream a
s'')

-- | \( O(\infty) \): like 'runStream', but advancing the stream indefinitely.
--
-- For infinite streams, this produces an infinite list. For finite streams,
-- this produces a finite list.
runStreamIndefinitely :: Stream a -> [Maybe a]
runStreamIndefinitely :: forall a. Stream a -> [Maybe a]
runStreamIndefinitely (UnsafeStream InternalInfo
_ [Maybe a]
as) = [Maybe a]
as [Maybe a] -> [Maybe a] -> [Maybe a]
forall a. [a] -> [a] -> [a]
++ Maybe a -> [Maybe a]
forall a. a -> [a]
repeat Maybe a
forall a. Maybe a
Nothing

{-------------------------------------------------------------------------------
  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
UnsafeStream 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
UnsafeStream 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
UnsafeStream InternalInfo
Infinite ([Maybe a] -> Stream a) -> [Maybe a] -> Stream a
forall a b. (a -> b) -> a -> b
$ [Maybe a] -> [Maybe a]
forall a. HasCallStack => [a] -> [a]
cycle [Maybe a]
xs

-- | UNSAFE: Make a 'Stream' that is marked as finite. It is the user's
-- responsibility to only pass in finite lists. See 'UnsafeStream' for more
-- information.
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
UnsafeStream InternalInfo
Finite

-- | UNSAFE: Make a 'Stream' that is marked as infinite. It is the user's
-- responsibility to only pass in infinite lists. See 'UnsafeStream' for more
-- information.
unsafeMkInfinite :: [Maybe a] -> Stream a
unsafeMkInfinite :: forall a. [Maybe a] -> Stream a
unsafeMkInfinite = InternalInfo -> [Maybe a] -> Stream a
forall a. InternalInfo -> [Maybe a] -> Stream a
UnsafeStream InternalInfo
Infinite

{-------------------------------------------------------------------------------
  Modify
-------------------------------------------------------------------------------}

-- | Filter a 'Stream', preserving finiteness.
filter :: (Maybe a -> Bool) -> Stream a -> Stream a
filter :: forall a. (Maybe a -> Bool) -> Stream a -> Stream a
filter Maybe a -> Bool
p (UnsafeStream InternalInfo
info [Maybe a]
xs) = InternalInfo -> [Maybe a] -> Stream a
forall a. InternalInfo -> [Maybe a] -> Stream a
UnsafeStream InternalInfo
info ((Maybe a -> Bool) -> [Maybe a] -> [Maybe a]
forall a. (a -> Bool) -> [a] -> [a]
Prelude.filter Maybe a -> Bool
p [Maybe a]
xs)

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

-- | Check that the stream is empty.
--
-- In general, a stream is only empty if the stream is equivalent to 'empty'.
--
-- A finite\/infinite stream consisting of only 'Nothing's is not considered to
-- be empty. In particular, @'null' ('always' Nothing) /= True@.
null :: Stream a -> Bool
null :: forall a. Stream a -> Bool
null (UnsafeStream InternalInfo
Finite []) = Bool
True
null Stream a
_                        = Bool
False

-- | Check that the stream is finite
isFinite :: Stream a -> Bool
isFinite :: forall a. Stream a -> Bool
isFinite (UnsafeStream InternalInfo
Finite [Maybe a]
_)   = Bool
True
isFinite (UnsafeStream InternalInfo
Infinite [Maybe a]
_) = Bool
False

-- | Check that the stream is infinite
isInfinite :: Stream a -> Bool
isInfinite :: forall a. Stream a -> Bool
isInfinite (UnsafeStream InternalInfo
Finite [Maybe a]
_)   = Bool
False
isInfinite (UnsafeStream InternalInfo
Infinite [Maybe a]
_) = Bool
True

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

-- | Shrink a stream like it is an 'QC.InfiniteList'.
--
-- Infinite streams are shrunk differently than lists that are finite, which is
-- to ensure that we shrink infinite lists towards finite lists.
--
-- * Infinite streams are shrunk by taking finite prefixes of the argument
--   stream. Note that there are an infinite number of finite prefixes, so even
--   though the *shrink list* is infinite, the individual *list elements* are
--   finite.
--
-- * Finite streams are shrunk like lists are shrunk normally, preserving
--   finiteness.
shrinkStream :: Stream a -> [Stream a]
shrinkStream :: forall a. Stream a -> [Stream a]
shrinkStream (UnsafeStream InternalInfo
info [Maybe a]
xs0) = case InternalInfo
info of
    InternalInfo
Infinite -> InternalInfo -> [Maybe a] -> Stream a
forall a. InternalInfo -> [Maybe a] -> Stream a
UnsafeStream 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
UnsafeStream 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

-- | Like 'shrinkStream', but with a custom shrinker for elements of the stream.
liftShrinkStream :: (Maybe a -> [Maybe a]) -> Stream a -> [Stream a]
liftShrinkStream :: forall a. (Maybe a -> [Maybe a]) -> Stream a -> [Stream a]
liftShrinkStream Maybe a -> [Maybe a]
shrinkOne (UnsafeStream InternalInfo
info [Maybe a]
xs0) = case InternalInfo
info of
    InternalInfo
Infinite -> InternalInfo -> [Maybe a] -> Stream a
forall a. InternalInfo -> [Maybe a] -> Stream a
UnsafeStream 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
UnsafeStream 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]
shrinkOne [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'). These /likelihoods/ are
-- passed to 'QC.frequency'.
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)
    ]

-- | Generate a finite 'Stream' of length @n@.
genFiniteN ::
     Int -- ^ Requested size of finite stream.
  -> Gen (Maybe a)
  -> Gen (Stream a)
genFiniteN :: forall a. Int -> Gen (Maybe a) -> Gen (Stream a)
genFiniteN Int
n Gen (Maybe a)
gen = InternalInfo -> [Maybe a] -> Stream a
forall a. InternalInfo -> [Maybe a] -> Stream a
UnsafeStream 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 a sized, finite 'Stream'.
genFinite ::
     Gen (Maybe a)
  -> Gen (Stream a)
genFinite :: forall a. Gen (Maybe a) -> Gen (Stream a)
genFinite Gen (Maybe a)
gen = InternalInfo -> [Maybe a] -> Stream a
forall a. InternalInfo -> [Maybe a] -> Stream a
UnsafeStream InternalInfo
Finite ([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

-- | Generate an infinite 'Stream'.
genInfinite ::
     Gen (Maybe a)
  -> 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
UnsafeStream 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.infiniteListOf Gen (Maybe a)
gen