-- |
-- Copyright   : Anders Claesson 2013-2016
-- Maintainer  : Anders Claesson <anders.claesson@gmail.com>
--
-- Convenience functions for dealing with arrays of 'CLong's.

module Sym.Internal.CLongArray
    (
    -- * Data type
    CLongArray

    -- * Conversions
    , fromList
    , toList
    , slices

    -- * Accessors
    , size
    , at
    , unsafeAt
    , elemIndices

    -- * Map
    , imap
    , izipWith

    -- * Low level functions
    , unsafeNew
    , unsafeWith
    ) where

import Data.Ord
import qualified Data.Vector.Storable as V
import qualified Data.Vector.Storable.Mutable as MV
import Sym.Internal.Size
import Foreign
import Foreign.C.Types

infixl 9 `at`
infixl 9 `unsafeAt`

-- Data type
-- ---------

-- | An array of 'CLong's
newtype CLongArray = CArr (V.Vector CLong) deriving (CLongArray -> CLongArray -> Bool
(CLongArray -> CLongArray -> Bool)
-> (CLongArray -> CLongArray -> Bool) -> Eq CLongArray
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: CLongArray -> CLongArray -> Bool
== :: CLongArray -> CLongArray -> Bool
$c/= :: CLongArray -> CLongArray -> Bool
/= :: CLongArray -> CLongArray -> Bool
Eq)

instance Ord CLongArray where
    compare :: CLongArray -> CLongArray -> Ordering
compare CLongArray
u CLongArray
v =
        case (CLongArray -> Int) -> CLongArray -> CLongArray -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing CLongArray -> Int
forall a. Size a => a -> Int
size CLongArray
u CLongArray
v of
          Ordering
EQ -> (CLongArray -> [Int]) -> CLongArray -> CLongArray -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing CLongArray -> [Int]
toList CLongArray
u CLongArray
v
          Ordering
x  -> Ordering
x

instance Size CLongArray where
    size :: CLongArray -> Int
size (CArr Vector CLong
w) = Vector CLong -> Int
forall a. Storable a => Vector a -> Int
V.length Vector CLong
w
    {-# INLINE size #-}

instance Show CLongArray where
    show :: CLongArray -> String
show CLongArray
w = String
"fromList " String -> ShowS
forall a. [a] -> [a] -> [a]
++ [Int] -> String
forall a. Show a => a -> String
show (CLongArray -> [Int]
toList CLongArray
w)


-- Conversions
-- -----------

-- | Construct an array from a list of elements.
fromList :: [Int] -> CLongArray
fromList :: [Int] -> CLongArray
fromList = Vector CLong -> CLongArray
CArr (Vector CLong -> CLongArray)
-> ([Int] -> Vector CLong) -> [Int] -> CLongArray
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [CLong] -> Vector CLong
forall a. Storable a => [a] -> Vector a
V.fromList ([CLong] -> Vector CLong)
-> ([Int] -> [CLong]) -> [Int] -> Vector CLong
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> CLong) -> [Int] -> [CLong]
forall a b. (a -> b) -> [a] -> [b]
map Int -> CLong
forall a b. (Integral a, Num b) => a -> b
fromIntegral

-- | The list of elements.
toList :: CLongArray -> [Int]
toList :: CLongArray -> [Int]
toList (CArr Vector CLong
w) = (CLong -> Int) -> [CLong] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map CLong -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([CLong] -> [Int]) -> [CLong] -> [Int]
forall a b. (a -> b) -> a -> b
$ Vector CLong -> [CLong]
forall a. Storable a => Vector a -> [a]
V.toList Vector CLong
w

-- | Slice a 'CLongArray' into contiguous segments of the given
-- sizes. Each segment size must be positive and they must sum to the
-- size of the array.
slices :: [Int] -> CLongArray -> [CLongArray]
slices :: [Int] -> CLongArray -> [CLongArray]
slices [] CLongArray
_ = []
slices (Int
k:[Int]
ks) (CArr Vector CLong
w) = let (Vector CLong
u,Vector CLong
v) = Int -> Vector CLong -> (Vector CLong, Vector CLong)
forall a. Storable a => Int -> Vector a -> (Vector a, Vector a)
V.splitAt Int
k Vector CLong
w in Vector CLong -> CLongArray
CArr Vector CLong
u CLongArray -> [CLongArray] -> [CLongArray]
forall a. a -> [a] -> [a]
: [Int] -> CLongArray -> [CLongArray]
slices [Int]
ks (Vector CLong -> CLongArray
CArr Vector CLong
v)


-- Accessors
-- ---------

-- | @w \`at\` i@ is the value of @w@ at @i@, where @i@ is in @[0..size w-1]@.
at :: CLongArray -> Int -> Int
at :: CLongArray -> Int -> Int
at (CArr Vector CLong
w) Int
i =
    case Vector CLong -> Int -> Maybe CLong
forall a. Storable a => Vector a -> Int -> Maybe a
(V.!?) Vector CLong
w Int
i of
      Maybe CLong
Nothing -> String -> Int
forall a. HasCallStack => String -> a
error String
"Sym.Internal.CLongArray.at: out of range"
      Just CLong
j  -> CLong -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral CLong
j

-- | Like 'at' but without range checking.
unsafeAt :: CLongArray -> Int -> Int
unsafeAt :: CLongArray -> Int -> Int
unsafeAt (CArr Vector CLong
w) = CLong -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (CLong -> Int) -> (Int -> CLong) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector CLong -> Int -> CLong
forall a. Storable a => Vector a -> Int -> a
(V.!) Vector CLong
w

-- | The indices of all elements equal to the query element, in
-- ascending order.
elemIndices :: CLong -> CLongArray -> V.Vector Int
elemIndices :: CLong -> CLongArray -> Vector Int
elemIndices CLong
x (CArr Vector CLong
w) = CLong -> Vector CLong -> Vector Int
forall a. (Storable a, Eq a) => a -> Vector a -> Vector Int
V.elemIndices CLong
x Vector CLong
w


-- Map and Zip
-- -----------

-- | Apply a function to every element of an array and its index.
imap :: (Int -> CLong -> CLong) -> CLongArray -> CLongArray
imap :: (Int -> CLong -> CLong) -> CLongArray -> CLongArray
imap Int -> CLong -> CLong
f (CArr Vector CLong
w) =  Vector CLong -> CLongArray
CArr ((Int -> CLong -> CLong) -> Vector CLong -> Vector CLong
forall a b.
(Storable a, Storable b) =>
(Int -> a -> b) -> Vector a -> Vector b
V.imap Int -> CLong -> CLong
f Vector CLong
w)

-- | Apply a function to corresponding pairs of elements and their (shared) index.
izipWith :: (Int -> CLong -> CLong -> CLong) -> CLongArray -> CLongArray -> CLongArray
izipWith :: (Int -> CLong -> CLong -> CLong)
-> CLongArray -> CLongArray -> CLongArray
izipWith Int -> CLong -> CLong -> CLong
f (CArr Vector CLong
u) (CArr Vector CLong
v) = Vector CLong -> CLongArray
CArr ((Int -> CLong -> CLong -> CLong)
-> Vector CLong -> Vector CLong -> Vector CLong
forall a b c.
(Storable a, Storable b, Storable c) =>
(Int -> a -> b -> c) -> Vector a -> Vector b -> Vector c
V.izipWith Int -> CLong -> CLong -> CLong
f Vector CLong
u Vector CLong
v)

-- Low level functions
-- -------------------

-- | Create a new array of the given size that is initialized through
-- an IO action.
unsafeNew :: Int -> (Ptr CLong -> IO ()) -> IO CLongArray
unsafeNew :: Int -> (Ptr CLong -> IO ()) -> IO CLongArray
unsafeNew Int
n Ptr CLong -> IO ()
act = do
    Vector CLong
v <- MVector RealWorld CLong -> IO (Vector CLong)
MVector (PrimState IO) CLong -> IO (Vector CLong)
forall a (m :: * -> *).
(Storable a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
V.unsafeFreeze (MVector RealWorld CLong -> IO (Vector CLong))
-> IO (MVector RealWorld CLong) -> IO (Vector CLong)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Int -> IO (MVector (PrimState IO) CLong)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
Int -> m (MVector (PrimState m) a)
MV.unsafeNew Int
n
    let (ForeignPtr CLong
ptr, Int
_) = Vector CLong -> (ForeignPtr CLong, Int)
forall a. Vector a -> (ForeignPtr a, Int)
V.unsafeToForeignPtr0 Vector CLong
v
    ForeignPtr CLong -> (Ptr CLong -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr CLong
ptr Ptr CLong -> IO ()
act
    CLongArray -> IO CLongArray
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (CLongArray -> IO CLongArray) -> CLongArray -> IO CLongArray
forall a b. (a -> b) -> a -> b
$ Vector CLong -> CLongArray
CArr (ForeignPtr CLong -> Int -> Vector CLong
forall a. ForeignPtr a -> Int -> Vector a
V.unsafeFromForeignPtr0 ForeignPtr CLong
ptr Int
n)

-- | Pass a pointer to the array to an IO action; the array may not be
-- modified through the pointer.
unsafeWith :: CLongArray -> (Ptr CLong -> IO a) -> IO a
unsafeWith :: forall a. CLongArray -> (Ptr CLong -> IO a) -> IO a
unsafeWith (CArr Vector CLong
w) = Vector CLong -> (Ptr CLong -> IO a) -> IO a
forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
V.unsafeWith Vector CLong
w