content: {-# LANGUAGE DataKinds, GADTs, OverloadedLabels, ScopedTypeVariables,
               TypeOperators #-}
  -- A benchmark where we initialize a 'V.Vector' of random vertices,
  -- each carrying 3D position, 2D texture coordinates, and a 3D normal
  -- vector. A calculation is carried out where we multiply the y
  -- coordinate of each point's normal vector by 2, then sum all normal
  -- vector coordinates over all vertices. This calculation is performed
  -- by interfacing the vertex data as a flat record, a traditional
  -- record of "Linear" finite dimensional vector types, and a vinyl
  -- record of linear fields.
  import Lens.Micro
  import Lens.Micro.Extras (view)
  import Control.Monad (when)
  import qualified Data.Foldable as F
  import qualified Data.Vector.Storable as V
  import qualified Data.Vector.Storable.Mutable as VM
  import Data.Vinyl
  import Foreign.Ptr (castPtr)
  import Foreign.Storable (Storable(..))
  import Linear (V2, V3, _y)
  import System.Random.MWC (withSystemRandom, Variate(..), GenIO)
  import Criterion.Main
  
  randVec :: (Storable a, Variate a) => Int -> GenIO -> IO (V.Vector a)
  randVec n g = VM.replicateM n (uniform g) >>=
                V.unsafeFreeze
  
  randVecStd :: (Storable a, Variate a) => Int -> IO (V.Vector a)
  randVecStd = withSystemRandom . randVec
  
  type MyFields a = [ '("pos", V3 a), '("tex", V2 a), '("normal", V3 a) ]
  type MyVertex a = FieldRec (MyFields a)
  
  (*~) :: Num a => ASetter s t a a -> a -> s -> t
  l *~ x = l %~ (* x)
  infixr 4 *~
  
  vinylNormSumLens :: (Num a, Storable a) => V.Vector (MyVertex a) -> a
  vinylNormSumLens = V.sum . V.map (F.sum . view (rlensf #normal))
  
  vinylNormSumLabel :: (Num a, Storable a) => V.Vector (MyVertex a) -> a
  vinylNormSumLabel = V.sum . V.map (F.sum . rvalf #normal)
  
  doubleNormYLens :: V.Vector (MyVertex Float) -> V.Vector (MyVertex Float)
  doubleNormYLens = V.map (rlensf #normal . _y *~ (2::Float))
  
  doubleNormY :: V.Vector (MyVertex Float) -> V.Vector (MyVertex Float)
  doubleNormY = V.map (\(p :& t :& Field n :& RNil) ->
                         p :& t :& Field (_y *~ (2::Float) $ n) :& RNil)
  
  vinylNormSum :: (Num a, Storable a) => V.Vector (MyVertex a) -> a
  vinylNormSum = V.sum . V.map (F.sum . (\(_ :& _ :& Field vn :& RNil) -> vn))
  
  main :: IO ()
  main = do vals <- randVecStd $ n * 8 :: IO (V.Vector Float)
            let vinylVerts = V.unsafeCast vals :: V.Vector (MyVertex Float)
                flatVerts = V.unsafeCast vals
                reasVerts = V.unsafeCast vals
                vinylAns = vinylNormSum $ doubleNormY vinylVerts
                vinylLans = vinylNormSumLens $ doubleNormYLens vinylVerts
                vinylLabAns = vinylNormSumLabel $ doubleNormYLens vinylVerts
                flatAns = flatNormSum $ doubleNormFlat flatVerts
                reasAns = reasNormSum $ doubleNormReas reasVerts
            when (any (/= vinylAns) [ vinylLans, flatAns, reasAns, vinylLabAns ])
                 (error "Not all versions compute the same answer")
            defaultMain [ bench "flat" $
                          whnf (flatNormSum . doubleNormFlat) flatVerts
                        , bench "vinyl" $
                          whnf (vinylNormSum . doubleNormY) vinylVerts
                        , bench "vinyl-lens" $
                          whnf (vinylNormSumLens . doubleNormYLens) vinylVerts
                        , bench "vinyl-label" $
                          whnf (vinylNormSumLabel . doubleNormYLens) vinylVerts
                        , bench "reasonable" $
                          whnf (reasNormSum . doubleNormReas) reasVerts ]
    where n = 1000
  
  --------------------------------------------------------------------------------
  -- Baseline data structures for comparison
  
  -- Don't trust data structures at all? Use a flat record where only
  -- the field prefixes denote their roles.
  data TotallyFlat a = Flat { px :: !a
                            , py :: !a
                            , pz :: !a
                            , tu :: !a
                            , tv :: !a
                            , nx :: !a
                            , ny :: !a
                            , nz :: !a }
  
  instance Storable a => Storable (TotallyFlat a) where
    sizeOf _ = sizeOf (undefined::a) * 8
    alignment _ = alignment (undefined::a)
    peek ptr = Flat <$> peek ptr' <*> peekElemOff ptr' 1 <*> peekElemOff ptr' 2
                    <*> peekElemOff ptr' 3 <*> peekElemOff ptr' 4
                    <*> peekElemOff ptr' 5 <*> peekElemOff ptr' 6
                    <*> peekElemOff ptr' 7
      where ptr' = castPtr ptr
    poke ptr (Flat px' py' pz' tu' tv' nx' ny' nz') = do poke ptr' px'
                                                         pokeElemOff ptr' 1 py'
                                                         pokeElemOff ptr' 2 pz'
                                                         pokeElemOff ptr' 3 tu'
                                                         pokeElemOff ptr' 4 tv'
                                                         pokeElemOff ptr' 5 nx'
                                                         pokeElemOff ptr' 6 ny'
                                                         pokeElemOff ptr' 7 nz'
      where ptr' = castPtr ptr
  
  flatNormSum :: (Num a, Storable a) => V.Vector (TotallyFlat a) -> a
  flatNormSum = V.sum . V.map (\v -> nx v + ny v + nz v)
  
  doubleNormFlat :: V.Vector (TotallyFlat Float) -> V.Vector (TotallyFlat Float)
  doubleNormFlat = V.map (\v -> v { ny = ny v * 2 })
  
  -- A more reasonable approach to a vertex record.
  data Reasonable a = Reasonable { rPos  :: V3 a
                                 , rTex  :: V2 a
                                 , rNorm :: V3 a }
  
  instance Storable a => Storable (Reasonable a) where
    sizeOf _ = sizeOf (undefined::V3 a)*2 + sizeOf (undefined::V2 a)
    alignment _ = alignment (undefined::V3 a)
    peek ptr = Reasonable <$> peek (castPtr ptr)
                          <*> peekByteOff (castPtr ptr) szx
                          <*> peekByteOff (castPtr ptr) (szx + szy)
      where szx = sizeOf (undefined::V3 a)
            szy = sizeOf (undefined::V2 a)
    poke ptr (Reasonable p t n) = do poke (castPtr ptr) p
                                     pokeByteOff (castPtr ptr) szx t
                                     pokeByteOff (castPtr ptr) (szx + szy) n
      where szx = sizeOf (undefined::V3 a)
            szy = sizeOf (undefined::V2 a)
  
  reasNormSum :: (Num a, Storable a) => V.Vector (Reasonable a) -> a
  reasNormSum = V.sum . V.map (F.sum . rNorm)
  
  doubleNormReas :: V.Vector (Reasonable Float) -> V.Vector (Reasonable Float)
  doubleNormReas = V.map (\v -> v { rNorm = (_y *~ 2) $ rNorm v })
  