Array.hs

{-# LANGUAGE CPP #-}
{-# OPTIONS_GHC -fno-warn-missing-methods #-}
{-# OPTIONS_GHC -fno-warn-incomplete-patterns #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module SubHask.Algebra.Array
    ( BArray (..)
    , UArray (..)
    , Unboxable
    )
    where

import Control.Monad
import Control.Monad.Primitive
import Data.Primitive as Prim
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Unboxed.Mutable as VUM
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Generic.Mutable as VGM

import qualified Prelude as P
import SubHask.Algebra
import SubHask.Algebra.Parallel
import SubHask.Algebra.Vector
import SubHask.Category
import SubHask.Internal.Prelude

-------------------------------------------------------------------------------
-- boxed arrays

newtype BArray e = BArray (V.Vector e)

type instance Index (BArray e) = Int
type instance Logic (BArray e) = Logic e
type instance Scalar (BArray e) = Int
type instance Elem (BArray e) = e
type instance SetElem (BArray e) e' = BArray e'

----------------------------------------
-- mutability

mkMutable [t| forall e. BArray e |]

----------------------------------------
-- misc instances

instance Arbitrary e => Arbitrary (BArray e) where
    arbitrary = fmap fromList arbitrary

instance NFData e => NFData (BArray e) where
    rnf (BArray v) = rnf v

instance Show e => Show (BArray e) where
    show (BArray v) = "BArray " ++ show (VG.toList v)

----------------------------------------
-- algebra

instance Semigroup (BArray e) where
    (BArray v1)+(BArray v2) = fromList $ VG.toList v1 ++ VG.toList v2

instance Monoid (BArray e) where
    zero = BArray VG.empty

instance Normed (BArray e) where
    size (BArray v) = VG.length v

----------------------------------------
-- comparison

instance (ValidLogic e, Eq_ e) => Eq_ (BArray e) where
    a1==a2 = toList a1==toList a2

instance (ClassicalLogic e, POrd_ e) => POrd_ (BArray e) where
    inf a1 a2 = fromList $ inf (toList a1) (toList a2)

instance (ClassicalLogic e, POrd_ e) => MinBound_ (BArray e) where
    minBound = zero

----------------------------------------
-- container

instance Constructible (BArray e) where
    fromList1 x xs = BArray $ VG.fromList (x:xs)

instance (ValidLogic e, Eq_ e) => Container (BArray e) where
    elem e arr = elem e $ toList arr

instance Foldable (BArray e) where

    {-# INLINE toList #-}
    toList (BArray v) = VG.toList v

    {-# INLINE uncons #-}
    uncons (BArray v) = if VG.null v
        then Nothing
        else Just (VG.head v, BArray $ VG.tail v)

    {-# INLINE unsnoc #-}
    unsnoc (BArray v) = if VG.null v
        then Nothing
        else Just (BArray $ VG.init v, VG.last v)

    {-# INLINE foldMap #-}
    foldMap f   (BArray v) = VG.foldl' (\a e -> a + f e) zero v

    {-# INLINE foldr #-}
    {-# INLINE foldr' #-}
    {-# INLINE foldr1 #-}
    {-# INLINE foldr1' #-}
    {-# INLINE foldl #-}
    {-# INLINE foldl' #-}
    {-# INLINE foldl1 #-}
    {-# INLINE foldl1' #-}
    foldr   f x (BArray v) = VG.foldr   f x v
    foldr'  f x (BArray v) = VG.foldr'  f x v
    foldr1  f   (BArray v) = VG.foldr1  f   v
    foldr1' f   (BArray v) = VG.foldr1' f   v
    foldl   f x (BArray v) = VG.foldl   f x v
    foldl'  f x (BArray v) = VG.foldl'  f x v
    foldl1  f   (BArray v) = VG.foldl1  f   v
    foldl1' f   (BArray v) = VG.foldl1' f   v

instance ValidLogic e => Sliceable (BArray e) where
    slice i n (BArray v) = BArray $ VG.slice i n v

instance ValidLogic e => IxContainer (BArray e) where
    lookup i (BArray v) = v VG.!? i
    (!) (BArray v) = VG.unsafeIndex v
    indices (BArray v) = [0..VG.length v-1]
    values (BArray v) = VG.toList v
    imap f (BArray v) = BArray $ VG.imap f v

instance ValidLogic e => Partitionable (BArray e) where
    partition n arr = go 0
        where
            go i = if i>=length arr
                then []
                else (slice i len arr):(go $ i+lenmax)
                where
                    len = if i+lenmax >= length arr
                        then (length arr)-i
                        else lenmax

            lenmax = length arr `quot` n

-------------------------------------------------------------------------------
-- unboxed arrays

data UArray e
    = UArray !(VU.Vector e)
--     | UArray_Zero

type instance Index (UArray e) = Int
type instance Logic (UArray e) = Logic e
type instance Scalar (UArray e) = Int
type instance Elem (UArray e) = e
type instance SetElem (UArray e) e' = UArray e'

----------------------------------------
-- mutability

mkMutable [t| forall e. UArray e |]

----------------------------------------
-- misc instances

instance (Unboxable e, Arbitrary e) => Arbitrary (UArray e) where
    arbitrary = fmap fromList arbitrary

instance (NFData e) => NFData (UArray e) where
    rnf (UArray v) = rnf v
--     rnf UArray_Zero = ()

instance (Unboxable e, Show e) => Show (UArray e) where
    show arr = "UArray " ++ show (toList arr)

----------------------------------------
-- algebra

instance Unboxable e => Semigroup (UArray e) where
--     UArray_Zero + a           = a
--     a           + UArray_Zero = a
    (UArray v1) + (UArray v2) = fromList $ VG.toList v1 ++ VG.toList v2

instance Unboxable e => Monoid (UArray e) where
    zero = UArray VG.empty
--     zero = UArray_Zero

instance Unbox e => Normed (UArray e) where
    size (UArray v) = VG.length v
--     size UArray_Zero = 0

----------------------------------------
-- comparison

instance (Unboxable e, Eq_ e) => Eq_ (UArray e) where
    a1==a2 = toList a1==toList a2

instance (Unboxable e, POrd_ e) => POrd_ (UArray e) where
    inf a1 a2 = fromList $ inf (toList a1) (toList a2)

instance (Unboxable e, POrd_ e) => MinBound_ (UArray e) where
    minBound = zero

----------------------------------------
-- container

type Unboxable e = (Constructible (UArray e), Eq e, Unbox e)

#define mkConstructible(e) \
instance Constructible (UArray e) where\
    { fromList1 x xs = UArray $ VG.fromList (x:xs) } ;

mkConstructible(Int)
mkConstructible(Float)
mkConstructible(Double)
mkConstructible(Char)
mkConstructible(Bool)

instance
    ( ClassicalLogic r
    , Eq_ r
    , Unbox r
    , Prim r
    , FreeModule r
    , IsScalar r
    , ValidUVector s r
    ) => Constructible (UArray (UVector (s::Symbol) r))
        where

    {-# INLINABLE fromList1 #-}
    fromList1 x xs = fromList1N (length $ x:xs) x xs

    {-# INLINABLE fromList1N #-}
    fromList1N n x xs = unsafeInlineIO $ do
        marr <- safeNewByteArray (n*size'*rbytes) 16
        let mv = UArray_MUVector marr 0 n size'

        let go [] (-1) = return ()
            go (x':xs') i = do
                VGM.unsafeWrite mv i x'
                go xs' (i-1)

        go (P.reverse $ x:xs) (n-1)
        v <- VG.basicUnsafeFreeze mv
        return $ UArray v
        where
            rbytes=Prim.sizeOf (undefined::r)
            size'=roundUpToNearest 4 $ dim x

instance Unboxable e => Container (UArray e) where
    elem e (UArray v) = elem e $ VG.toList v

instance Unboxable e => Foldable (UArray e) where

    {-# INLINE toList #-}
    toList (UArray v) = VG.toList v

    {-# INLINE uncons #-}
    uncons (UArray v) = if VG.null v
        then Nothing
        else Just (VG.head v, UArray $ VG.tail v)

    {-# INLINE unsnoc #-}
    unsnoc (UArray v) = if VG.null v
        then Nothing
        else Just (UArray $ VG.init v, VG.last v)

    {-# INLINE foldMap #-}
    foldMap f   (UArray v) = VG.foldl' (\a e -> a + f e) zero v

    {-# INLINE foldr #-}
    {-# INLINE foldr' #-}
    {-# INLINE foldr1 #-}
    {-# INLINE foldr1' #-}
    {-# INLINE foldl #-}
    {-# INLINE foldl' #-}
    {-# INLINE foldl1 #-}
    {-# INLINE foldl1' #-}
    foldr   f x (UArray v) = VG.foldr   f x v
    foldr'  f x (UArray v) = VG.foldr'  f x v
    foldr1  f   (UArray v) = VG.foldr1  f   v
    foldr1' f   (UArray v) = VG.foldr1' f   v
    foldl   f x (UArray v) = VG.foldl   f x v
    foldl'  f x (UArray v) = VG.foldl'  f x v
    foldl1  f   (UArray v) = VG.foldl1  f   v
    foldl1' f   (UArray v) = VG.foldl1' f   v

instance Unboxable e => Sliceable (UArray e) where
    slice i n (UArray v) = UArray $ VG.slice i n v

instance Unboxable e => IxContainer (UArray e) where
    type ValidElem (UArray e) e = Unboxable e

    lookup i (UArray v) = v VG.!? i
    (!) (UArray v) = VG.unsafeIndex v
    indices (UArray v) = [0..VG.length v-1]
    values (UArray v) = VG.toList v

instance Unboxable e => Partitionable (UArray e) where
    partition n arr = go 0
        where
            go i = if i>=length arr
                then []
                else (slice i len arr):(go $ i+lenmax)
                where
                    len = if i+lenmax >= length arr
                        then (length arr)-i
                        else lenmax

            lenmax = length arr `quot` n

-------------------------------------------------------------------------------
-- UVector

instance
    ( IsScalar elem
    , ClassicalLogic elem
    , Unbox elem
    , Prim elem
    ) => Unbox (UVector (n::Symbol) elem)

data instance VU.Vector (UVector (n::Symbol) elem) = UArray_UVector
    {-#UNPACK#-}!ByteArray
    {-#UNPACK#-}!Int -- offset
    {-#UNPACK#-}!Int -- length of container
    {-#UNPACK#-}!Int -- length of element vectors

instance
    ( IsScalar elem
    , Unbox elem
    , Prim elem
    ) => VG.Vector VU.Vector (UVector (n::Symbol) elem)
        where

    {-# INLINABLE basicLength #-}
    basicLength (UArray_UVector _ _ n _) = n

    {-# INLINABLE basicUnsafeSlice #-}
    basicUnsafeSlice i len' (UArray_UVector arr off _ size') = UArray_UVector arr (off+i*size') len' size'

    {-# INLINABLE basicUnsafeFreeze #-}
    basicUnsafeFreeze (UArray_MUVector marr off n size') = do
        arr <- unsafeFreezeByteArray marr
        return $ UArray_UVector arr off n size'

    {-# INLINABLE basicUnsafeThaw #-}
    basicUnsafeThaw (UArray_UVector arr off n size')= do
        marr <- unsafeThawByteArray arr
        return $ UArray_MUVector marr off n size'

    {-# INLINABLE basicUnsafeIndexM #-}
    basicUnsafeIndexM (UArray_UVector arr off _ size') i =
        return $ UVector_Dynamic arr (off+i*size') size'

data instance VUM.MVector s (UVector (n::Symbol) elem) = UArray_MUVector
    {-#UNPACK#-}!(MutableByteArray s)
    {-#UNPACK#-}!Int -- offset in number of elem
    {-#UNPACK#-}!Int -- length of container
    {-#UNPACK#-}!Int -- length of element vectors

instance
    ( ClassicalLogic elem
    , IsScalar elem
    , Unbox elem
    , Prim elem
    ) => VGM.MVector VUM.MVector (UVector (n::Symbol) elem)
        where

    {-# INLINABLE basicLength #-}
    basicLength (UArray_MUVector _ _ n _) = n

    {-# INLINABLE basicUnsafeSlice #-}
    basicUnsafeSlice i lenM' (UArray_MUVector marr off _ size')
        = UArray_MUVector marr (off+i*size') lenM' size'

    {-# INLINABLE basicOverlaps #-}
    basicOverlaps (UArray_MUVector marr1 _ _ _) (UArray_MUVector marr2 _ _ _)
        = sameMutableByteArray marr1 marr2

    {-# INLINABLE basicUnsafeNew #-}
    basicUnsafeNew 0 = do
        marr <- newByteArray 0
        return $ UArray_MUVector marr 0 0 0
    basicUnsafeNew _ = error "basicUnsafeNew not supported on UArray_MUVector with nonzero size"

    {-# INLINABLE basicUnsafeRead #-}
    basicUnsafeRead (UArray_MUVector marr off _ size') i = do
        let b=Prim.sizeOf (undefined::elem)
        marr' <- safeNewByteArray (size'*b) 16
        copyMutableByteArray marr' 0 marr ((off+i*size')*b) (size'*b)
        arr <- unsafeFreezeByteArray marr'
        return $ UVector_Dynamic arr 0 size'

    {-# INLINABLE basicUnsafeWrite #-}
    basicUnsafeWrite (UArray_MUVector marr1 off1 _ size1) loc (UVector_Dynamic arr2 off2 _) =
        copyByteArray marr1 ((off1+size1*loc)*b) arr2 (off2*b) (size1*b)
        where
            b=Prim.sizeOf (undefined::elem)

    {-# INLINABLE basicUnsafeCopy #-}
    basicUnsafeCopy (UArray_MUVector marr1 off1 _ size1) (UArray_MUVector marr2 off2 n2 _) =
        copyMutableByteArray marr1 (off1*b) marr2 (off2*b) (n2*b)
        where
            b = size1*Prim.sizeOf (undefined::elem)

    {-# INLINABLE basicUnsafeMove #-}
    basicUnsafeMove (UArray_MUVector marr1 off1 _ size1) (UArray_MUVector marr2 off2 n2 _) =
        moveByteArray marr1 (off1*b) marr2 (off2*b) (n2*b)
        where
            b = size1*Prim.sizeOf (undefined::elem)

--------------------------------------------------------------------------------
-- Labeled'

instance
    ( Unbox y
    , Prim y
    , ClassicalLogic a
    , IsScalar a
    , Unbox a
    , Prim a
    ) => Unbox (Labeled' (UVector (s::Symbol) a) y)

data instance VUM.MVector s (Labeled' (UVector (n::Symbol) elem) y) = UArray_Labeled'_MUVector
    {-#UNPACK#-}!(MutableByteArray s)
    {-#UNPACK#-}!Int -- offset in number of elem
    {-#UNPACK#-}!Int -- length of container
    {-#UNPACK#-}!Int -- length of element vectors

instance
    ( ClassicalLogic elem
    , IsScalar elem
    , Unbox elem
    , Prim elem
    , Prim y
    ) => VGM.MVector VUM.MVector (Labeled' (UVector (n::Symbol) elem) y)
        where

    {-# INLINABLE basicLength #-}
    basicLength (UArray_Labeled'_MUVector _ _ n _) = n

    {-# INLINABLE basicUnsafeSlice #-}
    basicUnsafeSlice i lenM' (UArray_Labeled'_MUVector marr off _ size')
        = UArray_Labeled'_MUVector marr (off+i*(size'+ysize)) lenM' size'
        where
            ysize=roundUpToNearest 4 $ Prim.sizeOf (undefined::y) `quot` Prim.sizeOf (undefined::elem)

    {-# INLINABLE basicOverlaps #-}
    basicOverlaps (UArray_Labeled'_MUVector marr1 _ _ _) (UArray_Labeled'_MUVector marr2 _ _ _)
        = sameMutableByteArray marr1 marr2

    {-# INLINABLE basicUnsafeNew #-}
    basicUnsafeNew 0 = do
        marr <- newByteArray 0
        return $ UArray_Labeled'_MUVector marr 0 0 0
    basicUnsafeNew _ = error "basicUnsafeNew not supported on UArray_MUVector with nonzero size"

    {-# INLINABLE basicUnsafeRead #-}
    basicUnsafeRead (UArray_Labeled'_MUVector marr off _ size') i = do
        marr' <- safeNewByteArray (size'*b) 16
        copyMutableByteArray marr' 0 marr ((off+i*(size'+ysize))*b) (size'*b)
        arr <- unsafeFreezeByteArray marr'
        let x=UVector_Dynamic arr 0 size'
        y <- readByteArray marr $ (off+i*(size'+ysize)+size') `quot` ysizereal
        return $ Labeled' x y
        where
            b=Prim.sizeOf (undefined::elem)
            ysizereal = Prim.sizeOf (undefined::y) `quot` Prim.sizeOf (undefined::elem)
            ysize=roundUpToNearest 4 $ ysizereal

    {-# INLINABLE basicUnsafeWrite #-}
    basicUnsafeWrite
        (UArray_Labeled'_MUVector marr1 off1 _ size')
        i
        (Labeled' (UVector_Dynamic arr2 off2 _) y)
        = do
            copyByteArray marr1 ((off1+i*(size'+ysize))*b) arr2 (off2*b) (size'*b)
            writeByteArray marr1 ((off1+i*(size'+ysize)+size') `quot` ysizereal) y
        where
            b=Prim.sizeOf (undefined::elem)
            ysizereal = Prim.sizeOf (undefined::y) `quot` Prim.sizeOf (undefined::elem)
            ysize=roundUpToNearest 4 $ ysizereal

    {-# INLINABLE basicUnsafeCopy #-}
    basicUnsafeCopy
        (UArray_Labeled'_MUVector marr1 off1 _ size1)
        (UArray_Labeled'_MUVector marr2 off2 n2 _)
        = copyMutableByteArray marr1 (off1*b) marr2 (off2*b) (n2*b)
        where
            b = (size1+ysize)*Prim.sizeOf (undefined::elem)
            ysize=roundUpToNearest 4 $ Prim.sizeOf (undefined::y) `quot` Prim.sizeOf (undefined::elem)

    {-# INLINABLE basicUnsafeMove #-}
    basicUnsafeMove
        (UArray_Labeled'_MUVector marr1 off1 _ size1)
        (UArray_Labeled'_MUVector marr2 off2 n2 _)
        = moveByteArray marr1 (off1*b) marr2 (off2*b) (n2*b)
        where
            b = (size1+ysize)*Prim.sizeOf (undefined::elem)
            ysize=roundUpToNearest 4 $ Prim.sizeOf (undefined::y) `quot` Prim.sizeOf (undefined::elem)

data instance VU.Vector (Labeled' (UVector (n::Symbol) elem) y) = UArray_Labeled'_UVector
    {-#UNPACK#-}!ByteArray
    {-#UNPACK#-}!Int -- offset
    {-#UNPACK#-}!Int -- length of container
    {-#UNPACK#-}!Int -- length of element vectors

instance
    ( IsScalar elem
    , Unbox elem
    , Prim elem
    , Prim y
    ) => VG.Vector VU.Vector (Labeled' (UVector (n::Symbol) elem) y)
        where

    {-# INLINABLE basicLength #-}
    basicLength (UArray_Labeled'_UVector _ _ n _) = n

    {-# INLINABLE basicUnsafeSlice #-}
    basicUnsafeSlice i len' (UArray_Labeled'_UVector arr off _ size')
        = UArray_Labeled'_UVector arr (off+i*(size'+ysize)) len' size'
        where
            ysize=roundUpToNearest 4 $ Prim.sizeOf (undefined::y) `quot` Prim.sizeOf (undefined::elem)

    {-# INLINABLE basicUnsafeFreeze #-}
    basicUnsafeFreeze (UArray_Labeled'_MUVector marr off n size') = do
        arr <- unsafeFreezeByteArray marr
        return $ UArray_Labeled'_UVector arr off n size'

    {-# INLINABLE basicUnsafeThaw #-}
    basicUnsafeThaw (UArray_Labeled'_UVector arr off n size')= do
        marr <- unsafeThawByteArray arr
        return $ UArray_Labeled'_MUVector marr off n size'

    {-# INLINE basicUnsafeIndexM #-}
    basicUnsafeIndexM (UArray_Labeled'_UVector arr off _ size') i =
--         trace ("off'="+show off') $
        return $ Labeled' x y
        where
            off' = off+i*(size'+ysize)
            x = UVector_Dynamic arr off' size'
            y = indexByteArray arr $ (off'+size') `quot` ysizereal
            ysizereal = Prim.sizeOf (undefined::y) `quot` Prim.sizeOf (undefined::elem)
            ysize=roundUpToNearest 4 $ ysizereal

instance
    ( ClassicalLogic r
    , Eq_ r
    , Unbox r
    , Prim r
    , FreeModule r
    , IsScalar r
    , Prim y
    , Unbox y
    , ValidUVector s r
    ) => Constructible (UArray (Labeled' (UVector (s::Symbol) r) y))
        where

    {-# INLINABLE fromList1 #-}
    fromList1 x xs = fromList1N (length $ x:xs) x xs

    {-# INLINABLE fromList1N #-}
    fromList1N n x xs = unsafeInlineIO $ do
        let arrlen = n*(xsize+ysize)
        marr <- safeNewByteArray (arrlen*rbytes) 16
        setByteArray marr 0 arrlen (0::r)
        let mv = UArray_Labeled'_MUVector marr 0 n xsize

        let go [] (-1) = return ()
            go (x':xs') i = do
                VGM.unsafeWrite mv i x'
                go xs' (i-1)

        go (P.reverse $ x:xs) (n-1)
        v <- VG.basicUnsafeFreeze mv
        return $ UArray v
        where
            rbytes=Prim.sizeOf (undefined::r)

            xsize=roundUpToNearest 4 $ dim $ xLabeled' x
            ysize=roundUpToNearest 4 $ Prim.sizeOf (undefined::y) `quot` rbytes