{-# LANGUAGE GHC2024 #-}
{-# LANGUAGE BlockArguments, ImpredicativeTypes, TemplateHaskell #-}

import Control.Monad.Trans.Class
import Data.Foldable
import Data.Functor
import Data.Generics
import Data.List
import Data.Traversable
import GHC.TypeLits
import Language.Haskell.TH hiding (ppr)
import Unsafe.Coerce

newtype Exists a = MkExists a

runExists :: Exists a -> a
runExists (MkExists a) = a
{-# NOINLINE runExists #-}


data Vec n a where
  VNil :: Vec 0 a
  VCons :: a -> Vec n a -> Vec (n + 1) a

data T a where
  Int  :: T Int
  Bool :: T Bool


type Doc = String

class Pretty a where
  ppr :: a -> Doc

instance Pretty (T a) where
  ppr = serialise

serialise :: T a -> String
serialise Int  = "int"
serialise Bool = "bool"

sep :: [Doc] -> Doc
sep = intercalate "," 

data c /\ a = c => MkCPair a


let
  -- Since we can't write `exists vars. stuff` in Haskell yet, as a hacky
  -- approximation, write `$(exists [t| forall vars. stuff |])` instead.
  exists = (>>= \(ForallT vars [] ty) -> [t| Exists (forall r. $(forallT vars (pure []) [t| $(pure ty) -> r |]) -> r) |])

  findExists = \case
    t@(ConT e `AppT` ForallT _ [] (ArrowT `AppT` body `AppT` _)) | e == ''Exists -> [(t, body)]
    other -> []

  addRules decsQ = do
    decs <- decsQ
    let pairs = everything (<>) (mempty `mkQ` findExists) decs
    rules <- flip foldMap pairs \(t, body) ->
      -- We can't simply write this rule once for all types and expect it to work,    
      -- because of the limitations on polymorphic variables discussed at
      -- https://gitlab.haskell.org/ghc/ghc/-/issues/21093. But if we give `x` and `y`
      -- type annotations specific to each existential type used, that makes GHC
      -- willing to use it.
      [d| {-# RULES "clever unique name here" forall (x :: $(pure t)) (y :: $(pure body)). runExists x y = y (unsafeCoerce x id) #-} |]
    pure (rules <> decs)
  in
  addRules [d|

  listToVec :: [a] -> $(exists [t| forall n. Vec n a |])
  listToVec = \case
    [] -> MkExists \k -> k VNil
    x : xs -> runExists (listToVec xs) \ys -> MkExists \k -> k (VCons x ys)

  vecToList :: $(exists [t| forall n. Vec n a |]) -> [a]
  vecToList e = runExists e go
    where
    go :: Vec n a -> [a]
    go = \case
      VNil -> []
      VCons x xs -> x : go xs

  vfilter :: (a -> Bool) -> Vec n a -> $(exists [t| forall n. Vec n a |])
  vfilter p = \case
    VNil -> MkExists \k -> k VNil
    VCons x xs
      | p x -> runExists (vfilter p xs) \ys -> MkExists \k -> k (VCons x ys)
      | otherwise -> vfilter p xs


  deserialise :: String -> $(exists [t| forall a. T a |])
  deserialise "int"  = MkExists \k -> k Int
  deserialise "bool" = MkExists \k -> k Bool

  pprList :: [$(exists [t| forall a. Pretty a /\ a |])] -> Doc
  pprList = sep . map \e -> runExists e \(MkCPair x) -> ppr x


  main :: IO ()
  main = do
    runExists (listToVec [1..]) \xs -> print $ take 10 $ vecToList $ vfilter even $ xs

    print $ pprList $ map @($(exists [t| forall a. T a |])) (\e -> runExists e \t -> MkExists \k -> k (MkCPair t)) $ map deserialise ["int", "bool", "int"]

  |]