brush-strokes-0.1.0.0: src/lib/Math/Interval/Internal.hs

{-# LANGUAGE CPP #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE TemplateHaskell      #-}
{-# LANGUAGE UndecidableInstances #-}

{-# OPTIONS_GHC -Wno-orphans #-}

module Math.Interval.Internal
  ( 𝕀(𝕀, inf, sup)
  , scaleInterval, width, singleton
  , hull
  , addWithErr, subWithErr, prodWithErr, divWithErr
  , 𝕀ℝ(..)
  )
  where

-- base
import Prelude hiding ( Num(..), Fractional(..), Floating(..), (^) )
import Data.Kind
  ( Type )
import Data.Monoid
  ( Sum(..) )
import GHC.Generics
  ( Generic )
import GHC.Show
  ( showCommaSpace )
import GHC.TypeNats
  ( Nat )
import GHC.Exts
  ( (+#), (*#) )

-- deepseq
import Control.DeepSeq
  ( NFData(..) )

-- primitive
import Data.Primitive
  ( Prim(..) )

-- brush-strokes
#if defined(USE_SIMD)
import Math.Interval.Internal.SIMD
#elif defined(USE_FMA)
import Math.Interval.Internal.FMA
#else
import Math.Interval.Internal.RoundedHW
#endif

import Math.Interval.Internal.FMA
  ( addI, subI, prodI, divI )

-- brush-strokes
import Math.Float.Utils
  ( succFP, prevFP )
import Math.Linear
  ( T(..)
  , RepDim, RepresentableQ(..), Representable(..)
  , Fin(..)
  )
import Math.Module
  ( Module(..) )
import Math.Ring

--------------------------------------------------------------------------------

-- | An interval reduced to a single point.
singleton :: Double -> 𝕀
singleton x = 𝕀 x x

-- | The width of an interval.
--
-- NB: assumes the endpoints are neither @NaN@ nor infinity.
width :: 𝕀 -> Double
width ( 𝕀 lo hi ) = hi - lo

-- | Union (hull) of two intervals.
hull :: 𝕀 -> 𝕀 -> 𝕀
hull ( 𝕀 lo1 hi1 ) ( 𝕀 lo2 hi2 ) = 𝕀 ( min lo1 lo2 ) ( max hi1 hi2 )

addWithErr, subWithErr, prodWithErr, divWithErr :: Double -> Double -> (# Double, 𝕀 #)
addWithErr  = addI withError2
subWithErr  = subI withError2
prodWithErr = prodI withError2
divWithErr  = divI withError2

-- TODO: not sure why but these don't work properly right now
--{-# SPECIALISE INLINE addI withError2 #-}
--{-# SPECIALISE INLINE subI withError2 #-}
--{-# SPECIALISE INLINE prodI withError2 #-}
--{-# SPECIALISE INLINE divI withError2 #-}

{-# INLINE withError2 #-}
withError2 :: Double -> Double -> (# Double, 𝕀 #)
withError2 f e = (# f , #)
  case compare e 0 of
    LT -> 𝕀 ( prevFP f - f ) 0
    EQ -> 𝕀 0 0
    GT -> 𝕀 0 ( succFP f - f)

--------------------------------------------------------------------------------
-- Instances for (1D) intervals.

instance Eq 𝕀 where
  𝕀 a b == 𝕀 c d =
    a == c && b == d

instance Show 𝕀 where
  showsPrec _ ( 𝕀 x y )
    = showString "["
    . showsPrec 0 x
    . showCommaSpace
    . showsPrec 0 y
    . showString "]"

deriving via ViaPrelude 𝕀
  instance AbelianGroup ( T 𝕀 )
deriving via Sum 𝕀
  instance Module 𝕀 ( T 𝕀 )


-- NB: the SIMD implementation of 𝕀 uses DoubleX2#, and has a custom 'Prim' instance.
#if !defined(USE_SIMD)
instance Prim 𝕀 where
  sizeOf# _ = 2# *# sizeOf# (undefined :: Double)

  alignment# _ = alignment# (undefined :: Double)

  indexByteArray# arr# i# =
    let infVal = indexByteArray# arr# (2# *# i#)
        supVal = indexByteArray# arr# (2# *# i# +# 1#)
    in 𝕀 infVal supVal

  readByteArray# arr# i# s0 =
    case readByteArray# arr# (2# *# i#) s0 of
      (# s1, infVal #) -> case readByteArray# arr# (2# *# i# +# 1#) s1 of
        (# s2, supVal #) -> (# s2, 𝕀 infVal supVal #)

  writeByteArray# arr# i# (𝕀 infVal supVal) s0 =
    case writeByteArray# arr# (2# *# i#) infVal s0 of
      s1 -> writeByteArray# arr# (2# *# i# +# 1#) supVal s1

  indexOffAddr# addr# i# =
    let infVal = indexOffAddr# addr# (2# *# i#)
        supVal = indexOffAddr# addr# (2# *# i# +# 1#)
    in 𝕀 infVal supVal

  readOffAddr# addr# i# s0 =
    case readOffAddr# addr# (2# *# i#) s0 of
      (# s1, infVal #) -> case readOffAddr# addr# (2# *# i# +# 1#) s1 of
        (# s2, supVal #) -> (# s2, 𝕀 infVal supVal #)

  writeOffAddr# addr# i# (𝕀 infVal supVal) s0 =
    case writeOffAddr# addr# (2# *# i#) infVal s0 of
      s1 -> writeOffAddr# addr# (2# *# i# +# 1#) supVal s1
#endif

--------------------------------------------------------------------------------

type 𝕀ℝ :: Nat -> Type
data family 𝕀ℝ n

data instance 𝕀ℝ 0 = 𝕀ℝ0
  deriving stock ( Show, Eq, Ord, Generic )
  deriving anyclass NFData
newtype instance 𝕀ℝ 1 = 𝕀ℝ1 { un𝕀ℝ1 :: 𝕀 }
  deriving stock   ( Show, Generic )
  deriving newtype ( Eq, NFData )
data    instance 𝕀ℝ 2 = 𝕀ℝ2 { _𝕀ℝ2_x, _𝕀ℝ2_y :: !𝕀 }
  deriving stock    Generic
  deriving anyclass NFData
  deriving stock    ( Show, Eq )
data    instance 𝕀ℝ 3 = 𝕀ℝ3 { _𝕀ℝ3_x, _𝕀ℝ3_y, _𝕀ℝ3_z :: !𝕀 }
  deriving stock    Generic
  deriving anyclass NFData
  deriving stock    ( Show, Eq )
data    instance 𝕀ℝ 4 = 𝕀ℝ4 { _𝕀ℝ4_x, _𝕀ℝ4_y, _𝕀ℝ4_z, _𝕀ℝ4_w :: !𝕀 }
  deriving stock    Generic
  deriving anyclass NFData
  deriving stock    ( Show, Eq )

type instance RepDim ( 𝕀ℝ n ) = n

instance RepresentableQ 𝕀 ( 𝕀ℝ 0 ) where
  tabulateQ _ = [|| 𝕀ℝ0 ||]
  indexQ _ _ = [|| 0 ||]

instance RepresentableQ 𝕀 ( 𝕀ℝ 1 ) where
  tabulateQ f = [|| 𝕀ℝ1 $$( f ( Fin 1 ) ) ||]
  indexQ p = \ case
    _ -> [|| un𝕀ℝ1 $$p ||]

instance RepresentableQ 𝕀 ( 𝕀ℝ 2 ) where
  tabulateQ f = [|| 𝕀ℝ2 $$( f ( Fin 1 ) ) $$( f ( Fin 2 ) ) ||]
  indexQ p = \ case
    Fin 1 -> [|| _𝕀ℝ2_x $$p ||]
    _     -> [|| _𝕀ℝ2_y $$p ||]

instance RepresentableQ 𝕀 ( 𝕀ℝ 3 ) where
  tabulateQ f = [|| 𝕀ℝ3 $$( f ( Fin 1 ) ) $$( f ( Fin 2 ) ) $$( f ( Fin 3 ) ) ||]
  indexQ p = \ case
    Fin 1 -> [|| _𝕀ℝ3_x $$p ||]
    Fin 2 -> [|| _𝕀ℝ3_y $$p ||]
    _     -> [|| _𝕀ℝ3_z $$p ||]

instance RepresentableQ 𝕀 ( 𝕀ℝ 4 ) where
  tabulateQ f = [|| 𝕀ℝ4 $$( f ( Fin 1 ) ) $$( f ( Fin 2 ) ) $$( f ( Fin 3 ) ) $$( f ( Fin 4 ) ) ||]
  indexQ p = \ case
    Fin 1 -> [|| _𝕀ℝ4_x $$p ||]
    Fin 2 -> [|| _𝕀ℝ4_y $$p ||]
    Fin 3 -> [|| _𝕀ℝ4_z $$p ||]
    _     -> [|| _𝕀ℝ4_w $$p ||]


instance Representable 𝕀 ( 𝕀ℝ 0 ) where
  tabulate _ = 𝕀ℝ0
  {-# INLINE tabulate #-}
  index _ _ = 0
  {-# INLINE index #-}

instance Representable 𝕀 ( 𝕀ℝ 1 ) where
  tabulate f = 𝕀ℝ1 ( f ( Fin 1 ) )
  {-# INLINE tabulate #-}
  index p = \ case
    _ -> un𝕀ℝ1 p
  {-# INLINE index #-}

instance Representable 𝕀 ( 𝕀ℝ 2 ) where
  tabulate f = 𝕀ℝ2 ( f ( Fin 1 ) ) ( f ( Fin 2 ) )
  {-# INLINE tabulate #-}
  index p = \ case
    Fin 1 -> _𝕀ℝ2_x p
    _     -> _𝕀ℝ2_y p
  {-# INLINE index #-}

instance Representable 𝕀 ( 𝕀ℝ 3 ) where
  tabulate f = 𝕀ℝ3 ( f ( Fin 1 ) ) ( f ( Fin 2 ) ) ( f ( Fin 3 ) )
  {-# INLINE tabulate #-}
  index p = \ case
    Fin 1 -> _𝕀ℝ3_x p
    Fin 2 -> _𝕀ℝ3_y p
    _     -> _𝕀ℝ3_z p
  {-# INLINE index #-}

instance Representable 𝕀 ( 𝕀ℝ 4 ) where
  tabulate f = 𝕀ℝ4 ( f ( Fin 1 ) ) ( f ( Fin 2 ) ) ( f ( Fin 3 ) ) ( f ( Fin 4 ) )
  {-# INLINE tabulate #-}
  index p = \ case
    Fin 1 -> _𝕀ℝ4_x p
    Fin 2 -> _𝕀ℝ4_y p
    Fin 3 -> _𝕀ℝ4_z p
    _     -> _𝕀ℝ4_w p
  {-# INLINE index #-}