{-# LANGUAGE FlexibleInstances, ScopedTypeVariables #-}
{-# OPTIONS -fno-warn-orphans #-}

-- | Reading and writing Repa arrays as binary files.
module Data.Array.Repa.IO.Binary
        ( readArrayFromStorableFile
        , writeArrayToStorableFile)
where
import Foreign.Storable
import Foreign.Ptr
import Foreign.ForeignPtr
import Foreign.Marshal.Alloc
import System.IO
import Data.Array.Repa                  as R
import Data.Array.Repa.Repr.ForeignPtr  as R
import Prelude                          as P
import Control.Monad


-- | Read an array from a file.
--   Data appears in host byte order.
--   If the file size does match the provided shape then `error`.
readArrayFromStorableFile 
        :: forall a sh
        .  (Shape sh, Storable a)
        => FilePath 
        -> sh
        -> IO (Array F sh a)

readArrayFromStorableFile :: FilePath -> sh -> IO (Array F sh a)
readArrayFromStorableFile FilePath
filePath sh
sh
 = do
        -- Determine number of bytes per element.
        let (a
fake   :: a)               = a
forall a. HasCallStack => a
undefined
        let (Integer
bytes1 :: Integer)         = Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Integer) -> Int -> Integer
forall a b. (a -> b) -> a -> b
$ a -> Int
forall a. Storable a => a -> Int
sizeOf a
fake

        -- Determine how many elements the whole file will give us.
        Handle
h :: Handle     <- FilePath -> IOMode -> IO Handle
openBinaryFile FilePath
filePath IOMode
ReadMode
        Integer
bytesTotal      <- Handle -> IO Integer
hFileSize Handle
h

        let lenTotal :: Integer
lenTotal      =  Integer
bytesTotal Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
bytes1
        let bytesExpected :: Integer
bytesExpected =  Integer
bytes1 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
lenTotal
        
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Integer
bytesTotal Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
/= Integer
bytesExpected)
         (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath -> IO ()
forall a. HasCallStack => FilePath -> a
error (FilePath -> IO ()) -> FilePath -> IO ()
forall a b. (a -> b) -> a -> b
$ [FilePath] -> FilePath
unlines
                [FilePath
"Data.Array.Repa.IO.Binary.readArrayFromStorableFile: not a whole number of elements in file"
                , FilePath
"element length = " FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
P.++ Integer -> FilePath
forall a. Show a => a -> FilePath
show Integer
bytes1
                , FilePath
"file size      = " FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
P.++ Integer -> FilePath
forall a. Show a => a -> FilePath
show Integer
bytesTotal
                , FilePath
"slack space    = " FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
P.++ Integer -> FilePath
forall a. Show a => a -> FilePath
show (Integer
bytesTotal Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
bytes1) ]
         
        let bytesTotal' :: Int
bytesTotal' = Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
bytesTotal
        Ptr a
buf :: Ptr a    <- Int -> IO (Ptr a)
forall a. Int -> IO (Ptr a)
mallocBytes Int
bytesTotal' 
        Int
bytesRead       <- Handle -> Ptr a -> Int -> IO Int
forall a. Handle -> Ptr a -> Int -> IO Int
hGetBuf Handle
h Ptr a
buf Int
bytesTotal'
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
bytesTotal' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
bytesRead)
         (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath -> IO ()
forall a. HasCallStack => FilePath -> a
error FilePath
"Data.Array.Repa.IO.Binary.readArrayFromStorableFile: read failed"

        Handle -> IO ()
hClose Handle
h
        ForeignPtr a
fptr     <- FinalizerPtr a -> Ptr a -> IO (ForeignPtr a)
forall a. FinalizerPtr a -> Ptr a -> IO (ForeignPtr a)
newForeignPtr FinalizerPtr a
forall a. FinalizerPtr a
finalizerFree Ptr a
buf        

        -- Converting the foreign ptr like this means that the array
        -- elements are used directly from the buffer, and not copied.
        let arr :: Array F sh a
arr  =  sh -> ForeignPtr a -> Array F sh a
forall sh e. Shape sh => sh -> ForeignPtr e -> Array F sh e
R.fromForeignPtr sh
sh ForeignPtr a
fptr

        Array F sh a -> IO (Array F sh a)
forall (m :: * -> *) a. Monad m => a -> m a
return   (Array F sh a -> IO (Array F sh a))
-> Array F sh a -> IO (Array F sh a)
forall a b. (a -> b) -> a -> b
$  Array F sh a
arr


-- | Write an array to a file.
--   Data appears in host byte order.
writeArrayToStorableFile
        :: forall sh a r
        .  (Shape sh, Source r a, Storable a)
        => FilePath 
        -> Array r sh a
        -> IO ()

writeArrayToStorableFile :: FilePath -> Array r sh a -> IO ()
writeArrayToStorableFile FilePath
filePath Array r sh a
arr
 = do   let bytes1 :: Int
bytes1      = a -> Int
forall a. Storable a => a -> Int
sizeOf (Array r sh a
arr Array r sh a -> sh -> a
forall sh r e. (Shape sh, Source r e) => Array r sh e -> sh -> e
R.! sh
forall sh. Shape sh => sh
R.zeroDim)
        let bytesTotal :: Int
bytesTotal  = Int
bytes1 Int -> Int -> Int
forall a. Num a => a -> a -> a
* (sh -> Int
forall sh. Shape sh => sh -> Int
R.size (sh -> Int) -> sh -> Int
forall a b. (a -> b) -> a -> b
$ Array r sh a -> sh
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
R.extent Array r sh a
arr)
        
        Ptr a
buf :: Ptr a    <- Int -> IO (Ptr a)
forall a. Int -> IO (Ptr a)
mallocBytes Int
bytesTotal
        ForeignPtr a
fptr            <- FinalizerPtr a -> Ptr a -> IO (ForeignPtr a)
forall a. FinalizerPtr a -> Ptr a -> IO (ForeignPtr a)
newForeignPtr FinalizerPtr a
forall a. FinalizerPtr a
finalizerFree Ptr a
buf        
        ForeignPtr a -> Array D sh a -> IO ()
forall r1 sh e.
(Load r1 sh e, Storable e) =>
ForeignPtr e -> Array r1 sh e -> IO ()
R.computeIntoP ForeignPtr a
fptr (Array r sh a -> Array D sh a
forall sh r e.
(Shape sh, Source r e) =>
Array r sh e -> Array D sh e
delay Array r sh a
arr)
        
        Handle
h <- FilePath -> IOMode -> IO Handle
openBinaryFile FilePath
filePath IOMode
WriteMode
        Handle -> Ptr a -> Int -> IO ()
forall a. Handle -> Ptr a -> Int -> IO ()
hPutBuf Handle
h Ptr a
buf Int
bytesTotal 
        Handle -> IO ()
hClose Handle
h