geomancy-0.4.0.0: src/Geomancy/DVec4.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE ViewPatterns #-}

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

module Geomancy.DVec4
  ( DVec4(..)
  , dvec4
  , withDVec4
  , pattern WithDVec4
  , convert
  , fromDVec2
  , fromDVec22
  , fromDVec3
  , fromTuple

  , (^*)
  , (^/)
  , lerp

  , dot
  , normalize

  , newDVec4

  , fromVec4
  , toVec4
  ) where

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

import Control.DeepSeq (NFData(rnf))
import Data.MonoTraversable (Element, MonoFunctor(..), MonoPointed(..))
import Data.VectorSpace (VectorSpace)
import Data.VectorSpace qualified as VectorSpace
import Foreign (Storable(..))
import Foreign.Ptr.Diff (peekDiffOff, pokeDiffOff)
import GHC.Float (double2Float, float2Double)
import GHC.IO (IO(..))
import GHC.OverloadedLabels (IsLabel(..))
import GHC.Records (HasField(..))
import Text.Printf (printf)
import WebColor.Labels (IsWebColorAlpha(..))

#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 Geomancy.DVec2 (DVec2, withDVec2)
import Geomancy.DVec3 (DVec3, withDVec3)
import Geomancy.Elementwise (Elementwise(..))
import Graphics.Gl.Block (Block(..))
import Geomancy.Vec4 (Vec4)
import Geomancy.Vec4 qualified as Vec4

data DVec4 = DVec4 ByteArray#

{-# INLINE [1] dvec4 #-}
dvec4 :: Double -> Double -> Double -> Double -> DVec4
dvec4 v0 v1 v2 v3 =
  runRW# \world ->
    let
      !(# world_, arr #) = newByteArray# 32# world

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

      world0 = write 0x0# v0 world_
      world1 = write 0x1# v1 world0
      world2 = write 0x2# v2 world1
      world3 = write 0x3# v3 world2
      !(# _world', arr' #) = unsafeFreezeByteArray# arr world3
    in
      DVec4 arr'

{-# INLINE [1] withDVec4 #-}
withDVec4
  :: DVec4
  -> (Double -> Double -> Double -> Double -> r)
  -> r
withDVec4 (DVec4 arr) f =
  f
    (D# (indexDoubleArray# arr 0x0#))
    (D# (indexDoubleArray# arr 0x1#))
    (D# (indexDoubleArray# arr 0x2#))
    (D# (indexDoubleArray# arr 0x3#))

{-# RULES
"withDVec4/dvec4" forall a b c d f.
  withDVec4 (dvec4 a b c d) f = f a b c d
"withDVec4/withDVec4" forall v k f.
  withDVec4 (withDVec4 v k) f = withDVec4 v \a b c d -> withDVec4 (k a b c d) f
  #-}

{-# INLINE convert #-}
convert :: Coercible v DVec4 => (Double -> a) -> (a -> a -> a -> a -> r) -> v -> r
convert f t v =
  withDVec4 (coerce v) \a b c d->
    t (f a) (f b) (f c) (f d)

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

instance Ord DVec4 where
  compare = lexicographic LT GT compare
  a < b = lexicographic True False (<) a b
  a <= b = not (b < a)
  a > b = b < a
  a >= b = not (a < b)

{-# INLINE lexicographic #-}
lexicographic :: r -> r -> (Double -> Double -> r) -> DVec4 -> DVec4 -> r
lexicographic lt gt lastElement (DVec4 a) (DVec4 b) = go 0#
  where
    go 3# = lastElement (D# (indexDoubleArray# a 3#)) (D# (indexDoubleArray# b 3#))
    go i =
      case compare (D# (indexDoubleArray# a i)) (D# (indexDoubleArray# b i)) of
        LT -> lt
        EQ -> go (i +# 1#)
        GT -> gt

instance Show DVec4 where
  show v =
    withDVec4 v $
      printf "DVec4 %.4f %.4f %.4f %.4f"

pattern WithDVec4 :: Double -> Double -> Double -> Double -> DVec4
pattern WithDVec4 a b c d <- ((`withDVec4` (,,,)) -> (a, b, c, d))
{-# COMPLETE WithDVec4 #-}

instance HasField "x" DVec4 Double where
  {-# INLINE getField #-}
  getField v = withDVec4 v \x _ _ _ -> x

instance HasField "y" DVec4 Double where
  {-# INLINE getField #-}
  getField v = withDVec4 v \_ y _ _ -> y

instance HasField "z" DVec4 Double where
  {-# INLINE getField #-}
  getField v = withDVec4 v \_ _ z _ -> z

instance HasField "w" DVec4 Double where
  {-# INLINE getField #-}
  getField v = withDVec4 v \_ _ _ w -> w

{-# INLINE fromDVec2 #-}
fromDVec2 :: DVec2 -> Double -> Double -> DVec4
fromDVec2 xy z w =
  withDVec2 xy \x y ->
    dvec4 x y z w

{-# INLINE fromDVec22 #-}
fromDVec22 :: DVec2 -> DVec2 -> DVec4
fromDVec22 xy zw =
  withDVec2 xy \x y ->
  withDVec2 zw \z w ->
    dvec4 x y z w

{-# INLINE fromDVec3 #-}
fromDVec3 :: Coercible a DVec3 => a -> Double -> DVec4
fromDVec3 xyz w =
  withDVec3 (coerce xyz) \x y z ->
    dvec4 x y z w

{-# INLINE fromTuple #-}
fromTuple :: (Double, Double, Double, Double) -> DVec4
fromTuple (x, y, z, w) = dvec4 x y z w

{-# INLINE fromVec4 #-}
fromVec4 :: Vec4 -> DVec4
fromVec4 = Vec4.convert float2Double dvec4

{-# INLINE toVec4 #-}
toVec4 :: DVec4 -> Vec4
toVec4 = convert double2Float Vec4.vec4

instance IsWebColorAlpha s => IsLabel s DVec4 where
  {-# INLINE fromLabel #-}
  fromLabel =
    webColorAlpha @s \r g b a ->
      dvec4
        (fromIntegral r / 255)
        (fromIntegral g / 255)
        (fromIntegral b / 255)
        (fromIntegral a / 255)

instance NFData DVec4 where
  rnf DVec4{} = ()

#ifdef TH_LIFT
instance Lift DVec4 where
  lift (WithDVec4 a b c d) = [| dvec4 $(lift a) $(lift b) $(lift c) $(lift d) |]
  liftTyped (WithDVec4 a b c d) = [|| dvec4 $$(liftTyped a) $$(liftTyped b) $$(liftTyped c) $$(liftTyped d) ||]
#endif

type instance Element DVec4 = Double

instance MonoFunctor DVec4 where
  {-# INLINE omap #-}
  omap f v =
    withDVec4 v \x y z w ->
      dvec4 (f x) (f y) (f z) (f w)

instance MonoPointed DVec4 where
  {-# INLINE opoint #-}
  opoint x = dvec4 x x x x

instance Elementwise DVec4 where
  {-# INLINE emap2 #-}
  emap2 f p0 p1 =
    withDVec4 p0 \x0 y0 z0 w0 ->
    withDVec4 p1 \x1 y1 z1 w1 ->
      dvec4
        (f x0 x1)
        (f y0 y1)
        (f z0 z1)
        (f w0 w1)

  {-# INLINE emap3 #-}
  emap3 f p0 p1 p2 =
    withDVec4 p0 \x0 y0 z0 w0 ->
    withDVec4 p1 \x1 y1 z1 w1 ->
    withDVec4 p2 \x2 y2 z2 w2 ->
      dvec4
        (f x0 x1 x2)
        (f y0 y1 y2)
        (f z0 z1 z2)
        (f w0 w1 w2)

  {-# INLINE emap4 #-}
  emap4 f p0 p1 p2 p3 =
    withDVec4 p0 \x0 y0 z0 w0 ->
    withDVec4 p1 \x1 y1 z1 w1 ->
    withDVec4 p2 \x2 y2 z2 w2 ->
    withDVec4 p3 \x3 y3 z3 w3 ->
      dvec4
        (f x0 x1 x2 x3)
        (f y0 y1 y2 y3)
        (f z0 z1 z2 z3)
        (f w0 w1 w2 w3)

  {-# INLINE emap5 #-}
  emap5 f p0 p1 p2 p3 p4 =
    withDVec4 p0 \x0 y0 z0 w0 ->
    withDVec4 p1 \x1 y1 z1 w1 ->
    withDVec4 p2 \x2 y2 z2 w2 ->
    withDVec4 p3 \x3 y3 z3 w3 ->
    withDVec4 p4 \x4 y4 z4 w4 ->
      dvec4
        (f x0 x1 x2 x3 x4)
        (f y0 y1 y2 y3 y4)
        (f z0 z1 z2 z3 z4)
        (f w0 w1 w2 w3 w4)

instance Num DVec4 where
  {-# INLINE (+) #-}
  (+) l r =
    withDVec4 l \l1 l2 l3 l4 ->
      withDVec4 r \r1 r2 r3 r4 ->
        dvec4
          (l1 + r1)
          (l2 + r2)
          (l3 + r3)
          (l4 + r4)

  {-# INLINE (-) #-}
  (-) l r =
    withDVec4 l \l1 l2 l3 l4 ->
      withDVec4 r \r1 r2 r3 r4 ->
        dvec4
          (l1 - r1)
          (l2 - r2)
          (l3 - r3)
          (l4 - r4)

  {-# INLINE (*) #-}
  (*) l r =
    withDVec4 l \l1 l2 l3 l4 ->
      withDVec4 r \r1 r2 r3 r4 ->
        dvec4
          (l1 * r1)
          (l2 * r2)
          (l3 * r3)
          (l4 * r4)

  {-# INLINE abs #-}
  abs v =
    withDVec4 v \a b c d ->
      dvec4 (abs a) (abs b) (abs c) (abs d)

  {-# INLINE signum #-}
  signum v =
    withDVec4 v \a b c d ->
      dvec4 (signum a) (signum b) (signum c) (signum d)

  {-# INLINE fromInteger #-}
  fromInteger x = dvec4 x' x' x' x'
    where
      x' = fromInteger x

instance Fractional DVec4 where
  {-# INLINE (/) #-}
  (/) l r =
    withDVec4 l \l1 l2 l3 l4 ->
      withDVec4 r \r1 r2 r3 r4 ->
        dvec4 (l1 / r1) (l2 / r2) (l3 / r3) (l4 / r4)

  {-# INLINE recip #-}
  recip v =
    withDVec4 v \a b c d ->
      dvec4 (recip a) (recip b) (recip c) (recip d)

  {-# INLINE fromRational #-}
  fromRational x = dvec4 x' x' x' x'
    where
      x' = fromRational x

instance Floating DVec4 where
  pi = opoint pi

  {-# INLINE exp #-}
  exp = omap exp
  {-# INLINE log #-}
  log = omap log
  {-# INLINE sqrt #-}
  sqrt = omap sqrt
  {-# INLINE sin #-}
  sin = omap sin
  {-# INLINE cos #-}
  cos = omap cos
  {-# INLINE asin #-}
  asin = omap asin
  {-# INLINE acos #-}
  acos = omap acos
  {-# INLINE atan #-}
  atan = omap atan
  {-# INLINE sinh #-}
  sinh = omap sinh
  {-# INLINE cosh #-}
  cosh = omap cosh
  {-# INLINE asinh #-}
  asinh = omap asinh
  {-# INLINE acosh #-}
  acosh = omap acosh
  {-# INLINE atanh #-}
  atanh = omap atanh

  {-# INLINE (**) #-}
  a ** b =
    withDVec4 a \ax ay az aw ->
    withDVec4 b \bx by bz bw ->
      dvec4
        (ax ** bx)
        (ay ** by)
        (az ** bz)
        (aw ** bw)

instance Storable DVec4 where
  {-# INLINE sizeOf #-}
  sizeOf _ = 32

  {-# INLINE alignment #-}
  alignment _ = 32

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

  {-# INLINE peek #-}
  peek (Ptr addr) = IO \world ->
    let
      !(# world0, v0 #) = readDoubleOffAddr# addr 0x0# world
      !(# world1, v1 #) = readDoubleOffAddr# addr 0x1# world0
      !(# world2, v2 #) = readDoubleOffAddr# addr 0x2# world1
      !(# world3, v3 #) = readDoubleOffAddr# addr 0x3# world2
      !v = dvec4 (D# v0) (D# v1) (D# v2) (D# v3)
    in
      (# world3, v #)

instance Block DVec4 where
  type PackedSize DVec4 = 32
  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 DVec4 where
  encode v = withDVec4 v \x y z w ->
    Serialise.encodeListLen 4 <> encode x <> encode y <> encode z <> encode w
  decode =
    Serialise.decodeListLenOf 4 *> (dvec4 <$> decode <*> decode <*> decode <*> decode)
#endif

instance VectorSpace DVec4 Double where
  zeroVector = epoint 0

  {-# INLINE (*^) #-}
  (*^) = flip (Geomancy.DVec4.^*)

  {-# INLINE (^/) #-}
  (^/) = (Geomancy.DVec4.^/)

  {-# INLINE (^+^) #-}
  (^+^) = emap2 (+)

  {-# INLINE (^-^) #-}
  (^-^) = emap2 (-)

  {-# INLINE negateVector #-}
  negateVector = emap negate

  {-# INLINE dot #-}
  dot = Geomancy.DVec4.dot

  {-# INLINE normalize #-}
  normalize = Geomancy.DVec4.normalize

{-# INLINE (^*) #-}
(^*) :: DVec4 -> Double -> DVec4
(^*) v x =
  withDVec4 v \a b c d ->
    dvec4
      (a * x)
      (b * x)
      (c * x)
      (d * x)

{-# INLINE (^/) #-}
(^/) :: DVec4 -> Double -> DVec4
v ^/ x = v ^* recip x

{-# INLINE lerp #-}
lerp :: Double -> DVec4 -> DVec4 -> DVec4
lerp alpha u v = u ^* alpha + v ^* (1 - alpha)

{-# INLINE dot #-}
dot :: DVec4 -> DVec4 -> Double
dot a b =
  withDVec4 a \a1 a2 a3 a4 ->
    withDVec4 b \b1 b2 b3 b4 ->
      a1 * b1 +
      a2 * b2 +
      a3 * b3 +
      a4 * b4

{-# INLINE normalize #-}
normalize :: DVec4 -> DVec4
normalize v =
  if nearZero q || nearZero (1-q) then
    v
  else
    v ^/ l

  where
    q = dot v v
    l = sqrt q

    nearZero a = abs a <= 1e-12

{-# INLINE newDVec4 #-}
newDVec4 :: IO DVec4
newDVec4 =
  IO \world ->
    let
      !(# world_, arr_ #) = newByteArray# 32# world
      !(# world', arr #) = unsafeFreezeByteArray# arr_ world_
    in
      (# world', DVec4 arr #)