Basics.hs

{-# LANGUAGE CPP           #-}
{-# LANGUAGE BangPatterns  #-}
{-# LANGUAGE TypeOperators #-}
-----------------------------------------------------------------------------
-- |
-- Module      : Data.Array.Vector.Strict.Basics
-- Copyright   : (c) [2001..2002] Manuel M T Chakravarty & Gabriele Keller
--               (c) 2006         Manuel M T Chakravarty & Roman Leshchinskiy
-- License     : see libraries/ndp/LICENSE
-- 
-- Maintainer  : Roman Leshchinskiy <rl@cse.unsw.edu.au>
-- Stability   : internal
-- Portability : portable
--
-- Description ---------------------------------------------------------------
--
--  Basic operations on flat unlifted arrays.
--
-- Todo ----------------------------------------------------------------------
--


#include "fusion-phases.h"

module Data.Array.Vector.Strict.Basics where

import Data.Array.Vector.Stream
import Data.Array.Vector.UArr hiding (lengthU, indexU)
import qualified Data.Array.Vector.UArr as Prim (lengthU, indexU)

import Data.Array.Vector.Strict.Stream

import Data.Array.Vector.Prim.Debug
import Data.Array.Vector.Prim.Hyperstrict
import GHC.ST

import Debug.Trace

------------------------------------------------------------------------

instance (Eq e, UA e) => Eq (UArr e) where (==) = eqU
   -- not really fusible

{-# INLINE_U eqU #-}
eqU :: (Eq e, UA e) => UArr e -> UArr e -> Bool
eqU a1 a2 = lengthU a1 == lengthU a2 && foldlU cmp True (zipU a1 a2)
   where
       cmp r (e1 :*: e2) = e1 == e2 && r

------------------------------------------------------------------------

-- XXX general rule:
-- If we have direct implementations with better complexity than streams:
--
--   length, index, take, drop
--
-- Then use the direct version


-- | /O(1)/, 'length' returns the length of a UArr as an 'Int'.
lengthU :: UA e => UArr e -> Int
lengthU = foldlU (const . (+1)) 0
-- lengthU = Prim.lengthU
{-# INLINE_U lengthU #-}

{-

Unfused version:

$ time ./length
100000000
./length  1.10s user 1.13s system 91% cpu 2.433 total

Fusible version:

100000000
./length  0.31s user 0.00s system 97% cpu 0.318 total

-}

-- lengthU_stream :: UA e => UArr e -> Int
-- {-# INLINE lengthU_stream #-}

-- lengthU is reexported from UArr

-- |Test whether the given array is empty
--
nullU :: UA e => UArr e -> Bool
nullU  = nullS . streamU  -- better code if we short circuit
{-# INLINE_U nullU #-}
-- nullU  = (== 0) . lengthU

-- |Yield an empty array
--
emptyU :: UA e => UArr e
emptyU = unstreamU emptyS
{-# INLINE_U emptyU #-}

-- emptyU = newU 0 (const $ return ())

-- |Yield a singleton array
--
singletonU :: UA e => e -> UArr e
{-# INLINE_U singletonU #-}
singletonU = unstreamU . singletonS

-- |Prepend an element to an array
--
consU :: UA e => e -> UArr e -> UArr e
{-# INLINE_U consU #-}
consU x = unstreamU . consS x . streamU

-- |Append an element to an array
--
snocU :: UA e => UArr e -> e -> UArr e
{-# INLINE_U snocU #-}
snocU s x = unstreamU (snocS (streamU s) x)

-- unitsU is reexported from Loop

-- |Yield an array where all elements contain the same value
--
replicateU :: UA e => Int -> e -> UArr e
{-# INLINE_U replicateU #-}
replicateU n e = unstreamU (replicateS n e)

replicateEachU :: UA e => Int -> UArr Int -> UArr e -> UArr e
{-# INLINE_U replicateEachU #-}
replicateEachU n ns es = unstreamU
                       . replicateEachS n
                       $ zipS (streamU ns) (streamU es)

-- |Array indexing
--
indexU :: UA e => UArr e -> Int -> e
indexU arr n = indexS (streamU arr) n
{-# INLINE_U indexU #-}

headU :: UA e => UArr e -> e
headU = headS . streamU
{-# INLINE_U headU #-}

lastU :: UA e => UArr e -> e
lastU = foldl1U (flip const)
{-# INLINE lastU #-}

-- |Concatenate two arrays
--
appendU :: UA e => UArr e -> UArr e -> UArr e
{-# INLINE_U appendU #-}
a1 `appendU` a2 = unstreamU (streamU a1 +++ streamU a2)

initU :: UA e => UArr e -> UArr e -- not unboxing
initU = unstreamU . initS . streamU
{-# INLINE initU #-}

-- |Repeat an array @n@ times
--
repeatU :: UA e => Int -> UArr e -> UArr e
repeatU n = unstreamU . repS n
{-# INLINE_U repeatU #-}

-- No work duplicated
repS :: UA e => Int -> UArr e -> Stream e
{-# INLINE_STREAM repS #-}
repS k xs = Stream next (0 :*: k) (k*n)
  where
    n = lengthU xs

    {-# INLINE next #-}
    next (i :*: 0) = Done
    next (i :*: k) | i == n    = Skip (0 :*: k-1)
                   | otherwise = Yield (xs `Prim.indexU` i) (i+1 :*: k)

-- |Indexing
-- ---------

-- |Associate each element of the array with its index
--
indexedU :: UA e => UArr e -> UArr (Int :*: e)
{-# INLINE_U indexedU #-}
indexedU = unstreamU . indexedS . streamU

-- |Conversion
-- -----------

-- |Turn a list into a parallel array
--
toU :: UA e => [e] -> UArr e
{-# INLINE_U toU #-}
toU = unstreamU . toStream

-- |Collect the elements of a parallel array in a list
--
fromU :: UA e => UArr e -> [e]
{-# INLINE_U fromU #-}
fromU a = [a `Prim.indexU` i | i <- [0..lengthU a - 1]]

------------------------------------------------------------------------


here s = "Data.Array.Vector.Strict.Combinators." ++ s

-- |Map a function over an array
--
mapU :: (UA e, UA e') => (e -> e') -> UArr e -> UArr e'
{-# INLINE_U mapU #-}
mapU f = unstreamU . mapS f . streamU

-- |Extract all elements from an array that meet the given predicate
--
filterU :: UA e => (e -> Bool) -> UArr e -> UArr e 
{-# INLINE_U filterU #-}
filterU p = unstreamU . filterS p . streamU

-- |Extract all elements from an array according to a given flag array
-- 
packU:: UA e => UArr e -> UArr Bool -> UArr e
{-# INLINE_U packU #-}
packU xs = fstU . filterU sndS . zipU xs



-- |Array reduction proceeding from the left
--
foldlU :: UA a => (b -> a -> b) -> b -> UArr a -> b
{-# INLINE_U foldlU #-}
foldlU f z = foldS f z . streamU

-- |Array reduction proceeding from the left for non-empty arrays
--
-- FIXME: Rewrite for 'Stream's.
--
-- foldl1U :: UA a => (a -> a -> a) -> UArr a -> a
-- {-# INLINE_U foldl1U #-}
-- foldl1U f arr = checkNotEmpty (here "foldl1U") (lengthU arr) $
--                 foldlU f (arr `Prim.indexU` 0) (sliceU arr 1 (lengthU arr - 1))

foldl1U :: UA a => (a -> a -> a) -> UArr a -> a
foldl1U f = foldl1S f . streamU
{-# INLINE foldl1U #-}

foldl1MaybeU :: UA a => (a -> a -> a) -> UArr a -> MaybeS a
{-# INLINE_U foldl1MaybeU #-}
foldl1MaybeU f = fold1MaybeS f . streamU

-- |Array reduction that requires an associative combination function with its
-- unit
--
foldU :: UA a => (a -> a -> a) -> a -> UArr a -> a
{-# INLINE_U foldU #-}
foldU = foldlU

fold1MaybeU :: UA a => (a -> a -> a) -> UArr a -> MaybeS a
{-# INLINE_U fold1MaybeU #-}
fold1MaybeU = foldl1MaybeU

-- |Reduction of a non-empty array which requires an associative combination
-- function
--
fold1U :: UA a => (a -> a -> a) -> UArr a -> a
{-# INLINE_U fold1U #-}
fold1U = foldl1U

-- |Prefix scan proceedings from left to right
--
scanlU :: (UA a, UA b) => (b -> a -> b) -> b -> UArr a -> UArr b
{-# INLINE_U scanlU #-}
scanlU f z = unstreamU . scanS f z . streamU

-- |Prefix scan of a non-empty array proceeding from left to right
--
scanl1U :: UA a => (a -> a -> a) -> UArr a -> UArr a
{-# INLINE_U scanl1U #-}
scanl1U f arr = checkNotEmpty (here "scanl1U") (lengthU arr) $
                unstreamU (scan1S f (streamU arr))

-- |Prefix scan proceeding from left to right that needs an associative
-- combination function with its unit
--
scanU :: UA a => (a -> a -> a) -> a -> UArr a -> UArr a
{-# INLINE_U scanU #-}
scanU = scanlU

-- |Prefix scan of a non-empty array proceeding from left to right that needs
-- an associative combination function
--
scan1U :: UA a => (a -> a -> a) -> UArr a -> UArr a
{-# INLINE_U scan1U #-}
scan1U = scanl1U

scanResU :: UA a => (a -> a -> a) -> a -> UArr a -> UArr a :*: a
{-# INLINE_U scanResU #-}
scanResU f z = unstreamScan f z . streamU

unstreamScan :: UA a => (a -> a -> a) -> a -> Stream a -> UArr a :*: a
{-# INLINE_STREAM unstreamScan #-}
unstreamScan f z st@(Stream _ _ n)
  = newDynResU n (\marr -> unstreamScanM marr f z st)

unstreamScanM :: UA a => MUArr a s -> (a -> a -> a) -> a -> Stream a
                      -> ST s (Int :*: a)
{-# INLINE_U unstreamScanM #-}
unstreamScanM marr f z (Stream next s n) = fill s z 0
  where
    fill s !z !i = case next s of
                     Done       -> return (i :*: z)
                     Skip    s' -> s' `seq` fill s' z i
                     Yield x s' -> s' `seq`
                                   do
                                     writeMU marr i z
                                     fill s' (f z x) (i+1)

-- |Accumulating map from left to right. Does not return the accumulator.
--
-- FIXME: Naming inconsistent with lists.
--
mapAccumLU :: (UA a, UA b) => (c -> a -> c :*: b) -> c -> UArr a -> UArr b
{-# INLINE_U mapAccumLU #-}
mapAccumLU f z = unstreamU . mapAccumS f z . streamU

-- zipU is re-exported from UArr

-- |
--
zip3U :: (UA e1, UA e2, UA e3) 
      => UArr e1 -> UArr e2 -> UArr e3 -> UArr (e1 :*: e2 :*: e3)
{-# INLINE_U zip3U #-}
zip3U a1 a2 a3 = (a1 `zipU` a2) `zipU` a3

-- |
zipWithU :: (UA a, UA b, UA c) 
         => (a -> b -> c) -> UArr a -> UArr b -> UArr c
{-# INLINE_U zipWithU #-}
zipWithU f a1 a2 = unstreamU (zipWithS f (streamU a1) (streamU a2))

-- |
zipWith3U :: (UA a, UA b, UA c, UA d) 
          => (a -> b -> c -> d) -> UArr a -> UArr b -> UArr c -> UArr d
{-# INLINE_U zipWith3U #-}
zipWith3U f a1 a2 a3 = unstreamU (zipWith3S f (streamU a1)
                                              (streamU a2)
                                              (streamU a3))

-- unzipU is re-exported from UArr

-- |
unzip3U :: (UA e1, UA e2, UA e3) 
        => UArr (e1 :*: e2 :*: e3) -> (UArr e1 :*: UArr e2 :*: UArr e3)
{-# INLINE_U unzip3U #-}
unzip3U a = let (a12 :*: a3) = unzipU a
                (a1  :*: a2) = unzipU a12
            in
            (a1 :*: a2 :*: a3)

-- fstU and sndU reexported from UArr
-- |
combineU :: UA a
         => UArr Bool -> UArr a -> UArr a -> UArr a
{-# INLINE_U combineU #-}
combineU f a1 a2 = checkEq (here "combineU") 
     ("flag length not equal to sum of arg length")
     (lengthU f) (lengthU a1 + lengthU a2) $ 
  trace ("combineU:\n\t"  ++ show (lengthU f)  ++ "\n\t" ++ show (lengthU a1) ++ "\n\t" ++ show (lengthU a2) ++ "\n")
  unstreamU (combineS (streamU f) (streamU a1) (streamU a2))

------------------------------------------------------------------------

-- sliceU reexported from UArr

-- {-# INLINE_U extractU #-}
-- extractU :: UA a => UArr a -> Int -> Int -> UArr a
-- extractU arr i n = newU n $ \marr -> copyMU marr 0 (sliceU arr i n)

tailU :: UA e => UArr e -> UArr e
tailU = unstreamU . tailS . streamU
{-# INLINE_U tailU #-}

dropU :: UA e=> Int -> UArr e -> UArr e
dropU n = unstreamU . dropS n . streamU
{-# INLINE dropU #-}

takeU :: UA e=> Int -> UArr e -> UArr e
takeU n = unstreamU . takeS n . streamU
{-# INLINE takeU #-}

-- |Split an array into two halves at the given index
--
splitAtU :: UA e => Int -> UArr e -> (UArr e, UArr e)
splitAtU n a = (takeU n a, dropU n a)
{-# INLINE splitAtU #-}

{-
-- |Yield the tail of an array
--
tailU :: UA e => UArr e -> UArr e
{-# INLINE_U tailU #-}
tailU = unstreamU . tailS . streamU

-- |Extract a prefix of an array
--
takeU :: UA e=> Int -> UArr e -> UArr e
{-# INLINE_U takeU #-}
takeU n a = extractU a 0 n

-- |Extract a suffix of an array
--
dropU :: UA e => Int -> UArr e -> UArr e
{-# INLINE_U dropU #-}
dropU n a = let len = lengthU a
            in
            extractU a n (len - n)

-- |Split an array into two halves at the given index
--
splitAtU :: UA e => Int -> UArr e -> (UArr e, UArr e)
{-# INLINE_U splitAtU #-}
splitAtU n a = (takeU n a, dropU n a)

-}

------------------------------------------------------------------------

-- | 'takeWhile', applied to a predicate @p@ and a UArr @xs@,
-- returns the longest prefix (possibly empty) of @xs@ of elements that satisfy @p@.
--
takeWhileU :: UA e => (e -> Bool) -> UArr e -> UArr e
takeWhileU f = unstreamU . takeWhileS f . streamU
{-# INLINE_U takeWhileU #-}

-- | 'dropWhile' @p xs@ returns the suffix remaining after 'takeWhile' @p xs@.
dropWhileU :: UA e => (e -> Bool) -> UArr e -> UArr e
dropWhileU f = unstreamU . dropWhileS f . streamU
{-# INLINE_U dropWhileU #-}

------------------------------------------------------------------------
-- ** Searching with a predicate

-- | /O(n)/,/fusion/. The 'find' function takes a predicate and an array
-- and returns the first element in the list matching the predicate, or
-- 'Nothing' if there is no such element.
findU :: UA a => (a -> Bool) -> UArr a -> Maybe a
{-# INLINE_U findU #-}
findU p = findS p . streamU

-- | /O(n)/, /fusion/, The 'findIndex' function takes a predicate and an array and returns
-- the index of the first element in the array satisfying the predicate,
-- or 'Nothing' if there is no such element.
--
findIndexU :: UA a => (a -> Bool) -> UArr a -> Maybe Int
{-# INLINE_U findIndexU #-}
findIndexU p = findIndexS p . streamU

-- | /O(n)/,/fusion/. 'lookup' @key assocs@ looks up a key in an array
-- of pairs treated as an association table.
--
lookupU :: (Eq a, UA a, UA b) => a -> UArr (a :*: b) -> Maybe b
{-# INLINE_U lookupU #-}
lookupU p = lookupS p . streamU

------------------------------------------------------------------------

unfoldU :: UA a => Int -> (b -> MaybeS (a :*: b)) -> b -> UArr a
{-# INLINE_U unfoldU #-}
unfoldU n f z = unstreamU (unfoldS n f z)