{-# language DataKinds #-}
{-# language ExplicitForAll #-}
{-# language KindSignatures #-}
{-# language MagicHash #-}
{-# language RankNTypes #-}
{-# language ScopedTypeVariables #-}
{-# language TypeOperators #-}

module Arithmetic.Nat
  ( -- * Addition
    plus
    -- * Subtraction
  , monus
    -- * Successor
  , succ
    -- * Compare
  , testEqual
  , testLessThan
  , testLessThanEqual
  , testZero
  , (=?)
  , (<?)
  , (<=?)
    -- * Constants
  , zero
  , one
  , two
  , three
  , constant
    -- * Convert
  , demote
  , with
  ) where

import Prelude hiding (succ)

import Arithmetic.Types
import Arithmetic.Unsafe ((:=:)(Eq), type (<=)(Lte))
import Arithmetic.Unsafe (Nat(Nat),type (<)(Lt))
import GHC.Exts (Proxy#,proxy#)
import GHC.TypeNats (type (+),KnownNat,natVal')

-- | Infix synonym of 'testLessThan'.
(<?) :: Nat a -> Nat b -> Maybe (a < b)
{-# inline (<?) #-}
<? :: Nat a -> Nat b -> Maybe (a < b)
(<?) = Nat a -> Nat b -> Maybe (a < b)
forall (a :: Nat) (b :: Nat). Nat a -> Nat b -> Maybe (a < b)
testLessThan

-- | Infix synonym of 'testLessThanEqual'.
(<=?) :: Nat a -> Nat b -> Maybe (a <= b)
{-# inline (<=?) #-}
<=? :: Nat a -> Nat b -> Maybe (a <= b)
(<=?) = Nat a -> Nat b -> Maybe (a <= b)
forall (a :: Nat) (b :: Nat). Nat a -> Nat b -> Maybe (a <= b)
testLessThanEqual

-- | Infix synonym of 'testEqual'.
(=?) :: Nat a -> Nat b -> Maybe (a :=: b)
{-# inline (=?) #-}
=? :: Nat a -> Nat b -> Maybe (a :=: b)
(=?) = Nat a -> Nat b -> Maybe (a :=: b)
forall (a :: Nat) (b :: Nat). Nat a -> Nat b -> Maybe (a :=: b)
testEqual

-- | Is the first argument strictly less than the second
-- argument?
testLessThan :: Nat a -> Nat b -> Maybe (a < b)
{-# inline testLessThan #-}
testLessThan :: Nat a -> Nat b -> Maybe (a < b)
testLessThan (Nat Int
x) (Nat Int
y) = if Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
y
  then (a < b) -> Maybe (a < b)
forall a. a -> Maybe a
Just a < b
forall (a :: Nat) (b :: Nat). a < b
Lt
  else Maybe (a < b)
forall a. Maybe a
Nothing

-- | Is the first argument less-than-or-equal-to the second
-- argument?
testLessThanEqual :: Nat a -> Nat b -> Maybe (a <= b)
{-# inline testLessThanEqual #-}
testLessThanEqual :: Nat a -> Nat b -> Maybe (a <= b)
testLessThanEqual (Nat Int
x) (Nat Int
y) = if Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
y
  then (a <= b) -> Maybe (a <= b)
forall a. a -> Maybe a
Just a <= b
forall (a :: Nat) (b :: Nat). a <= b
Lte
  else Maybe (a <= b)
forall a. Maybe a
Nothing

-- | Are the two arguments equal to one another?
testEqual :: Nat a -> Nat b -> Maybe (a :=: b)
{-# inline testEqual #-}
testEqual :: Nat a -> Nat b -> Maybe (a :=: b)
testEqual (Nat Int
x) (Nat Int
y) = if Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
y
  then (a :=: b) -> Maybe (a :=: b)
forall a. a -> Maybe a
Just a :=: b
forall (a :: Nat) (b :: Nat). a :=: b
Eq
  else Maybe (a :=: b)
forall a. Maybe a
Nothing

-- | Is zero equal to this number or less than it?
testZero :: Nat a -> Either (0 :=: a) (0 < a)
{-# inline testZero #-}
testZero :: Nat a -> Either (0 :=: a) (0 < a)
testZero (Nat Int
x) = case Int
x of
  Int
0 -> (0 :=: a) -> Either (0 :=: a) (0 < a)
forall a b. a -> Either a b
Left 0 :=: a
forall (a :: Nat) (b :: Nat). a :=: b
Eq
  Int
_ -> (0 < a) -> Either (0 :=: a) (0 < a)
forall a b. b -> Either a b
Right 0 < a
forall (a :: Nat) (b :: Nat). a < b
Lt

-- | Add two numbers.
plus :: Nat a -> Nat b -> Nat (a + b)
{-# inline plus #-}
plus :: Nat a -> Nat b -> Nat (a + b)
plus (Nat Int
x) (Nat Int
y) = Int -> Nat (a + b)
forall (n :: Nat). Int -> Nat n
Nat (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
y)

-- | The successor of a number.
succ :: Nat a -> Nat (a + 1)
{-# inline succ #-}
succ :: Nat a -> Nat (a + 1)
succ Nat a
n = Nat a -> Nat 1 -> Nat (a + 1)
forall (a :: Nat) (b :: Nat). Nat a -> Nat b -> Nat (a + b)
plus Nat a
n Nat 1
one

-- | Subtract the second argument from the first argument.
monus :: Nat a -> Nat b -> Maybe (Difference a b)
{-# inline monus #-}
monus :: Nat a -> Nat b -> Maybe (Difference a b)
monus (Nat Int
a) (Nat Int
b) = let c :: Int
c = Int
a Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
b in if Int
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0
  then Difference a b -> Maybe (Difference a b)
forall a. a -> Maybe a
Just (Nat Any -> ((Any + b) :=: a) -> Difference a b
forall (a :: Nat) (b :: Nat) (c :: Nat).
Nat c -> ((c + b) :=: a) -> Difference a b
Difference (Int -> Nat Any
forall (n :: Nat). Int -> Nat n
Nat Int
c) (Any + b) :=: a
forall (a :: Nat) (b :: Nat). a :=: b
Eq)
  else Maybe (Difference a b)
forall a. Maybe a
Nothing

-- | The number zero.
zero :: Nat 0
{-# inline zero #-}
zero :: Nat 0
zero = Int -> Nat 0
forall (n :: Nat). Int -> Nat n
Nat Int
0

-- | The number one.
one :: Nat 1
{-# inline one #-}
one :: Nat 1
one = Int -> Nat 1
forall (n :: Nat). Int -> Nat n
Nat Int
1

-- | The number two.
two :: Nat 2
{-# inline two #-}
two :: Nat 2
two = Int -> Nat 2
forall (n :: Nat). Int -> Nat n
Nat Int
2

-- | The number three.
three :: Nat 3
{-# inline three #-}
three :: Nat 3
three = Int -> Nat 3
forall (n :: Nat). Int -> Nat n
Nat Int
3

-- | Use GHC's built-in type-level arithmetic to create a witness
-- of a type-level number. This only reduces if the number is a
-- constant.
constant :: forall n. KnownNat n => Nat n
{-# inline constant #-}
constant :: Nat n
constant = Int -> Nat n
forall (n :: Nat). Int -> Nat n
Nat (Natural -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Proxy# n -> Natural
forall (n :: Nat). KnownNat n => Proxy# n -> Natural
natVal' (Proxy# n
forall k (a :: k). Proxy# a
proxy# :: Proxy# n)))

-- | Extract the 'Int' from a 'Nat'. This is intended to be used
-- at a boundary where a safe interface meets the unsafe primitives
-- on top of which it is built.
demote :: Nat n -> Int
{-# inline demote #-}
demote :: Nat n -> Int
demote (Nat Int
n) = Int
n

-- | Run a computation on a witness of a type-level number. The
-- argument 'Int' must be greater than or equal to zero. This is
-- not checked. Failure to upload this invariant will lead to a
-- segfault.
with :: Int -> (forall n. Nat n -> a) -> a
{-# inline with #-}
with :: Int -> (forall (n :: Nat). Nat n -> a) -> a
with Int
i forall (n :: Nat). Nat n -> a
f = Nat Any -> a
forall (n :: Nat). Nat n -> a
f (Int -> Nat Any
forall (n :: Nat). Int -> Nat n
Nat Int
i)