Countable sets, finally tagless implementation of Fin

IsCountable.hs

{-# LANGUAGE GHC2021 #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ViewPatterns #-}

import Data.Kind (Type)
import Data.Void (Void)
import GHC.TypeNats (type (+))
import Prelude hiding (pred)
import Numeric.Natural (Natural)

class IsCountable s where
  -- | A countable set is either empty: @Pred s ~ 'Nothing@ or one bigger than another countable set, its predecessor: @Pred s ~ 'Just p@
  type Pred s :: Maybe Type
  mkZ :: Pred s ~ 'Just p => s -- the 'extra' element
  mkS :: Pred s ~ 'Just p => p -> s -- lift elements from the predecessor set
  -- | pattern match on elements of @s@
  pred :: (forall p. (Pred s ~ 'Just p, IsCountable p) => r) -- ^ the extra element case
       -> (forall p. (Pred s ~ 'Just p, IsCountable p) => p -> r) -- ^ the lifted elements case
       -> s -> r
  -- | An "if" based on the type of @Pred s@.
  tIf :: (Pred s ~ 'Nothing => r)
      -> (forall p. (Pred s ~ 'Just p, IsCountable p) => r)
      -> r

absurd :: (IsCountable s, Pred s ~ 'Nothing) => s -> x
absurd = pred undefined undefined -- TODO: How to get rid of these undefineds

enumAll :: forall s. IsCountable s => [s]
enumAll = tIf @s [] (Z : (S <$> enumAll))

-- | Type level size, this is SLOW on larger types.
type Size s = SIZE (Pred s)
type family SIZE ps where
  SIZE 'Nothing = 0
  SIZE ('Just p) = 1 + Size p


data PredCase s where
  PredZ :: (Pred s ~ 'Just p, IsCountable p) => PredCase s
  PredS :: (Pred s ~ 'Just p, IsCountable p) => p -> PredCase s

toPredCase :: IsCountable s => s -> PredCase s
toPredCase = pred PredZ PredS

pattern Z :: IsCountable s => (Pred s ~ 'Just p, IsCountable p) => s
pattern Z <- (toPredCase -> PredZ) where Z = mkZ

pattern S :: IsCountable s => (Pred s ~ 'Just p, IsCountable p) => p -> s
pattern S p <- (toPredCase -> PredS p) where S = mkS

{-# COMPLETE Z, S #-}


data Nat = NZ | NS Nat
data Fin n where
  Fz :: Fin (NS n)
  Fs :: Fin n -> Fin (NS n)

instance IsCountable (Fin NZ) where
  type Pred (Fin NZ) = 'Nothing
  mkZ = undefined
  mkS = undefined
  pred _ _ = \case {}
  tIf z _ = z

instance IsCountable (Fin n) => IsCountable (Fin (NS n)) where
  type Pred (Fin (NS n)) = 'Just (Fin n)
  mkZ = Fz
  mkS = Fs
  pred z _ Fz = z
  pred _ s (Fs n) = s n
  tIf _ s = s


instance IsCountable Void where
  type Pred Void = 'Nothing
  mkZ = undefined
  mkS = undefined
  pred _ _ = \case {}
  tIf z _ = z

instance IsCountable () where
  type Pred () = 'Just Void
  mkZ = ()
  mkS = \case {}
  pred z _ () = z
  tIf _ s = s

instance IsCountable Bool where
  type Pred Bool = 'Just ()
  mkZ = False
  mkS () = True
  pred z _ False = z
  pred _ s True = s ()
  tIf _ s = s

instance IsCountable Natural where
  type Pred Natural = 'Just Natural
  mkZ = 0
  mkS = (+ 1)
  pred z _ 0 = z
  pred _ s n = s (n - 1)
  tIf _ s = s

instance IsCountable a => IsCountable (Maybe a) where
  type Pred (Maybe a) = 'Just a
  mkZ = Nothing
  mkS = Just
  pred z _ Nothing = z
  pred _ s (Just a) = s a
  tIf _ s = s


type family SUM pa b where
  SUM 'Nothing b = Pred b
  SUM ('Just a) b = 'Just (Either a b)

instance (IsCountable a, IsCountable b) => IsCountable (Either a b) where
  type Pred (Either a b) = SUM (Pred a) b
  mkZ = tIf @a (Right mkZ) (Left mkZ)
  mkS = tIf @a (Right . mkS) (either (Left . mkS) Right)
  pred z s (Left a) = pred z (s . Left) a
  pred z s (Right b) = tIf @a (pred z s b) (s (Right b))
  tIf z s = tIf @a (tIf @b z s) s


type ProdStep a b = Either b (a, b)

fromProdStep :: (IsCountable a, Pred a ~ 'Just a') => ProdStep a' b -> (a, b)
fromProdStep (Left b) = (mkZ, b)
fromProdStep (Right (a, b)) = (mkS a, b)

toProdStep :: (IsCountable a, Pred a ~ 'Just a') => (a, b) -> ProdStep a' b
toProdStep (a, b) = pred (Left b) (\a' -> Right (a', b)) a

type family PROD pa b where
  PROD 'Nothing b = 'Nothing
  PROD ('Just a) b = Pred (ProdStep a b)

instance (IsCountable a, IsCountable b) => IsCountable (a, b) where
  type Pred (a, b) = PROD (Pred a) b
  mkZ = tIf @a undefined $ fromProdStep mkZ
  mkS p = tIf @a undefined $ fromProdStep (mkS p)
  pred z s = tIf @a (\(a, _) -> absurd a) $ pred z s . toProdStep
  tIf z s = tIf @a z (tIf' z s)
    where
      tIf' :: forall a' r. (Pred a ~ 'Just a', IsCountable a') => (Pred (a, b) ~ 'Nothing => r) -> (forall p. (Pred (a, b) ~ 'Just p, IsCountable p) => r) -> r
      tIf' = tIf @(ProdStep a' b)


type ExpStep a b = (b, a -> b)

fromExpStep :: (IsCountable a, Pred a ~ 'Just a') => ExpStep a' b -> (a -> b)
fromExpStep (b, _) Z = b
fromExpStep (_, a'b) (S a') = a'b a'

toExpStep :: (IsCountable a, Pred a ~ 'Just a', IsCountable a') => (a -> b) -> ExpStep a' b
toExpStep ab = (ab Z, ab . S)

type family EXP pa b where
  EXP 'Nothing b = 'Just Void
  EXP ('Just a) b = Pred (ExpStep a b)

instance (IsCountable a, IsCountable b) => IsCountable (a -> b) where
  type Pred (a -> b) = EXP (Pred a) b
  mkZ = tIf @a absurd $ fromExpStep mkZ
  mkS p = tIf @a (\case {}) $ fromExpStep (mkS p)
  pred z s = tIf @a (const z) $ pred z s . toExpStep
  tIf z s = tIf @a s (tIf' z s)
    where
      tIf' :: forall a' r. (Pred a ~ 'Just a', IsCountable a') => (Pred (a -> b) ~ 'Nothing => r) -> (forall p. (Pred (a -> b) ~ 'Just p, IsCountable p) => r) -> r
      tIf' = tIf @(ExpStep a' b)


type Fin2 = Bool
type Fin3 = Maybe Bool
type Fin4 = (Fin2, Fin2)
type Fin16 = (Fin4, Fin4)
type Fin256 = Fin4 -> Fin4
type Fin65536 = (Fin256, Fin256)
type Fin4294967296 = (Fin65536, Fin65536)