{-# LANGUAGE NoStarIsType #-}
{-# LANGUAGE ExistentialQuantification #-}

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

-- |

--

-- 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)