    -- |
    -- Module: Utils
    -- Description: Utility bounded-list functions (e.g., folds, scans, etc.)
    -- Copyright: (c) 2011 National Institute of Aerospace / Galois, Inc.
    --
    -- Utility bounded-list functions (e.g., folds, scans, etc.)

    module Copilot.Library.Utils
           ( -- * Functions similar to the Prelude functions on lists
             take, tails, cycle,
             -- ** Folds
             nfoldl, nfoldl1, nfoldr, nfoldr1,
             -- ** Scans
             nscanl, nscanr, nscanl1, nscanr1,
             -- ** Indexing
             case', (!!!))
    where

    import Copilot.Language
    import qualified Prelude as P

    -- | Given a stream, produce an infinite list of streams dropping an increasing
    -- number of elements of the given stream. For example, for a given stream @s@,
    -- the expression @tails s@ is equal to @[ drop 0 s, drop 1 s, drop 2 s, ...]@.
    --
    tails :: ( Typed a )
             => Stream a -> [ Stream a ]
    tails :: forall a. Typed a => Stream a -> [Stream a]
    tails Stream a
    s = [ Int -> Stream a -> Stream a
    forall a. Typed a => Int -> Stream a -> Stream a
    drop Int
    x Stream a
    s | Int
    x <- [ Int
    0 .. ] ]

    -- | Given a stream and a number, produce a finite list of streams dropping an
    -- increasing number of elements of the given stream, up to that number. For
    -- example, for a given stream @s@, the expression @take 2 s@ is equal to
    -- @[ drop 0 s, drop 1 s]@.
    take :: ( Integral a, Typed b )
            => a -> Stream b -> [ Stream b ]
    take :: forall a b. (Integral a, Typed b) => a -> Stream b -> [Stream b]
    take a
    n Stream b
    s = Int -> [Stream b] -> [Stream b]
    forall a. Int -> [a] -> [a]
    P.take ( a -> Int
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral a
    n ) ([Stream b] -> [Stream b]) -> [Stream b] -> [Stream b]
    forall a b. (a -> b) -> a -> b
    $ Stream b -> [Stream b]
    forall a. Typed a => Stream a -> [Stream a]
    tails Stream b
    s

    -- | Given a number, a function on streams, and two streams, fold from the left
    -- the function over the finite list of tails of the second stream (up to the
    -- given number).
    nfoldl :: ( Typed a, Typed b )
              => Int -> ( Stream a -> Stream b -> Stream a )
                     ->   Stream a -> Stream b -> Stream a
    nfoldl :: forall a b.
    (Typed a, Typed b) =>
    Int
    -> (Stream a -> Stream b -> Stream a)
    -> Stream a
    -> Stream b
    -> Stream a
    nfoldl Int
    n Stream a -> Stream b -> Stream a
    f Stream a
    e Stream b
    s = (Stream a -> Stream b -> Stream a)
    -> Stream a -> [Stream b] -> Stream a
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl Stream a -> Stream b -> Stream a
    f Stream a
    e ([Stream b] -> Stream a) -> [Stream b] -> Stream a
    forall a b. (a -> b) -> a -> b
    $ Int -> Stream b -> [Stream b]
    forall a b. (Integral a, Typed b) => a -> Stream b -> [Stream b]
    take Int
    n Stream b
    s

    -- | Given a number, a function on streams, and two streams, fold from the left
    -- the function over the finite list of tails of the second stream (up to the
    -- given number).
    --
    -- This function differs from 'nfoldl' in that it does not require an initial
    -- accumulator and it assumes the argument number @n@ is positive.
    nfoldl1 :: ( Typed a )
               => Int -> ( Stream a -> Stream a -> Stream a )
                      ->   Stream a -> Stream a
    nfoldl1 :: forall a.
    Typed a =>
    Int -> (Stream a -> Stream a -> Stream a) -> Stream a -> Stream a
    nfoldl1 Int
    n Stream a -> Stream a -> Stream a
    f Stream a
    s = (Stream a -> Stream a -> Stream a) -> [Stream a] -> Stream a
    forall a. (a -> a -> a) -> [a] -> a
    forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
    foldl1 Stream a -> Stream a -> Stream a
    f ([Stream a] -> Stream a) -> [Stream a] -> Stream a
    forall a b. (a -> b) -> a -> b
    $ Int -> Stream a -> [Stream a]
    forall a b. (Integral a, Typed b) => a -> Stream b -> [Stream b]
    take Int
    n Stream a
    s

    -- | Given a number, a function on streams, and two streams, fold from the
    -- right the function over the finite list of tails of the second stream (up to
    -- the given number).
    nfoldr :: ( Typed a, Typed b )
              => Int -> ( Stream a -> Stream b -> Stream b )
                     ->   Stream b -> Stream a -> Stream b
    nfoldr :: forall a b.
    (Typed a, Typed b) =>
    Int
    -> (Stream a -> Stream b -> Stream b)
    -> Stream b
    -> Stream a
    -> Stream b
    nfoldr Int
    n Stream a -> Stream b -> Stream b
    f Stream b
    e Stream a
    s = (Stream a -> Stream b -> Stream b)
    -> Stream b -> [Stream a] -> Stream b
    forall a b. (a -> b -> b) -> b -> [a] -> b
    forall (t :: * -> *) a b.
    Foldable t =>
    (a -> b -> b) -> b -> t a -> b
    foldr Stream a -> Stream b -> Stream b
    f Stream b
    e ([Stream a] -> Stream b) -> [Stream a] -> Stream b
    forall a b. (a -> b) -> a -> b
    $ Int -> Stream a -> [Stream a]
    forall a b. (Integral a, Typed b) => a -> Stream b -> [Stream b]
    take Int
    n Stream a
    s

    -- | Given a number, a function on streams, and two streams, fold from the
    -- right the function over the finite list of tails of the second stream (up to
    -- the given number).
    --
    -- This function differs from 'nfoldr' in that it does not require an initial
    -- accumulator and it assumes the argument number @n@ is positive.
    nfoldr1 :: ( Typed a )
               => Int -> ( Stream a -> Stream a -> Stream a )
                      ->   Stream a -> Stream a
    nfoldr1 :: forall a.
    Typed a =>
    Int -> (Stream a -> Stream a -> Stream a) -> Stream a -> Stream a
    nfoldr1 Int
    n Stream a -> Stream a -> Stream a
    f Stream a
    s = (Stream a -> Stream a -> Stream a) -> [Stream a] -> Stream a
    forall a. (a -> a -> a) -> [a] -> a
    forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
    foldr1 Stream a -> Stream a -> Stream a
    f ([Stream a] -> Stream a) -> [Stream a] -> Stream a
    forall a b. (a -> b) -> a -> b
    $ Int -> Stream a -> [Stream a]
    forall a b. (Integral a, Typed b) => a -> Stream b -> [Stream b]
    take Int
    n Stream a
    s

    -- | Given a number, a function on streams, and two streams, fold from the left
    -- the function over the finite list of tails of the second stream (up to the
    -- given number).
    --
    -- This function differs from 'nfoldl' in that it returns the intermediate
    -- results as well.
    nscanl :: ( Typed a, Typed b )
              => Int -> ( Stream a -> Stream b -> Stream a )
              -> Stream a -> Stream b -> [ Stream a ]
    nscanl :: forall a b.
    (Typed a, Typed b) =>
    Int
    -> (Stream a -> Stream b -> Stream a)
    -> Stream a
    -> Stream b
    -> [Stream a]
    nscanl Int
    n Stream a -> Stream b -> Stream a
    f Stream a
    e Stream b
    s = (Stream a -> Stream b -> Stream a)
    -> Stream a -> [Stream b] -> [Stream a]
    forall b a. (b -> a -> b) -> b -> [a] -> [b]
    scanl Stream a -> Stream b -> Stream a
    f Stream a
    e ([Stream b] -> [Stream a]) -> [Stream b] -> [Stream a]
    forall a b. (a -> b) -> a -> b
    $ Int -> Stream b -> [Stream b]
    forall a b. (Integral a, Typed b) => a -> Stream b -> [Stream b]
    take Int
    n Stream b
    s

    -- | Given a number, a function on streams, and two streams, fold from the
    -- right the function over the finite list of tails of the second stream (up to
    -- the given number).
    --
    -- This function differs from 'nfoldr' in that it returns the intermediate
    -- results as well.
    nscanr :: ( Typed a )
              => Int -> ( Stream a -> Stream b -> Stream b )
              -> Stream b -> Stream a -> [ Stream b ]
    nscanr :: forall a b.
    Typed a =>
    Int
    -> (Stream a -> Stream b -> Stream b)
    -> Stream b
    -> Stream a
    -> [Stream b]
    nscanr Int
    n Stream a -> Stream b -> Stream b
    f Stream b
    e Stream a
    s = (Stream a -> Stream b -> Stream b)
    -> Stream b -> [Stream a] -> [Stream b]
    forall a b. (a -> b -> b) -> b -> [a] -> [b]
    scanr Stream a -> Stream b -> Stream b
    f Stream b
    e ([Stream a] -> [Stream b]) -> [Stream a] -> [Stream b]
    forall a b. (a -> b) -> a -> b
    $ Int -> Stream a -> [Stream a]
    forall a b. (Integral a, Typed b) => a -> Stream b -> [Stream b]
    take Int
    n Stream a
    s

    -- | Given a number, a function on streams, and two streams, fold from the left
    -- the function over the finite list of tails of the second stream (up to the
    -- given number).
    --
    -- This function assumes the number of elements to scan is positive, and it
    -- also returns the intermediate results.
    nscanl1 :: ( Typed a )
               => Int -> ( Stream a -> Stream a -> Stream a )
               -> Stream a -> [ Stream a ]
    nscanl1 :: forall a.
    Typed a =>
    Int -> (Stream a -> Stream a -> Stream a) -> Stream a -> [Stream a]
    nscanl1 Int
    n Stream a -> Stream a -> Stream a
    f Stream a
    s = (Stream a -> Stream a -> Stream a) -> [Stream a] -> [Stream a]
    forall a. (a -> a -> a) -> [a] -> [a]
    scanl1 Stream a -> Stream a -> Stream a
    f ([Stream a] -> [Stream a]) -> [Stream a] -> [Stream a]
    forall a b. (a -> b) -> a -> b
    $ Int -> Stream a -> [Stream a]
    forall a b. (Integral a, Typed b) => a -> Stream b -> [Stream b]
    take Int
    n Stream a
    s

    -- | Given a number, a function on streams, and two streams, fold from the
    -- right the function over the finite list of tails of the second stream (up to
    -- the given number).
    --
    -- This function assumes the number of elements to scan is positive, and it
    -- also returns the intermediate results.
    nscanr1 :: ( Typed a )
               => Int -> ( Stream a -> Stream a -> Stream a )
               -> Stream a -> [ Stream a ]
    nscanr1 :: forall a.
    Typed a =>
    Int -> (Stream a -> Stream a -> Stream a) -> Stream a -> [Stream a]
    nscanr1 Int
    n Stream a -> Stream a -> Stream a
    f Stream a
    s = (Stream a -> Stream a -> Stream a) -> [Stream a] -> [Stream a]
    forall a. (a -> a -> a) -> [a] -> [a]
    scanr1 Stream a -> Stream a -> Stream a
    f ([Stream a] -> [Stream a]) -> [Stream a] -> [Stream a]
    forall a b. (a -> b) -> a -> b
    $ Int -> Stream a -> [Stream a]
    forall a b. (Integral a, Typed b) => a -> Stream b -> [Stream b]
    take Int
    n Stream a
    s

    -- | Case-like function: The index of the first predicate that is true
    -- in the predicate list selects the stream result. If no predicate
    -- is true, the last element is chosen (default element)
    case' :: ( Typed a )
             => [ Stream Bool ] -> [ Stream a ] -> Stream a
    case' :: forall a. Typed a => [Stream Bool] -> [Stream a] -> Stream a
    case' [Stream Bool]
    predicates [Stream a]
    alternatives =
      let case'' :: [Stream Bool] -> [Stream a] -> Stream a
    case'' []         ( Stream a
    default' : [Stream a]
    _ ) = Stream a
    default'
          case'' ( Stream Bool
    p : [Stream Bool]
    ps ) ( Stream a
    a : [Stream a]
    as )       = Stream Bool -> Stream a -> Stream a -> Stream a
    forall a.
    Typed a =>
    Stream Bool -> Stream a -> Stream a -> Stream a
    mux Stream Bool
    p Stream a
    a ( [Stream Bool] -> [Stream a] -> Stream a
    case'' [Stream Bool]
    ps [Stream a]
    as )
          case'' [Stream Bool]
    _          [Stream a]
    _                =
            String -> Stream a
    forall a. String -> a
    badUsage (String -> Stream a) -> String -> Stream a
    forall a b. (a -> b) -> a -> b
    $ String
    "in case' in Utils library: "
                       String -> String -> String
    forall a. [a] -> [a] -> [a]
    P.++ String
    "length of alternatives list is not "
                       String -> String -> String
    forall a. [a] -> [a] -> [a]
    P.++ String
    "greater by one than the length of predicates list"
      in [Stream Bool] -> [Stream a] -> Stream a
    forall a. Typed a => [Stream Bool] -> [Stream a] -> Stream a
    case'' [Stream Bool]
    predicates [Stream a]
    alternatives

    -- | Index.
    --
    -- WARNING: Very expensive! Consider using this only for very short lists.
    (!!!) :: (Typed a, Eq b, Num b, Typed b) => [Stream a] -> Stream b -> Stream a
    [Stream a]
    ls !!! :: forall a b.
    (Typed a, Eq b, Num b, Typed b) =>
    [Stream a] -> Stream b -> Stream a
    !!! Stream b
    n = let indices :: [Stream b]
    indices      = (Int -> Stream b) -> [Int] -> [Stream b]
    forall a b. (a -> b) -> [a] -> [b]
    map
                                 ( b -> Stream b
    forall a. Typed a => a -> Stream a
    constant (b -> Stream b) -> (Int -> b) -> Int -> Stream b
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Int -> b
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral )
                                 [ Int
    0 .. [Stream a] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    P.length [Stream a]
    ls Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - Int
    1 ]
                   select :: [Stream b] -> [Stream a] -> Stream a
    select [] [Stream a]
    _  = [Stream a] -> Stream a
    forall a. HasCallStack => [a] -> a
    last [Stream a]
    ls
                   select
                     ( Stream b
    i : [Stream b]
    is )
                     ( Stream a
    x : [Stream a]
    xs ) = Stream Bool -> Stream a -> Stream a -> Stream a
    forall a.
    Typed a =>
    Stream Bool -> Stream a -> Stream a -> Stream a
    mux ( Stream b
    i Stream b -> Stream b -> Stream Bool
    forall a. (Eq a, Typed a) => Stream a -> Stream a -> Stream Bool
    == Stream b
    n ) Stream a
    x ( [Stream b] -> [Stream a] -> Stream a
    select [Stream b]
    is [Stream a]
    xs )
                                  -- should not happen
                   select [Stream b]
    _ []  = String -> Stream a
    forall a. String -> a
    badUsage (String
    "in (!!) defined in Utils.hs " String -> String -> String
    forall a. [a] -> [a] -> [a]
    P.++
                                    String
    "in copilot-libraries")
               in if [Stream a] -> Bool
    forall a. [a] -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null [Stream a]
    ls then
                    String -> Stream a
    forall a. String -> a
    badUsage (String
    "in (!!) defined in Utils.hs " String -> String -> String
    forall a. [a] -> [a] -> [a]
    P.++
                                 String
    "indexing the empty list with !! is not defined")
                  else
                    [Stream b] -> [Stream a] -> Stream a
    select [Stream b]
    indices [Stream a]
    ls

    -- | Cycle a list to form an infinite stream.
    cycle :: ( Typed a ) => [ a ] -> Stream a
    cycle :: forall a. Typed a => [a] -> Stream a
    cycle [a]
    ls = Stream a
    cycle'
      where
        cycle' :: Stream a
    cycle' = [a]
    ls [a] -> Stream a -> Stream a
    forall a. Typed a => [a] -> Stream a -> Stream a
    ++ Stream a
    cycle'