{- |
Provide allocation functions with correct alignment.
GHC's functions are wrong until at least GHC-6.12.3.
-}
module Synthesizer.LLVM.Alloc (
   malloc, free,
   mallocArray, freeArray,
   F.alloca, F.with,
   ) where
{- reexport alloca and with, since they work correctly in GHC-6.12.3 -}

import qualified Foreign as F
import Foreign.Storable.FixedArray (sizeOfArray, roundUp, )
import Foreign.Storable (Storable, sizeOf, alignment, )
import Foreign.Ptr (Ptr, minusPtr, plusPtr, nullPtr, )

import qualified Synthesizer.LLVM.Debug.Storable as Debug


{-
diffClip :: Int -> Int -> Int
-- this implementation also works for unsigned types
diffClip x y = x - min x y
-- diffClip x y = max 0 (x-y)
-}

defaultMallocAlign :: Int
defaultMallocAlign = 8


seekPointer :: Int -> Ptr a -> Ptr b
seekPointer size ptr =
   alignPointer (alignment nullPtr) (plusPtr ptr size)

alignPointer :: Int -> Ptr a -> Ptr b
alignPointer align ptr =
   plusPtr nullPtr $
   roundUp align $
   minusPtr ptr nullPtr

{- |
> mallocAligned align size
-}
mallocAligned :: Int -> Int -> IO (Ptr a)
mallocAligned align size =
   if mod defaultMallocAlign align == 0
     then F.mallocBytes size
     else do
        let {- pessimistic but safe -}
            ptrOffset = size + alignment nullPtr
            {- This should be optimal and I think it is also correct,
               but better safe than sorry.
            common = gcd align (alignment nullPtr)
            ptrOffset = roundUp common size + alignment nullPtr - common
            -}
            padSize = align - gcd defaultMallocAlign align
        allocPtr <- F.mallocBytes (padSize + ptrOffset + sizeOf nullPtr)
        let alignedPtr = alignPointer align allocPtr
        F.poke (seekPointer size alignedPtr) allocPtr
        return alignedPtr

{- |
> mallocBytes align size
-}
mallocBytes :: Storable a => a -> Int -> IO (Ptr a)
mallocBytes a size =
   Debug.traceMalloc a size =<< mallocAligned (alignment a) size

peekMock :: Ptr a -> a
peekMock _ = error "auxiliary object for free functions"

freeAligned :: Int -> Int -> Ptr a -> IO ()
freeAligned align size ptr =
   F.free =<<
   if mod defaultMallocAlign align == 0
     then return ptr
     else F.peek $ seekPointer size ptr

freeBytes :: Storable a => a -> Int -> Ptr a -> IO ()
freeBytes a size ptr = freeAligned (alignment a) size ptr

malloc :: (Storable a) => IO (Ptr a)
malloc =
   case error "auxiliary object for LLVM.Alloc.malloc" of
      mock -> mallocBytes mock (sizeOf mock)

free :: (Storable a) => Ptr a -> IO ()
free ptr =
   freeBytes (peekMock ptr) (sizeOf (peekMock ptr)) ptr


mallocArray :: (Storable a) => Int -> IO (Ptr a)
mallocArray n =
   case error "auxiliary object for LLVM.Alloc.mallocArray" of
      mock -> mallocBytes mock (sizeOfArray n mock)

freeArray :: (Storable a) => Int -> Ptr a -> IO ()
freeArray n ptr =
   freeBytes (peekMock ptr) (sizeOfArray n (peekMock ptr)) ptr