--------------------------------------------------------------------------------

-- |

--

-- Module      :  Data.Type.Int

-- Description :  Type-level integers

-- Copyright   :  (c) Alice Rixte 2025

-- License     :  BSD 3

-- Maintainer  :  alice.rixte@u-bordeaux.fr

-- Stability   :  unstable

-- Portability :  non-portable (GHC extensions)

--

-- Type level integers.

--

--------------------------------------------------------------------------------



module Data.Type.Int
  ( module Data.Type.Int
  ) where

import GHC.TypeLits
import Data.Type.Ord
import Data.Type.Equality
import Data.Type.Bool

import Data.Proxy

-- | Add a sign to any type

--

data Signed a = Pos a | Neg a | Zero

-- | Type integers

--

-- ZZ represents the mathematical font for the set of integers

type ZZ = Signed Nat



type instance Compare (a :: Signed k) (b :: Signed k) = CmpSigned a b

type family IsPos (a :: ZZ) :: Bool where
  IsPos (Pos a) = 'True
  IsPos b       = 'False

-- | Compare Signed kinds when those kinds are comparable.

type family CmpSigned a b where
  CmpSigned (Neg a) (Neg b) = FlipOrdering (Compare a b)
  CmpSigned (Neg a) b = 'LT
  CmpSigned Zero (Neg a) = 'GT
  CmpSigned Zero Zero = 'EQ
  CmpSigned Zero (Pos a) = 'GT
  CmpSigned (Pos a) (Pos b) = Compare a b
  CmpSigned (Pos a) b = 'GT

-- | Always use @Zero@ instead of @Pos 0@ or @Neg 0@.

type family NormalizeInt (a :: ZZ) :: ZZ where
  NormalizeInt (Pos 0) = Zero
  NormalizeInt (Neg 0) = Zero
  NormalizeInt n = n


-- | Reverse the order of an Ordering

--

-- This should be declared in to Data.Type.Ord in base

type family FlipOrdering (o :: Ordering) :: Ordering where
  FlipOrdering 'LT = 'GT
  FlipOrdering 'EQ = 'EQ
  FlipOrdering 'GT = 'LT


-- | Absolute value

--

type family Abs (a :: ZZ) :: Nat where
  Abs (Pos a) = a
  Abs (Neg a) = a
  Abs Zero = 0

-- | Unary negation

--

type family Negate (a :: ZZ) :: ZZ where
  Negate (Pos a) = Neg a
  Negate (Neg a) = Pos a
  Negate Zero = Zero


-- | Utility family for Add

--

type family AddCmp (cmp :: Ordering) (a :: ZZ) (b :: ZZ) where
  AddCmp _ a Zero = a
  AddCmp _ Zero b = b
  AddCmp _ (Pos a) (Pos b) = Pos (a + b)
  AddCmp _ (Neg a) (Neg b) = Neg (a + b)
  AddCmp EQ _ _ = Zero
  AddCmp LT (Pos a) (Neg b) = Neg (b - a)
  AddCmp GT (Pos a) (Neg b) = Pos (a - b)
  AddCmp LT (Neg a) (Pos b) = Pos (b - a)
  AddCmp GT (Neg a) (Pos b) = Neg (a - b)

-- | Addition

type family Add (a :: ZZ) (b :: ZZ) :: ZZ where
  Add a b = AddCmp (Compare (Abs a) (Abs b)) a b

-- | Subtraction

type family Sub (a :: ZZ) (b :: ZZ) :: ZZ where
  Sub a b = Add a (Negate b)

-- | Multiplication

type family Mul (a :: ZZ) (b :: ZZ) :: ZZ where
  Mul a Zero = a
  Mul Zero b = b
  Mul (Pos a) (Pos b) = Pos (a * b)
  Mul (Pos a) (Neg b) = Neg (a * b)
  Mul (Neg a) (Pos b) = Neg (a * b)
  Mul (Neg a) (Neg b) = Pos (a * b)


-- | Exponentiation

type family Pow (a :: ZZ) (n :: Nat) :: ZZ where
  Pow Zero 0 = Pos 1 -- Following Nat from Base : 0^0 :: Natural = 1

  Pow Zero n = Zero
  Pow (Pos a) n = Pos (a ^ n)
  Pow (Neg a) n = If (Mod n 2 == 0) (Pos (a^n)) (Neg (a^n))


-- | Gives the integer associated to a type-level integer.

class KnownInt (r :: ZZ) where
  -- | Reify a type integer to an integer.

  intVal :: proxy r -> Integer

instance KnownInt Zero where
  intVal :: forall (proxy :: ZZ -> Type). proxy 'Zero -> Integer
intVal proxy 'Zero
_  = Integer
0

instance KnownNat n => KnownInt (Pos n) where
  intVal :: forall (proxy :: ZZ -> Type). proxy ('Pos n) -> Integer
intVal proxy ('Pos n)
_  = Proxy n -> Integer
forall (n :: Nat) (proxy :: Nat -> Type).
KnownNat n =>
proxy n -> Integer
natVal (Proxy n
forall {k} (t :: k). Proxy t
Proxy :: Proxy n)

instance KnownNat n => KnownInt (Neg n) where
  intVal :: forall (proxy :: ZZ -> Type). proxy ('Neg n) -> Integer
intVal proxy ('Neg n)
_  = -Proxy n -> Integer
forall (n :: Nat) (proxy :: Nat -> Type).
KnownNat n =>
proxy n -> Integer
natVal (Proxy n
forall {k} (t :: k). Proxy t
Proxy :: Proxy n)

-- | Singleton type for type-level integers.

data SZZ (z :: ZZ) where
  SPos :: KnownNat n => SNat n -> SZZ ('Pos n)
  SNeg :: KnownNat n => SNat n -> SZZ ('Neg n)
  SZero :: SZZ 'Zero

-- | Singleton for zero.

zero :: SZZ 'Zero
zero :: SZZ 'Zero
zero = SZZ 'Zero
SZero

-- | Singleton for positive integers.

--

-- >>> :t pos @3

--  pos @3 :: KnownNat 3 => SZZ ('Pos 3)

--

pos :: KnownNat n => SZZ ('Pos n)
pos :: forall (n :: Nat). KnownNat n => SZZ ('Pos n)
pos = SNat n -> SZZ ('Pos n)
forall (n :: Nat). KnownNat n => SNat n -> SZZ ('Pos n)
SPos SNat n
forall (n :: Nat). KnownNat n => SNat n
SNat

-- | Integer singleton for 1.

pos1 :: SZZ ('Pos 1)
pos1 :: SZZ ('Pos 1)
pos1 = forall (n :: Nat). KnownNat n => SZZ ('Pos n)
pos @1

-- | Integer singleton for 2.

pos2 :: SZZ ('Pos 2)
pos2 :: SZZ ('Pos 2)
pos2 = forall (n :: Nat). KnownNat n => SZZ ('Pos n)
pos @2

-- | Integer singleton for 3.

pos3 :: SZZ ('Pos 3)
pos3 :: SZZ ('Pos 3)
pos3 = forall (n :: Nat). KnownNat n => SZZ ('Pos n)
pos @3

-- | Integer singleton for 4.

pos4 :: SZZ ('Pos 4)
pos4 :: SZZ ('Pos 4)
pos4 = forall (n :: Nat). KnownNat n => SZZ ('Pos n)
pos @4

-- | Singleton for negative integers.

--

-- >>> :t neg @3

-- neg @3 :: KnownNat 3 => SZZ ('Neg 3)

--

neg :: KnownNat n => SZZ ('Neg n)
neg :: forall (n :: Nat). KnownNat n => SZZ ('Neg n)
neg = SNat n -> SZZ ('Neg n)
forall (n :: Nat). KnownNat n => SNat n -> SZZ ('Neg n)
SNeg SNat n
forall (n :: Nat). KnownNat n => SNat n
SNat

-- | Integer singleton for -1.

neg1 :: SZZ ('Neg 1)
neg1 :: SZZ ('Neg 1)
neg1 = forall (n :: Nat). KnownNat n => SZZ ('Neg n)
neg @1

-- | Integer singleton for -2.

neg2 :: SZZ ('Neg 2)
neg2 :: SZZ ('Neg 2)
neg2 = forall (n :: Nat). KnownNat n => SZZ ('Neg n)
neg @2

-- | Integer singleton for -3.

neg3 :: SZZ ('Neg 3)
neg3 :: SZZ ('Neg 3)
neg3 = forall (n :: Nat). KnownNat n => SZZ ('Neg n)
neg @3

-- | Integer singleton for -4.

neg4 :: SZZ ('Neg 4)
neg4 :: SZZ ('Neg 4)
neg4 = forall (n :: Nat). KnownNat n => SZZ ('Neg n)
neg @4


-- | Return the 'Integer' corresponding to @n@ in an @SZZ n@ value.

fromSZZ :: SZZ n -> Integer
fromSZZ :: forall (n :: ZZ). SZZ n -> Integer
fromSZZ SZZ n
SZero      = Integer
0
fromSZZ (SPos SNat n
sn)  = SNat n -> Integer
forall (n :: Nat) (proxy :: Nat -> Type).
KnownNat n =>
proxy n -> Integer
natVal SNat n
sn
fromSZZ (SNeg SNat n
sn)  = -SNat n -> Integer
forall (n :: Nat) (proxy :: Nat -> Type).
KnownNat n =>
proxy n -> Integer
natVal SNat n
sn

-- | Convert an 'Integer' into an @'SZZ' n@ value, where @n@ is a fresh

-- type-level symbol.

withSomeSZZ :: Integer -> (forall (n :: ZZ). SZZ n -> r) -> r
withSomeSZZ :: forall r. Integer -> (forall (n :: ZZ). SZZ n -> r) -> r
withSomeSZZ Integer
0 forall (n :: ZZ). SZZ n -> r
f = SZZ 'Zero -> r
forall (n :: ZZ). SZZ n -> r
f SZZ 'Zero
SZero
withSomeSZZ Integer
i forall (n :: ZZ). SZZ n -> r
f
  | Integer
i Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0    = SZZ 'Zero -> r
forall (n :: ZZ). SZZ n -> r
f SZZ 'Zero
SZero
  | Integer
i Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
0     = Integer -> (forall (n :: Nat). Maybe (SNat n) -> r) -> r
forall r. Integer -> (forall (n :: Nat). Maybe (SNat n) -> r) -> r
withSomeSNat Integer
i ((SZZ ('Pos n) -> r) -> Maybe (SNat n) -> r
forall {n :: Nat} {p}. (SZZ ('Pos n) -> p) -> Maybe (SNat n) -> p
fpos SZZ ('Pos n) -> r
forall (n :: ZZ). SZZ n -> r
f)
  | Bool
otherwise = Integer -> (forall (n :: Nat). Maybe (SNat n) -> r) -> r
forall r. Integer -> (forall (n :: Nat). Maybe (SNat n) -> r) -> r
withSomeSNat (Integer -> Integer
forall a. Num a => a -> a
negate Integer
i) ((SZZ ('Neg n) -> r) -> Maybe (SNat n) -> r
forall {n :: Nat} {p}. (SZZ ('Neg n) -> p) -> Maybe (SNat n) -> p
fneg SZZ ('Neg n) -> r
forall (n :: ZZ). SZZ n -> r
f)
  where
    fpos :: (SZZ ('Pos n) -> p) -> Maybe (SNat n) -> p
fpos SZZ ('Pos n) -> p
f' (Just (SNat n
SNat :: SNat n)) = SZZ ('Pos n) -> p
f' (SNat n -> SZZ ('Pos n)
forall (n :: Nat). KnownNat n => SNat n -> SZZ ('Pos n)
SPos (forall (n :: Nat). KnownNat n => SNat n
SNat @n))
    fpos SZZ ('Pos n) -> p
_ Maybe (SNat n)
Nothing = [Char] -> p
forall a. HasCallStack => [Char] -> a
error [Char]
"withSomeSZZ: This should never happen.\
      \ A bug report would be appreciated."

    fneg :: (SZZ ('Neg n) -> p) -> Maybe (SNat n) -> p
fneg SZZ ('Neg n) -> p
f' (Just (SNat n
SNat :: SNat n)) = SZZ ('Neg n) -> p
f' (SNat n -> SZZ ('Neg n)
forall (n :: Nat). KnownNat n => SNat n -> SZZ ('Neg n)
SNeg (forall (n :: Nat). KnownNat n => SNat n
SNat @n))
    fneg SZZ ('Neg n) -> p
_ Maybe (SNat n)
Nothing = [Char] -> p
forall a. HasCallStack => [Char] -> a
error [Char]
"withSomeSZZ: This should never happen.\
      \ A bug report would be appreciated."

-- | Convert an explicit @'SZZ' n@ value into an implicit @'KnownInt' n@

-- constraint.

withKnownInt :: SZZ n -> (KnownInt n => r) -> r
withKnownInt :: forall (n :: ZZ) r. SZZ n -> (KnownInt n => r) -> r
withKnownInt SZZ n
SZero      KnownInt n => r
r = r
KnownInt n => r
r
withKnownInt (SPos SNat n
sn)  KnownInt n => r
r = SNat n -> (KnownNat n => r) -> r
forall (n :: Nat) r. SNat n -> (KnownNat n => r) -> r
withKnownNat SNat n
sn r
KnownNat n => r
KnownInt n => r
r
withKnownInt (SNeg SNat n
sn)  KnownInt n => r
r = SNat n -> (KnownNat n => r) -> r
forall (n :: Nat) r. SNat n -> (KnownNat n => r) -> r
withKnownNat SNat n
sn r
KnownNat n => r
KnownInt n => r
r