Towards poly-kinded optics (fusing multiplate and lens)

PolyLens.hs

{-# language DataKinds #-}
{-# language PolyKinds #-}
{-# language RankNTypes #-}
{-# language ConstraintKinds #-}
{-# language DeriveFunctor #-}
{-# language GADTs #-}
{-# language TypeApplications #-}
{-# language ScopedTypeVariables #-}
{-# language FlexibleInstances #-}
{-# language FlexibleContexts #-}
{-# language DefaultSignatures #-}
{-# language DeriveAnyClass #-}
{-# language DerivingStrategies #-}
{-# language TypeOperators #-}
{-# language StandaloneKindSignatures #-}
{-# language BlockArguments #-}
{-# language LambdaCase #-}
{-# language RankNTypes #-}
{-# language TypeFamilies #-}
{-# language ImpredicativeTypes #-}
{-# language PatternSynonyms #-}
{-# language ViewPatterns #-}
{-# language ImportQualifiedPost #-}
{-# language TypeSynonymInstances #-}

-- polykinded lenses?

import Control.Applicative
import Control.Lens qualified as Lens
import Data.Bitraversable
import Data.Constraint
import Data.Functor
import Data.Functor.Compose
import Data.Functor.Identity
import Data.Functor.Product
import Data.Functor.Sum
import Data.Kind
import Data.Profunctor
import Data.Profunctor.Rep
import Data.Profunctor.Sieve
import Data.Profunctor.Traversing
import Data.Proxy
import Data.Traversable
import Data.Type.Equality
import GHC.Generics qualified as G
import GHC.Generics.Lens

--
-- It would be nice to be able to handle "multi-plate" style optics using the
-- same machinery as normal lenses.
-- 
-- The obvious first stab at this is something like:
--
-- @
-- type PLens s t a b = forall f. Functor f
--   => (forall i. (a i -> f (b i))) 
--   -> forall j. s j -> f (t j)
-- @
--
-- which is what I've used in the past, but then normal lenses, and 'poly' lenses
-- can't compose. You can of course then define Lens01 or Lens10 to go back and forth
-- but this invites the question of whether we can do better.
--
-- I really want to share 'f' between both kinds of lenses.

--------------------------------------------------------------------------------
-- * lifting profunctors to handle polymorphic functions and dictionaries
--------------------------------------------------------------------------------

-- |
-- The next idea is blocked by https://gitlab.haskell.org/ghc/ghc/-/issues/9269
--
-- @
-- type family BadK (f :: Type -> Type) (a :: k) (b :: k) :: Type where
--   BadK f a b = a -> f b
--   BadK f a b = forall i. a i -> f (b i)
-- @
--
-- which would let me use
--
-- @
-- type BadLens s t a b = forall f. Functor f => BadK f a b -> BadK f s t 
-- @
--
-- and this might be a bit lighter for lenses/traversals using VL style
-- but unlike lens, we need custom data types all over, so it probably
-- should maybe go for profunctor optics as well, so let's do that:

data family K (p :: Type -> Type -> Type) :: k -> k -> Type
newtype instance K p a b = K0 { runK0 :: p a b }
newtype instance K p a b = KS { runKS :: forall i. K p (a i) (b i) }
newtype instance K p a b = KD { runKD :: p (Dict a) (Dict b) }

runK1 :: K p a b -> forall i. p (a i) (b i)
runK1 f = runK0 (runKS f)

pattern K1 :: (forall i. p (a i) (b i)) -> K p a b
pattern K1 f <- (runK1 -> f) where
  K1 f = KS (K0 f)

runK2 :: K p a b -> forall i j. p (a i j) (b i j)
runK2 f = runK1 (runKS f)

pattern K2 :: (forall i j. p (a i j) (b i j)) -> K p a b
pattern K2 f <- (runK2 -> f) where
  K2 f = KS (K1 f)

--------------------------------------------------------------------------------
-- * polykinded optics
--------------------------------------------------------------------------------

type KOptic :: forall i j. (Type -> Type -> Type) -> i -> i -> j -> j -> Type
type KOptic p s t a b = K p a b -> K p s t
type KOptic' p s a = KOptic p s s a a

o1 :: KOptic p s t a b -> (forall i. p (a i) (b i)) -> (forall j. p (s j) (t j))
o1 l f = runK1 $ l $ K1 f

-- convert to/from a multi-plate style profunctor optic
pattern O1 :: ((forall i. p (a i) (b i)) -> (forall j. p (s j) (t j))) -> KOptic p s t a b
pattern O1 l <- (o1 -> l) where
  O1 l = \(K1 f) -> K1 (l f)
{-# complete O1 #-}

-- convert to/from normal profunctor optics
pattern O0 :: (p a b -> p s t) -> KOptic p s t a b
pattern O0 l' <- (\l -> runK0 . l . K0 -> l') where
  O0 l = \(K0 f) -> K0 (l f)
{-# complete O0 #-}

type KIso :: forall i j. i -> i -> j -> j -> Type
type KIso s t a b = forall p. Profunctor p => KOptic p s t a b 
type KIso' s a = KIso s s a a

type KLens :: forall i j. i -> i -> j -> j -> Type
type KLens s t a b = forall p. Strong p => KOptic p s t a b
type KLens' s a = KLens s s a a

type KPrism :: forall i j. i -> i -> j -> j -> Type
type KPrism s t a b = forall p. Choice p => KOptic p s t a b
type KPrism' s a = KPrism s s a a

type KTraversal :: forall i j. i -> i -> j -> j -> Type
type KTraversal s t a b = forall p. Traversing p => KOptic p s t a b
type KTraversal' s a = KTraversal s s a a

-- using this for now, because it is a little easier to think about than wander.
-- (Representable p, Applicative (Rep p))
type KTraversingish p = (Representable p, Applicative (Rep p))
type KTraversalish :: forall i j. i -> i -> j -> j -> Type
type KTraversalish s t a b = forall p. KTraversingish p => KOptic p s t a b
type KTraversalish' s a = KTraversalish s s a a

-- and this would imply that instead of Strong I should maybe use
type KLensish :: forall i j. i -> i -> j -> j -> Type
type KLensish s t a b = forall p. Representable p => KOptic p s t a b
type KLensish' s a = KLensish s s a a

-- this requires me to have "representable by an applicative functor" form KTraversals
vlo1 :: Representable p => ((forall i. p (a i) (b i)) -> (forall j. p (s j) (t j))) ->
       (forall i. a i -> Rep p (b i)) -> (forall j. s j -> Rep p (t j))
vlo1 l f = sieve $ l (tabulate f)

-- construct a van laarhoven-style (multi)-plate optic, given a representable profunctor
pattern VL1 :: Representable p =>
  ( (forall i. a i -> Rep p (b i)) ->
    (forall j. s j -> Rep p (t j))
  ) -> KOptic p s t a b
pattern VL1 f <- O1 (vlo1 -> f) where
  VL1 l = O1 \f -> tabulate $ l (sieve f)

--------------------------------------------------------------------------------
-- * tools for inspecting poly-kinds optics
--------------------------------------------------------------------------------

type N :: Type -> Type
data N n where
  N0 :: N Type
  NS :: N j -> N (i -> j) 
  ND :: N Constraint

class Poly n where
  poly :: N n

instance Poly Type where
  poly = N0

instance Poly j => Poly (i -> j) where
  poly = NS poly

instance Poly Constraint where
  poly = ND

kdimap :: forall k p (a::k) (b::k) (c::k) (d::k). (Poly k, Profunctor p) => K (->) a b -> K (->) c d -> K p b c -> K p a d
kdimap = kdimap' (poly @k) where
  kdimap' :: forall k p (a::k) (b::k) (c::k) (d::k). Profunctor p => N k -> K (->) a b -> K (->) c d -> K p b c -> K p a d
  kdimap' N0 (K0 f) (K0 g) (K0 h) = K0 $ dimap f g h
  kdimap' (NS k) (KS f) (KS g) (KS h) = KS $ kdimap' k f g h
  kdimap' ND (KD f) (KD g) (KD h) = KD $ dimap f g h

--------------------------------------------------------------------------------
-- * KFunctor
--------------------------------------------------------------------------------

class KFunctor (t :: i -> j) where
  kmap :: Mapping p => K p a b -> K p (t a) (t b)
  default kmap :: (t ~~ (s::Type -> Type), Functor s, Mapping p) => K p a b -> K p (t a) (t b)
  kmap (K0 f) = K0 $ map' f

instance KFunctor []

instance KFunctor Maybe

_KConst :: KIso (Const a) (Const b) a b
_KConst (K0 f) = K1 $ dimap getConst Const f

instance KFunctor Const where
  kmap = _KConst

kfirst :: KLens ((,) a) ((,) b) a b
kfirst (K0 f) = K1 $ first' f

instance KFunctor (,) where
  kmap = kfirst

instance KFunctor ((,) e)

kleft :: KPrism (Either a) (Either b) a b
kleft (K0 f) = K1 $ left' f

kright :: KPrism (Either c a) (Either c b) a b
kright (K0 f) = K0 $ right' f

instance KFunctor Either where
  kmap = kleft

instance KFunctor (Either e) where
  kmap = kright

-- composeF :: KIso (Compose f) (Compose g) f g
-- composeF (K1 f) = K2 $ dimap getCompose Compose f

_Compose :: KIso (Compose f g a) (Compose h i b) (f (g a)) (h (i b))
_Compose (K0 f) = K0 $ dimap getCompose Compose f

instance KFunctor Compose where
  -- kmap = composeF
  kmap (KS f) = KS $ KS $ _Compose f

instance Functor f => KFunctor (Compose f) where
  kmap (K1 f) = KS $ _Compose $ K0 $ map' f
  -- kmap (K1 f) = K1 $ dimap getCompose Compose $ map' f

_GCompose :: KIso ((G.:.:) f g a) ((G.:.:) h i b) (f (g a)) (h (i b))
_GCompose (K0 f) = K0 $ dimap G.unComp1 G.Comp1 f

instance KFunctor (G.:.:) where
  -- kmap (K1 f) = K2 $ dimap G.unComp1 G.Comp1 f
  kmap (KS f) = KS $ KS $ _GCompose f

instance Functor f => KFunctor ((G.:.:) f) where
  -- kmap (K1 f) = K1 $ dimap G.unComp1 G.Comp1 $ map' f
  kmap (K1 f) = KS $ _GCompose $ K0 $ map' f

_Pair :: KIso (Product f g a) (Product h i b) (f a, g a) (h b, i b)
_Pair (K0 f) = K0 $ dimap (\(Pair a b) -> (a,b)) (\(a,b) -> (Pair a b)) f

instance KFunctor Product where
  -- kmap (K1 f) = K2 $ dimap (\(Pair a b) -> (a,b)) (\(a,b) -> (Pair a b)) $ first' f
  kmap (K1 f) = KS $ KS $ _Pair $ K0 $ first' f

instance KFunctor (Product f) where
  -- kmap (K1 f) = K1 $ dimap (\(Pair a b) -> (a,b)) (\(a,b) -> (Pair a b)) $ second' f
  kmap (K1 f) = KS $ _Pair $ K0 $ second' f

_Sum :: KIso (Sum f g a) (Sum h i b) (Either (f a) (g a)) (Either (h b) (i b))
_Sum (K0 f) = K0 $ dimap hither yon f where
  hither (InL f) = Left f
  hither (InR g) = Right g
  yon (Left h) = InL h
  yon (Right i) = InR i

instance KFunctor Sum where
  kmap (K1 f) = KS $ KS $ _Sum $ K0 $ left' f

instance KFunctor (Sum f) where
  kmap (K1 f) = KS $ _Sum $ K0 $ right' f

_GSum :: KIso ((f G.:+: g) a) ((h G.:+: i) b) (Either (f a) (g a)) (Either (h b) (i b))
_GSum (K0 f) = K0 $ dimap hither yon f where
  hither (G.L1 f) = Left f
  hither (G.R1 g) = Right g
  yon (Left h) = G.L1 h
  yon (Right i) = G.R1 i

instance KFunctor (G.:+:) where
  kmap (K1 f) = KS $ KS $ _GSum $ K0 $ left' f

instance KFunctor ((G.:+:) f) where
  kmap (K1 f) = KS $ _GSum $ K0 $ right' f

instance KFunctor (G.:*:) where
  kmap (K1 f) = K2 $ dimap (\(a G.:*:b) -> (a,b)) (\(a,b) -> (a G.:*:b)) $ first' f

instance KFunctor ((G.:*:) f) where
  kmap (K1 f) = K1 $ dimap (\(a G.:*:b) -> (a,b)) (\(a,b) -> (a G.:*:b)) $ second' f

{-
instance Poly k => KFunctor (Proxy :: k -> Type) where
  kmap = kmap' (poly @k) where
    kmap' :: forall i (a::i) (b::i) p. Mapping p => N i -> K p a b -> K p (Proxy a) (Proxy b)
    kmap' N0 (K0 p) = K0 (map' p)
    -- kmap' (NS k) (KS p) = ...
-}


-- instance KFunctor (Const a :: i -> Type) -- generalize cleanly using Poly?
-- instance KFunctor (Const a :: (i -> Type) -> Type) where
{-
instance Poly k => KFunctor (Const a :: k -> Type) where
  kmap = go (poly @k) where
    go :: forall i (b::i) (c::i) (p::Type->Type->Type). Mapping p => N i -> K p b c -> K p (Const a b) (Const a c)
    go N0 (K0 f) = K0 (map' f)
    -- go NS 
    -- go ND (KD f) = K0 ...
-}
    

-- TODO: KFoldable

--------------------------------------------------------------------------------
-- * KTraversable
--------------------------------------------------------------------------------

class KFunctor t => KTraversable (t :: i -> j) where
  ktraverse :: Traversing p => K p a b -> K p (t a) (t b)
  -- GHC seems unwilling to let me learn i ~ Type, j ~ Type, and then use t directly?
  -- default ktraverse :: (i ~ Type, j ~ Type, Traversable t, Traversing p) => K p a b -> K p (t a) (t b)
  default ktraverse :: (t ~~ (s::Type -> Type), Traversable s, Traversing p) => K p a b -> K p (t a) (t b)
  ktraverse = K0 . traverse' . runK0

instance KTraversable []

instance KTraversable Maybe

instance KTraversable (,) where
  ktraverse (K0 f) = K1 $ wander Lens._1 f
  
instance KTraversable ((,) e)

instance KTraversable Either where
  ktraverse (K0 f) = K1 $ wander Lens._Left f

instance KTraversable (Either e)

instance KTraversable Compose where 
  ktraverse (KS f) = KS $ KS $ _Compose f

instance Traversable f => KTraversable (Compose f) where
  -- ktraverse (K1 f) = K1 $ dimap getCompose Compose $ traverse' f
  ktraverse (K1 f) = KS $ _Compose $ K0 $ traverse' f

instance KTraversable Sum where
  -- ktraverse (KS f) = KS $ _Sum . kleft f
  ktraverse (K1 f) = KS $ KS $ _Sum $ K0 $ left' f

instance KTraversable (Sum f) where
  ktraverse (K1 f) = KS $ _Sum $ K0 $ right' f

instance KTraversable Product where
  ktraverse (K1 f) = KS $ KS $ _Pair $ K0 $ first' f

instance KTraversable (Product f) where
  ktraverse (K1 f) = KS $ _Pair $ K0 $ second' f

instance KTraversable (G.:.:) where
  ktraverse (KS f) = KS $ KS $ _GCompose f

instance Traversable f => KTraversable ((G.:.:) f) where
  ktraverse (K1 f) = KS $ _GCompose $ K0 $ traverse' f

--------------------------------------------------------------------------------
-- * Flailing aound trying to get a nice version of Bazaar
--------------------------------------------------------------------------------

type KBazaar :: k -> k -> Type -> Type
newtype KBazaar a b t = KBazaar 
  { runKBazaar :: forall f. Applicative f => K (Star f) a b -> f t 
  }
  deriving stock Functor
  deriving anyclass KFunctor

--class Corepresentable p => Sellable p w | w -> p where
--  sell :: p a (w a b b)

sell :: a -> KBazaar a b b
sell a = KBazaar \(K0 (Star k)) -> k a

instance Applicative (KBazaar a b) where
  pure a = KBazaar $ \_ -> pure a
  mf <*> ma = KBazaar $ \k -> runKBazaar mf k <*> runKBazaar ma k
  
type KBaz :: Type -> k -> k -> Type
newtype KBaz t b a = KBaz { runKBaz :: forall f. Applicative f => K (Star f) a b -> f t }

instance Foldable (KBaz t b) where
  foldMap = foldMapDefault

instance Functor (KBaz t b) where
  fmap = fmapDefault

instance Traversable (KBaz t b) where
  traverse f bz = fmap (\m -> KBaz (runKBazaar m)) . getCompose . runKBaz bz $ K0 $ Star \x -> Compose $ sell <$> f x

sold1 :: forall i (t :: Type) (a :: i -> Type). KBaz t a a -> t
sold1 m = runIdentity $ runKBaz m $ K1 (Star Identity)

{-
-- dead end?
kwander :: forall k p (s::k) (t::k) (a::k) (b::k). (Poly k, Traversing p) => (forall f. Applicative f => K (Star f) a b -> K (Star f) s t) -> K p a b -> K p s t
kwander = kwander' (poly @k) where
  kwander' :: forall k p (s::k) (t::k) (a::k) (b::k). Traversing p => N k -> (forall f. Applicative f => K (Star f) a b -> K (Star f) s t) -> K p a b -> K p s t
  kwander' N0 l (K0 p) = K0 $ wander (\f s -> runStar (runK0 (l $ K0 $ Star f)) s) p
  kwander' ND l (KD p) = KD $ wander (\f s -> runStar (runKD (l $ KD $ Star f)) s) p
  -- stuck:
  --kwander' (NS k) l (KS p) = KS $ kwander' k (go l) p where
  --  go :: (K (Star f) a b -> K (Star f) s t) -> (forall i. K (Star f) (a i) (b i)) -> forall j. K (Star f) (s j) (t j)
  --  go = undefined
-}

--wander1 :: Traversing p => (forall f. Applicative f => (forall i. a i -> f (b i)) -> forall j. s j -> f (t j)) -> K p a b -> K p s t
--wander1 t (K1 f) = K1 $ undefined

--------------------------------------------------------------------------------
-- * Polykinded Plate
--------------------------------------------------------------------------------

--instance Poly k => KFunctor (KBaz @k t b) where

data Expr a
  = Var a
  | App (Expr a) (Expr a)
  | Lam (Expr (Maybe a))

plate :: KTraversalish' Expr Expr
plate = VL1 \f -> \case -- works if we go with the 'representable by Applicative' notion of traversals
-- plate = wander1 \f -> \case -- needs me to figure out a good wander1 equivalent
  Var a -> pure $ Var a
  App l r -> App <$> f l <*> f r
  Lam b -> Lam <$> f b

-- note: by fixing ourselves to profunctors on Hask, and lifting via K, we can
-- have arguments with different index structures and still compose, e.g.
plate0 :: KTraversalish' (Expr i) Expr
plate0 = runKS . plate

type A :: forall k. k -> Type
data A a where
  A0 :: { runA0 :: x } -> A x
  AD :: Dict x -> A x
  AS :: forall i (x :: i) j (a :: i -> j). A (a x) -> A a

pattern A1 :: a i -> A a
pattern A1 x = AS (A0 x)

pattern A2 :: a i j -> A a
pattern A2 x = AS (A1 x)

type An = A

-- these should allow for a generic implementation of view

-- view :: KOptic' (Star (Const (An a))) s a -> An s -> An a -- ?