-- SPDX-FileCopyrightText: 2021 Tocqueville Group
--
-- SPDX-License-Identifier: LicenseRef-MIT-TQ

-- | Utility for 'PeanoNatural'
--
-- At the moment, we have no other libriaries with that would
-- provide effective implementation of type-level natural
-- numbers. So we define our own data type called @PeanoNatural@.
-- Using this type one may extract a term-level natural number
-- from a type-level one quite simply.
module Util.PeanoNatural
  ( PeanoNatural (Zero, Succ, One, Two)
  , toPeanoNatural
  , toPeanoNatural'
  , fromPeanoNatural
  , singPeanoVal
  ) where

import Data.Singletons (SingI (..))
import Text.PrettyPrint.Leijen.Text (integer)
import Util.Peano (Nat (Z, S), SingNat(..), ToPeano)
import qualified GHC.TypeNats as GHC (Nat)

import Michelson.Printer.Util

-- | PeanoNatural data type
--
-- The @PN@ constructor stores @s :: SingNat n@ and @k :: Natural@
-- with the following invariant:
-- if @PN s k :: PeanoNatural n@, then @k == n@.
-- This definition allows extracting values of Natural
-- without O(n) conversion from @SingNat n@.
data PeanoNatural (n :: Nat) = PN !(SingNat n) !Natural

deriving stock instance Show (PeanoNatural n)
deriving stock instance Eq (PeanoNatural n)

instance RenderDoc (PeanoNatural n) where
  renderDoc :: RenderContext -> PeanoNatural n -> Doc
renderDoc RenderContext
_ = Integer -> Doc
integer (Integer -> Doc)
-> (PeanoNatural n -> Integer) -> PeanoNatural n -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Natural -> Integer
forall a. Integral a => a -> Integer
toInteger (Natural -> Integer)
-> (PeanoNatural n -> Natural) -> PeanoNatural n -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PeanoNatural n -> Natural
forall (n :: Nat). PeanoNatural n -> Natural
fromPeanoNatural

instance NFData (PeanoNatural n) where
  rnf :: PeanoNatural n -> ()
rnf (PN SingNat n
s Natural
_) = SingNat n -> ()
forall a. NFData a => a -> ()
rnf SingNat n
s

data MatchPS n where
  PS_Match :: PeanoNatural n -> MatchPS ('S n)
  PS_Mismatch :: MatchPS n

matchPS :: PeanoNatural n -> MatchPS n
matchPS :: PeanoNatural n -> MatchPS n
matchPS (PN (SS SingNat n
m) Natural
k) = PeanoNatural n -> MatchPS ('S n)
forall (n :: Nat). PeanoNatural n -> MatchPS ('S n)
PS_Match (SingNat n -> Natural -> PeanoNatural n
forall (n :: Nat). SingNat n -> Natural -> PeanoNatural n
PN SingNat n
m (Natural
k Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
- Natural
1))
matchPS PeanoNatural n
_ = MatchPS n
forall (n :: Nat). MatchPS n
PS_Mismatch

-- | Patterns 'Zero' and 'Succ'
-- We introduce pattern synonyms 'Zero' and 'Succ' assuming that
-- 'Zero' and 'Succ' cover all possible cases that satisfy the invariant.
-- Using these patterns, we also avoid cases when `k /= peanoValSing @n s`
pattern Zero :: () => (n ~ 'Z) => PeanoNatural n
pattern $bZero :: PeanoNatural n
$mZero :: forall r (n :: Nat).
PeanoNatural n -> ((n ~ 'Z) => r) -> (Void# -> r) -> r
Zero = PN SZ 0

pattern Succ :: () => (n ~ 'S m) => PeanoNatural m -> PeanoNatural n
pattern $bSucc :: PeanoNatural m -> PeanoNatural n
$mSucc :: forall r (n :: Nat).
PeanoNatural n
-> (forall (m :: Nat). (n ~ 'S m) => PeanoNatural m -> r)
-> (Void# -> r)
-> r
Succ s <- (matchPS -> PS_Match s) where
  Succ (PN SingNat m
n Natural
k) = SingNat ('S m) -> Natural -> PeanoNatural ('S m)
forall (n :: Nat). SingNat n -> Natural -> PeanoNatural n
PN (SingNat m -> SingNat ('S m)
forall (n :: Nat). SingNat n -> SingNat ('S n)
SS SingNat m
n) (Natural
kNatural -> Natural -> Natural
forall a. Num a => a -> a -> a
+Natural
1)
{-# COMPLETE Zero, Succ #-}

-- | The following patterns are introduced for convenience.
-- This allow us to avoid writing @Succ (Succ Zero)@ in
-- several places.
pattern One :: () => (n ~ ('S 'Z)) => PeanoNatural n
pattern $bOne :: PeanoNatural n
$mOne :: forall r (n :: Nat).
PeanoNatural n -> ((n ~ 'S 'Z) => r) -> (Void# -> r) -> r
One = Succ Zero

pattern Two :: () => (n ~ ('S ('S 'Z))) => PeanoNatural n
pattern $bTwo :: PeanoNatural n
$mTwo :: forall r (n :: Nat).
PeanoNatural n -> ((n ~ 'S ('S 'Z)) => r) -> (Void# -> r) -> r
Two = Succ One

fromPeanoNatural :: forall n. PeanoNatural n -> Natural
fromPeanoNatural :: PeanoNatural n -> Natural
fromPeanoNatural (PN SingNat n
_ Natural
n) = Natural
n

-- | toPeanoNatural and toPeanoNatural' connect PeanoNatural with
-- natural numbers known at run-time.
singPeanoVal :: forall (n :: Nat). SingNat n -> Natural
singPeanoVal :: SingNat n -> Natural
singPeanoVal = \case
  SingNat n
SZ   -> Natural
0
  SS SingNat n
a -> Natural
1 Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
+ SingNat n -> Natural
forall (n :: Nat). SingNat n -> Natural
singPeanoVal SingNat n
a

toPeanoNatural :: forall n. SingI n => PeanoNatural n
toPeanoNatural :: PeanoNatural n
toPeanoNatural = let pSing :: Sing n
pSing = SingI n => Sing n
forall k (a :: k). SingI a => Sing a
sing @n in
  SingNat n -> Natural -> PeanoNatural n
forall (n :: Nat). SingNat n -> Natural -> PeanoNatural n
PN Sing n
SingNat n
pSing (Natural -> PeanoNatural n) -> Natural -> PeanoNatural n
forall a b. (a -> b) -> a -> b
$ SingNat n -> Natural
forall (n :: Nat). SingNat n -> Natural
singPeanoVal Sing n
SingNat n
pSing

toPeanoNatural'
  :: forall (n :: GHC.Nat). SingI (ToPeano n) => PeanoNatural (ToPeano n)
toPeanoNatural' :: PeanoNatural (ToPeano n)
toPeanoNatural' = SingI (ToPeano n) => PeanoNatural (ToPeano n)
forall (n :: Nat). SingI n => PeanoNatural n
toPeanoNatural @(ToPeano n)