Records.hs

{-# LANGUAGE AllowAmbiguousTypes  #-}
{-# LANGUAGE ParallelListComp     #-}
{-# LANGUAGE PolyKinds            #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE UndecidableInstances #-}

module BrushStroking.Records where

-- base
import Data.Functor
  ( (<&>) )
import Data.Kind
  ( Type, Constraint )
import Data.List
  ( findIndex, intersperse, sortBy )
import Data.Ord
  ( comparing )
import Data.Typeable
  ( Typeable, eqT )
import Data.Type.Equality
  ( (:~:)(Refl) )
import GHC.Exts
  ( Proxy#, proxy# )
import GHC.Show
  ( showCommaSpace )
import GHC.TypeLits
  ( Symbol, KnownSymbol, symbolVal'
  , SomeSymbol(..), someSymbolVal
  )
import GHC.TypeNats
  ( Nat, KnownNat
  , type (<=), type (<=?), type (+)
  )
import Unsafe.Coerce
  ( unsafeCoerce )

-- acts
import Data.Act
  ( Act(..), Torsor(..) )

-- containers
import qualified Data.Map.Strict as Map

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

-- groups
import Data.Group
  ( Group(..) )

-- text
import Data.Text
  ( Text )
import qualified Data.Text as Text
  ( pack, unpack )

-- MetaBrush
import Math.Differentiable
import Math.Linear
import Math.Module
  ( Module )

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

-- | A record of 'Double' values.
type Record :: [ k ] -> Type
newtype Record ks = MkR { recordKeyVals :: ℝ ( Length ks ) }

deriving newtype
  instance Eq ( ℝ ( Length ks ) )
        => Eq ( Record ks )
deriving newtype
  instance Ord ( ℝ ( Length ks ) )
        => Ord ( Record ks )
deriving newtype
  instance NFData ( ℝ ( Length ks ) )
        => NFData ( Record ks )

-- | Show a record, using the given type-level field names.
instance ( KnownSymbols ks, Representable Double ( ℝ ( Length ks ) ) )
       => Show ( Record ks ) where
  showsPrec p ( MkR r )
    = showParen ( p >= 11 )
    $ showString "{"
    . foldr (.) id ( intersperse showCommaSpace fields )
    . showString "}"
      where
        fields :: [ ShowS ]
        fields =
          zip [ 1.. ] ( knownSymbols @ks ) <&> \ ( i, fld ) ->
            let v = index r ( Fin i )
            in showString ( Text.unpack fld ) . showString " = " . showsPrec 0 v

deriving via ( T ( ℝ ( Length ks ) ) )
  instance Semigroup ( T ( ℝ ( Length ks ) ) )
        => Semigroup ( T ( Record ks ) )
deriving via ( T ( ℝ ( Length ks ) ) )
  instance Monoid    ( T ( ℝ ( Length ks ) ) )
        => Monoid    ( T ( Record ks ) )
deriving via ( T ( ℝ ( Length ks ) ) )
  instance Group     ( T ( ℝ ( Length ks ) ) )
        => Group     ( T ( Record ks ) )
deriving via ( T ( ℝ ( Length ks ) ) )
  instance Module Double ( T ( ℝ ( Length ks ) ) )
        => Module Double ( T ( Record ks ) )

instance ( Act       ( T ( ℝ ( Length ks ) ) ) ( ℝ ( Length ks ) )
         , Semigroup ( T ( ℝ ( Length ks ) ) ) )
        => Act ( T ( Record ks ) ) ( Record ks ) where
  T ( MkR g ) • MkR a = MkR ( T g • a )
instance ( Torsor ( T ( ℝ ( Length ks ) ) ) ( ℝ ( Length ks ) )
         , Group  ( T ( ℝ ( Length ks ) ) ) )
        => Torsor ( T ( Record ks ) ) ( Record ks ) where
  MkR g --> MkR a = T $ MkR $ unT $ g --> a

type instance RepDim ( Record ks ) = Length ks
deriving newtype instance ( KnownNat (Length ks)
                          , Representable r ( ℝ ( Length ks ) ) )
                       => Representable r ( Record ks )

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

type Length :: [ k ] -> Nat
type family Length xs where
  Length '[]        = 0
  Length ( _ : xs ) = 1 + Length xs

type KnownSymbols :: [ Symbol ] -> Constraint
class Typeable ks => KnownSymbols ks where
  knownSymbols :: [ Text ]
instance KnownSymbols '[] where
  knownSymbols = []
  {-# INLINE knownSymbols #-}
instance ( KnownSymbol k, KnownSymbols ks ) => KnownSymbols ( k ': ks ) where
  knownSymbols = Text.pack ( symbolVal' @k proxy# ) : knownSymbols @ks
  {-# INLINE knownSymbols #-}

--------------------------------------------------------------------------------
-- Intersection of two records.

{- Note [intersect and specialisation]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The "intersect" function does two things:

  1. Given a row of "path" parameters (the parametrs that vary along the stroke path),
     and a row of "brush" parameters (used to parametrise the brush shape),
     it computes an intersection row, which is the row for which the main
     brush stroking algorithms is going to be used with.

  2. Additionally, to ensure good performance, we dispatch and then call
     the continuation with specific arguments. This requires:

       - INLINE pragmas on 'intersection, 'doIntersection' etc
       - explicit type applications of the continuation argument to specific types
       - usage of oneShot

     This allows us to dispatch and call the appropriately specialised function,
     instead of the overloaded function.

     It's essentially the following pattern:

        type C :: Type -> Constraint
        type C t = ( Show t, NFData t, Num t, ... )
          -- a big bunch of constraints

        bigOverloadedFn :: forall t. C t => t -> Bool

        callOfBigFn :: Typeable t => t -> Bool
        callOfBigFn x = ensureSpecialisation @t $ bigOverloadedFn ( 3 * x + 1 )

        {-# INLINE ensureSpecialisation #-}
        ensureSpecialisation :: forall x r. KnownNat x => ( forall y. C y => r ) -> r
        ensureSpecialisation f
          | Just Refl <- eqT @x @1
          = f @Int
          | Just Refl <- eqT @x @2
          = f @Word
          | Just Refl <- eqT @x @3
          = f @Double
          ...
          | otherwise
          = error "ensureSpecialisation: unsupported type"
-}

{-# INLINE intersect #-}
intersect :: forall r1 r2 kont
          .  ( KnownSymbols r1, KnownSymbols r2 )
          => ( Intersection r1 r2 -> kont ) -> kont
intersect f =
  doIntersection @r1 @r2 \ ( lg1 :: Proxy# l1 ) ( lg2 :: Proxy# l2 ) ( _ :: Proxy# r1r2 ) ( lg12 :: Proxy# l12 ) r1_idxs r2_idxs ->
    case ( eqT @r1 @r2, eqT @r2 @r1r2 ) of
      ( Just Refl, Just Refl ) ->
        -- Shortcut when the two rows are equal
        f $ Intersection lg1 lg2 lg12 id id ( \ _ -> id ) ( \ _ -> id )
      _ ->
        let
          project1 :: Record r1 -> Record r1r2
          project1 = \ ( MkR r1 ) -> MkR $ projection ( (!) r1_idxs ) r1
          {-# INLINE project1 #-}

          project2 :: Record r2 -> Record r1r2
          project2 = \ ( MkR r2 ) -> MkR $ projection ( (!) r2_idxs ) r2
          {-# INLINE project2 #-}

          inject1 :: Record r1 -> Record r1r2 -> Record r1
          inject1 = \ ( MkR r1 ) -> \ ( MkR r1r2 ) -> MkR $ injection ( \ i -> find ( == i ) r1_idxs ) r1r2 r1
          {-# INLINE inject1 #-}

          inject2 :: Record r2 -> Record r1r2 -> Record r2
          inject2 = \ ( MkR r2 ) -> \ ( MkR r1r2 ) -> MkR $ injection ( \ i -> find ( == i ) r2_idxs ) r1r2 r2
          {-# INLINE inject2 #-}
        in
          f $ Intersection { lg1, lg2, lg12, project1, project2, inject1, inject2 }

data Intersection r1 r2 where
  Intersection
    :: forall r1r2 r1 r2 l12 l1 l2
    .  ( l12 ~ Length r1r2, l1 ~ Length r1, l2 ~ Length r2
       , l1 <= 4, l2 <= 4, l12 <= 4
       , l12 <= l1, l12 <= l2
       , KnownSymbols r1r2
       , Representable Double ( ℝ l1 )
       , Show ( ℝ l1 ), NFData ( ℝ l1 )
       , Show ( ℝ l2 ), NFData ( ℝ l2 )
       , Representable Double ( ℝ l2 )
       , Representable Double ( ℝ l12 )
       , Show ( ℝ l12 ), NFData ( ℝ l12 )
       )
    => { lg1 :: Proxy# l1
       , lg2 :: Proxy# l2
       , lg12 :: Proxy# l12
       , project1 :: Record r1 -> Record r1r2
          -- ^ project out fields present in both rows
          -- (linear non-decreasing mapping)
       , project2 :: Record r2 -> Record r1r2
          -- ^ project out fields present in both rows
          -- (linear non-decreasing mapping)
       , inject1  :: Record r1 -> Record r1r2 -> Record r1
          -- ^ overrides the components of the first record with the second
          -- (linear non-decreasing mapping in its second argument)
       , inject2  :: Record r2 -> Record r1r2 -> Record r2
          -- ^ overrides the components of the first record with the second
          -- (linear non-decreasing mapping in its second argument)
       } -> Intersection r1 r2

{-# INLINE doIntersection #-}
doIntersection
  :: forall r1 r2 kont
  .  ( KnownSymbols r1, KnownSymbols r2
     )
  => ( forall r1r2 l12 l1 l2.
        ( l1 ~ Length r1, l2 ~ Length r2
        , l1 <= 4, l2 <= 4, l12 <= 4
        , r1r2 ~ Intersect r1 r2
        , KnownSymbols r1r2, l12 ~ Length r1r2
        , l12 <= l1, l12 <= l2
        , Representable Double ( ℝ l1 )
        , Show ( ℝ l1 ), NFData ( ℝ l1 )
        , Show ( ℝ l2 ), NFData ( ℝ l2 )
        , Representable Double ( ℝ l2 )
        , Representable Double ( ℝ l12 )
        , Show ( ℝ l12 ), NFData ( ℝ l12 )
        )
      => Proxy# l1 -> Proxy# l2 -> Proxy# r1r2 -> Proxy# l12 -> Vec l12 ( Fin l1 ) -> Vec l12 ( Fin l2 ) -> kont )
  -> kont
doIntersection k =
  proveDicts @r1 $ \ ( px_l1 :: Proxy# l1 ) ->
  proveDicts @r2 $ \ ( px_l2 :: Proxy# l2 ) ->
  case knownSymbols @r1 `intersectLists` knownSymbols @r2 of

    [ ]
      | Refl <- ( unsafeCoerce Refl :: '[ ] :~: Intersect r1 r2 )
      , Refl <- ( unsafeCoerce Refl :: ( 0 <=? l1 ) :~: True )
      , Refl <- ( unsafeCoerce Refl :: ( 0 <=? l2 ) :~: True )
      -> k @'[ ] px_l1 px_l2 proxy# ( proxy# :: Proxy# 0 )
           ( Vec [] )
           ( Vec [] )

    [   ( f1, r1_i1, r2_i1 ) ]
      | SomeSymbol @f1 _     <- someSymbolVal ( Text.unpack f1 )
      , Refl <- ( unsafeCoerce Refl :: '[ f1 ] :~: Intersect r1 r2 )
      , Refl <- ( unsafeCoerce Refl :: ( 1 <=? l1 ) :~: True )
      , Refl <- ( unsafeCoerce Refl :: ( 1 <=? l2 ) :~: True )
      -> k @'[ f1 ] px_l1 px_l2 proxy# ( proxy# :: Proxy# 1 )
           ( Vec [ Fin r1_i1 ] )
           ( Vec [ Fin r2_i1 ] )

    [   ( f1, r1_i1, r2_i1 )
      , ( f2, r1_i2, r2_i2 ) ]
      | SomeSymbol @f1 _     <- someSymbolVal ( Text.unpack f1 )
      , SomeSymbol @f2 _     <- someSymbolVal ( Text.unpack f2 )
      , Refl <- ( unsafeCoerce Refl :: '[ f1, f2 ] :~: Intersect r1 r2 )
      , Refl <- ( unsafeCoerce Refl :: ( 2 <=? l1 ) :~: True )
      , Refl <- ( unsafeCoerce Refl :: ( 2 <=? l2 ) :~: True )
      -> k @'[ f1, f2 ] px_l1 px_l2 proxy# ( proxy# :: Proxy# 2 )
           ( Vec [ Fin r1_i1, Fin r1_i2 ] )
           ( Vec [ Fin r2_i1, Fin r2_i2 ] )

    [   ( f1, r1_i1, r2_i1 )
      , ( f2, r1_i2, r2_i2 )
      , ( f3, r1_i3, r2_i3 ) ]
      | SomeSymbol @f1 _     <- someSymbolVal ( Text.unpack f1 )
      , SomeSymbol @f2 _     <- someSymbolVal ( Text.unpack f2 )
      , SomeSymbol @f3 _     <- someSymbolVal ( Text.unpack f3 )
      , Refl <- ( unsafeCoerce Refl :: '[ f1, f2, f3 ] :~: Intersect r1 r2 )
      , Refl <- ( unsafeCoerce Refl :: ( 3 <=? l1 ) :~: True )
      , Refl <- ( unsafeCoerce Refl :: ( 3 <=? l2 ) :~: True )
      -> k @'[ f1, f2, f3 ] px_l1 px_l2 proxy# ( proxy# :: Proxy# 3 )
           ( Vec [ Fin r1_i1, Fin r1_i2, Fin r1_i3 ] )
           ( Vec [ Fin r2_i1, Fin r2_i2, Fin r2_i3 ] )

    [   ( f1, r1_i1, r2_i1 )
      , ( f2, r1_i2, r2_i2 )
      , ( f3, r1_i3, r2_i3 )
      , ( f4, r1_i4, r2_i4 ) ]
      | SomeSymbol @f1 _     <- someSymbolVal ( Text.unpack f1 )
      , SomeSymbol @f2 _     <- someSymbolVal ( Text.unpack f2 )
      , SomeSymbol @f3 _     <- someSymbolVal ( Text.unpack f3 )
      , SomeSymbol @f4 _     <- someSymbolVal ( Text.unpack f4 )
      , Refl <- ( unsafeCoerce Refl :: [ f1, f2, f3, f4 ] :~: Intersect r1 r2 )
      , Refl <- ( unsafeCoerce Refl :: ( 4 <=? l1 ) :~: True )
      , Refl <- ( unsafeCoerce Refl :: ( 4 <=? l2 ) :~: True )
      -> k @'[ f1, f2, f3, f4 ] px_l1 px_l2 proxy# ( proxy# :: Proxy# 4 )
           ( Vec [ Fin r1_i1, Fin r1_i2, Fin r1_i3, Fin r1_i4 ] )
           ( Vec [ Fin r2_i1, Fin r2_i2, Fin r2_i3, Fin r2_i4 ] )

    other ->
      error $ "Intersection not defined in dimension " ++ show ( length other )

{-# INLINE proveDicts #-}
proveDicts
  :: forall r kont
  . ( KnownSymbols r ) =>
    ( forall l
      . ( Length r ~ l, l <= 4
        , Representable Double ( ℝ l )
        , Show ( ℝ l ), NFData ( ℝ l )
        )
      => Proxy# l -> kont
      )
  -> kont
proveDicts k =
  case knownSymbols @r of
    [ ]
      | Refl <- ( unsafeCoerce Refl :: Length r :~: 0 )
      -> k @0 proxy#
    [ _ ]
      | Refl <- ( unsafeCoerce Refl :: Length r :~: 1 )
      -> k @1 proxy#
    [ _, _ ]
      | Refl <- ( unsafeCoerce Refl :: Length r :~: 2 )
      -> k @2 proxy#
    [ _, _, _ ]
      | Refl <- ( unsafeCoerce Refl :: Length r :~: 3 )
      -> k @3 proxy#
    [ _, _, _, _ ]
      | Refl <- ( unsafeCoerce Refl :: Length r :~: 4 )
      -> k @4 proxy#
    _ ->
      error "proveDicts: instances not defined in dimension >= 5"

------
-- Functions for intersection.

intersectLists :: forall k. Eq k => [ k ] -> [ k ] -> [ ( k, Word, Word ) ]
intersectLists = go 1
  where
    go :: Word -> [ k ] -> [ k ] -> [ ( k, Word, Word ) ]
    go _ [] _
      = []
    go i ( k : ks ) r
      | Just j <- findIndex ( k == ) r
      = ( k, i, fromIntegral j + 1 ) : go ( i + 1 ) ks r
      | otherwise
      = go ( i + 1 ) ks r

type Intersect :: [ k ] -> [ k ] -> [ k ]
type family Intersect r1 r2 where
  Intersect '[]         _ = '[]
  Intersect ( k ': ks ) r = DoIntersection k ks r ( Elem k r )

type DoIntersection :: k -> [ k ] -> [ k ] -> Bool -> [ k ]
type family DoIntersection k ks r mb_j where
  DoIntersection _ ks r False = Intersect ks r
  DoIntersection k ks r True  = k ': Intersect ks r

type Elem :: k -> [ k ] -> Bool
type family Elem k ks where
  Elem _ '[]         = False
  Elem k ( k ': _  ) = True
  Elem k ( _ ': ks ) = Elem k ks

--------------------------------------------------------------------------------
-- Union of two records.

data Union r1 r2 where
  Union
    :: forall r1r2 r1 r2 l12
    .  ( l12 ~ Length r1r2
       , KnownSymbols r1r2
       , Representable Double ( ℝ l12 )
       , Show ( ℝ l12 )
       , NFData ( ℝ l12 )
       , DiffInterp 2 NonIV l12
       , DiffInterp 3 IsIV l12
       )
    => { unionWith :: ( Double -> Double -> Double )
                   -> Record r1 -> Record r2 -> Record r1r2
         -- ^ union of two records
       } -> Union r1 r2

{-# INLINE union #-}
union :: forall r1 r2 l1 l2
          . ( KnownSymbols r1, KnownSymbols r2
            , l1 ~ Length r1, l2 ~ Length r2
            , Representable Double ( ℝ l1 )
            , Representable Double ( ℝ l2 )
            , Show ( ℝ l1 )
            , Show ( ℝ l2 )
            , NFData ( ℝ l2 )
            , DiffInterp 2 NonIV l2
            , DiffInterp 3 IsIV  l2
            )
          => Union r1 r2
union
  -- Shortcut when the two rows are equal.
  | Just Refl <- eqT @r1 @r2
  = Union { unionWith = \ f ( MkR l ) ( MkR r ) -> MkR @Symbol @r1 ( tabulate $ \ i -> f ( index l i ) ( index r i ) ) }
  | otherwise
  = doUnion @r1 @r2 \ ( _ :: Proxy# r1r2 ) idxs ->
      let
        unionWith :: ( Double -> Double -> Double )
                  -> Record r1 -> Record r2 -> Record r1r2
        unionWith f = \ ( MkR r1 ) ( MkR r2 ) -> MkR $ tabulate $ \ i ->
          case idxs ! i of
            InBoth i1 i2 -> f ( index r1 i1 ) ( index r2 i2 )
            InL    i1    -> index r1 i1
            InR    i2    -> index r2 i2
      in Union { unionWith }


data LR l r
  = InBoth !l !r
  | InL !l
  | InR !r
  deriving stock ( Eq, Show )

bimapLR :: ( l1 -> l2 ) -> ( r1 -> r2 ) -> LR l1 r1 -> LR l2 r2
bimapLR f g = \case
  InL    l   -> InL    ( f l )
  InR      r -> InR            ( g r )
  InBoth l r -> InBoth ( f l ) ( g r )

instance ( Ord l, Ord r ) => Ord ( LR l r ) where
  compare a b = compare ( getL a, getR a ) ( getL b, getR b )
    where
      getL :: LR l r -> Maybe l
      getL ( InL l ) = Just l
      getL ( InBoth l _ ) = Just l
      getL ( InR {} ) = Nothing
      getR :: LR l r -> Maybe r
      getR ( InL {} ) = Nothing
      getR ( InBoth _ r ) = Just r
      getR ( InR r ) = Just r

{-# INLINE doUnion #-}
doUnion
  :: forall r1 r2 l1 l2 kont
  .  ( KnownSymbols r1, KnownSymbols r2
     , l1 ~ Length r1, l2 ~ Length r2
     )
  => ( forall r1r2 l12.
        ( KnownSymbols r1r2, l12 ~ Length r1r2
        , DiffInterp 2 NonIV l12
        , DiffInterp 3 IsIV  l12
        , Representable Double ( ℝ l12 )
        , Show ( ℝ l12 )
        , NFData ( ℝ l12 )
        )
      => Proxy# r1r2 -> Vec l12 ( LR ( Fin l1 ) ( Fin l2 ) ) -> kont )
  -> kont
doUnion k =
  case knownSymbols @r1 `unionLists` knownSymbols @r2 of

    [ ]
      -> k @'[ ] proxy#
           ( Vec [] )

    [ ( f1, i1 ) ]
      | SomeSymbol @f1 _     <- someSymbolVal ( Text.unpack f1 )
      -> k @'[ f1 ] proxy#
           ( Vec $ map ( bimapLR Fin Fin ) [ i1 ] )

    [ ( f1, i1 ), ( f2, i2 ) ]
      | SomeSymbol @f1 _     <- someSymbolVal ( Text.unpack f1 )
      , SomeSymbol @f2 _     <- someSymbolVal ( Text.unpack f2 )
      -> k @'[ f1, f2 ]proxy#
           ( Vec $ map ( bimapLR Fin Fin ) [ i1, i2 ] )

    [ ( f1, i1 ), ( f2, i2 ), ( f3, i3 ) ]
      | SomeSymbol @f1 _     <- someSymbolVal ( Text.unpack f1 )
      , SomeSymbol @f2 _     <- someSymbolVal ( Text.unpack f2 )
      , SomeSymbol @f3 _     <- someSymbolVal ( Text.unpack f3 )
      -> k @'[ f1, f2, f3 ] proxy#
           ( Vec $ map ( bimapLR Fin Fin ) [ i1, i2, i3 ] )

    [ ( f1, i1 ), ( f2, i2 ), ( f3, i3 ), ( f4, i4 ) ]
      | SomeSymbol @f1 _     <- someSymbolVal ( Text.unpack f1 )
      , SomeSymbol @f2 _     <- someSymbolVal ( Text.unpack f2 )
      , SomeSymbol @f3 _     <- someSymbolVal ( Text.unpack f3 )
      , SomeSymbol @f4 _     <- someSymbolVal ( Text.unpack f4 )
      -> k @'[ f1, f2, f3, f4 ] proxy#
           ( Vec $ map ( bimapLR Fin Fin ) [ i1, i2, i3, i4 ] )

    other -> error $ "Union not defined in dimension " ++ show ( length other )

unionLists :: forall k. Ord k => [ k ] -> [ k ] -> [ ( k, LR Word Word ) ]
unionLists l r = sortBy ( comparing snd )
               $ Map.toList
               $ Map.unionWith f
                  ( Map.fromList [ ( k, InL i ) | k <- l | i <- [ 1 .. ] ] )
                  ( Map.fromList [ ( k, InR i ) | k <- r | i <- [ 1 .. ] ] )
  where
    f :: LR Word Word -> LR Word Word -> LR Word Word
    f ( InL i ) ( InR j ) = InBoth i j
    f _ _ = error "unionList: internal error"