Storable.hs

{-# LANGUAGE TypeFamilyDependencies, FlexibleContexts, GADTs, FlexibleInstances, TypeOperators #-}

{-|
Storable arrays for primitive values
-}

module Futhask.Array.Storable where
import Futhask.PrimTypes
import Futhask.Array.Element
import Futhask.Array.Boxed
import System.IO.Unsafe
import Foreign.Marshal.Utils (fromBool, toBool)
import qualified Foreign.ForeignPtr as F
import qualified Foreign.Storable as S

data SArray s where
    SArray :: S.Storable s => Int -> F.ForeignPtr s -> SArray s
data MSArray s where
    MSArray :: S.Storable s => Int -> F.ForeignPtr s -> MSArray s

instance Element Int where
    type Array Int = SArray Int
    type MArray Int = MSArray Int
    size = sizeS
    unsafeRead = unsafeReadS
    msize = msizeS
    unsafePeek = unsafePeekS
    unsafePoke = unsafePokeS
    scratch = scratchS
    unsafeThaw = unsafeThawS
    unsafeFreeze = unsafeFreezeS
    unsafeSlice = unsafeSliceS

instance Element F16 where
    type Array F16 = SArray F16
    type MArray F16 = MSArray F16
    size = sizeS
    unsafeRead = unsafeReadS
    msize = msizeS
    unsafePeek = unsafePeekS
    unsafePoke = unsafePokeS
    scratch = scratchS
    unsafeThaw = unsafeThawS
    unsafeFreeze = unsafeFreezeS
    unsafeSlice = unsafeSliceS

instance Element F32 where
    type Array F32 = SArray F32
    type MArray F32 = MSArray F32
    size = sizeS
    unsafeRead = unsafeReadS
    msize = msizeS
    unsafePeek = unsafePeekS
    unsafePoke = unsafePokeS
    scratch = scratchS
    unsafeThaw = unsafeThawS
    unsafeFreeze = unsafeFreezeS
    unsafeSlice = unsafeSliceS

instance Element F64 where
    type Array F64 = SArray F64
    type MArray F64 = MSArray F64
    size = sizeS
    unsafeRead = unsafeReadS
    msize = msizeS
    unsafePeek = unsafePeekS
    unsafePoke = unsafePokeS
    scratch = scratchS
    unsafeThaw = unsafeThawS
    unsafeFreeze = unsafeFreezeS
    unsafeSlice = unsafeSliceS

instance Element I8 where
    type Array I8 = SArray I8
    type MArray I8 = MSArray I8
    size = sizeS
    unsafeRead = unsafeReadS
    msize = msizeS
    unsafePeek = unsafePeekS
    unsafePoke = unsafePokeS
    scratch = scratchS
    unsafeThaw = unsafeThawS
    unsafeFreeze = unsafeFreezeS
    unsafeSlice = unsafeSliceS

instance Element I16 where
    type Array I16 = SArray I16
    type MArray I16 = MSArray I16
    size = sizeS
    unsafeRead = unsafeReadS
    msize = msizeS
    unsafePeek = unsafePeekS
    unsafePoke = unsafePokeS
    scratch = scratchS
    unsafeThaw = unsafeThawS
    unsafeFreeze = unsafeFreezeS
    unsafeSlice = unsafeSliceS

instance Element I32 where
    type Array I32 = SArray I32
    type MArray I32 = MSArray I32
    size = sizeS
    unsafeRead = unsafeReadS
    msize = msizeS
    unsafePeek = unsafePeekS
    unsafePoke = unsafePokeS
    scratch = scratchS
    unsafeThaw = unsafeThawS
    unsafeFreeze = unsafeFreezeS
    unsafeSlice = unsafeSliceS

instance Element I64 where
    type Array I64 = SArray I64
    type MArray I64 = MSArray I64
    size = sizeS
    unsafeRead = unsafeReadS
    msize = msizeS
    unsafePeek = unsafePeekS
    unsafePoke = unsafePokeS
    scratch = scratchS
    unsafeThaw = unsafeThawS
    unsafeFreeze = unsafeFreezeS
    unsafeSlice = unsafeSliceS

instance Element U8 where
    type Array U8 = SArray U8
    type MArray U8 = MSArray U8
    size = sizeS
    unsafeRead = unsafeReadS
    msize = msizeS
    unsafePeek = unsafePeekS
    unsafePoke = unsafePokeS
    scratch = scratchS
    unsafeThaw = unsafeThawS
    unsafeFreeze = unsafeFreezeS
    unsafeSlice = unsafeSliceS

instance Element U16 where
    type Array U16 = SArray U16
    type MArray U16 = MSArray U16
    size = sizeS
    unsafeRead = unsafeReadS
    msize = msizeS
    unsafePeek = unsafePeekS
    unsafePoke = unsafePokeS
    scratch = scratchS
    unsafeThaw = unsafeThawS
    unsafeFreeze = unsafeFreezeS
    unsafeSlice = unsafeSliceS

instance Element U32 where
    type Array U32 = SArray U32
    type MArray U32 = MSArray U32
    size = sizeS
    unsafeRead = unsafeReadS
    msize = msizeS
    unsafePeek = unsafePeekS
    unsafePoke = unsafePokeS
    scratch = scratchS
    unsafeThaw = unsafeThawS
    unsafeFreeze = unsafeFreezeS
    unsafeSlice = unsafeSliceS

instance Element U64 where
    type Array U64 = SArray U64
    type MArray U64 = MSArray U64
    size = sizeS
    unsafeRead = unsafeReadS
    msize = msizeS
    unsafePeek = unsafePeekS
    unsafePoke = unsafePokeS
    scratch = scratchS
    unsafeThaw = unsafeThawS
    unsafeFreeze = unsafeFreezeS
    unsafeSlice = unsafeSliceS

instance Element Bool where
    type Array Bool = SArray CBool
    type MArray Bool = MSArray CBool
    size = sizeS
    unsafeRead a i = toBool $ unsafeReadS a i
    msize = msizeS
    unsafePeek a i = fmap toBool $ unsafePeekS a i
    unsafePoke a i e = unsafePokeS a i $ fromBool e
    scratch = scratchS
    unsafeThaw = unsafeThawS
    unsafeFreeze = unsafeFreezeS
    unsafeSlice = unsafeSliceS

instance Element (SArray elem) where
    type Array (SArray elem) = BArray (SArray elem)
    type MArray (SArray elem) = MBArray (SArray elem)
    size = sizeB
    unsafeRead = unsafeReadB
    msize = msizeB
    unsafePeek = unsafePeekB
    unsafePoke = unsafePokeB
    scratch = scratchB
    unsafeThaw = unsafeThawB
    unsafeFreeze = unsafeFreezeB
    unsafeSlice = unsafeSliceB

instance (Show a, Element a, SArray a ~ Array a) => Show (SArray a) where
    show = show . toList

sizeS :: SArray a -> Int
sizeS (SArray sz _) = sz 
unsafeReadS :: S.Storable a => SArray a -> Int -> a
unsafeReadS (SArray _ ptr) idx = unsafePerformIO $ F.withForeignPtr ptr (\ptr -> S.peekElemOff ptr idx)
msizeS :: S.Storable a => MSArray a -> Int
msizeS (MSArray sz _) = sz
unsafePeekS :: MSArray a -> Int -> IO a 
unsafePeekS (MSArray _ ptr) idx = F.withForeignPtr ptr (\ptr -> S.peekElemOff ptr idx)
unsafePokeS :: MSArray a -> Int -> a -> IO () 
unsafePokeS (MSArray _ ptr) idx elem = F.withForeignPtr ptr (\ptr -> S.pokeElemOff ptr idx elem) 
scratchS :: S.Storable a => Int -> IO (MSArray a)
scratchS sz = fmap (MSArray sz) (F.mallocForeignPtrArray sz)
unsafeThawS :: S.Storable a => SArray a -> IO (MSArray a)
unsafeThawS (SArray sz ptr) = pure (MSArray sz ptr)
unsafeFreezeS :: S.Storable a => MSArray a -> IO (SArray a)
unsafeFreezeS (MSArray sz ptr) = pure (SArray sz ptr)
unsafeSliceS (SArray _ ptr) i0 sz = SArray sz (F.plusForeignPtr ptr i0)