gmap.hs

gmap.hs

{-#  LANGUAGE AllowAmbiguousTypes    #-}
{-#  LANGUAGE DefaultSignatures      #-}
{-#  LANGUAGE DeriveAnyClass         #-}
{-#  LANGUAGE DeriveGeneric          #-}
{-#  LANGUAGE FlexibleContexts       #-}
{-#  LANGUAGE FlexibleInstances      #-}
{-#  LANGUAGE FunctionalDependencies #-}
{-#  LANGUAGE GADTs                  #-}
{-#  LANGUAGE InstanceSigs           #-}
{-#  LANGUAGE KindSignatures         #-}
{-#  LANGUAGE Rank2Types             #-}
{-#  LANGUAGE ScopedTypeVariables    #-}
{-#  LANGUAGE StandaloneDeriving     #-}
{-#  LANGUAGE TypeApplications       #-}
{-#  LANGUAGE TypeFamilies           #-}
{-#  LANGUAGE TypeInType             #-}
{-#  LANGUAGE TypeOperators          #-}
{-#  LANGUAGE UndecidableInstances   #-}

module GenericN where

import Data.Kind
import GHC.Generics
import GHC.TypeLits
import Data.Coerce


mmap :: forall a. Maybe a -> Either String a
mmap Nothing = Left "sad"
mmap (Just x) = Right x

test1 = map_k' @'[Maybe ~> Either String, Int -> String] @Either' (NatTrans mmap :> show :> HNil) (Right' 20)
--Right' "20"

test2 = map_k' @'[Maybe ~> Either String, Int -> String] @Either' (NatTrans mmap :> show :> HNil) (Left' (Just 20))
--Left' (Right "20")

type Test2 = (Doms_k '[Type -> Type, Type] '[Maybe ~> Maybe, Int -> String] :: HList '[Type -> Type, Type])
-- = Maybe ':> (Int ':> 'HNil)


data IntList1 = Empty1 | Cons1 Int IntList1
data IntPair = IntPair Int Int

data IntList2 a = Empty2 | Cons2 Int a (IntList2 a)

data T a = T a Int
  deriving Generic

newtype Param1 (n :: Nat) a = Param1 a

type family Indexed1 (t :: k) (i :: Nat) :: k where
  Indexed1 (t a) i = Indexed1 t (i + 1) (Param1 i a)
  Indexed1 t _ = t

data Foo f = Foo (f Int)
  deriving Generic

type family Param :: Nat -> k where

type family Extract (p :: k) :: Nat where
  Extract (_ n) = n

type family Indexed (t :: k) (i :: Nat) :: k where
  Indexed (t a) i = Indexed t (i + 1) (Param i)
  Indexed t _     = t

newtype Rec (p :: Type) a x = Rec (K1 R a x)
  deriving Show

type family Zip (a :: Type -> Type) (b :: Type -> Type) :: Type -> Type where
  Zip (M1 mt m s) (M1 mt m t)
    = M1 mt m (Zip s t)
  Zip (l :+: r) (l' :+: r')
    = Zip l l' :+: Zip r r'
  Zip (l :*: r) (l' :*: r')
    = Zip l l' :*: Zip r r'
  Zip (Rec0 p) (Rec0 a)
    = Rec p a

class
  ( Coercible (Rep a) (RepN a)
  , Generic a
  ) => GenericN (a :: Type) where
  type family RepN (a :: Type) :: Type -> Type
  type instance RepN a = Zip (Rep (Indexed a 0)) (Rep a)
  toN :: RepN a x -> a
  fromN :: a -> RepN a x

instance
  ( Coercible (Rep a) (RepN a)
  , Generic a
  ) => GenericN a where

  toN :: forall x. RepN a x -> a
  toN   = coerce (to :: Rep a x -> a)
  fromN :: forall x. a -> RepN a x
  fromN = coerce (from :: a -> Rep a x)

class Map (funs :: [Type]) f where
  map' :: HList funs -> App f (Doms funs) -> App f (Cods funs)

data HList :: [Type] -> Type where
  HNil :: HList '[]
  (:>) :: a -> HList as -> HList (a ': as)

infixr 5 :>

class Lookup (i :: Nat) (p :: [Type]) (a :: Type) | p i -> a where
  hlookup :: HList p -> a

instance Lookup 0 (t ': ts) t where
  hlookup (c :> _) = c

instance {-#  OVERLAPPABLE  #-}
  Lookup (n - 1) ts a => Lookup n (t ': ts) a where
  hlookup (_ :> cs) = hlookup @(n - 1) cs

type family Doms (t :: [Type]) :: [Type] where
  Doms '[d -> _] = '[d]
  Doms ((d -> _) ': xs) = d ': Doms xs

type family Cods (t :: [Type]) :: [Type] where
  Cods '[_ -> c] = '[c]
  Cods ((_ -> c) ': xs) = c ': Cods xs

type family App (t :: k) (ts :: [Type]) :: Type where
  App t '[x] = t x
  App t (x ': xs) = App (t x) xs

class Reverse xs ys | xs -> ys where
  hreverse :: HList xs -> HList ys

instance Reverse '[] '[] where
  hreverse = id

instance
  ( Reverse xs ys
  , ys' ~ (ys ++ '[x])
  , Append x ys
  ) => Reverse (x ': xs) ys' where
  hreverse (x :> xs) = happend x (hreverse @xs @ys xs)

class Append x xs where
  happend :: x -> HList xs -> HList (xs ++ '[x])

instance Append x '[] where
  happend x _ = x :> HNil

instance Append x xs => Append x (y ': xs) where
  happend x (y :> xs) = y :> (happend x xs)

type family xs ++ ys where
  '[] ++ ys = ys
  (x ': xs) ++ ys = x ': (xs ++ ys)

class GMap (s :: Type -> Type) (t :: Type -> Type) (ts :: [Type]) where
  gmap :: HList ts -> s x -> t x

instance GMap s t ts
  => GMap (M1 mt m s) (M1 mt m t) ts where
  gmap fs (M1 m) = M1 (gmap fs m)

instance
  ( GMap l l' ts
  , GMap r r' ts
  ) => GMap (l :+: r) (l' :+: r') ts where
  gmap fs (L1 l) = L1 (gmap fs l)
  gmap fs (R1 r) = R1 (gmap fs r)

instance
  ( GMap l l' ts
  , GMap r r' ts
  ) => GMap (l :*: r) (l' :*: r') ts where
  gmap fs (l :*: r) = gmap fs l :*: gmap fs r

instance Lookup n ts (a -> b)
  => GMap (Rec (p (n :: Nat)) a) (Rec (p n) b) ts where
  gmap fs (Rec (K1 x)) = Rec (K1 (hlookup @n fs x))

instance {-#  OVERLAPPABLE  #-}
  ( GenericN b
  , GenericN a
  , is ~ Collect s
  , Adapt is ts ts'
  , GMap (RepN a) (RepN b) ts'
  ) => GMap (Rec s a) (Rec t b) ts where
  gmap fs (Rec (K1 x))
    = Rec (K1 (toN $ gmap (adapt @is fs) (fromN x)))

type family Collect (a :: k) :: [Nat] where
  Collect (c (p n)) = n ': Collect c
  Collect (c a) = Collect c
  Collect c = '[]

class Adapt (is :: [Nat]) xs ys | is xs -> ys where
  adapt :: HList xs -> HList ys

instance Adapt '[] xs '[] where
  adapt _ = HNil

instance
  ( Lookup n xs y
  , Adapt ns xs ys
  ) => Adapt (n ': ns) xs (y ': ys) where
  adapt xs = hlookup @n xs :> adapt @ns xs

instance
  ( s ~ App f (Doms funs)
  , t ~ App f (Cods funs)
  , GenericN s
  , GenericN t
  , Reverse funs funs'
  , GMap (RepN s) (RepN t) funs'
  ) => Map funs f where
  map' funs s = toN (gmap (hreverse funs) (fromN s))

class Bifunctor p where
  bimap :: (a -> b) -> (c -> d) -> p a c -> p b d

  default bimap ::
    Map '[a -> b, c -> d] p
     => (a -> b) -> (c -> d) -> p a c -> p b d
  bimap f g = map' @_ @p (f :> g :> HNil)

deriving instance Bifunctor Either

data WTree a b
  = Leaf a
  | Leafo (a, a)
  | Fork (WTree a b) (WTree a b)
  | WithWeight (WTree a b) b
  deriving (Generic, Bifunctor)

data Weird a b
  = W1 a
  | W2 a b
  | W3 (Weird a b) (Weird b a)
  | W4 (Weird a a) (Weird b b)
  deriving (Generic, Bifunctor)
  --  W4 (Weird (Weird a b) (Weird a b)) <- this diverges now

data GTree f a = GTree a (f (GTree f a))
  deriving (Generic)

deriving instance (Show a, Show (f (GTree f a))) => Show (GTree f a)

gtreemap :: Functor f => (forall a. f a -> g a) -> GTree f a -> GTree g a
gtreemap f (GTree a ts) = GTree a (f (fmap (gtreemap f) ts))

type family Params' (a :: k) :: HList x where
  Params' (c a) = a ':> Params' c
  Params' _ = 'HNil

type family ParamKinds (a :: k) :: [Type] where
  ParamKinds (f (_ :: k)) = k ': ParamKinds f
  ParamKinds _     = '[]

type Params a = (Params' a :: HList (ParamKinds a))

class GMap_k (s :: Type -> Type) (t :: Type -> Type) (ts :: [Type]) where
  gmap_k :: HList ts -> s x -> t x

instance GMap_k s t ts
  => GMap_k (M1 mt m s) (M1 mt m t) ts where
  gmap_k fs (M1 m) = M1 (gmap_k fs m)

instance
  ( GMap_k l l' ts
  , GMap_k r r' ts
  ) => GMap_k (l :+: r) (l' :+: r') ts where
  gmap_k fs (L1 l) = L1 (gmap_k fs l)
  gmap_k fs (R1 r) = R1 (gmap_k fs r)

instance
  ( GMap_k l l' ts
  , GMap_k r r' ts
  ) => GMap_k (l :*: r) (l' :*: r') ts where
  gmap_k fs (l :*: r) = gmap_k fs l :*: gmap_k fs r

instance Lookup n ts (f ~> g)
  => GMap_k (Rec (p n a) (f a)) (Rec (p n a) (g a)) ts where
  gmap_k fs (Rec (K1 x)) = Rec (K1 (unFun (hlookup @n fs) x))

instance Lookup n ts (a -> b)
  => GMap_k (Rec (p n) a) (Rec (p n) b) ts where
  gmap_k fs (Rec (K1 x)) = Rec (K1 (hlookup @n fs x))

unFun :: (f ~> g) -> f a -> g a
unFun (NatTrans t) = t

newtype f ~> g
  = NatTrans (forall a. FunAt (f a) (g a))

type family FunAt (a :: k) (b :: k) :: Type where
  FunAt a b = a -> b
  FunAt (a :: Type -> k) b = a ~> b

-- 
data MyEither a b = MyLeft a | MyRight b

emap :: forall a b. Either a b -> MyEither a b
emap (Left a) = MyLeft a
emap (Right a) = MyRight a

-- etrans :: Either ~> MyEither -- inferred
etrans = NatTrans @Either (NatTrans emap)

data T1 a b c = T1 a b c
data T2 a b c = T2 a b c

tmap :: forall a b c. T1 a b c -> T2 a b c
tmap (T1 a b c) = T2 a b c

-- ttrans :: T1 ~> T2 -- inferred
ttrans = NatTrans @T1 (NatTrans (NatTrans tmap))

class Map_k (funs :: [Type]) f where
  map_k'
    :: HList funs
    -> App_k f (Doms_k (ParamList (KindOf f)) funs)
    -> App_k f (Cods_k (ParamList (KindOf f)) funs)

instance
  ( s ~ App_k f (Doms_k (ParamList (KindOf f)) funs)
  , t ~ App_k f (Cods_k (ParamList (KindOf f)) funs)
  , GenericN s
  , GenericN t
  , Reverse funs funs'
  , GMap_k (RepN s) (RepN t) funs'
  ) => Map_k funs f where
  map_k' funs s = toN (gmap_k (hreverse funs) (fromN s))

type family ParamList (k :: Type) :: [Type] where
  ParamList (k -> Type) = '[k]
  ParamList (k -> j) = k ': ParamList j

type family KindOf (a :: k) :: Type where
  KindOf (a :: k) = k

data Either' f a = Left' (f Int) | Right' a
  deriving Generic
deriving instance (Show a, Show (f Int)) => Show (Either' f a)

type family Doms_k (k :: [Type]) (fs :: [Type]) :: HList k where
  Doms_k '[Type] '[d -> _] = d ':> 'HNil
  Doms_k '[Type -> k] '[d ~> _] = d ':> 'HNil
  Doms_k (Type ': ks) ((d -> _) ': xs) = d ':> Doms_k ks xs
  Doms_k ((Type -> k) ': ks) ((d ~> _) ': xs) = d ':> Doms_k ks xs

type family Cods_k (k :: [Type]) (fs :: [Type]) :: HList k where
  Cods_k '[Type] '[_ -> c] = c ':> 'HNil
  Cods_k '[Type -> k] '[_ ~> c] = c ':> 'HNil
  Cods_k (Type ': ks) ((_ -> c) ': xs) = c ':> Cods_k ks xs
  Cods_k ((Type -> k) ': ks) ((_ ~> c) ': xs) = c ':> Cods_k ks xs

type family App_k (t :: k) (ts :: HList ks) :: Type where
  App_k t (x ':> 'HNil) = t x
  App_k t (x ':> xs) = App_k (t x) xs