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

tails :: ( Typed a )
         => Stream a -> [ Stream a ]
tails :: 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 .. ] ]


take :: ( Integral a, Typed b )
        => a -> Stream b -> [ Stream b ]
take :: 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

nfoldl :: ( Typed a, Typed b )
          => Int -> ( Stream a -> Stream b -> Stream a )
                 ->   Stream a -> Stream b -> Stream a
nfoldl :: 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 (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

nfoldl1 :: ( Typed a )
           => Int -> ( Stream a -> Stream a -> Stream a )
                  ->   Stream a -> Stream a
nfoldl1 :: 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 (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

nfoldr :: ( Typed a, Typed b )
          => Int -> ( Stream a -> Stream b -> Stream b )
                 ->   Stream b -> Stream a -> Stream b
nfoldr :: 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 (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

nfoldr1 :: ( Typed a )
           => Int -> ( Stream a -> Stream a -> Stream a )
                  ->   Stream a -> Stream a
nfoldr1 :: 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 (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

nscanl :: ( Typed a, Typed b )
          => Int -> ( Stream a -> Stream b -> Stream a )
          -> Stream a -> Stream b -> [ Stream a ]
nscanl :: 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

nscanr :: ( Typed a )
          => Int -> ( Stream a -> Stream b -> Stream b )
          -> Stream b -> Stream a -> [ Stream b ]
nscanr :: 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

nscanl1 :: ( Typed a )
           => Int -> ( Stream a -> Stream a -> Stream a )
           -> Stream a -> [ Stream a ]
nscanl1 :: 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

nscanr1 :: ( Typed a )
           => Int -> ( Stream a -> Stream a -> Stream a )
           -> Stream a -> [ Stream a ]
nscanr1 :: 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' :: [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 !! :: [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 (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. [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 (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 :: ( Typed a ) => [ a ] -> Stream a
cycle :: [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'