-- | A properly aligned ForeignPtr type.
-- This can be currently achieved only by wasting some bytes.

module Foreign.ForeignPtr.Aligned 
  ( AlignedForeignPtr 
  , Alignment, align, fromAlignment
  , mallocAlignedForeignPtr
  , mallocAlignedForeignPtrArray
  , mallocAlignedForeignPtrBytes
  , withAlignedForeignPtr
  , finalizeAlignedForeignPtr
  , touchAlignedForeignPtr
  , castAlignedForeignPtr
  ) 
  where

--------------------------------------------------------------------------------

import Data.Bits
import Foreign.Ptr
import Foreign.ForeignPtr
import Foreign.Storable
import System.IO.Unsafe

--------------------------------------------------------------------------------

-- | Should be a power of two.
newtype Alignment = Alignment { fromAlignment :: Int } deriving Show

isPowerOfTwo :: Int -> Bool
isPowerOfTwo k 
  | k< 1 = False
  | k==1 = True
  | True = if odd k 
      then False
      else isPowerOfTwo (k `div` 2)

-- | A \"smart\" constructor which checks whether the input 
-- is a power of two.
align :: Int -> Alignment
align k = if isPowerOfTwo k 
  then Alignment k
  else error "invalid alignment"

-- | The aligned ForeignPtr type
data AlignedForeignPtr a = AFPtr
  { _fptr   :: {-# UNPACK #-} !(ForeignPtr a)
  , _offset :: {-# UNPACK #-} !Int
  } 
 
instance Show (AlignedForeignPtr a) where
  show afptr = unsafePerformIO $ withAlignedForeignPtr afptr $ \p -> return (show p)
  
mallocAlignedForeignPtr :: Storable a => Alignment -> IO (AlignedForeignPtr a)
mallocAlignedForeignPtr align = mallocAlignedForeignPtrArray align 1
  
mallocAlignedForeignPtrArray :: Storable a => Alignment -> Int -> IO (AlignedForeignPtr a)
mallocAlignedForeignPtrArray align n = worker where
  worker = mallocAlignedForeignPtrBytes align (n * sizeOf undef)
  undef = ioAFPtrUndefined worker 

-- | This is here only to avoid the @ScopedTypeVariables@ language extension
-- for more portability.
ioAFPtrUndefined :: IO (AlignedForeignPtr a) -> a
ioAFPtrUndefined _ = undefined

mallocAlignedForeignPtrBytes :: Alignment -> Int -> IO (AlignedForeignPtr a)
mallocAlignedForeignPtrBytes (Alignment align) n = do
  fptr <- mallocForeignPtrBytes (n+align)    -- (n+align-1) would be enough, to be precise. But usually that's an odd number.
  ofs <- withForeignPtr fptr $ \p -> do
    let j = ptrToIntPtr p
        a = fromIntegral align :: IntPtr
        k = (j+a-1) .&. (complement (a-1))
    return $ fromIntegral (k-j)
  return (AFPtr fptr ofs)  

withAlignedForeignPtr :: AlignedForeignPtr a -> (Ptr a -> IO b) -> IO b
withAlignedForeignPtr (AFPtr fptr ofs) action = 
  withForeignPtr fptr $ \p -> action (p `plusPtr` ofs)

finalizeAlignedForeignPtr :: AlignedForeignPtr a -> IO ()
finalizeAlignedForeignPtr (AFPtr fptr ofs) = finalizeForeignPtr fptr

touchAlignedForeignPtr :: AlignedForeignPtr a -> IO ()
touchAlignedForeignPtr (AFPtr fptr ofs) = touchForeignPtr fptr

castAlignedForeignPtr :: AlignedForeignPtr a -> AlignedForeignPtr b
castAlignedForeignPtr (AFPtr fptr ofs) = AFPtr (castForeignPtr fptr) ofs

--------------------------------------------------------------------------------