{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
module Data.INat where

import Data.Proxy
import Data.Type.Equality ((:~:) (Refl))
import Numeric.Natural
import GHC.TypeLits
import Unsafe.Coerce (unsafeCoerce)

-- | Evidence for the constraint @c a@.
data Dict c a where
  Dict :: c a => Dict c a

-- | An inductive peano natural number. Intended to be used at the type level.
data INat = Z | S INat
  deriving (Show)

-- | Singleton for a 'INat'.
data SINat n where
  SZ :: SINat Z
  SS :: SINat n -> SINat (S n)
deriving instance Show (SINat n)

-- | A singleton 'SINat' corresponding to @n@.
class KnownINat n where inatSing :: SINat n
instance KnownINat Z where inatSing = SZ
instance KnownINat n => KnownINat (S n) where inatSing = SS inatSing

-- | Explicitly bidirectional pattern synonym that converts between a singleton
-- 'SINat' and evidence of a 'KnownINat' constraint. Analogous to 'GHC.SNat'.
pattern SINat' :: () => KnownINat n => SINat n
pattern SINat' <- (snatKnown -> Dict)
  where SINat' = inatSing

-- | A 'KnownINat' dictionary is just a singleton natural, so we can create
-- evidence of 'KnownINat' given an 'SINat'.
snatKnown :: SINat n -> Dict KnownINat n
snatKnown SZ = Dict
snatKnown (SS n) | Dict <- snatKnown n = Dict

-- | Convert a 'INat' to a normal number.
fromINat :: INat -> Natural
fromINat Z = 0
fromINat (S n) = 1 + fromINat n

-- | Convert an 'SINat' to a normal number.
fromSINat :: SINat n -> Natural
fromSINat SZ = 0
fromSINat (SS n) = 1 + fromSINat n

-- | The value of a known inductive natural as a value-level integer.
inatVal :: forall n. KnownINat n => Proxy n -> Natural
inatVal _ = fromSINat (inatSing @n)

-- | Add two 'INat's
type family n +! m where
  Z +! m = m
  S n +! m = S (n +! m)

-- | Convert a 'INat' to a "GHC.TypeLits" 'G.Nat'.
type family FromINat n where
  FromINat Z = 0
  FromINat (S n) = 1 + FromINat n

-- | Convert a "GHC.TypeLits" 'G.Nat' to a 'INat'.
type family ToINat (n :: Nat) where
  ToINat 0 = Z
  ToINat n = S (ToINat (n - 1))

lemInjectiveFromINat :: n :~: ToINat (FromINat n)
lemInjectiveFromINat = unsafeCoerce Refl

lemSuccFromINat :: Proxy n -> 1 + FromINat n :~: FromINat (S n)
lemSuccFromINat _ = unsafeCoerce Refl

lemAddFromINat :: Proxy m -> Proxy n
               -> FromINat m + FromINat n :~: FromINat (m +! n)
lemAddFromINat _ = unsafeCoerce Refl

lemInjectiveToINat :: n :~: FromINat (ToINat n)
lemInjectiveToINat = unsafeCoerce Refl

lemSuccToINat :: Proxy n -> ToINat (1 + n) :~: S (ToINat n)
lemSuccToINat _ = unsafeCoerce Refl

lemAddToINat :: Proxy m -> Proxy n -> ToINat (m + n) :~: ToINat m +! ToINat n
lemAddToINat _ _ = unsafeCoerce Refl

-- | If an inductive 'INat' is known, then the corresponding "GHC.TypeLits"
-- 'G.Nat' is also known.
knownNatFromINat :: KnownINat n => Proxy n -> Dict KnownNat (FromINat n)
knownNatFromINat (Proxy @n) = go (SINat' @n)
  where
    go :: SINat m -> Dict KnownNat (FromINat m)
    go SZ = Dict
    go (SS n) | Dict <- go n = Dict

-- * Some type-level inductive naturals

type I0 = Z
type I1 = S I0
type I2 = S I1
type I3 = S I2
type I4 = S I3
type I5 = S I4
type I6 = S I5
type I7 = S I6
type I8 = S I7
type I9 = S I8