-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

-- | 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 Morley.Util.PeanoNatural
  ( PeanoNatural (Zero, Succ, One, Two)
  , toPeanoNatural
  , toPeanoNatural'
  , fromPeanoNatural
  , singPeanoVal
  , eqPeanoNat
  , singPeanoNat
  , addPeanoNat
  ) where

import Data.Singletons (SingI(..))
import Data.Type.Equality (TestEquality(..), (:~:))
import Fmt (build)
import GHC.TypeNats (Nat)

import Morley.Michelson.Printer.Util
import Morley.Util.Peano

-- | 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 :: Peano) = PN !(SingNat n) !Natural

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

instance RenderDoc (PeanoNatural n) where
  renderDoc _ = build . fromPeanoNatural

instance NFData (PeanoNatural n) where
  rnf (PN s _) = rnf s

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

matchPS :: PeanoNatural n -> MatchPS n
matchPS (PN (SS m) k) = PS_Match (PN m (k - 1))
matchPS _ = 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 Zero = PN SZ 0

pattern Succ :: () => (n ~ 'S m) => PeanoNatural m -> PeanoNatural n
pattern Succ s <- (matchPS -> PS_Match s) where
  Succ (PN n k) = PN (SS n) (k + 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 One = Succ Zero

pattern Two :: () => (n ~ ('S ('S 'Z))) => PeanoNatural n
pattern Two = Succ One

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

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

toPeanoNatural :: forall n. SingI n => PeanoNatural n
toPeanoNatural = let pSing = sing @n in
  PN pSing $ singPeanoVal pSing

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

-- | Check if two 'PeanoNatural's are equal, and return evidence of type equality
-- between the corresponding 'Nat's.
eqPeanoNat :: PeanoNatural a -> PeanoNatural b -> Maybe (a :~: b)
eqPeanoNat (PN a _) (PN b _) = testEquality a b

-- | Get a corresponding 'SingNat' from a 'PeanoNatural'.
singPeanoNat :: PeanoNatural a -> SingNat a
singPeanoNat (PN n _) = n

addPeanoNat :: PeanoNatural a -> PeanoNatural b -> PeanoNatural (AddPeano a b)
addPeanoNat (PN a n) (PN b m) = PN (peanoSingAdd a b) (n + m)