module Foreign.Storable.Record (
   Dictionary, Access,
   element, run,

   alignment, sizeOf,
   peek, poke,
   ) where

import Control.Monad.Trans.Reader
          (ReaderT(ReaderT), runReaderT,
           Reader, reader, runReader, )
import Control.Monad.Trans.Writer
          (Writer, writer, runWriter, )
import Control.Monad.Trans.State
          (State, modify, state, runState, )
import Control.Applicative (Applicative(..), liftA2, )
import Data.Monoid (Monoid(mempty, mappend), )

import Foreign.Storable.FixedArray (roundUp, )
import qualified Foreign.Storable as St

import Foreign.Ptr (Ptr, )
import Foreign.Storable (Storable, )


data Dictionary r =
   Dictionary {
      sizeOf_ :: Int,
      alignment_ :: Alignment,
      ptrBox :: Reader (Ptr r) (Box r r)
   }

newtype Access r a =
   Access
      (Compose (Writer Alignment)
        (Compose (State Int)
          (Compose (Reader (Ptr r))
              (Box r)))
        a)

instance Functor (Access r) where
   {-# INLINE fmap #-}
   fmap f (Access m) = Access (fmap f m)

instance Applicative (Access r) where
   {-# INLINE pure #-}
   {-# INLINE (<*>) #-}
   pure a = Access (pure a)
   Access f <*> Access x = Access (f <*> x)


{- |
See (.:) in TypeCompose library.
However I find this library too heavy weight
with respect to type extensions in order to depend on it.
-}
newtype Compose f g a =
   Compose (f (g a))

instance (Functor f, Functor g) => Functor (Compose f g) where
   {-# INLINE fmap #-}
   fmap f (Compose x) = Compose (fmap (fmap f) x)

instance (Applicative f, Applicative g) => Applicative (Compose f g) where
   {-# INLINE pure #-}
   {-# INLINE (<*>) #-}
   pure x = Compose (pure (pure x))
   Compose f <*> Compose x =
      Compose (liftA2 (<*>) f x)


data Box r a =
   Box {
      peek_ :: IO a,
      poke_ :: ReaderT r IO ()
   }

instance Functor (Box r) where
   {-# INLINE fmap #-}
   fmap f (Box pe po) =
      Box (fmap f pe) po

instance Applicative (Box r) where
   {-# INLINE pure #-}
   {-# INLINE (<*>) #-}
   pure a = Box (pure a) (pure ())
   f <*> x = Box (peek_ f <*> peek_ x) (poke_ f >> poke_ x)


newtype Alignment = Alignment Int

instance Monoid Alignment where
   {-# INLINE mempty #-}
   {-# INLINE mappend #-}
   mempty = Alignment 1
   mappend (Alignment x) (Alignment y) = Alignment (lcm x y)


{-# INLINE element #-}
element :: Storable a => (r -> a) -> Access r a
element f =
   let align = St.alignment (f undefined)
       size  = St.sizeOf (f undefined)
   in  Access $
       Compose $ writer $ flip (,) (Alignment align) $
       Compose $
       modify (roundUp align) >>
       state (\offset ->
          (Compose $ reader $ \ptr ->
           Box
             (St.peekByteOff ptr offset)
             (ReaderT $ St.pokeByteOff ptr offset . f),
           offset+size))

{-# INLINE run #-}
run :: Access r r -> Dictionary r
run (Access (Compose m)) =
   let (Compose s, align) = runWriter m
       (Compose r, size) = runState s 0
   in  Dictionary size align r


{-# INLINE alignment #-}
alignment :: Dictionary r -> r -> Int
alignment dict _ =
   let (Alignment align) = alignment_ dict
   in  align

{-# INLINE sizeOf #-}
sizeOf :: Dictionary r -> r -> Int
sizeOf dict _ =
   sizeOf_ dict

{-# INLINE peek #-}
peek :: Dictionary r -> Ptr r -> IO r
peek dict ptr =
   peek_ $ runReader (ptrBox dict) ptr

{-# INLINE poke #-}
poke :: Dictionary r -> Ptr r -> r -> IO ()
poke dict ptr =
   runReaderT (poke_ $ runReader (ptrBox dict) ptr)