Internal.hs

{-# LANGUAGE PolyKinds           #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell     #-}

module Math.Linear.Internal
  ( ℝ(..)
  , Fin(..), MFin(..)
  , RepDim
  , RepresentableQ(..)
  , Representable(..)
  , set
  , projection, injection
  )
  where

-- base
import Data.Kind
  ( Type, Constraint )
import Data.List
  ( intersperse )
import GHC.Generics
  ( Generic )
import GHC.Show
  ( showSpace )
import GHC.TypeNats
  ( Nat, KnownNat )

-- template-haskell
import Language.Haskell.TH
  ( CodeQ )

-- deepseq
import Control.DeepSeq
  ( NFData)

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

-- | Euclidean space \( \mathbb{R}^n \).
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 :: Double }
  deriving stock   ( Generic )
  deriving newtype ( Eq, Ord, NFData )
data    instance ℝ 2 = ℝ2 { _ℝ2_x, _ℝ2_y :: Double }
  deriving stock    Generic
  deriving anyclass NFData
  deriving stock    ( Eq, Ord )
data    instance ℝ 3 = ℝ3 { _ℝ3_x, _ℝ3_y, _ℝ3_z :: Double }
  deriving stock    Generic
  deriving anyclass NFData
  deriving stock    ( Eq, Ord )
data    instance ℝ 4 = ℝ4 { _ℝ4_x, _ℝ4_y, _ℝ4_z, _ℝ4_w :: Double }
  deriving stock    Generic
  deriving anyclass NFData
  deriving stock    ( Eq, Ord )

instance Show ( ℝ 1 ) where
  showsPrec p ( ℝ1 x )
    = showParen ( p >= 11 )
    $ showString "ℝ1"
    . foldr (.) id ( showSpace : intersperse showSpace coords )
      where
        coords = map ( showsPrec 0 ) [ x ]

instance Show ( ℝ 2 ) where
  showsPrec p ( ℝ2 x y )
    = showParen ( p >= 11 )
    $ showString "ℝ2"
    . foldr (.) id ( showSpace : intersperse showSpace coords )
      where
        coords = map ( showsPrec 0 ) [ x, y ]

instance Show ( ℝ 3 ) where
  showsPrec p ( ℝ3 x y z )
    = showParen ( p >= 11 )
    $ showString "ℝ3"
    . foldr (.) id ( showSpace : intersperse showSpace coords )
      where
        coords = map ( showsPrec 0 ) [ x, y, z ]

instance Show ( ℝ 4 ) where
  showsPrec p ( ℝ4 x y z w )
    = showParen ( p >= 11 )
    $ showString "ℝ4"
    . foldr (.) id ( showSpace : intersperse showSpace coords )
      where
        coords = map ( showsPrec 0 ) [ x, y, z, w ]

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

-- | 1, ..., n
type Fin :: Nat -> Type
newtype Fin n = Fin Word
  deriving stock ( Eq, Ord, Show )

-- | 0, ..., n
type MFin :: Nat -> Type
newtype MFin n = MFin Word

type RepDim :: k -> Nat
type family RepDim v

type RepresentableQ :: Type -> Type -> Constraint
class RepresentableQ r v | v -> r where
  tabulateQ :: ( Fin ( RepDim v ) -> CodeQ r ) -> CodeQ v
  indexQ    :: CodeQ v -> Fin ( RepDim v ) -> CodeQ r

type Representable :: Type -> Type -> Constraint
class KnownNat ( RepDim v ) => Representable r v | v -> r where
  tabulate :: ( Fin ( RepDim v ) -> r ) -> v
  index    :: v -> Fin ( RepDim v ) -> r

set :: Representable r v => Fin ( RepDim v ) -> r -> v -> v
set i r u = tabulate \ j ->
  if i == j
  then r
  else index u j
{-# INLINEABLE set #-}

projection :: ( Representable r u, Representable r v )
           => ( Fin ( RepDim v ) -> Fin ( RepDim u ) )
           -> u -> v
projection f = \ u ->
  tabulate \ i -> index u ( f i )
{-# INLINEABLE projection #-}

injection :: ( Representable r u, Representable r v )
          => ( Fin ( RepDim v ) -> MFin ( RepDim u ) )
          -> u -> v -> v
injection f = \ u v ->
  tabulate \ i -> case f i of
    MFin 0 -> index v i
    MFin j -> index u ( Fin j )
{-# INLINEABLE injection #-}

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

type instance RepDim ( ℝ n ) = n

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

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

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

instance RepresentableQ Double ( ℝ 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 Double ( ℝ 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 Double ( ℝ 0 ) where
  tabulate _ = ℝ0
  {-# INLINE tabulate #-}
  index _ _ = 0
  {-# INLINE index #-}

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

instance Representable Double ( ℝ 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 Double ( ℝ 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 Double ( ℝ 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 #-}