Stream.hs

{-# 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 Test.QuickCheck (Gen)
import qualified Test.QuickCheck as QC
import Prelude hiding (filter, isInfinite, null)
import qualified Prelude

{-------------------------------------------------------------------------------
  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
    { unsafeStreamInternalInfo :: InternalInfo
    -- ^ UNSAFE: see 'UnsafeStream' for more information.
    --
    -- Info about the finiteness of the stream. It is used for 'QC.shrink'ing
    -- and the 'Show' instance.
    , unsafeStreamList :: [Maybe a]
    -- ^ UNSAFE: see 'UnsafeStream' for more information.
    --
    -- The internal list underlying the stream.
    }
  deriving 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 n (UnsafeStream info xs) = case info of
    Infinite -> ("<infinite stream>" ++)
    Finite ->
      (if n > 10 then ('(' :) else id)
        . shows xs
        . (" ++ ..." ++)
        . (if n > 10 then (')' :) else 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 s@(UnsafeStream _ []) = (Nothing, s)
runStream (UnsafeStream info (a : as)) = (a, UnsafeStream info 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 n s
  | n <= 0 = ([], s)
  | otherwise =
      let (x, s') = runStream s
          (xs, s'') = runStreamN (n - 1) s'
       in (x : xs, 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 (UnsafeStream _ as) = as ++ repeat Nothing

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

-- | Make an empty 'Stream'.
empty :: Stream a
empty = UnsafeStream Finite []

-- | Make a 'Stream' that always generates the given @a@.
always :: a -> Stream a
always x = UnsafeStream Infinite (repeat (Just x))

-- | Make a 'Stream' that infinitely repeats the given list.
repeating :: [Maybe a] -> Stream a
repeating xs = UnsafeStream Infinite $ cycle 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 = UnsafeStream 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 = UnsafeStream Infinite

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

-- | Filter a 'Stream', preserving finiteness.
filter :: (Maybe a -> Bool) -> Stream a -> Stream a
filter p (UnsafeStream info xs) = UnsafeStream info (Prelude.filter p 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 (UnsafeStream Finite []) = True
null _ = False

-- | Check that the stream is finite
isFinite :: Stream a -> Bool
isFinite (UnsafeStream Finite _) = True
isFinite (UnsafeStream Infinite _) = False

-- | Check that the stream is infinite
isInfinite :: Stream a -> Bool
isInfinite (UnsafeStream Finite _) = False
isInfinite (UnsafeStream Infinite _) = 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 (UnsafeStream info xs0) = case info of
  Infinite -> UnsafeStream Finite <$> [take n xs0 | n <- map (2 ^) [0 :: Int ..]]
  Finite -> UnsafeStream Finite <$> QC.shrinkList (const []) xs0

-- | Like 'shrinkStream', but with a custom shrinker for elements of the stream.
liftShrinkStream :: (Maybe a -> [Maybe a]) -> Stream a -> [Stream a]
liftShrinkStream shrinkOne (UnsafeStream info xs0) = case info of
  Infinite -> UnsafeStream Finite <$> [take n xs0 | n <- map (2 ^) [0 :: Int ..]]
  Finite -> UnsafeStream Finite <$> QC.shrinkList shrinkOne 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 ::
  -- | Likelihood of 'Nothing'
  Int ->
  -- | Likelihood of @'Just' a@
  Int ->
  Gen a ->
  Gen (Maybe a)
genMaybe nLi jLi genA =
  QC.frequency
    [ (nLi, return Nothing)
    , (jLi, Just <$> genA)
    ]

-- | Generate a finite 'Stream' of length @n@.
genFiniteN ::
  -- | Requested size of finite stream.
  Int ->
  Gen (Maybe a) ->
  Gen (Stream a)
genFiniteN n gen = UnsafeStream Finite <$> replicateM n gen

-- | Generate a sized, finite 'Stream'.
genFinite ::
  Gen (Maybe a) ->
  Gen (Stream a)
genFinite gen = UnsafeStream Finite <$> QC.listOf gen

-- | Generate an infinite 'Stream'.
genInfinite ::
  Gen (Maybe a) ->
  Gen (Stream a)
genInfinite gen = UnsafeStream Infinite <$> QC.infiniteListOf gen