{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE EmptyCase #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE NoStarIsType #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Presburger #-}

module Data.Type.Natural.Core
  ( SNat (Zero, Succ),
#if !MIN_VERSION_base(4,18,0)
    fromSNat,
    withKnownNat,
    withSomeSNat,
#endif
    unsafeLiftSBin,
    ZeroOrSucc (..),
    viewNat,
    sNat,
    (%+),
    (%-),
    (%*),
    (%^),
    sDiv,
    sMod,
    sLog2,
    (%<=?),
    sCmpNat,
    sCompare,
    Succ,
    S,
    sSucc,
    sS,
    Pred,
    sPred,
    Zero,
    One,
    sZero,
    sOne,
    Equality (..),
    equalAbsurdFromBool,
    type (===),
    (%~),
    sFlipOrdering,
    FlipOrdering,
    SOrdering (..),
    SBool (..),
    Natural,
    OrderingI(..),
    fromOrderingI,
    toOrderingI,
    -- Re-exports
    module GHC.TypeNats,
  )
where

import Data.Type.Equality
  ( type (:~:) (..),
    type (==),
  )
import GHC.TypeNats
import Math.NumberTheory.Logarithms (naturalLog2)
import Type.Reflection (Typeable)
import Unsafe.Coerce (unsafeCoerce)
import Numeric.Natural

#if MIN_VERSION_base(4,16,0)
import Data.Type.Ord (OrderingI(..))
#endif

#if !MIN_VERSION_base(4,18,0)
import Data.Proxy
import Data.Type.Equality
import GHC.Exts
#endif

#if !MIN_VERSION_base(4,18,0)
-- | A singleton for type-level naturals
newtype SNat (n :: Nat) = UnsafeSNat Natural
  deriving newtype (Show, Eq, Ord)

fromSNat :: SNat n -> Natural
fromSNat = coerce

withKnownNat :: forall n rep (r :: TYPE rep). SNat n -> (KnownNat n => r) -> r
withKnownNat (UnsafeSNat n) act =
  case someNatVal n of
    SomeNat (_ :: Proxy m) ->
      case unsafeCoerce (Refl @()) :: n :~: m of
        Refl -> act

data KnownNatInstance (n :: Nat) where
  KnownNatInstance :: KnownNat n => KnownNatInstance n

-- An internal function that is only used for defining the SNat pattern
-- synonym.
knownNatInstance :: SNat n -> KnownNatInstance n
knownNatInstance sn = withKnownNat sn KnownNatInstance

pattern SNat :: forall n. () => KnownNat n => SNat n
pattern SNat <- (knownNatInstance -> KnownNatInstance) 
  where SNat = sNat

withSomeSNat :: forall rep (r :: TYPE rep). Natural -> (forall n. SNat n -> r) -> r
withSomeSNat n f = f (UnsafeSNat n)
#endif

unsafeLiftSBin :: (Natural -> Natural -> Natural) -> SNat n -> SNat m -> SNat k
{-# INLINE unsafeLiftSBin #-}
unsafeLiftSBin :: forall (n :: Natural) (m :: Natural) (k :: Natural).
(Natural -> Natural -> Natural) -> SNat n -> SNat m -> SNat k
unsafeLiftSBin Natural -> Natural -> Natural
f = \SNat n
l SNat m
r -> Natural -> (forall (n :: Natural). SNat n -> SNat k) -> SNat k
forall r. Natural -> (forall (n :: Natural). SNat n -> r) -> r
withSomeSNat (SNat n -> Natural
forall (n :: Natural). SNat n -> Natural
fromSNat SNat n
l Natural -> Natural -> Natural
`f` SNat m -> Natural
forall (n :: Natural). SNat n -> Natural
fromSNat SNat m
r) SNat n -> SNat k
forall (n :: Natural). SNat n -> SNat k
forall a b. a -> b
unsafeCoerce

unsafeLiftSUnary :: (Natural -> Natural) -> SNat n -> SNat k
{-# INLINE unsafeLiftSUnary #-}
unsafeLiftSUnary :: forall (n :: Natural) (k :: Natural).
(Natural -> Natural) -> SNat n -> SNat k
unsafeLiftSUnary Natural -> Natural
f = \SNat n
l -> Natural -> (forall (n :: Natural). SNat n -> SNat k) -> SNat k
forall r. Natural -> (forall (n :: Natural). SNat n -> r) -> r
withSomeSNat (Natural -> Natural
f (Natural -> Natural) -> Natural -> Natural
forall a b. (a -> b) -> a -> b
$ SNat n -> Natural
forall (n :: Natural). SNat n -> Natural
fromSNat SNat n
l) SNat n -> SNat k
forall (n :: Natural). SNat n -> SNat k
forall a b. a -> b
unsafeCoerce

(%+) :: SNat n -> SNat m -> SNat (n + m)
{-# INLINE (%+) #-}
%+ :: forall (n :: Natural) (m :: Natural).
SNat n -> SNat m -> SNat (n + m)
(%+) = (Natural -> Natural -> Natural) -> SNat n -> SNat m -> SNat (n + m)
forall (n :: Natural) (m :: Natural) (k :: Natural).
(Natural -> Natural -> Natural) -> SNat n -> SNat m -> SNat k
unsafeLiftSBin Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
(+)

(%-) :: SNat n -> SNat m -> SNat (n - m)
%- :: forall (n :: Natural) (m :: Natural).
SNat n -> SNat m -> SNat (n - m)
(%-) = (Natural -> Natural -> Natural) -> SNat n -> SNat m -> SNat (n - m)
forall (n :: Natural) (m :: Natural) (k :: Natural).
(Natural -> Natural -> Natural) -> SNat n -> SNat m -> SNat k
unsafeLiftSBin (-)

(%*) :: SNat n -> SNat m -> SNat (n * m)
%* :: forall (n :: Natural) (m :: Natural).
SNat n -> SNat m -> SNat (n * m)
(%*) = (Natural -> Natural -> Natural) -> SNat n -> SNat m -> SNat (n * m)
forall (n :: Natural) (m :: Natural) (k :: Natural).
(Natural -> Natural -> Natural) -> SNat n -> SNat m -> SNat k
unsafeLiftSBin Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
(*)

sDiv :: SNat n -> SNat m -> SNat (Div n m)
sDiv :: forall (n :: Natural) (m :: Natural).
SNat n -> SNat m -> SNat (Div n m)
sDiv = (Natural -> Natural -> Natural)
-> SNat n -> SNat m -> SNat (Div n m)
forall (n :: Natural) (m :: Natural) (k :: Natural).
(Natural -> Natural -> Natural) -> SNat n -> SNat m -> SNat k
unsafeLiftSBin Natural -> Natural -> Natural
forall a. Integral a => a -> a -> a
quot

sMod :: SNat n -> SNat m -> SNat (Mod n m)
sMod :: forall (n :: Natural) (m :: Natural).
SNat n -> SNat m -> SNat (Mod n m)
sMod = (Natural -> Natural -> Natural)
-> SNat n -> SNat m -> SNat (Mod n m)
forall (n :: Natural) (m :: Natural) (k :: Natural).
(Natural -> Natural -> Natural) -> SNat n -> SNat m -> SNat k
unsafeLiftSBin Natural -> Natural -> Natural
forall a. Integral a => a -> a -> a
rem

(%^) :: SNat n -> SNat m -> SNat (n ^ m)
%^ :: forall (n :: Natural) (m :: Natural).
SNat n -> SNat m -> SNat (n ^ m)
(%^) = (Natural -> Natural -> Natural) -> SNat n -> SNat m -> SNat (n ^ m)
forall (n :: Natural) (m :: Natural) (k :: Natural).
(Natural -> Natural -> Natural) -> SNat n -> SNat m -> SNat k
unsafeLiftSBin Natural -> Natural -> Natural
forall a b. (Num a, Integral b) => a -> b -> a
(^)

sLog2 :: SNat n -> SNat (Log2 n)
sLog2 :: forall (n :: Natural). SNat n -> SNat (Log2 n)
sLog2 = (Natural -> Natural) -> SNat n -> SNat (Log2 n)
forall (n :: Natural) (k :: Natural).
(Natural -> Natural) -> SNat n -> SNat k
unsafeLiftSUnary ((Natural -> Natural) -> SNat n -> SNat (Log2 n))
-> (Natural -> Natural) -> SNat n -> SNat (Log2 n)
forall a b. (a -> b) -> a -> b
$ Int -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Natural) -> (Natural -> Int) -> Natural -> Natural
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Natural -> Int
naturalLog2

sNat :: forall n. KnownNat n => SNat n
#if MIN_VERSION_base(4,18,0)
sNat :: forall (n :: Natural). KnownNat n => SNat n
sNat = SNat n
forall (n :: Natural). KnownNat n => SNat n
SNat
#else
sNat = UnsafeSNat $ natVal' (proxy# :: Proxy# n)
#endif


infixl 6 %+, %-

infixl 7 %*, `sDiv`, `sMod`

infixr 8 %^

#if !MIN_VERSION_ghc(4,18,0)
instance TestEquality SNat where
  testEquality (UnsafeSNat l) (UnsafeSNat r) =
    if l == r
      then Just trustMe
      else Nothing
#endif


-- | Since 1.1.0.0 (Type changed)
data Equality n m where
  Equal :: ((n == n) ~ 'True) => Equality n n
  NonEqual ::
    ((n === m) ~ 'False, (n == m) ~ 'False) =>
    Equality n m

equalAbsurdFromBool ::
  (x === y) ~ 'False => x :~: y -> a
equalAbsurdFromBool :: forall {k} (x :: k) (y :: k) a.
((x === y) ~ 'False) =>
(x :~: y) -> a
equalAbsurdFromBool = \case {}

type family a === b where
  a === a = 'True
  _ === _ = 'False

infix 4 ===, %~

(%~) :: SNat l -> SNat r -> Equality l r
SNat l
l %~ :: forall (l :: Natural) (r :: Natural).
SNat l -> SNat r -> Equality l r
%~ SNat r
r =
  if SNat l -> Natural
forall (n :: Natural). SNat n -> Natural
fromSNat SNat l
l Natural -> Natural -> Bool
forall a. Eq a => a -> a -> Bool
== SNat r -> Natural
forall (n :: Natural). SNat n -> Natural
fromSNat SNat r
r
    then Equality () () -> Equality l r
forall a b. a -> b
unsafeCoerce (forall n. ((n == n) ~ 'True) => Equality n n
forall {k} (n :: k). ((n == n) ~ 'True) => Equality n n
Equal @())
    else Equality 0 1 -> Equality l r
forall a b. a -> b
unsafeCoerce (forall (n :: Natural) (m :: Natural).
((n === m) ~ 'False, (n == m) ~ 'False) =>
Equality n m
forall {k} (n :: k) (m :: k).
((n === m) ~ 'False, (n == m) ~ 'False) =>
Equality n m
NonEqual @0 @1)

type Zero = 0

type One = 1

sZero :: SNat 0
sZero :: SNat 0
sZero = SNat 0
forall (n :: Natural). KnownNat n => SNat n
sNat

sOne :: SNat 1
sOne :: SNat 1
sOne = SNat 1
forall (n :: Natural). KnownNat n => SNat n
sNat

type Succ n = n + 1

type S n = Succ n

sSucc, sS :: SNat n -> SNat (Succ n)
sS :: forall (n :: Natural). SNat n -> SNat (Succ n)
sS = (SNat n -> SNat 1 -> SNat (n + 1)
forall (n :: Natural) (m :: Natural).
SNat n -> SNat m -> SNat (n + m)
%+ SNat 1
sOne)
sSucc :: forall (n :: Natural). SNat n -> SNat (Succ n)
sSucc = SNat n -> SNat (Succ n)
forall (n :: Natural). SNat n -> SNat (Succ n)
sS

sPred :: SNat n -> SNat (Pred n)
sPred :: forall (n :: Natural). SNat n -> SNat (Pred n)
sPred = (SNat n -> SNat 1 -> SNat (n - 1)
forall (n :: Natural) (m :: Natural).
SNat n -> SNat m -> SNat (n - m)
%- SNat 1
sOne)

type Pred n = n - 1

data ZeroOrSucc n where
  IsZero :: ZeroOrSucc 0
  IsSucc ::
    SNat n ->
    ZeroOrSucc (n + 1)

pattern Zero :: forall n. () => n ~ 0 => SNat n
pattern $mZero :: forall {r} {n :: Natural}.
SNat n -> ((n ~ 0) => r) -> ((# #) -> r) -> r
$bZero :: forall (n :: Natural). (n ~ 0) => SNat n
Zero <-
  (viewNat -> IsZero)
  where
    Zero = SNat n
SNat 0
sZero

pattern Succ :: forall n. () => forall n1. n ~ Succ n1 => SNat n1 -> SNat n
pattern $mSucc :: forall {r} {n :: Natural}.
SNat n
-> (forall {n1 :: Natural}. (n ~ Succ n1) => SNat n1 -> r)
-> ((# #) -> r)
-> r
$bSucc :: forall (n :: Natural) (n1 :: Natural).
(n ~ Succ n1) =>
SNat n1 -> SNat n
Succ n <-
  (viewNat -> IsSucc n)
  where
    Succ SNat n1
n = SNat n1 -> SNat (Succ n1)
forall (n :: Natural). SNat n -> SNat (Succ n)
sSucc SNat n1
n

{-# COMPLETE Zero, Succ #-}

viewNat :: forall n. SNat n -> ZeroOrSucc n
viewNat :: forall (n :: Natural). SNat n -> ZeroOrSucc n
viewNat SNat n
n =
  case SNat n
n SNat n -> SNat 0 -> Equality n 0
forall (l :: Natural) (r :: Natural).
SNat l -> SNat r -> Equality l r
%~ forall (n :: Natural). KnownNat n => SNat n
sNat @0 of
    Equality n 0
Equal -> ZeroOrSucc n
ZeroOrSucc 0
IsZero
    Equality n 0
NonEqual -> SNat (n - 1) -> ZeroOrSucc ((n - 1) + 1)
forall (n :: Natural). SNat n -> ZeroOrSucc (n + 1)
IsSucc (SNat n -> SNat (n - 1)
forall (n :: Natural). SNat n -> SNat (Pred n)
sPred SNat n
n)


#if !MIN_VERSION_base(4,16,0)
data OrderingI (a :: Nat) (b :: Nat) where
  LTI :: CmpNat a b ~ 'LT => OrderingI a b
  EQI :: CmpNat a b ~ 'EQ => OrderingI a b
  GTI :: CmpNat a b ~ 'GT => OrderingI a b
#endif

type family FlipOrdering ord where
  FlipOrdering 'LT = 'GT
  FlipOrdering 'GT = 'LT
  FlipOrdering 'EQ = 'EQ

data SOrdering (ord :: Ordering) where
  SLT :: SOrdering 'LT
  SEQ :: SOrdering 'EQ
  SGT :: SOrdering 'GT

fromOrderingI :: OrderingI n m -> SOrdering (CmpNat n m)
fromOrderingI :: forall (n :: Natural) (m :: Natural).
OrderingI n m -> SOrdering (CmpNat n m)
fromOrderingI OrderingI n m
LTI = SOrdering (CmpNat n m)
SOrdering 'LT
SLT
fromOrderingI OrderingI n m
EQI = SOrdering (CmpNat n m)
SOrdering 'EQ
SEQ
fromOrderingI OrderingI n m
GTI = SOrdering (CmpNat n m)
SOrdering 'GT
SGT

toOrderingI :: SOrdering (CmpNat n m) -> OrderingI n m
toOrderingI :: forall (n :: Natural) (m :: Natural).
SOrdering (CmpNat n m) -> OrderingI n m
toOrderingI SOrdering (CmpNat n m)
SLT = OrderingI n m
forall {k} (a :: k) (b :: k). (Compare a b ~ 'LT) => OrderingI a b
LTI
toOrderingI SOrdering (CmpNat n m)
SEQ = OrderingI n n
OrderingI n m
forall {k} (a :: k). (Compare a a ~ 'EQ) => OrderingI a a
EQI
toOrderingI SOrdering (CmpNat n m)
SGT = OrderingI n m
forall {k} (a :: k) (b :: k). (Compare a b ~ 'GT) => OrderingI a b
GTI

deriving instance Show (SOrdering ord)

deriving instance Eq (SOrdering ord)

deriving instance Typeable SOrdering

sFlipOrdering :: SOrdering ord -> SOrdering (FlipOrdering ord)
sFlipOrdering :: forall (ord :: Ordering).
SOrdering ord -> SOrdering (FlipOrdering ord)
sFlipOrdering SOrdering ord
SLT = SOrdering 'GT
SOrdering (FlipOrdering ord)
SGT
sFlipOrdering SOrdering ord
SEQ = SOrdering 'EQ
SOrdering (FlipOrdering ord)
SEQ
sFlipOrdering SOrdering ord
SGT = SOrdering 'LT
SOrdering (FlipOrdering ord)
SLT

data SBool (b :: Bool) where
  SFalse :: SBool 'False
  STrue :: SBool 'True

deriving instance Show (SBool ord)

deriving instance Eq (SBool ord)

deriving instance Typeable SBool

infix 4 %<=?

(%<=?) :: SNat n -> SNat m -> SBool (n <=? m)
SNat n
n %<=? :: forall (n :: Natural) (m :: Natural).
SNat n -> SNat m -> SBool (n <=? m)
%<=? SNat m
m =
  if SNat n -> Natural
forall (n :: Natural). SNat n -> Natural
fromSNat SNat n
n Natural -> Natural -> Bool
forall a. Ord a => a -> a -> Bool
<= SNat m -> Natural
forall (n :: Natural). SNat n -> Natural
fromSNat SNat m
m
    then SBool 'True -> SBool (OrdCond (CmpNat n m) 'True 'True 'False)
forall a b. a -> b
unsafeCoerce SBool 'True
STrue
    else SBool 'False -> SBool (OrdCond (CmpNat n m) 'True 'True 'False)
forall a b. a -> b
unsafeCoerce SBool 'False
SFalse

sCmpNat, sCompare :: SNat n -> SNat m -> SOrdering (CmpNat n m)
sCompare :: forall (n :: Natural) (m :: Natural).
SNat n -> SNat m -> SOrdering (CmpNat n m)
sCompare = SNat n -> SNat m -> SOrdering (CmpNat n m)
forall (n :: Natural) (m :: Natural).
SNat n -> SNat m -> SOrdering (CmpNat n m)
sCmpNat
sCmpNat :: forall (n :: Natural) (m :: Natural).
SNat n -> SNat m -> SOrdering (CmpNat n m)
sCmpNat SNat n
n SNat m
m =
  case Natural -> Natural -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (SNat n -> Natural
forall (n :: Natural). SNat n -> Natural
fromSNat SNat n
n) (SNat m -> Natural
forall (n :: Natural). SNat n -> Natural
fromSNat SNat m
m) of
    Ordering
LT -> SOrdering 'LT -> SOrdering (CmpNat n m)
forall a b. a -> b
unsafeCoerce SOrdering 'LT
SLT
    Ordering
EQ -> SOrdering 'EQ -> SOrdering (CmpNat n m)
forall a b. a -> b
unsafeCoerce SOrdering 'EQ
SEQ
    Ordering
GT -> SOrdering 'GT -> SOrdering (CmpNat n m)
forall a b. a -> b
unsafeCoerce SOrdering 'GT
SGT