{-# OPTIONS -fglasgow-exts -fbang-patterns -cpp -fallow-undecidable-instances #-}

import Foreign
import Control.Monad


#ifdef POLY_OTHER
import OtherP
version = "Polymorphic version, other file"
#endif

#ifdef POLY_SAME

{- POLYMORPHIC TYPE 
 - when these definitions are in this file, it runs in constant space 
 - when you move them to another module, it uses heap proportional to n 
 -}

{- to see heap explosion, snip form HERE ... -}
-- data C a b = C {-# UNPACK #-} !a {-# UNPACK #-} !b deriving (Eq,Show)
data C a b = C !a !b deriving (Eq,Show)

type Vec3 = C Double (C Double (C Double ())) 

vec3 !a !b !c = C a (C b (C c ()))

instance (Show a, Eq a, Num a) => Num (C a ()) where
  (C a ()) + (C b ()) = C (a+b) ()

instance (Show (C a (C a b)), Eq (C a (C a b)), Num (C a b), Num a) => Num (C a (C a b)) where
  (C i (C j k)) + (C x (C y z)) = C (i+x) ((C j k)+(C y z))

instance Storable a => Storable (C a ()) where
  sizeOf _ = sizeOf (undefined::a)
  alignment _ = alignment (undefined::a)
  peek p = peek (castPtr p) >>= \a -> return (C a ())
  poke p (C a _) = poke (castPtr p) a

instance (Storable a, Storable (C a v)) => Storable (C a (C a v)) 
  where
  sizeOf _ = sizeOf (undefined::a) + sizeOf (undefined::(C a v))
  alignment _ = alignment (undefined::a)
  peek !p = 
    peek (castPtr p) >>= \a -> 
    peek (castPtr (p`plusPtr`sizeOf(undefined::a))) >>= \v -> 
    return (C a v)
  poke !p !(C a v) = 
    poke (castPtr p) a >> 
    poke (castPtr (p`plusPtr`sizeOf(undefined::a))) v

--  also, inlining these seem to defeat unboxing
--  {-# INLINE peek #-}
--  {-# INLINE poke #-}
--  {-# INLINE sizeOf #-}
--  {-# INLINE alignment #-}

{- ... TO HERE and move to another file -}
version = "Polymorphic version, same file"
#endif

#ifdef MONO_OTHER
import OtherM
version = "Monomorphic version, other file"
#endif

#ifdef MONO_SAME

{- MONOMORPHIC TYPE -}

data Vec3 = Vec3 {-# UNPACK #-} !Double
                 {-# UNPACK #-} !Double
                 {-# UNPACK #-} !Double
  deriving (Eq,Show)

vec3 !x !y !z = Vec3 x y z

instance Num Vec3 where
  (Vec3 a b c) + (Vec3 x y z) = Vec3 (a+x) (b+y) (c+z)

instance Storable Vec3 where
  sizeOf _ = 3 * sizeOf(undefined::Double)
  alignment _ = alignment (undefined::Double)
  peek !ptr = peek  (castPtr ptr)               >>= \x -> 
              peek ((castPtr ptr)`advancePtr`1) >>= \y -> 
              peek ((castPtr ptr)`advancePtr`2) >>= \z -> return (Vec3 x y z)

  poke !ptr !(Vec3 x y z) = 
             poke  (castPtr ptr)               x >>
             poke ((castPtr ptr)`advancePtr`1) y >>
             poke ((castPtr ptr)`advancePtr`2) z 


version = "Monomorphic version, same file"

#endif 

n = 10000000

main = 
  do
  a <- mallocArray n  :: IO( Ptr Vec3 )
  b <- mallocArray n  
  c <- mallocArray n  

  forM_ [0..n-1] $ \i ->
    pokeElemOff a i (vec3 1 2 3) >> pokeElemOff b i (vec3 4 5 6)

  forM_ [0..n-1] $ \i ->
    peekElemOff a i >>= \x -> peekElemOff b i >>= \y -> pokeElemOff c i (x+y)

  forM_ [1..n-1] $ \i ->
    peekElemOff c (i-1) >>= \x -> peekElemOff c i >>= \y -> pokeElemOff c i (x+y)
  
  print =<< peekElemOff c (n-1)

  putStrLn version