elim.hs

{-# LANGUAGE LambdaCase, BlockArguments, EmptyCase #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE TypeSynonymInstances, TypeOperators, FlexibleInstances, FlexibleContexts, UndecidableInstances #-}
{-# LANGUAGE TypeFamilies, GADTs, MultiParamTypeClasses, FunctionalDependencies #-}
{-# LANGUAGE DataKinds, PolyKinds, AllowAmbiguousTypes, TypeApplications, ScopedTypeVariables #-}

import GHC.Generics
import Unsafe.Coerce (unsafeCoerce)
import Data.Type.Equality


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

type family Map f xs where
  Map f '[] = '[]
  Map f (x:xs) = f x : Map f xs

mapIsLinear :: forall f g a b. f (Map g a :++: Map g b) -> f (Map g (a :++: b))
mapIsLinear = unsafeCoerce


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


data Prod xs where
  ProdNil :: Prod '[]
  ProdCons :: x -> Prod xs -> Prod (x:xs)

concatProd :: Prod xs -> Prod ys -> Prod (xs :++: ys)
concatProd ProdNil ys = ys
concatProd (ProdCons x xs) ys = ProdCons x (concatProd xs ys)

reverseProd :: Prod xs -> Prod (Reverse xs)
reverseProd ProdNil = ProdNil
reverseProd (ProdCons x xs) = concatProd (reverseProd xs) (ProdCons x ProdNil)


data Sum xs where
  SumHere :: x -> Sum (x:ys)
  SumThere :: Sum xs -> Sum (x:xs)

expandSumR :: forall xs ys. Sum xs -> Sum (xs :++: ys)
expandSumR (SumHere x) = SumHere x
expandSumR (SumThere x) = SumThere (expandSumR @_ @ys x)

class ExpandSumL xs ys where  
  expandSumL :: Sum ys -> Sum (xs :++: ys)
instance ExpandSumL '[] ys where
  expandSumL x = x
instance ExpandSumL xs ys => ExpandSumL (x:xs) ys where
  expandSumL x = SumThere (expandSumL @xs x)


class CTreeToSum rep where
  type CTreeSum rep :: [[*]]
  ctreeToSum :: rep x -> Sum (Map Prod (CTreeSum rep))

instance (CTreeToSum l, CTreeToSum r, ExpandSumL (Map Prod (CTreeSum l)) (Map Prod (CTreeSum r))) => CTreeToSum (l :+: r) where
  type CTreeSum (l :+: r) = CTreeSum l :++: CTreeSum r
  ctreeToSum (L1 x) = mapIsLinear @Sum @Prod @(CTreeSum l) @(CTreeSum r) $ expandSumR @(Map Prod (CTreeSum l)) @(Map Prod (CTreeSum r)) (ctreeToSum x)
  ctreeToSum (R1 x) = mapIsLinear @Sum @Prod @(CTreeSum l) @(CTreeSum r) $ expandSumL @(Map Prod (CTreeSum l)) @(Map Prod (CTreeSum r)) (ctreeToSum x)

instance STreeToProd sels => CTreeToSum (C1 meta sels) where
  type CTreeSum (M1 C meta sels) = '[STreeProd sels]
  ctreeToSum (M1 x) = SumHere (streeToProd x)


class STreeToProd rep where
  type STreeProd rep :: [*]
  streeToProd :: rep x -> Prod (STreeProd rep)

instance (STreeToProd l, STreeToProd r) => STreeToProd (l :*: r) where
  type STreeProd (l :*: r) = STreeProd l :++: STreeProd r
  streeToProd (l :*: r) = concatProd (streeToProd l) (streeToProd r)

instance STreeToProd U1 where
  type STreeProd U1 = '[]
  streeToProd U1 = ProdNil

instance STreeToProd (S1 meta (Rec0 a)) where
  type STreeProd (S1 meta (Rec0 a)) = '[a]
  streeToProd (M1 (K1 x)) = ProdCons x ProdNil
  

structure :: forall a meta cons. (Generic a, Rep a ~ D1 meta cons, CTreeToSum cons) => a -> Sum (Map Prod (CTreeSum cons))
structure = ctreeToSum . unM1 . from


class ProdElim (srcT :: *) xs acc where
  type ProdElimF srcT xs acc
  prodElim :: (srcT -> acc) -> ProdElimF srcT xs acc -> Prod xs -> acc

instance ProdElim srcT '[] acc where
  type ProdElimF srcT '[] acc = acc
  prodElim r acc ProdNil = acc

class ProdElim srcT xs acc => ProdElimCase (eq :: Bool) srcT x xs acc where
  type ProdElimCaseF eq srcT x xs acc
  prodElimCase :: (srcT -> acc) -> ProdElimCaseF eq srcT x xs acc -> Prod (x:xs) -> acc

instance ProdElim srcT xs acc => ProdElimCase True srcT x xs acc where
  type ProdElimCaseF True srcT x xs acc = acc -> ProdElimF srcT xs acc
  prodElimCase r f (ProdCons x xs) = prodElim @srcT r (f (r (unsafeCoerce x))) xs

instance ProdElim srcT xs acc => ProdElimCase False srcT x xs acc where
  type ProdElimCaseF False srcT x xs acc = x -> ProdElimF srcT xs acc
  prodElimCase r f (ProdCons x xs) = prodElim @srcT r (f x) xs

instance (ProdElimCase (srcT == x) srcT x xs acc, ProdElim srcT xs acc) => ProdElim srcT (x:xs) acc where
  type ProdElimF srcT (x:xs) acc = ProdElimCaseF (srcT == x) srcT x xs acc
  prodElim = prodElimCase @(srcT == x) @srcT @x @xs @acc


newtype SingleElim r x = SingleElim { runSingleElim :: x -> r }

elimSumByProdS :: Prod (Map (SingleElim acc) ts) -> Sum ts -> acc
elimSumByProdS ProdNil x = case x of {}
elimSumByProdS (ProdCons f _) (SumHere x) = runSingleElim f x
elimSumByProdS (ProdCons _ fs) (SumThere xs) = elimSumByProdS fs xs

class SumElim srcT (xs :: [[*]]) (prev :: [[*]]) (src :: [[*]]) (acc :: *) where
  type SumElimF srcT xs prev src acc
  sumElim :: (srcT -> acc) -> Sum (Map Prod src) -> Prod (Map (SingleElim acc) (Map Prod prev)) -> SumElimF srcT xs prev src acc

instance (Generic srcT, Rep srcT ~ D1 meta cons, CTreeToSum cons, CTreeSum cons ~ src) => SumElim srcT '[] prev src acc where
  type SumElimF srcT '[] prev src acc = (acc, srcT -> acc)
  sumElim r x elims = (r' x, r' . structure)
    where
      r' = elimSumByProdS (mapReverse (reverseProd elims))
      mapReverse :: Prod (Reverse (Map (SingleElim acc) (Map Prod prev)))
                 -> Prod (Map (SingleElim acc) (Map Prod src))
      mapReverse = unsafeCoerce

instance (SumElim srcT xs (x:prev) src acc, ProdElim srcT x acc) => SumElim srcT (x:xs) prev src acc where
  type SumElimF srcT (x:xs) prev src acc = ProdElimF srcT x acc -> SumElimF srcT xs (x:prev) src acc
  sumElim r x elims elim = sumElim @srcT @xs @(x:prev) @src @acc r x $ ProdCons (SingleElim (prodElim @srcT r elim)) elims


elim ::
  forall a acc src cons meta.
  Generic a =>
  Rep a ~ D1 meta cons =>
  CTreeToSum cons =>
  SumElim a (CTreeSum cons) '[] (CTreeSum cons) acc =>
  (a -> acc) -> a -> SumElimF a (CTreeSum cons) '[] (CTreeSum cons) acc
elim r x = sumElim @a @(CTreeSum cons) @'[] @(CTreeSum cons) @acc r (structure x) ProdNil

data A = A | B Int | C Int String deriving Generic
data List a = Cons a (List a) | Nil deriving Generic

x :: List Int
x = Cons 1 (Cons 2 (Cons 3 Nil))
(y, r) = elim r x (:) []

-- y
-- [1, 2, 3]