Internal.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UnboxedTuples #-}

#ifdef TH_LIFT
{-# LANGUAGE TemplateHaskell #-}
#endif

module Geomancy.DMat4.Internal
  ( DMat4(..)
  , fromMemory
  , withMemory
  , newDMat4
    -- * Order-independent operations
  , identity
  , transpose
  , pointwise
  , scalarMultiply
    -- * List operations in memory order
  , toListMemory
  , toList2dMemory
  , toListTrans
  , toList2dTrans
  , zipWith
  ) where

import Prelude hiding (zipWith)
import GHC.Exts hiding (VecCount(..), toList)

#ifdef TH_LIFT
import Language.Haskell.TH.Syntax (Lift(..))
#endif

#ifdef SERIALISE
import Codec.Serialise (Serialise(..))
import qualified Codec.Serialise.Encoding as Serialise
import qualified Codec.Serialise.Decoding as Serialise
#endif

import Control.DeepSeq (NFData(rnf))
import Foreign (Storable(..))
import Foreign.Ptr.Diff (peekDiffOff, pokeDiffOff)
import GHC.IO (IO(..))

import qualified Data.List as List

import Graphics.Gl.Block (Block(..))

data DMat4 = DMat4 ByteArray#

instance Eq DMat4 where
  DMat4 a == DMat4 b = go 0#
    where
      go i
        | isTrue# (i ==# 16#) = True
        | isTrue# (indexDoubleArray# a i ==## indexDoubleArray# b i) = go (i +# 1#)
        | otherwise = False

instance NFData DMat4 where
  rnf DMat4{} = ()

{-# INLINE fromMemory #-}
fromMemory
  :: Double -> Double -> Double -> Double
  -> Double -> Double -> Double -> Double
  -> Double -> Double -> Double -> Double
  -> Double -> Double -> Double -> Double
  -> DMat4
fromMemory
  e0 e1 e2 e3
  e4 e5 e6 e7
  e8 e9 ea eb
  ec ed ee ef =
  runRW# \world ->
    let
      !(# world_, arr #) = newByteArray# 128# world

      write i (D# e) = writeDoubleArray# arr i e

      world0 = write 0x0# e0 world_
      world1 = write 0x1# e1 world0
      world2 = write 0x2# e2 world1
      world3 = write 0x3# e3 world2

      world4 = write 0x4# e4 world3
      world5 = write 0x5# e5 world4
      world6 = write 0x6# e6 world5
      world7 = write 0x7# e7 world6

      world8 = write 0x8# e8 world7
      world9 = write 0x9# e9 world8
      worldA = write 0xA# ea world9
      worldB = write 0xB# eb worldA

      worldC = write 0xC# ec worldB
      worldD = write 0xD# ed worldC
      worldE = write 0xE# ee worldD
      worldF = write 0xF# ef worldE

      !(# _world', arr' #) = unsafeFreezeByteArray# arr worldF
    in
      DMat4 arr'

{-# INLINE withMemory #-}
withMemory
  :: DMat4
  ->
    ( Double -> Double -> Double -> Double ->
      Double -> Double -> Double -> Double ->
      Double -> Double -> Double -> Double ->
      Double -> Double -> Double -> Double ->
      r
    )
  -> r
withMemory (DMat4 arr) f =
  f
    (D# (indexDoubleArray# arr 0x0#))
    (D# (indexDoubleArray# arr 0x1#))
    (D# (indexDoubleArray# arr 0x2#))
    (D# (indexDoubleArray# arr 0x3#))

    (D# (indexDoubleArray# arr 0x4#))
    (D# (indexDoubleArray# arr 0x5#))
    (D# (indexDoubleArray# arr 0x6#))
    (D# (indexDoubleArray# arr 0x7#))

    (D# (indexDoubleArray# arr 0x8#))
    (D# (indexDoubleArray# arr 0x9#))
    (D# (indexDoubleArray# arr 0xA#))
    (D# (indexDoubleArray# arr 0xB#))

    (D# (indexDoubleArray# arr 0xC#))
    (D# (indexDoubleArray# arr 0xD#))
    (D# (indexDoubleArray# arr 0xE#))
    (D# (indexDoubleArray# arr 0xF#))

#ifdef TH_LIFT
instance Lift DMat4 where
  lift m = withMemory m
    \ e0 e1 e2 e3
      e4 e5 e6 e7
      e8 e9 eA eB
      eC eD eE eF ->
      [|
        fromMemory
          $(lift e0) $(lift e1) $(lift e2) $(lift e3)
          $(lift e4) $(lift e5) $(lift e6) $(lift e7)
          $(lift e8) $(lift e9) $(lift eA) $(lift eB)
          $(lift eC) $(lift eD) $(lift eE) $(lift eF)
      |]
  liftTyped m = withMemory m
    \ e0 e1 e2 e3
      e4 e5 e6 e7
      e8 e9 eA eB
      eC eD eE eF ->
      [||
        fromMemory
          $$(liftTyped e0) $$(liftTyped e1) $$(liftTyped e2) $$(liftTyped e3)
          $$(liftTyped e4) $$(liftTyped e5) $$(liftTyped e6) $$(liftTyped e7)
          $$(liftTyped e8) $$(liftTyped e9) $$(liftTyped eA) $$(liftTyped eB)
          $$(liftTyped eC) $$(liftTyped eD) $$(liftTyped eE) $$(liftTyped eF)
      ||]
#endif


{-# INLINE identity #-}
identity :: DMat4
identity = fromMemory
  1 0 0 0
  0 1 0 0
  0 0 1 0
  0 0 0 1

{-# INLINE transpose #-}
transpose :: DMat4 -> DMat4
transpose =
  flip withMemory
    \ e0 e1 e2 e3
      e4 e5 e6 e7
      e8 e9 eA eB
      eC eD eE eF ->
    fromMemory
      e0 e4 e8 eC
      e1 e5 e9 eD
      e2 e6 eA eE
      e3 e7 eB eF

pointwise :: DMat4 -> DMat4 -> (Double -> Double -> Double) -> DMat4
pointwise a b f =
  withMemory a
    \ a0 a1 a2 a3
      a4 a5 a6 a7
      a8 a9 aA aB
      aC aD aE aF ->
  withMemory b
    \ b0 b1 b2 b3
      b4 b5 b6 b7
      b8 b9 bA bB
      bC bD bE bF ->
  fromMemory
    (f a0 b0) (f a1 b1) (f a2 b2) (f a3 b3)
    (f a4 b4) (f a5 b5) (f a6 b6) (f a7 b7)
    (f a8 b8) (f a9 b9) (f aA bA) (f aB bB)
    (f aC bC) (f aD bD) (f aE bE) (f aF bF)

{-# INLINE newDMat4 #-}
newDMat4 :: IO DMat4
newDMat4 =
  IO \world ->
    let
      !(# world_, arr_ #) = newByteArray# 128# world
      !(# world', arr #) = unsafeFreezeByteArray# arr_ world_
    in
      (# world', DMat4 arr #)

{-# INLINE scalarMultiply #-}
scalarMultiply :: Double -> DMat4 -> DMat4
scalarMultiply x m =
  withMemory m
    \ e0 e1 e2 e3
      e4 e5 e6 e7
      e8 e9 eA eB
      eC eD eE eF ->
      fromMemory
      (e0 * x) (e1 * x) (e2 * x) (e3 * x)
      (e4 * x) (e5 * x) (e6 * x) (e7 * x)
      (e8 * x) (e9 * x) (eA * x) (eB * x)
      (eC * x) (eD * x) (eE * x) (eF * x)

instance Storable DMat4 where
  sizeOf _mat4 = 128

  alignment _mat4 = 32

  {-# INLINE poke #-}
  poke (Ptr addr) (DMat4 arr) = IO \world ->
    let
      world' = copyByteArrayToAddr# arr 0# addr 128# world
    in
      (# world', () #)

  {-# INLINE peek #-}
  peek (Ptr addr) = IO \world ->
    let
      load i = readDoubleOffAddr# addr i

      !(# world0, e0 #) = load 0x0# world
      !(# world1, e1 #) = load 0x1# world0
      !(# world2, e2 #) = load 0x2# world1
      !(# world3, e3 #) = load 0x3# world2

      !(# world4, e4 #) = load 0x4# world3
      !(# world5, e5 #) = load 0x5# world4
      !(# world6, e6 #) = load 0x6# world5
      !(# world7, e7 #) = load 0x7# world6

      !(# world8, e8 #) = load 0x8# world7
      !(# world9, e9 #) = load 0x9# world8
      !(# worldA, eA #) = load 0xA# world9
      !(# worldB, eB #) = load 0xB# worldA

      !(# worldC, eC #) = load 0xC# worldB
      !(# worldD, eD #) = load 0xD# worldC
      !(# worldE, eE #) = load 0xE# worldD
      !(# world', eF #) = load 0xF# worldE

      !mat = fromMemory
        (D# e0) (D# e1) (D# e2) (D# e3)
        (D# e4) (D# e5) (D# e6) (D# e7)
        (D# e8) (D# e9) (D# eA) (D# eB)
        (D# eC) (D# eD) (D# eE) (D# eF)
    in
      (# world', mat #)

instance Block DMat4 where
  type PackedSize DMat4 = 128
  alignment140 _  = 32
  sizeOf140       = sizeOfPacked
  alignment430    = alignment140
  sizeOf430       = sizeOf140
  isStruct _      = False
  read140     = peekDiffOff
  write140    = pokeDiffOff
  read430     = read140
  write430    = write140
  readPacked  = read140
  writePacked = write140
  {-# INLINE alignment140 #-}
  {-# INLINE sizeOf140 #-}
  {-# INLINE alignment430 #-}
  {-# INLINE sizeOf430 #-}
  {-# INLINE isStruct #-}
  {-# INLINE read140 #-}
  {-# INLINE write140 #-}
  {-# INLINE read430 #-}
  {-# INLINE write430 #-}
  {-# INLINE readPacked #-}
  {-# INLINE writePacked #-}

#ifdef SERIALISE
instance Serialise DMat4 where
  encode m =
    withMemory m
      \ e0 e1 e2 e3
        e4 e5 e6 e7
        e8 e9 eA eB
        eC eD eE eF ->
          Serialise.encodeListLen 16 <>
          encode e0 <>
          encode e1 <> encode e2 <> encode e3 <>
          encode e4 <> encode e5 <> encode e6  <> encode e7 <>
          encode e8 <> encode e9 <> encode eA  <> encode eB <>
          encode eC <> encode eD <> encode eE  <> encode eF
  decode = do
    Serialise.decodeListLenOf 16
    e0 <- decode
    e1 <- decode
    e2 <- decode
    e3 <- decode
    e4 <- decode
    e5 <- decode
    e6 <- decode
    e7 <- decode
    e8 <- decode
    e9 <- decode
    eA <- decode
    eB <- decode
    eC <- decode
    eD <- decode
    eE <- decode
    eF <- decode
    pure $! fromMemory
      e0 e1 e2 e3
      e4 e5 e6 e7
      e8 e9 eA eB
      eC eD eE eF
#endif

toListMemory :: DMat4 -> [Double]
toListMemory = flip withMemory
    \ e0 e1 e2 e3
      e4 e5 e6 e7
      e8 e9 eA eB
      eC eD eE eF ->
    [ e0, e1, e2, e3
    , e4, e5, e6, e7
    , e8, e9, eA, eB
    , eC, eD, eE, eF
    ]

toList2dMemory :: DMat4 -> [[Double]]
toList2dMemory = flip withMemory
    \ e0 e1 e2 e3
      e4 e5 e6 e7
      e8 e9 eA eB
      eC eD eE eF ->
    [ [e0, e1, e2, e3]
    , [e4, e5, e6, e7]
    , [e8, e9, eA, eB]
    , [eC, eD, eE, eF]
    ]

toListTrans :: DMat4 -> [Double]
toListTrans = flip withMemory
    \ e0 e1 e2 e3
      e4 e5 e6 e7
      e8 e9 eA eB
      eC eD eE eF ->
    [ e0, e4, e8, eC
    , e1, e5, e9, eD
    , e2, e6, eA, eE
    , e3, e7, eB, eF
    ]

toList2dTrans :: DMat4 -> [[Double]]
toList2dTrans = flip withMemory
    \ e0 e1 e2 e3
      e4 e5 e6 e7
      e8 e9 eA eB
      eC eD eE eF ->
    [ [e0, e4, e8, eC]
    , [e1, e5, e9, eD]
    , [e2, e6, eA, eE]
    , [e3, e7, eB, eF]
    ]

zipWith :: (Double -> Double -> c) -> DMat4 -> DMat4 -> [c]
zipWith f a b = List.zipWith f (toListMemory a) (toListMemory b)