-- |
-- 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 = [ 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 = forall a. Int -> [a] -> [a]
P.take ( forall a b. (Integral a, Num b) => a -> b
fromIntegral a
n ) forall a b. (a -> b) -> a -> 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 = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Stream a -> Stream b -> Stream a
f Stream a
e forall a b. (a -> b) -> a -> 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 = forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldl1 Stream a -> Stream a -> Stream a
f forall a b. (a -> b) -> a -> b
$ 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 = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Stream a -> Stream b -> Stream b
f Stream b
e forall a b. (a -> b) -> a -> b
$ 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 = forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 Stream a -> Stream a -> Stream a
f forall a b. (a -> b) -> a -> b
$ 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 = forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl Stream a -> Stream b -> Stream a
f Stream a
e forall a b. (a -> b) -> a -> 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 = forall a b. (a -> b -> b) -> b -> [a] -> [b]
scanr Stream a -> Stream b -> Stream b
f Stream b
e forall a b. (a -> b) -> a -> b
$ 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 = forall a. (a -> a -> a) -> [a] -> [a]
scanl1 Stream a -> Stream a -> Stream a
f forall a b. (a -> b) -> a -> b
$ 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 = forall a. (a -> a -> a) -> [a] -> [a]
scanr1 Stream a -> Stream a -> Stream a
f forall a b. (a -> b) -> a -> b
$ 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 )       = 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]
_                =
        forall a. String -> a
badUsage forall a b. (a -> b) -> a -> b
$ String
"in case' in Utils library: "
                   forall a. [a] -> [a] -> [a]
P.++ String
"length of alternatives list is not "
                   forall a. [a] -> [a] -> [a]
P.++ String
"greater by one than the length of predicates list"
  in 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      = forall a b. (a -> b) -> [a] -> [b]
map
                             ( forall a. Typed a => a -> Stream a
constant forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (Integral a, Num b) => a -> b
fromIntegral )
                             [ Int
0 .. forall (t :: * -> *) a. Foldable t => t a -> Int
P.length [Stream a]
ls forall a. Num a => a -> a -> a
- Int
1 ]
              select :: [Stream b] -> [Stream a] -> Stream a
select [] [Stream a]
_  = forall a. [a] -> a
last [Stream a]
ls
              select
                ( Stream b
i : [Stream b]
is )
                ( Stream a
x : [Stream a]
xs ) = forall a.
Typed a =>
Stream Bool -> Stream a -> Stream a -> Stream a
mux ( Stream b
i 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]
_ []  = forall a. String -> a
badUsage (String
"in (!!) defined in Utils.hs " forall a. [a] -> [a] -> [a]
P.++
                               String
"in copilot-libraries")
          in if forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Stream a]
ls then
               forall a. String -> a
badUsage (String
"in (!!) defined in Utils.hs " 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 forall a. Typed a => [a] -> Stream a -> Stream a
++ Stream a
cycle'