-- |
-- 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 hiding ((!!))
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
!! :: forall a b.
(Typed a, Eq b, Num b, Typed b) =>
[Stream a] -> Stream b -> Stream a
(!!) = [Stream a] -> Stream b -> Stream a
forall a b.
(Typed a, Eq b, Num b, Typed b) =>
[Stream a] -> Stream b -> Stream a
(!!!)
{-# DEPRECATED (!!) "This function is deprecated in Copilot 4. Use (!!!)." #-}

-- | 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'