{-# LANGUAGE ScopedTypeVariables, StandaloneDeriving, TypeOperators, FlexibleContexts #-}
{-# LANGUAGE Trustworthy #-}
-- | Some instances making compound structures of Storables easier to work with. For instance, you can poke
--   two adjacent Storables into memory by poking a pair.
--
--   These work with an equivalent memory layout as the Swiz instances.
module Data.Storables (first, second, toPairs, toPairs4, toPairs5, toPairs6, toPairs7, toEither) where

import Foreign.Ptr
import Foreign.Storable
import Control.Monad
import Control.Compose
import Data.Word
import Data.ByteString (ByteString, unpack, pack)
import Data.Bits

align x y = ((sizeOf x - 1) `div` alignment y + 1) * alignment y

first :: Ptr (a, b) -> Ptr a
first = castPtr

second :: forall a b. (Storable a, Storable b) => Ptr (a, b) -> Ptr b
second = castPtr . (`plusPtr` align (undefined :: a) (undefined :: b))

toPairs :: Ptr (a, b, c) -> Ptr (a, (b, c))
toPairs = castPtr

toPairs4 :: Ptr (a, b, c, d) -> Ptr (a, (b, (c, d)))
toPairs4 = castPtr

toPairs5 :: Ptr (a, b, c, d, e) -> Ptr (a, (b, (c, (d, e))))
toPairs5 = castPtr

toPairs6 :: Ptr (a, b, c, d, e, f) -> Ptr (a, (b, (c, (d, (e, f)))))
toPairs6 = castPtr

toPairs7 :: Ptr (a, b, c, d, e, f, g) -> Ptr (a, (b, (c, (d, (e, (f, g))))))
toPairs7 = castPtr

toEither :: Ptr (Maybe t) -> Ptr (Either () t)
toEither = castPtr

instance (Storable a, Storable b) => Storable (a, b) where
	sizeOf _ = align (undefined :: a) (undefined :: b) + sizeOf (undefined :: b)
	alignment _ = alignment (undefined :: a)
		`lcm` alignment (undefined :: b)
	peek p = liftM2 (,) (peek (first p)) (peek (second p))
	poke p (x, y) = do
		poke (first p) x
		poke (second p) y

instance (Storable a, Storable b, Storable c) => Storable (a, b, c) where
	sizeOf _ = sizeOf (undefined :: (a, (b, c)))
	alignment _ = alignment (undefined :: (a, (b, c)))
	peek = liftM (\(x, (y, z)) -> (x, y, z)) . peek . toPairs
	poke p (x, y, z) = poke (toPairs p) (x, (y, z))

instance (Storable a, Storable b, Storable c, Storable d) => Storable (a, b, c, d) where
	sizeOf _ = sizeOf (undefined :: (a, (b, (c, d))))
	alignment _ = alignment (undefined :: (a, (b, (c, d))))
	peek = liftM (\(x, (y, (z, a))) -> (x, y, z, a)) . peek . toPairs4
	poke p (x, y, z, a) = poke (toPairs4 p) (x, (y, (z, a)))

instance (Storable a, Storable b, Storable c, Storable d, Storable e) => Storable (a, b, c, d, e) where
	sizeOf _ = sizeOf (undefined :: (a, (b, (c, (d, e)))))
	alignment _ = alignment (undefined :: (a, (b, (c, (d, e)))))
	peek = liftM (\(x, (y, (z, (a, b)))) -> (x, y, z, a, b)) . peek . toPairs5
	poke p (x, y, z, a, b) = poke (toPairs5 p) (x, (y, (z, (a, b))))

instance (Storable a, Storable b, Storable c, Storable d, Storable e, Storable f) => Storable (a, b, c, d, e, f) where
	sizeOf _ = sizeOf (undefined :: (a, (b, (c, (d, (e, f))))))
	alignment _ = alignment (undefined :: (a, (b, (c, (d, (e, f))))))
	peek = liftM (\(x, (y, (z, (a, (b, c))))) -> (x, y, z, a, b, c)) . peek . toPairs6
	poke p (x, y, z, a, b, c) = poke (toPairs6 p) (x, (y, (z, (a, (b, c)))))

instance (Storable a, Storable b, Storable c, Storable d, Storable e, Storable f, Storable g) => Storable (a, b, c, d, e, f, g) where
	sizeOf _ = sizeOf (undefined :: (a, (b, (c, (d, (e, (f, g)))))))
	alignment _ = alignment (undefined :: (a, (b, (c, (d, (e, (f, g)))))))
	peek = liftM (\(x, (y, (z, (a, (b, (c, d)))))) -> (x, y, z, a, b, c, d)) . peek . toPairs7
	poke p (x, y, z, a, b, c, d) = poke (toPairs7 p) (x, (y, (z, (a, (b, (c, d))))))

instance (Storable t, Storable u) => Storable (Either t u) where
	sizeOf ei = sizeOf (undefined :: t) `max` sizeOf (undefined :: u) + 1
	alignment _ = alignment (undefined :: t)
		`lcm` alignment (undefined :: u)
	peek p = do
		n :: Word8 <- peekByteOff p (sizeOf (undefined :: Either t u) - 1)
		if n == 0 then
				liftM Left (peekByteOff p 0)
			else
				liftM Right (peekByteOff p 0)
	poke p ei@(Left x) = do
		pokeByteOff p 0 x
		pokeByteOff p (sizeOf ei - 1) (0 :: Word8)
	poke p ei@(Right x) = do
		pokeByteOff p 0 x
		pokeByteOff p (sizeOf ei - 1) (1 :: Word8)

instance (Storable t) => Storable (Maybe t) where
	sizeOf _ = sizeOf (undefined :: t) + 1
	alignment _ = alignment (undefined :: t)
	peek = liftM (either (const Nothing) Just) . peek . toEither
	poke p = poke (toEither p) . maybe (Left ()) Right