DVec3.hs

{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ViewPatterns #-}

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

module Geomancy.DVec3
  ( DVec3
  , dvec3
  , withDVec3
  , pattern WithDVec3
  , convert
  , fromDVec2
  , fromTuple

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

  , cross
  , dot
  , normalize

  , Packed(..)
  , packed

  , emap2
  , emap3
  , emap4

  , fromVec3
  , toVec3
  ) where

import Control.DeepSeq (NFData(rnf))
import Data.Coerce (Coercible, coerce)
import Data.MonoTraversable (Element, MonoFunctor(..), MonoPointed(..))
import Data.VectorSpace (VectorSpace)
import Data.VectorSpace qualified as VectorSpace
import Foreign (Storable(..), castPtr)
import Foreign.Ptr.Diff (peekDiffOff, pokeDiffOff)
import GHC.Float (double2Float, float2Double)
import GHC.OverloadedLabels (IsLabel(..))
import GHC.Records (HasField(..))
import WebColor.Labels (IsWebColor(..))

#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.Elementwise (Elementwise(..))
import Graphics.Gl.Block (Block(..))
import Geomancy.Vec3 (Vec3)
import Geomancy.Vec3 qualified as Vec3

data DVec3 = DVec3
  {-# UNPACK #-} !Double
  {-# UNPACK #-} !Double
  {-# UNPACK #-} !Double
  deriving (Eq, Ord, Show)
#ifdef TH_LIFT
  deriving Lift
#endif

{-# INLINE dvec3 #-}
dvec3 :: Double -> Double -> Double -> DVec3
dvec3 = DVec3

{-# INLINE withDVec3 #-}
withDVec3
  :: DVec3
  -> (Double -> Double -> Double -> r)
  -> r
withDVec3 (DVec3 a b c) f = f a b c

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

pattern WithDVec3 :: Double -> Double -> Double -> DVec3
pattern WithDVec3 a b c <- ((`withDVec3` (,,)) -> (a, b, c))
{-# COMPLETE WithDVec3 #-}

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

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

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

{-# INLINE fromDVec2 #-}
fromDVec2 :: Coercible DVec3 a => DVec2 -> Double -> a
fromDVec2 xy z =
  withDVec2 xy \x y ->
    coerce (dvec3 x y z)

{-# INLINE fromTuple #-}
fromTuple :: Coercible DVec3 a => (Double, Double, Double) -> a
fromTuple (x, y, z) = coerce (dvec3 x y z)

{-# INLINE fromVec3 #-}
fromVec3 :: Vec3 -> DVec3
fromVec3 = Vec3.convert float2Double dvec3

{-# INLINE toVec3 #-}
toVec3 :: DVec3 -> Vec3
toVec3 = convert double2Float Vec3.vec3

instance IsWebColor s => IsLabel s DVec3 where
  {-# INLINE fromLabel #-}
  fromLabel =
    webColor @s \r g b ->
      dvec3
        (fromIntegral r / 255)
        (fromIntegral g / 255)
        (fromIntegral b / 255)

instance NFData DVec3 where
  rnf DVec3{} = ()

type instance Element DVec3 = Double

instance MonoFunctor DVec3 where
  {-# INLINE omap #-}
  omap f v =
    withDVec3 v \x y z ->
      dvec3 (f x) (f y) (f z)

instance MonoPointed DVec3 where
  opoint x = dvec3 x x x

instance Elementwise DVec3 where
  {-# INLINE emap2 #-}
  emap2 f p0 p1 =
    withDVec3 p0 \x0 y0 z0 ->
    withDVec3 p1 \x1 y1 z1 ->
      dvec3
        (f x0 x1)
        (f y0 y1)
        (f z0 z1)

  {-# INLINE emap3 #-}
  emap3 f p0 p1 p2 =
    withDVec3 p0 \x0 y0 z0 ->
    withDVec3 p1 \x1 y1 z1 ->
    withDVec3 p2 \x2 y2 z2 ->
      dvec3
        (f x0 x1 x2)
        (f y0 y1 y2)
        (f z0 z1 z2)

  {-# INLINE emap4 #-}
  emap4 f p0 p1 p2 p3 =
    withDVec3 p0 \x0 y0 z0 ->
    withDVec3 p1 \x1 y1 z1 ->
    withDVec3 p2 \x2 y2 z2 ->
    withDVec3 p3 \x3 y3 z3 ->
      dvec3
        (f x0 x1 x2 x3)
        (f y0 y1 y2 y3)
        (f z0 z1 z2 z3)

  {-# INLINE emap5 #-}
  emap5 f p0 p1 p2 p3 p4 =
    withDVec3 p0 \x0 y0 z0 ->
    withDVec3 p1 \x1 y1 z1 ->
    withDVec3 p2 \x2 y2 z2 ->
    withDVec3 p3 \x3 y3 z3 ->
    withDVec3 p4 \x4 y4 z4 ->
      dvec3
        (f x0 x1 x2 x3 x4)
        (f y0 y1 y2 y3 y4)
        (f z0 z1 z2 z3 z4)

instance Num DVec3 where
  {-# INLINE (+) #-}
  DVec3 a b c + DVec3 d e f =
    DVec3
      (a + d)
      (b + e)
      (c + f)

  {-# INLINE (-) #-}
  DVec3 a b c - DVec3 d e f =
    DVec3
      (a - d)
      (b - e)
      (c - f)

  {-# INLINE (*) #-}
  DVec3 a b c * DVec3 d e f =
    DVec3
      (a * d)
      (b * e)
      (c * f)

  {-# INLINE abs #-}
  abs (DVec3 a b c) =
    DVec3 (abs a) (abs b) (abs c)

  {-# INLINE signum #-}
  signum (DVec3 a b c) =
    DVec3 (signum a) (signum b) (signum c)

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

instance Fractional DVec3 where
  {-# INLINE (/) #-}
  DVec3 l1 l2 l3 / DVec3 r1 r2 r3 =
    DVec3 (l1 / r1) (l2 / r2) (l3 / r3)

  {-# INLINE recip #-}
  recip (DVec3 a b c) =
    DVec3 (recip a) (recip b) (recip c)

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

instance Floating DVec3 where
  pi = opoint pi

  exp = omap exp
  log = omap log
  sqrt = omap sqrt
  sin = omap sin
  cos = omap cos
  asin = omap asin
  acos = omap acos
  atan = omap atan
  sinh = omap sinh
  cosh = omap cosh
  asinh = omap asinh
  acosh = omap acosh
  atanh = omap atanh

  a ** b =
    withDVec3 a \ax ay az ->
    withDVec3 b \bx by bz ->
      dvec3
        (ax ** bx)
        (ay ** by)
        (az ** bz)

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

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

  {-# INLINE poke #-}
  poke ptr v3 =
    withDVec3 v3 \a b c -> do
      poke ptr' a
      pokeElemOff ptr' 1 b
      pokeElemOff ptr' 2 c
      pokeElemOff ptr' 3 (1.0 :: Double)
    where
      ptr' = castPtr ptr

  {-# INLINE peek #-}
  peek ptr =
    dvec3 <$> peek ptr' <*> peekElemOff ptr' 1 <*> peekElemOff ptr' 2
    where
      ptr' = castPtr ptr

instance Block DVec3 where
  type PackedSize DVec3 = 24
  alignment140 _  = 32
  sizeOf140 _     = 32
  alignment430    = alignment140
  sizeOf430       = sizeOf140
  isStruct _      = False
  read140     = peekDiffOff
  write140    = pokeDiffOff
  read430     = read140
  write430    = write140
  readPacked  = read140
  writePacked ptr diff = pokeDiffOff ptr (coerce diff) . Packed
  {-# 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 DVec3 where
  encode v = withDVec3 v \x y z ->
    Serialise.encodeListLen 3 <> encode x <> encode y <> encode z
  decode =
    Serialise.decodeListLenOf 3 *> (dvec3 <$> decode <*> decode <*> decode)
#endif

{-# INLINE (^*) #-}
(^*) :: DVec3 -> Double -> DVec3
DVec3 a b c ^* x =
  DVec3
    (a * x)
    (b * x)
    (c * x)

{-# INLINE (^/) #-}
(^/) :: DVec3 -> Double -> DVec3
DVec3 a b c ^/ x =
  DVec3
    (a / x)
    (b / x)
    (c / x)

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

{-# INLINE cross #-}
cross :: DVec3 -> DVec3 -> DVec3
cross (DVec3 a b c) (DVec3 d e f) =
  DVec3
    (b * f - c * e)
    (c * d - a * f)
    (a * e - b * d)

{-# INLINE dot #-}
dot :: DVec3 -> DVec3 -> Double
dot (DVec3 a b c) (DVec3 d e f) =
  a * d +
  b * e +
  c * f

{-# INLINE normalize #-}
normalize :: DVec3 -> DVec3
normalize v =
  if nearZero q || nearZero (1-q) then
    v
  else
    let
      DVec3 x y z = v
    in
      DVec3 (x / l) (y / l) (z / l)

  where
    q = dot v v
    l = sqrt q

    nearZero a = abs a <= 1e-12

instance VectorSpace DVec3 Double where
  zeroVector = epoint 0

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

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

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

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

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

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

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

type instance Element Packed = Double

newtype Packed = Packed { unPacked :: DVec3 }
  deriving stock
    ( Eq, Ord, Show
    )
  deriving newtype
    ( NFData, Num, Fractional, Floating
    , MonoFunctor, MonoPointed
    , Elementwise
    )
#ifdef TH_LIFT
  deriving Lift
#endif

{-# INLINE packed #-}
packed :: Double -> Double -> Double -> Packed
packed x y z = Packed (dvec3 x y z)

instance HasField "x" Packed Double where
  {-# INLINE getField #-}
  getField (Packed v) = withDVec3 v \x _ _ -> x

instance HasField "y" Packed Double where
  {-# INLINE getField #-}
  getField (Packed v) = withDVec3 v \_ y _ -> y

instance HasField "z" Packed Double where
  {-# INLINE getField #-}
  getField (Packed v) = withDVec3 v \_ _ z -> z

instance Storable Packed where
  {-# INLINE sizeOf #-}
  sizeOf _ = 24

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

  {-# INLINE poke #-}
  poke ptr (Packed v3) =
    withDVec3 v3 \a b c -> do
      poke ptr' a
      pokeElemOff ptr' 1 b
      pokeElemOff ptr' 2 c
    where
      ptr' = castPtr ptr

  {-# INLINE peek #-}
  peek ptr = packed
    <$> peek ptr'
    <*> peekElemOff ptr' 1
    <*> peekElemOff ptr' 2
    where
      ptr' = castPtr ptr

instance Block Packed where
  type PackedSize Packed = 24
  alignment140 _  = 32
  sizeOf140 _     = 24
  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 #-}

instance VectorSpace Packed Double where
  zeroVector = epoint 0

  {-# INLINE (*^) #-}
  (*^) = flip $ coerce (Geomancy.DVec3.^*)

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

  {-# INLINE (^+^) #-}
  v1 ^+^ v2 = v1 + v2

  {-# INLINE (^-^) #-}
  v1 ^-^ v2 = v1 - v2

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

  {-# INLINE dot #-}
  dot = coerce Geomancy.DVec3.dot

  {-# INLINE normalize #-}
  normalize = coerce Geomancy.DVec3.normalize