oneliner-repr.hs

{-# LANGUAGE RankNTypes, TypeOperators, DefaultSignatures, FlexibleContexts, DeriveGeneric, DeriveAnyClass #-}
import Generics.OneLiner
import Data.Profunctor
import Data.Functor.Contravariant
import GHC.Generics
import Control.Applicative
import Unsafe.Coerce (unsafeCoerce)

type Lens s t a b = forall f. Functor f => (a -> f b) -> s -> f t
type Key t = forall x. Lens (t x) (t x) x x

constLens :: x -> Lens s t a b -> x
constLens x _ = x

newtype Keyed s t a b = Keyed { getKeyed :: Lens s t a b -> b }
instance Profunctor (Keyed s t) where
  dimap f g (Keyed ix) = Keyed $ \l -> g (ix (l . (fmap g .) . (. f)))
instance GenericRecordProfunctor (Keyed s t) where
  unit = Keyed (constLens U1)
  mult (Keyed a) (Keyed b) = Keyed (\l -> a (l . fstl) :*: b (l . sndl))

fstl :: Lens ((a :*: b) x) ((c :*: b') x') (a x) (c x')
fstl f (a :*: b) = (\c -> c :*: unsafeCoerce b) <$> f a
sndl :: Lens ((a :*: b) x) ((a' :*: c) x') (b x) (c x')
sndl f (a :*: b) = (\c -> unsafeCoerce a :*: c) <$> f b

class Repr f where
  keyed :: (Lens s t a b -> b) -> Lens s t (f a) (f b) -> f b
  default keyed :: (ADTRecord1 f, Constraints1 f Repr) => (Lens s t a b -> b) -> Lens s t (f a) (f b) -> f b
  keyed f = getKeyed $ record1 (For :: For Repr) (\(Keyed g) -> Keyed $ keyed g) (Keyed f)

newtype Wrapped s t a b = Lens { runLens :: Lens s t a b }
newtype WrappedKey t = Key { runKey :: Key t }

repr :: Repr f => (Key f -> a) -> f a
repr f = keyed (\l -> f (runKey (unsafeCoerce (Lens l)))) id

index :: f a -> Key f -> a
index f l = getConst $ l Const f

data V3 a = V3 a a a deriving (Show, Generic1, Repr)

scale :: Num a => a -> V3 a -> V3 a
scale a v = repr (\l -> a * index v l)

diag :: V3 (V3 a) -> V3 a
diag m = repr (\l -> m `index` l `index` l)