{-# LANGUAGE GADTs #-}
module Data.IDX.Internal where
import Control.Monad (replicateM)
import Data.Binary
import Data.Int
import GHC.Float ( castFloatToWord32
, castDoubleToWord64
, castWord64ToDouble
, castWord32ToFloat
)
import qualified Data.Vector.Unboxed as V
import Data.Vector.Unboxed ((!))
import Data.Functor ((<$>))
data IDXContentType where
IDXUnsignedByte :: IDXContentType
IDXSignedByte :: IDXContentType
IDXShort :: IDXContentType
IDXInt :: IDXContentType
IDXFloat :: IDXContentType
IDXDouble :: IDXContentType
deriving (Int -> IDXContentType -> ShowS
[IDXContentType] -> ShowS
IDXContentType -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [IDXContentType] -> ShowS
$cshowList :: [IDXContentType] -> ShowS
show :: IDXContentType -> String
$cshow :: IDXContentType -> String
showsPrec :: Int -> IDXContentType -> ShowS
$cshowsPrec :: Int -> IDXContentType -> ShowS
Show, IDXContentType -> IDXContentType -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: IDXContentType -> IDXContentType -> Bool
$c/= :: IDXContentType -> IDXContentType -> Bool
== :: IDXContentType -> IDXContentType -> Bool
$c== :: IDXContentType -> IDXContentType -> Bool
Eq)
instance Binary IDXContentType where
get :: Get IDXContentType
get = do
Word8
w <- Get Word8
getWord8
case Word8
w of
Word8
0x08 -> forall (m :: * -> *) a. Monad m => a -> m a
return IDXContentType
IDXUnsignedByte
Word8
0x09 -> forall (m :: * -> *) a. Monad m => a -> m a
return IDXContentType
IDXSignedByte
Word8
0x0B -> forall (m :: * -> *) a. Monad m => a -> m a
return IDXContentType
IDXShort
Word8
0x0C -> forall (m :: * -> *) a. Monad m => a -> m a
return IDXContentType
IDXInt
Word8
0x0D -> forall (m :: * -> *) a. Monad m => a -> m a
return IDXContentType
IDXFloat
Word8
0x0E -> forall (m :: * -> *) a. Monad m => a -> m a
return IDXContentType
IDXDouble
Word8
_ -> forall (m :: * -> *) a. MonadFail m => String -> m a
fail forall a b. (a -> b) -> a -> b
$ String
"Unrecognized IDX content type: " forall a. [a] -> [a] -> [a]
++ (forall a. Show a => a -> String
show Word8
w)
put :: IDXContentType -> Put
put IDXContentType
IDXUnsignedByte = Word8 -> Put
putWord8 Word8
0x08
put IDXContentType
IDXSignedByte = Word8 -> Put
putWord8 Word8
0x09
put IDXContentType
IDXShort = Word8 -> Put
putWord8 Word8
0x0B
put IDXContentType
IDXInt = Word8 -> Put
putWord8 Word8
0x0C
put IDXContentType
IDXFloat = Word8 -> Put
putWord8 Word8
0x0D
put IDXContentType
IDXDouble = Word8 -> Put
putWord8 Word8
0x0E
data IDXData = IDXInts IDXContentType (V.Vector Int) (V.Vector Int )
| IDXDoubles IDXContentType (V.Vector Int) (V.Vector Double)
deriving (Int -> IDXData -> ShowS
[IDXData] -> ShowS
IDXData -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [IDXData] -> ShowS
$cshowList :: [IDXData] -> ShowS
show :: IDXData -> String
$cshow :: IDXData -> String
showsPrec :: Int -> IDXData -> ShowS
$cshowsPrec :: Int -> IDXData -> ShowS
Show, IDXData -> IDXData -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: IDXData -> IDXData -> Bool
$c/= :: IDXData -> IDXData -> Bool
== :: IDXData -> IDXData -> Bool
$c== :: IDXData -> IDXData -> Bool
Eq)
instance Binary IDXData where
get :: Get IDXData
get = do
Get Word8
getWord8
Get Word8
getWord8
IDXContentType
idxType <- forall t. Binary t => Get t
get :: Get IDXContentType
Int
nDimensions <- forall a b. (Integral a, Num b) => a -> b
fromIntegral forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Word8
getWord8
[Int]
dimensionSizes <- forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
nDimensions (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Int32
getInt32)
let nEntries :: Int
nEntries = forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
dimensionSizes
dimV :: Vector Int
dimV = forall a. Unbox a => [a] -> Vector a
V.fromList [Int]
dimensionSizes
case IDXContentType
idxType of
t :: IDXContentType
t@IDXContentType
IDXUnsignedByte -> forall a.
Integral a =>
Int -> IDXContentType -> Vector Int -> Get a -> Get IDXData
buildIntResult Int
nEntries IDXContentType
t Vector Int
dimV Get Word8
getWord8
t :: IDXContentType
t@IDXContentType
IDXSignedByte -> forall a.
Integral a =>
Int -> IDXContentType -> Vector Int -> Get a -> Get IDXData
buildIntResult Int
nEntries IDXContentType
t Vector Int
dimV Get Int8
getInt8
t :: IDXContentType
t@IDXContentType
IDXShort -> forall a.
Integral a =>
Int -> IDXContentType -> Vector Int -> Get a -> Get IDXData
buildIntResult Int
nEntries IDXContentType
t Vector Int
dimV Get Int16
getInt16
t :: IDXContentType
t@IDXContentType
IDXInt -> forall a.
Integral a =>
Int -> IDXContentType -> Vector Int -> Get a -> Get IDXData
buildIntResult Int
nEntries IDXContentType
t Vector Int
dimV Get Int32
getInt32
t :: IDXContentType
t@IDXContentType
IDXFloat -> forall a.
Real a =>
Int -> IDXContentType -> Vector Int -> Get a -> Get IDXData
buildDoubleResult Int
nEntries IDXContentType
t Vector Int
dimV Get Float
getFloat
t :: IDXContentType
t@IDXContentType
IDXDouble -> forall a.
Real a =>
Int -> IDXContentType -> Vector Int -> Get a -> Get IDXData
buildDoubleResult Int
nEntries IDXContentType
t Vector Int
dimV Get Double
getDouble
put :: IDXData -> Put
put IDXData
d = do
Word8 -> Put
putWord8 Word8
0
Word8 -> Put
putWord8 Word8
0
forall t. Binary t => t -> Put
put forall a b. (a -> b) -> a -> b
$ IDXData -> IDXContentType
idxType IDXData
d
let dimensions :: Vector Int
dimensions = IDXData -> Vector Int
idxDimensions IDXData
d
forall t. Binary t => t -> Put
put forall a b. (a -> b) -> a -> b
$ (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ forall a. Unbox a => Vector a -> Int
V.length Vector Int
dimensions :: Word8)
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
V.forM_ Vector Int
dimensions forall a b. (a -> b) -> a -> b
$ (\Int
x -> forall t. Binary t => t -> Put
put forall a b. (a -> b) -> a -> b
$! (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
x :: Int32))
case IDXData
d of
IDXDoubles IDXContentType
t Vector Int
_ Vector Double
content -> forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
V.forM_ Vector Double
content forall a b. (a -> b) -> a -> b
$ IDXContentType -> Double -> Put
putReal IDXContentType
t
IDXInts IDXContentType
t Vector Int
_ Vector Int
content -> forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
V.forM_ Vector Int
content forall a b. (a -> b) -> a -> b
$ IDXContentType -> Int -> Put
putIntegral IDXContentType
t
newtype IDXLabels = IDXLabels (V.Vector Int)
instance Binary IDXLabels where
get :: Get IDXLabels
get = do
Get Int32
getInt32
Int
nItems <- forall a b. (Integral a, Num b) => a -> b
fromIntegral forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Int32
getInt32
let readEntries :: Int -> Get (Vector a)
readEntries Int
n = forall (m :: * -> *) a.
(Monad m, Unbox a) =>
Int -> m a -> m (Vector a)
V.replicateM Int
n forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Word8
getWord8 forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$!)
Vector Int
v <- forall a.
Unbox a =>
(Int -> Get (Vector a)) -> Int -> Int -> Get (Vector a)
readContent forall {a}. (Unbox a, Num a) => Int -> Get (Vector a)
readEntries Int
500 Int
nItems
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Vector Int -> IDXLabels
IDXLabels Vector Int
v
put :: IDXLabels -> Put
put (IDXLabels Vector Int
v) = do
forall t. Binary t => t -> Put
put (Int32
0 :: Int32)
let len :: Int
len = forall a. Unbox a => Vector a -> Int
V.length Vector Int
v
forall t. Binary t => t -> Put
put (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len :: Int32)
forall (m :: * -> *) a b.
(Monad m, Unbox a) =>
Vector a -> (a -> m b) -> m ()
V.forM_ Vector Int
v (\Int
x -> forall t. Binary t => t -> Put
put forall a b. (a -> b) -> a -> b
$! (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
x :: Word8))
idxType :: IDXData -> IDXContentType
idxType :: IDXData -> IDXContentType
idxType (IDXInts IDXContentType
t Vector Int
_ Vector Int
_) = IDXContentType
t
idxType (IDXDoubles IDXContentType
t Vector Int
_ Vector Double
_) = IDXContentType
t
idxDimensions :: IDXData -> V.Vector Int
idxDimensions :: IDXData -> Vector Int
idxDimensions (IDXInts IDXContentType
_ Vector Int
ds Vector Int
_) = Vector Int
ds
idxDimensions (IDXDoubles IDXContentType
_ Vector Int
ds Vector Double
_) = Vector Int
ds
isIDXIntegral :: IDXData -> Bool
isIDXIntegral :: IDXData -> Bool
isIDXIntegral (IDXInts IDXContentType
_ Vector Int
_ Vector Int
_) = Bool
True
isIDXIntegral (IDXData
_ ) = Bool
False
isIDXReal :: IDXData -> Bool
isIDXReal :: IDXData -> Bool
isIDXReal (IDXDoubles IDXContentType
_ Vector Int
_ Vector Double
_) = Bool
True
isIDXReal (IDXData
_ ) = Bool
False
idxIntContent :: IDXData -> V.Vector Int
idxIntContent :: IDXData -> Vector Int
idxIntContent (IDXInts IDXContentType
_ Vector Int
_ Vector Int
v) = Vector Int
v
idxIntContent (IDXDoubles IDXContentType
_ Vector Int
_ Vector Double
v) =
forall a. Unbox a => [a] -> Vector a
V.fromList forall a b. (a -> b) -> a -> b
$ [forall a b. (RealFrac a, Integral b) => a -> b
round forall a b. (a -> b) -> a -> b
$ (Vector Double
v forall a. Unbox a => Vector a -> Int -> a
! Int
i) | Int
i <- [Int
0.. ((forall a. Unbox a => Vector a -> Int
V.length Vector Double
v)forall a. Num a => a -> a -> a
-Int
1)]]
idxDoubleContent :: IDXData -> V.Vector Double
idxDoubleContent :: IDXData -> Vector Double
idxDoubleContent (IDXDoubles IDXContentType
_ Vector Int
_ Vector Double
v) = Vector Double
v
idxDoubleContent (IDXInts IDXContentType
_ Vector Int
_ Vector Int
v) =
forall a. Unbox a => [a] -> Vector a
V.fromList forall a b. (a -> b) -> a -> b
$ [forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ (Vector Int
v forall a. Unbox a => Vector a -> Int -> a
! Int
i) | Int
i <- [Int
0.. ((forall a. Unbox a => Vector a -> Int
V.length Vector Int
v) forall a. Num a => a -> a -> a
- Int
1)]]
readContent :: (V.Unbox a)
=> (Int -> Get (V.Vector a))
-> Int
-> Int
-> Get (V.Vector a)
readContent :: forall a.
Unbox a =>
(Int -> Get (Vector a)) -> Int -> Int -> Get (Vector a)
readContent Int -> Get (Vector a)
readEntries Int
chunkSize Int
n =
if Int
n forall a. Ord a => a -> a -> Bool
> Int
chunkSize
then do
Vector a
headChunk <- Int -> Get (Vector a)
readEntries (Int
n forall a. Integral a => a -> a -> a
`mod` Int
chunkSize)
let nChunks :: Int
nChunks = Int
n forall a. Integral a => a -> a -> a
`div` Int
chunkSize
[Vector a]
chunkList <- forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
nChunks (forall a.
Unbox a =>
(Int -> Get (Vector a)) -> Int -> Int -> Get (Vector a)
readContent Int -> Get (Vector a)
readEntries Int
chunkSize Int
chunkSize)
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! forall a. Unbox a => [Vector a] -> Vector a
V.concat forall a b. (a -> b) -> a -> b
$ Vector a
headChunkforall a. a -> [a] -> [a]
:[Vector a]
chunkList
else do
Vector a
rest <- Int -> Get (Vector a)
readEntries Int
n
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! Vector a
rest
getInt8 :: Get Int8
getInt8 :: Get Int8
getInt8 = forall t. Binary t => Get t
get
getInt16 :: Get Int16
getInt16 :: Get Int16
getInt16 = forall t. Binary t => Get t
get
getInt32 :: Get Int32
getInt32 :: Get Int32
getInt32 = forall t. Binary t => Get t
get
getFloat :: Get Float
getFloat :: Get Float
getFloat = Word32 -> Float
castWord32ToFloat forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall t. Binary t => Get t
get
getDouble :: Get Double
getDouble :: Get Double
getDouble = Word64 -> Double
castWord64ToDouble forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall t. Binary t => Get t
get
buildIntResult :: Integral a
=> Int
-> IDXContentType
-> V.Vector Int
-> Get a
-> Get IDXData
buildIntResult :: forall a.
Integral a =>
Int -> IDXContentType -> Vector Int -> Get a -> Get IDXData
buildIntResult Int
nEntries IDXContentType
typ Vector Int
dimV Get a
getContent = do
Vector Int
content <- forall a.
Unbox a =>
(Int -> Get (Vector a)) -> Int -> Int -> Get (Vector a)
readContent Int -> Get (Vector Int)
readEntries Int
500 Int
nEntries
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ IDXContentType -> Vector Int -> Vector Int -> IDXData
IDXInts IDXContentType
typ Vector Int
dimV Vector Int
content
where
readEntries :: Int -> Get (Vector Int)
readEntries Int
n = forall (m :: * -> *) a.
(Monad m, Unbox a) =>
Int -> m a -> m (Vector a)
V.replicateM Int
n forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get a
getContent forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$!)
buildDoubleResult :: Real a
=> Int
-> IDXContentType
-> V.Vector Int
-> Get a
-> Get IDXData
buildDoubleResult :: forall a.
Real a =>
Int -> IDXContentType -> Vector Int -> Get a -> Get IDXData
buildDoubleResult Int
nEntries IDXContentType
typ Vector Int
dimV Get a
getContent = do
Vector Double
content <- forall a.
Unbox a =>
(Int -> Get (Vector a)) -> Int -> Int -> Get (Vector a)
readContent Int -> Get (Vector Double)
readEntries Int
500 Int
nEntries
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ IDXContentType -> Vector Int -> Vector Double -> IDXData
IDXDoubles IDXContentType
typ Vector Int
dimV Vector Double
content
where
readEntries :: Int -> Get (Vector Double)
readEntries Int
n = forall (m :: * -> *) a.
(Monad m, Unbox a) =>
Int -> m a -> m (Vector a)
V.replicateM Int
n forall a b. (a -> b) -> a -> b
$ forall a b. (Real a, Fractional b) => a -> b
realToFrac forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get a
getContent forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$!)
putIntegral :: IDXContentType -> Int -> Put
putIntegral :: IDXContentType -> Int -> Put
putIntegral IDXContentType
IDXUnsignedByte Int
n = forall t. Binary t => t -> Put
put forall a b. (a -> b) -> a -> b
$! (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n :: Word8)
putIntegral IDXContentType
IDXSignedByte Int
n = forall t. Binary t => t -> Put
put forall a b. (a -> b) -> a -> b
$! (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n :: Int8 )
putIntegral IDXContentType
IDXShort Int
n = forall t. Binary t => t -> Put
put forall a b. (a -> b) -> a -> b
$! (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n :: Int16)
putIntegral IDXContentType
IDXInt Int
n = forall t. Binary t => t -> Put
put forall a b. (a -> b) -> a -> b
$! (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n :: Int32)
putIntegral IDXContentType
t Int
_ = forall a. HasCallStack => String -> a
error forall a b. (a -> b) -> a -> b
$ String
"IDX.putIntegral " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show IDXContentType
t
putReal :: IDXContentType -> Double -> Put
putReal :: IDXContentType -> Double -> Put
putReal IDXContentType
IDXDouble Double
n = forall t. Binary t => t -> Put
put forall a b. (a -> b) -> a -> b
$ Double -> Word64
castDoubleToWord64 Double
n
putReal IDXContentType
IDXFloat Double
n = forall t. Binary t => t -> Put
put forall a b. (a -> b) -> a -> b
$! Float -> Word32
castFloatToWord32 (forall a b. (Real a, Fractional b) => a -> b
realToFrac Double
n :: Float )
partitionedData :: V.Unbox a => (IDXData -> V.Vector a) -> IDXData -> [V.Vector a]
partitionedData :: forall a. Unbox a => (IDXData -> Vector a) -> IDXData -> [Vector a]
partitionedData IDXData -> Vector a
getContent IDXData
idxData = do
Int
i <- [Int
0 .. Int
dim0 forall a. Num a => a -> a -> a
- Int
1]
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ (forall a. Unbox a => Int -> Int -> Vector a -> Vector a
V.slice (Int
iforall a. Num a => a -> a -> a
*Int
entrySize) Int
entrySize Vector a
content)
where
dim0 :: Int
dim0 = forall a. Unbox a => Vector a -> a
V.head forall a b. (a -> b) -> a -> b
$ IDXData -> Vector Int
idxDimensions IDXData
idxData
content :: Vector a
content = IDXData -> Vector a
getContent IDXData
idxData
entrySize :: Int
entrySize = (forall a. (Unbox a, Num a) => Vector a -> a
V.product forall a b. (a -> b) -> a -> b
$ IDXData -> Vector Int
idxDimensions IDXData
idxData) forall a. Integral a => a -> a -> a
`div` Int
dim0