{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE InstanceSigs      #-}

module HaskellWorks.Data.FromByteString
  ( FromByteString(..)
  ) where

import Data.Bits
import Data.Word
import Foreign.ForeignPtr

import qualified Data.ByteString              as BS
import qualified Data.ByteString.Internal     as BS
import qualified Data.Vector.Storable         as DVS
import qualified Data.Vector.Storable.Mutable as DVSM

-- | Class for byte-string-like datastructures
class FromByteString a where
  -- | Convert a byte string to a value of type @a
  fromByteString :: BS.ByteString -> a

instance FromByteString (DVS.Vector Word8) where
  fromByteString :: BS.ByteString -> DVS.Vector Word8
  fromByteString :: ByteString -> Vector Word8
fromByteString ByteString
bs = Int
-> (ByteString -> Maybe (Word8, ByteString))
-> ByteString
-> Vector Word8
forall a b.
Storable a =>
Int -> (b -> Maybe (a, b)) -> b -> Vector a
DVS.unfoldrN (ByteString -> Int
BS.length ByteString
bs) ByteString -> Maybe (Word8, ByteString)
gen ByteString
bs
    where gen :: BS.ByteString -> Maybe (Word8, BS.ByteString)
          gen :: ByteString -> Maybe (Word8, ByteString)
gen ByteString
cs = case ByteString -> Maybe (Word8, ByteString)
BS.uncons ByteString
cs of
            Just (Word8
d, ByteString
ds) -> (Word8, ByteString) -> Maybe (Word8, ByteString)
forall a. a -> Maybe a
Just (Word8
d, ByteString
ds)
            Maybe (Word8, ByteString)
Nothing      -> Maybe (Word8, ByteString)
forall a. Maybe a
Nothing
  {-# INLINE fromByteString #-}

instance FromByteString (DVS.Vector Word16) where
  fromByteString :: BS.ByteString -> DVS.Vector Word16
  fromByteString :: ByteString -> Vector Word16
fromByteString ByteString
bs = Int
-> (ByteString -> Maybe (Word16, ByteString))
-> ByteString
-> Vector Word16
forall a b.
Storable a =>
Int -> (b -> Maybe (a, b)) -> b -> Vector a
DVS.unfoldrN (ByteString -> Int
BS.length ByteString
bs Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2) ByteString -> Maybe (Word16, ByteString)
gen ByteString
bs
    where gen :: BS.ByteString -> Maybe (Word16, BS.ByteString)
          gen :: ByteString -> Maybe (Word16, ByteString)
gen ByteString
cs = case ByteString -> Maybe (Word8, ByteString)
BS.uncons ByteString
cs of
            Just (Word8
d, ByteString
ds) -> Int -> Word16 -> ByteString -> Maybe (Word16, ByteString)
gen' Int
8 (Word8 -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
d) ByteString
ds
            Maybe (Word8, ByteString)
Nothing      -> Maybe (Word16, ByteString)
forall a. Maybe a
Nothing
          gen' :: Int -> Word16 -> BS.ByteString -> Maybe (Word16, BS.ByteString)
          gen' :: Int -> Word16 -> ByteString -> Maybe (Word16, ByteString)
gen' Int
n Word16
w ByteString
cs
            | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
16   = (Word16, ByteString) -> Maybe (Word16, ByteString)
forall a. a -> Maybe a
Just (Word16
w, ByteString
cs)
            | Bool
otherwise = case ByteString -> Maybe (Word8, ByteString)
BS.uncons ByteString
cs of
                Just (Word8
d, ByteString
ds) -> Int -> Word16 -> ByteString -> Maybe (Word16, ByteString)
gen' (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
8) (Word16
w Word16 -> Word16 -> Word16
forall a. Bits a => a -> a -> a
.|. (Word8 -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
d Word16 -> Int -> Word16
forall a. Bits a => a -> Int -> a
`shiftL` Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n)) ByteString
ds
                Maybe (Word8, ByteString)
Nothing      -> (Word16, ByteString) -> Maybe (Word16, ByteString)
forall a. a -> Maybe a
Just (Word16
w, ByteString
cs)
  {-# INLINE fromByteString #-}

instance FromByteString (DVS.Vector Word32) where
  fromByteString :: BS.ByteString -> DVS.Vector Word32
  fromByteString :: ByteString -> Vector Word32
fromByteString ByteString
bs = Int
-> (ByteString -> Maybe (Word32, ByteString))
-> ByteString
-> Vector Word32
forall a b.
Storable a =>
Int -> (b -> Maybe (a, b)) -> b -> Vector a
DVS.unfoldrN (ByteString -> Int
BS.length ByteString
bs Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
4) ByteString -> Maybe (Word32, ByteString)
gen ByteString
bs
    where gen :: BS.ByteString -> Maybe (Word32, BS.ByteString)
          gen :: ByteString -> Maybe (Word32, ByteString)
gen ByteString
cs = case ByteString -> Maybe (Word8, ByteString)
BS.uncons ByteString
cs of
            Just (Word8
d, ByteString
ds) -> Int -> Word32 -> ByteString -> Maybe (Word32, ByteString)
gen' Int
8 (Word8 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
d) ByteString
ds
            Maybe (Word8, ByteString)
Nothing      -> Maybe (Word32, ByteString)
forall a. Maybe a
Nothing
          gen' :: Int -> Word32 -> BS.ByteString -> Maybe (Word32, BS.ByteString)
          gen' :: Int -> Word32 -> ByteString -> Maybe (Word32, ByteString)
gen' Int
n Word32
w ByteString
cs
            | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
32   = (Word32, ByteString) -> Maybe (Word32, ByteString)
forall a. a -> Maybe a
Just (Word32
w, ByteString
cs)
            | Bool
otherwise = case ByteString -> Maybe (Word8, ByteString)
BS.uncons ByteString
cs of
                Just (Word8
d, ByteString
ds) -> Int -> Word32 -> ByteString -> Maybe (Word32, ByteString)
gen' (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
8) (Word32
w Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.|. (Word8 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
d Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`shiftL` Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n)) ByteString
ds
                Maybe (Word8, ByteString)
Nothing      -> (Word32, ByteString) -> Maybe (Word32, ByteString)
forall a. a -> Maybe a
Just (Word32
w, ByteString
cs)
  {-# INLINE fromByteString #-}

instance FromByteString (DVS.Vector Word64) where
  fromByteString :: BS.ByteString -> DVS.Vector Word64
  fromByteString :: ByteString -> Vector Word64
fromByteString ByteString
bs = Int
-> (ByteString -> Maybe (Word64, ByteString))
-> ByteString
-> Vector Word64
forall a b.
Storable a =>
Int -> (b -> Maybe (a, b)) -> b -> Vector a
DVS.unfoldrN (ByteString -> Int
BS.length ByteString
bs Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
8 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
8) ByteString -> Maybe (Word64, ByteString)
gen ByteString
bs
    where gen :: BS.ByteString -> Maybe (Word64, BS.ByteString)
          gen :: ByteString -> Maybe (Word64, ByteString)
gen ByteString
cs = case ByteString -> Maybe (Word8, ByteString)
BS.uncons ByteString
cs of
            Just (Word8
d, ByteString
ds) -> Int -> Word64 -> ByteString -> Maybe (Word64, ByteString)
gen' Int
8 (Word8 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
d) ByteString
ds
            Maybe (Word8, ByteString)
Nothing      -> Maybe (Word64, ByteString)
forall a. Maybe a
Nothing
          gen' :: Int -> Word64 -> BS.ByteString -> Maybe (Word64, BS.ByteString)
          gen' :: Int -> Word64 -> ByteString -> Maybe (Word64, ByteString)
gen' Int
n Word64
w ByteString
cs
            | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
64   = (Word64, ByteString) -> Maybe (Word64, ByteString)
forall a. a -> Maybe a
Just (Word64
w, ByteString
cs)
            | Bool
otherwise = case ByteString -> Maybe (Word8, ByteString)
BS.uncons ByteString
cs of
                Just (Word8
d, ByteString
ds) -> Int -> Word64 -> ByteString -> Maybe (Word64, ByteString)
gen' (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
8) (Word64
w Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. (Word8 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
d Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftL` Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n)) ByteString
ds
                Maybe (Word8, ByteString)
Nothing      -> (Word64, ByteString) -> Maybe (Word64, ByteString)
forall a. a -> Maybe a
Just (Word64
w, ByteString
cs)
  {-# INLINE fromByteString #-}

instance FromByteString (DVSM.MVector s Word8) where
  fromByteString :: ByteString -> MVector s Word8
fromByteString ByteString
bs = case ByteString -> (ForeignPtr Word8, Int, Int)
BS.toForeignPtr ByteString
bs of
    (ForeignPtr Word8
fptr, Int
off, Int
len) -> ForeignPtr Word8 -> Int -> Int -> MVector s Word8
forall a s. Storable a => ForeignPtr a -> Int -> Int -> MVector s a
DVSM.unsafeFromForeignPtr (ForeignPtr Word8 -> ForeignPtr Word8
forall a b. ForeignPtr a -> ForeignPtr b
castForeignPtr ForeignPtr Word8
fptr) Int
off Int
len
  {-# INLINE fromByteString #-}