Type-level functions to find the cardinality of various types

Cardinality.hs

{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE NoStarIsType #-}
{-# LANGUAGE DeriveGeneric #-}

module Cardinality where

import Data.Void
import GHC.Generics
import GHC.TypeLits

data TrafficLight = Red | Yellow | Green
  deriving Generic

data TwoBit = TwoBit Bool Bool
  deriving Generic

type family CardinalityNat a :: Nat where
  CardinalityNat Void         = 0
  CardinalityNat ()           = 1
  CardinalityNat Bool         = 2
  CardinalityNat (Maybe a)    = 1 + CardinalityNat a
  CardinalityNat (Either a b) = CardinalityNat a + CardinalityNat b
  CardinalityNat (a, b)       = CardinalityNat a * CardinalityNat b
  CardinalityNat (a -> b)     = CardinalityNat b ^ CardinalityNat a

data NatOrInf = N Nat | Inf

type family (a :: NatOrInf) +% (b :: NatOrInf) :: NatOrInf where
  'N a +% 'N b = 'N (a + b)
  'Inf +% _    = 'Inf
  _    +% 'Inf = 'Inf

type family (a :: NatOrInf) *% (b :: NatOrInf) :: NatOrInf where
  'N a *% 'N b = 'N (a GHC.TypeLits.* b)
  'N 0 *% _    = 'N 0
  _    *% 'N 0 = 'N 0
  'Inf *% _    = 'Inf
  _    *% 'Inf = 'Inf

type family (a :: NatOrInf) ^% (b :: NatOrInf) :: NatOrInf where
  'N a ^% 'N b = 'N (a ^ b)
  'N 0 ^% _    = 'N 0
  _    ^% 'N 0 = 'N 1
  'Inf ^% _    = 'Inf
  _    ^% 'Inf = 'Inf

type family Cardinality a :: NatOrInf where
  Cardinality Void         = 'N 0
  Cardinality ()           = 'N 1
  Cardinality Bool         = 'N 2
  Cardinality (Maybe a)    = 'N 1 +% Cardinality a
  Cardinality (Either a b) = Cardinality a +% Cardinality b
  Cardinality (a, b)       = Cardinality a *% Cardinality b
  Cardinality (a -> b)     = Cardinality b ^% Cardinality a
  Cardinality Integer      = 'Inf

-- :k! Cardinality (Maybe Bool)

data Const = Const ()
  deriving Generic

type family CardGeneric a :: NatOrInf where
  CardGeneric V1        = 'N 0
  CardGeneric U1        = 'N 1
  CardGeneric (K1 _ x)  = Cardinality x
  CardGeneric (D1 _ x)  = CardGeneric x
  CardGeneric (C1 _ x)  = CardGeneric x
  CardGeneric (S1 _ x)  = CardGeneric x
  CardGeneric (a :+: b) = CardGeneric a +% CardGeneric b
  CardGeneric (a :*: b) = CardGeneric a *% CardGeneric b

-- :k! CardGeneric (Rep (Maybe Bool))