{-# LANGUAGE ConstraintKinds       #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DefaultSignatures     #-}
{-# LANGUAGE EmptyCase             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds             #-}
{-# LANGUAGE Safe                  #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}
-- |
--
-- This module is designed to be imported qualified:
--
-- @
-- import qualified Data.Fin.Enum as E
-- @
--
module Data.Fin.Enum (
    Enum (..),
    -- * Generic implementation
    gfrom, GFrom,
    gto, GTo,
    GEnumSize,
    ) where

import Prelude hiding (Enum (..))

import Data.Bifunctor (bimap)
import Data.Fin       (Fin (..))
import Data.Nat       (Nat (..))
import Data.Proxy     (Proxy (..))
import GHC.Generics   (M1 (..), U1 (..), V1, (:+:) (..))

import qualified Data.Fin      as F
import qualified Data.Type.Nat as N
import qualified Data.Void     as V
import qualified GHC.Generics  as G

-- $setup
-- >>> import qualified Data.Fin as F

-- | Generic enumerations.
--
-- /Examples:/
--
-- >>> from ()
-- 0
--
-- >>> to 0 :: ()
-- ()
--
-- >>> to 0 :: Bool
-- False
--
-- >>> map to F.universe :: [Bool]
-- [False,True]
--
-- >>> map (to . (+1) . from) [LT, EQ, GT] :: [Ordering] -- Num Fin is modulo arithmetic
-- [EQ,GT,LT]
--
class Enum a where
    -- | The size of an enumeration.
    type EnumSize a :: Nat
    type EnumSize a = GEnumSize a

    -- | Converts a value to its index.
    from :: a -> Fin (EnumSize a)
    default from :: (G.Generic a, GFrom a, EnumSize a ~ GEnumSize a) => a -> Fin (EnumSize a)
    from = a -> Fin (EnumSize a)
forall a. (Generic a, GFrom a) => a -> Fin (GEnumSize a)
gfrom

    -- | Converts from index to the original value.
    to :: Fin (EnumSize a) -> a
    default to :: (G.Generic a, GTo a, EnumSize a ~ GEnumSize a) => Fin (EnumSize a) -> a
    to = Fin (EnumSize a) -> a
forall a. (Generic a, GTo a) => Fin (GEnumSize a) -> a
gto

-- | 'Void' ~ 0
instance Enum V.Void where
    -- this should be written by hand to work with all @base@
    type EnumSize V.Void = N.Nat0
    from :: Void -> Fin (EnumSize Void)
from = Void -> Fin (EnumSize Void)
forall a. Void -> a
V.absurd
    to :: Fin (EnumSize Void) -> Void
to   = Fin (EnumSize Void) -> Void
forall b. Fin Nat0 -> b
F.absurd

-- | () ~ 1
instance Enum ()

-- | 'Bool' ~ 2
instance Enum Bool

-- | 'Ordering' ~ 3
instance Enum Ordering

-- | 'Either' ~ @+@
instance (Enum a, Enum b, N.SNatI (EnumSize a)) => Enum (Either a b) where
    type EnumSize (Either a b) = N.Plus (EnumSize a) (EnumSize b)

    to :: Fin (EnumSize (Either a b)) -> Either a b
to   = (Fin (EnumSize a) -> a)
-> (Fin (EnumSize b) -> b)
-> Either (Fin (EnumSize a)) (Fin (EnumSize b))
-> Either a b
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap Fin (EnumSize a) -> a
forall a. Enum a => Fin (EnumSize a) -> a
to Fin (EnumSize b) -> b
forall a. Enum a => Fin (EnumSize a) -> a
to (Either (Fin (EnumSize a)) (Fin (EnumSize b)) -> Either a b)
-> (Fin (Plus (EnumSize a) (EnumSize b))
    -> Either (Fin (EnumSize a)) (Fin (EnumSize b)))
-> Fin (Plus (EnumSize a) (EnumSize b))
-> Either a b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Fin (Plus (EnumSize a) (EnumSize b))
-> Either (Fin (EnumSize a)) (Fin (EnumSize b))
forall (n :: Nat) (m :: Nat).
SNatI n =>
Fin (Plus n m) -> Either (Fin n) (Fin m)
F.split
    from :: Either a b -> Fin (EnumSize (Either a b))
from = Either (Fin (EnumSize a)) (Fin (EnumSize b))
-> Fin (Plus (EnumSize a) (EnumSize b))
forall (n :: Nat) (m :: Nat).
SNatI n =>
Either (Fin n) (Fin m) -> Fin (Plus n m)
F.append (Either (Fin (EnumSize a)) (Fin (EnumSize b))
 -> Fin (Plus (EnumSize a) (EnumSize b)))
-> (Either a b -> Either (Fin (EnumSize a)) (Fin (EnumSize b)))
-> Either a b
-> Fin (Plus (EnumSize a) (EnumSize b))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Fin (EnumSize a))
-> (b -> Fin (EnumSize b))
-> Either a b
-> Either (Fin (EnumSize a)) (Fin (EnumSize b))
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap a -> Fin (EnumSize a)
forall a. Enum a => a -> Fin (EnumSize a)
from b -> Fin (EnumSize b)
forall a. Enum a => a -> Fin (EnumSize a)
from

-------------------------------------------------------------------------------
-- EnumSize
-------------------------------------------------------------------------------

-- | Compute the size from the type.
type GEnumSize a = EnumSizeRep (G.Rep a) N.Nat0

type family EnumSizeRep (a :: * -> *) (n :: Nat) :: Nat where
    EnumSizeRep (a :+: b )   n = EnumSizeRep a (EnumSizeRep b n)
    EnumSizeRep V1           n = n
    EnumSizeRep (M1 _d _c a) n = EnumSizeRep a n
    EnumSizeRep U1           n = 'S n
    -- No instance for K1 or :*:

-------------------------------------------------------------------------------
-- From
-------------------------------------------------------------------------------

-- | Generic version of 'from'.
gfrom :: (G.Generic a, GFrom a) => a -> Fin (GEnumSize a)
gfrom :: a -> Fin (GEnumSize a)
gfrom = \a
x -> Rep a (Any @*) -> Proxy @Nat Nat0 -> Fin (GEnumSize a)
forall (a :: * -> *) x (n :: Nat).
GFromRep a =>
a x -> Proxy @Nat n -> Fin (EnumSizeRep a n)
gfromRep (a -> Rep a (Any @*)
forall a x. Generic a => a -> Rep a x
G.from a
x) (Proxy @Nat Nat0
forall k (t :: k). Proxy @k t
Proxy :: Proxy N.Nat0)

-- | Constraint for the class that computes 'gfrom'.
type GFrom a = GFromRep (G.Rep a)

class GFromRep (a :: * -> *)  where
    gfromRep  :: a x     -> Proxy n -> Fin (EnumSizeRep a n)
    gfromSkip :: Proxy a -> Fin n   -> Fin (EnumSizeRep a n)

instance (GFromRep a, GFromRep b) => GFromRep (a :+: b) where
    gfromRep :: (:+:) @* a b x
-> Proxy @Nat n -> Fin (EnumSizeRep ((:+:) @* a b) n)
gfromRep (L1 a x
a) Proxy @Nat n
n = a x
-> Proxy @Nat (EnumSizeRep b n)
-> Fin (EnumSizeRep a (EnumSizeRep b n))
forall (a :: * -> *) x (n :: Nat).
GFromRep a =>
a x -> Proxy @Nat n -> Fin (EnumSizeRep a n)
gfromRep a x
a (Proxy @Nat n -> Proxy @Nat (EnumSizeRep b n)
forall (n :: Nat). Proxy @Nat n -> Proxy @Nat (EnumSizeRep b n)
prSkip Proxy @Nat n
n) where
        prSkip :: Proxy n -> Proxy (EnumSizeRep b n)
        prSkip :: Proxy @Nat n -> Proxy @Nat (EnumSizeRep b n)
prSkip  Proxy @Nat n
_ = Proxy @Nat (EnumSizeRep b n)
forall k (t :: k). Proxy @k t
Proxy
    gfromRep (R1 b x
b) Proxy @Nat n
n = Proxy @(* -> *) a
-> Fin (EnumSizeRep b n) -> Fin (EnumSizeRep a (EnumSizeRep b n))
forall (a :: * -> *) (n :: Nat).
GFromRep a =>
Proxy @(* -> *) a -> Fin n -> Fin (EnumSizeRep a n)
gfromSkip (Proxy @(* -> *) a
forall k (t :: k). Proxy @k t
Proxy :: Proxy a) (b x -> Proxy @Nat n -> Fin (EnumSizeRep b n)
forall (a :: * -> *) x (n :: Nat).
GFromRep a =>
a x -> Proxy @Nat n -> Fin (EnumSizeRep a n)
gfromRep b x
b Proxy @Nat n
n)

    gfromSkip :: Proxy @(* -> *) ((:+:) @* a b)
-> Fin n -> Fin (EnumSizeRep ((:+:) @* a b) n)
gfromSkip Proxy @(* -> *) ((:+:) @* a b)
_ Fin n
n = Proxy @(* -> *) a
-> Fin (EnumSizeRep b n) -> Fin (EnumSizeRep a (EnumSizeRep b n))
forall (a :: * -> *) (n :: Nat).
GFromRep a =>
Proxy @(* -> *) a -> Fin n -> Fin (EnumSizeRep a n)
gfromSkip (Proxy @(* -> *) a
forall k (t :: k). Proxy @k t
Proxy :: Proxy a) (Proxy @(* -> *) b -> Fin n -> Fin (EnumSizeRep b n)
forall (a :: * -> *) (n :: Nat).
GFromRep a =>
Proxy @(* -> *) a -> Fin n -> Fin (EnumSizeRep a n)
gfromSkip (Proxy @(* -> *) b
forall k (t :: k). Proxy @k t
Proxy :: Proxy b) Fin n
n)

instance GFromRep a => GFromRep (M1 d c a) where
    gfromRep :: M1 @* d c a x -> Proxy @Nat n -> Fin (EnumSizeRep (M1 @* d c a) n)
gfromRep (M1 a x
a) Proxy @Nat n
n = a x -> Proxy @Nat n -> Fin (EnumSizeRep a n)
forall (a :: * -> *) x (n :: Nat).
GFromRep a =>
a x -> Proxy @Nat n -> Fin (EnumSizeRep a n)
gfromRep a x
a Proxy @Nat n
n
    gfromSkip :: Proxy @(* -> *) (M1 @* d c a)
-> Fin n -> Fin (EnumSizeRep (M1 @* d c a) n)
gfromSkip Proxy @(* -> *) (M1 @* d c a)
_     Fin n
n = Proxy @(* -> *) a -> Fin n -> Fin (EnumSizeRep a n)
forall (a :: * -> *) (n :: Nat).
GFromRep a =>
Proxy @(* -> *) a -> Fin n -> Fin (EnumSizeRep a n)
gfromSkip (Proxy @(* -> *) a
forall k (t :: k). Proxy @k t
Proxy :: Proxy a) Fin n
n

instance GFromRep V1 where
    gfromRep :: V1 @* x -> Proxy @Nat n -> Fin (EnumSizeRep (V1 @*) n)
gfromRep  V1 @* x
v Proxy @Nat n
_ = case V1 @* x
v of {}
    gfromSkip :: Proxy @(* -> *) (V1 @*) -> Fin n -> Fin (EnumSizeRep (V1 @*) n)
gfromSkip Proxy @(* -> *) (V1 @*)
_ Fin n
n = Fin n
Fin (EnumSizeRep (V1 @*) n)
n

instance GFromRep U1 where
    gfromRep :: U1 @* x -> Proxy @Nat n -> Fin (EnumSizeRep (U1 @*) n)
gfromRep U1 @* x
U1 Proxy @Nat n
_ = Fin (EnumSizeRep (U1 @*) n)
forall (n :: Nat). Fin ('S n)
FZ
    gfromSkip :: Proxy @(* -> *) (U1 @*) -> Fin n -> Fin (EnumSizeRep (U1 @*) n)
gfromSkip Proxy @(* -> *) (U1 @*)
_ Fin n
n = Fin n -> Fin ('S n)
forall (n :: Nat). Fin n -> Fin ('S n)
FS Fin n
n

-------------------------------------------------------------------------------
-- To
-------------------------------------------------------------------------------

-- | Generic version of 'to'.
gto :: (G.Generic a, GTo a) => Fin (GEnumSize a) -> a
gto :: Fin (GEnumSize a) -> a
gto = \Fin (GEnumSize a)
x -> Rep a (Any @*) -> a
forall a x. Generic a => Rep a x -> a
G.to (Rep a (Any @*) -> a) -> Rep a (Any @*) -> a
forall a b. (a -> b) -> a -> b
$ Fin (GEnumSize a)
-> (Rep a (Any @*) -> Rep a (Any @*))
-> (Fin Nat0 -> Rep a (Any @*))
-> Rep a (Any @*)
forall (a :: * -> *) (n :: Nat) x r.
GToRep a =>
Fin (EnumSizeRep a n) -> (a x -> r) -> (Fin n -> r) -> r
gtoRep Fin (GEnumSize a)
x Rep a (Any @*) -> Rep a (Any @*)
forall a. a -> a
id Fin Nat0 -> Rep a (Any @*)
forall b. Fin Nat0 -> b
F.absurd

-- | Constraint for the class that computes 'gto'.
type GTo a = GToRep (G.Rep a)

class GToRep (a :: * -> *) where
    gtoRep :: Fin (EnumSizeRep a n) -> (a x -> r) -> (Fin n -> r) -> r

instance (GToRep a, GToRep b) => GToRep (a :+: b) where
    gtoRep :: Fin (EnumSizeRep ((:+:) @* a b) n)
-> ((:+:) @* a b x -> r) -> (Fin n -> r) -> r
gtoRep Fin (EnumSizeRep ((:+:) @* a b) n)
n (:+:) @* a b x -> r
s Fin n -> r
k = Fin (EnumSizeRep a (EnumSizeRep b n))
-> (a x -> r) -> (Fin (EnumSizeRep b n) -> r) -> r
forall (a :: * -> *) (n :: Nat) x r.
GToRep a =>
Fin (EnumSizeRep a n) -> (a x -> r) -> (Fin n -> r) -> r
gtoRep Fin (EnumSizeRep a (EnumSizeRep b n))
Fin (EnumSizeRep ((:+:) @* a b) n)
n ((:+:) @* a b x -> r
s ((:+:) @* a b x -> r) -> (a x -> (:+:) @* a b x) -> a x -> r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a x -> (:+:) @* a b x
forall k (f :: k -> *) (g :: k -> *) (p :: k).
f p -> (:+:) @k f g p
L1) ((Fin (EnumSizeRep b n) -> r) -> r)
-> (Fin (EnumSizeRep b n) -> r) -> r
forall a b. (a -> b) -> a -> b
$ \Fin (EnumSizeRep b n)
r -> Fin (EnumSizeRep b n) -> (b x -> r) -> (Fin n -> r) -> r
forall (a :: * -> *) (n :: Nat) x r.
GToRep a =>
Fin (EnumSizeRep a n) -> (a x -> r) -> (Fin n -> r) -> r
gtoRep Fin (EnumSizeRep b n)
r ((:+:) @* a b x -> r
s ((:+:) @* a b x -> r) -> (b x -> (:+:) @* a b x) -> b x -> r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b x -> (:+:) @* a b x
forall k (f :: k -> *) (g :: k -> *) (p :: k).
g p -> (:+:) @k f g p
R1) Fin n -> r
k

instance GToRep a => GToRep (M1 d c a) where
    gtoRep :: Fin (EnumSizeRep (M1 @* d c a) n)
-> (M1 @* d c a x -> r) -> (Fin n -> r) -> r
gtoRep Fin (EnumSizeRep (M1 @* d c a) n)
n M1 @* d c a x -> r
s = Fin (EnumSizeRep a n) -> (a x -> r) -> (Fin n -> r) -> r
forall (a :: * -> *) (n :: Nat) x r.
GToRep a =>
Fin (EnumSizeRep a n) -> (a x -> r) -> (Fin n -> r) -> r
gtoRep Fin (EnumSizeRep a n)
Fin (EnumSizeRep (M1 @* d c a) n)
n (M1 @* d c a x -> r
s (M1 @* d c a x -> r) -> (a x -> M1 @* d c a x) -> a x -> r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a x -> M1 @* d c a x
forall k i (c :: Meta) (f :: k -> *) (p :: k). f p -> M1 @k i c f p
M1)

instance GToRep V1 where
    gtoRep :: Fin (EnumSizeRep (V1 @*) n) -> (V1 @* x -> r) -> (Fin n -> r) -> r
gtoRep Fin (EnumSizeRep (V1 @*) n)
n V1 @* x -> r
_ Fin n -> r
k = Fin n -> r
k Fin n
Fin (EnumSizeRep (V1 @*) n)
n

instance GToRep U1 where
    gtoRep :: Fin (EnumSizeRep (U1 @*) n) -> (U1 @* x -> r) -> (Fin n -> r) -> r
gtoRep Fin (EnumSizeRep (U1 @*) n)
FZ     U1 @* x -> r
s Fin n -> r
_ = U1 @* x -> r
s U1 @* x
forall k (p :: k). U1 @k p
U1
    gtoRep (FS Fin n
n) U1 @* x -> r
_ Fin n -> r
k = Fin n -> r
k Fin n
Fin n
n