{-# 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 = 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 = 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 = forall a. Void -> a
V.absurd
    to :: Fin (EnumSize Void) -> Void
to   = forall b. Fin 'Z -> 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   = forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap forall a. Enum a => Fin (EnumSize a) -> a
to forall a. Enum a => Fin (EnumSize a) -> a
to forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 = forall (n :: Nat) (m :: Nat).
SNatI n =>
Either (Fin n) (Fin m) -> Fin (Plus n m)
F.append forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap forall a. Enum a => a -> Fin (EnumSize a)
from 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 :: forall a. (Generic a, GFrom a) => a -> Fin (GEnumSize a)
gfrom = \a
x -> forall (a :: * -> *) x (n :: Nat).
GFromRep a =>
a x -> Proxy @Nat n -> Fin (EnumSizeRep a n)
gfromRep (forall a x. Generic a => a -> Rep a x
G.from a
x) (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 :: forall x (n :: Nat).
(:+:) @(*) a b x
-> Proxy @Nat n -> Fin (EnumSizeRep ((:+:) @(*) a b) n)
gfromRep (L1 a x
a) Proxy @Nat n
n = forall (a :: * -> *) x (n :: Nat).
GFromRep a =>
a x -> Proxy @Nat n -> Fin (EnumSizeRep a n)
gfromRep a x
a (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 :: forall (n :: Nat). Proxy @Nat n -> Proxy @Nat (EnumSizeRep b n)
prSkip  Proxy @Nat n
_ = forall {k} (t :: k). Proxy @k t
Proxy
    gfromRep (R1 b x
b) Proxy @Nat n
n = forall (a :: * -> *) (n :: Nat).
GFromRep a =>
Proxy @(* -> *) a -> Fin n -> Fin (EnumSizeRep a n)
gfromSkip (forall {k} (t :: k). Proxy @k t
Proxy :: Proxy a) (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 :: forall (n :: Nat).
Proxy @(* -> *) ((:+:) @(*) a b)
-> Fin n -> Fin (EnumSizeRep ((:+:) @(*) a b) n)
gfromSkip Proxy @(* -> *) ((:+:) @(*) a b)
_ Fin n
n = forall (a :: * -> *) (n :: Nat).
GFromRep a =>
Proxy @(* -> *) a -> Fin n -> Fin (EnumSizeRep a n)
gfromSkip (forall {k} (t :: k). Proxy @k t
Proxy :: Proxy a) (forall (a :: * -> *) (n :: Nat).
GFromRep a =>
Proxy @(* -> *) a -> Fin n -> Fin (EnumSizeRep a n)
gfromSkip (forall {k} (t :: k). Proxy @k t
Proxy :: Proxy b) Fin n
n)

instance GFromRep a => GFromRep (M1 d c a) where
    gfromRep :: forall x (n :: Nat).
M1 @(*) d c a x
-> Proxy @Nat n -> Fin (EnumSizeRep (M1 @(*) d c a) n)
gfromRep (M1 a x
a) Proxy @Nat n
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 :: forall (n :: Nat).
Proxy @(* -> *) (M1 @(*) d c a)
-> Fin n -> Fin (EnumSizeRep (M1 @(*) d c a) n)
gfromSkip Proxy @(* -> *) (M1 @(*) d c a)
_     Fin n
n = forall (a :: * -> *) (n :: Nat).
GFromRep a =>
Proxy @(* -> *) a -> Fin n -> Fin (EnumSizeRep a n)
gfromSkip (forall {k} (t :: k). Proxy @k t
Proxy :: Proxy a) Fin n
n

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

instance GFromRep U1 where
    gfromRep :: forall x (n :: Nat).
U1 @(*) x -> Proxy @Nat n -> Fin (EnumSizeRep (U1 @(*)) n)
gfromRep U1 @(*) x
U1 Proxy @Nat n
_ = forall (n :: Nat). Fin ('S n)
FZ
    gfromSkip :: forall (n :: Nat).
Proxy @(* -> *) (U1 @(*)) -> Fin n -> Fin (EnumSizeRep (U1 @(*)) n)
gfromSkip Proxy @(* -> *) (U1 @(*))
_ Fin n
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 :: forall a. (Generic a, GTo a) => Fin (GEnumSize a) -> a
gto = \Fin (GEnumSize a)
x -> forall a x. Generic a => Rep a x -> a
G.to forall a b. (a -> b) -> a -> b
$ forall (a :: * -> *) (n :: Nat) x r.
GToRep a =>
Fin (EnumSizeRep a n) -> (a x -> r) -> (Fin n -> r) -> r
gtoRep Fin (GEnumSize a)
x forall a. a -> a
id forall b. Fin 'Z -> 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 :: forall (n :: Nat) x r.
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 = forall (a :: * -> *) (n :: Nat) x r.
GToRep a =>
Fin (EnumSizeRep a n) -> (a x -> r) -> (Fin n -> r) -> r
gtoRep Fin (EnumSizeRep ((:+:) @(*) a b) n)
n ((:+:) @(*) a b x -> r
s forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k (f :: k -> *) (g :: k -> *) (p :: k).
f p -> (:+:) @k f g p
L1) forall a b. (a -> b) -> a -> b
$ \Fin (EnumSizeRep b n)
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 forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 :: forall (n :: Nat) x r.
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 = forall (a :: * -> *) (n :: Nat) x r.
GToRep a =>
Fin (EnumSizeRep a n) -> (a x -> r) -> (Fin n -> r) -> r
gtoRep Fin (EnumSizeRep (M1 @(*) d c a) n)
n (M1 @(*) d c a x -> r
s forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k i (c :: Meta) (f :: k -> *) (p :: k). f p -> M1 @k i c f p
M1)

instance GToRep V1 where
    gtoRep :: forall (n :: Nat) x r.
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 (EnumSizeRep (V1 @(*)) n)
n

instance GToRep U1 where
    gtoRep :: forall (n :: Nat) x r.
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 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
n