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

{-# OPTIONS_GHC -Wno-orphans #-}
{-# LANGUAGE UndecidableSuperClasses #-}
{-# OPTIONS_GHC -Wno-redundant-constraints #-}

-- | Type-nat utilities.
--
-- We take Peano numbers as base for operations because they make it
-- much easer to prove things to compiler. Their performance does not
-- seem to introduce a problem, because we use nats primarily along with
-- stack which is a linked list with similar performance characteristics.
--
-- Many of things we introduce here are covered in @type-natural@ package,
-- but unfortunatelly it does not work with GHC 8.6 at the moment of writing
-- this module. We use 'Data.Vinyl' as source of Peano @Nat@ for now.
module Util.Peano
  ( -- * General
    Peano
    -- @gromak: I think we should not export Nat, but we want to
    -- export `Z` and `S` and I do not know how to do it without
    -- exporting `Nat`.
  , Nat (Z, S)
  , ToPeano
  , FromPeano
  , KnownPeano (..)
  , SingNat (SZ, SS)
  , peanoVal'
  , peanoValSing

  -- * Lists
  , Length
  , At
  , Drop
  , Take

    -- * Morley-specific utils
  , IsLongerThan
  , LongerThan
  , RequireLongerThan
  , IsLongerOrSameLength
  , LongerOrSameLength
  , RequireLongerOrSameLength

  -- * Length constraints 'Dict'ionaries
  , requireLongerThan
  , requireLongerOrSameLength
  ) where

import Data.Constraint (Dict(..))
import Data.Singletons (Sing, SingI(..))
import Data.Vinyl (Rec(..))
import Data.Vinyl.TypeLevel (Nat(..), RLength)
import GHC.TypeLits (ErrorMessage(..), TypeError)
import GHC.TypeNats (type (+), type (-))
import qualified GHC.TypeNats as GHC (Nat)
import Unsafe.Coerce (unsafeCoerce)

import Util.Type

----------------------------------------------------------------------------
-- General
----------------------------------------------------------------------------

-- | A convenient alias.
--
-- We are going to use 'Peano' numbers for type-dependent logic and
-- normal 'Nat's in user API, need to distinguish them somehow.
type Peano = Nat

type family ToPeano (n :: GHC.Nat) :: Peano where
  ToPeano 0 = 'Z
  ToPeano a = 'S (ToPeano (a - 1))

type family FromPeano (n :: Peano) :: GHC.Nat where
  FromPeano 'Z = 0
  FromPeano ('S a) = 1 + FromPeano a

class KnownPeano (n :: Peano) where
  peanoVal :: proxy n -> Natural

instance KnownPeano 'Z where
  peanoVal :: proxy 'Z -> Natural
peanoVal _ = 0
instance KnownPeano a => KnownPeano ('S a) where
  peanoVal :: proxy ('S a) -> Natural
peanoVal _ = KnownPeano a => Natural
forall (n :: Peano). KnownPeano n => Natural
peanoVal' @a Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
+ 1

peanoVal' :: forall n. KnownPeano n => Natural
peanoVal' :: Natural
peanoVal' = Proxy n -> Natural
forall (n :: Peano) (proxy :: Peano -> *).
KnownPeano n =>
proxy n -> Natural
peanoVal (Proxy n
forall k (t :: k). Proxy t
Proxy @n)

-- | Get runtime value from singleton.
peanoValSing :: forall n. KnownPeano n => Sing n -> Natural
peanoValSing :: Sing n -> Natural
peanoValSing _ = KnownPeano n => Natural
forall (n :: Peano). KnownPeano n => Natural
peanoVal' @n

instance KnownPeano a => MockableConstraint (KnownPeano a) where
  provideConstraintUnsafe :: Dict (KnownPeano a)
provideConstraintUnsafe = Dict (KnownPeano a)
forall (a :: Constraint). a => Dict a
Dict

data SingNat (n :: Nat) where
  SZ :: SingNat 'Z
  SS :: (SingI n, KnownPeano n) => SingNat n -> SingNat ('S n)

deriving stock instance Show (SingNat (n :: Nat))
deriving stock instance Eq (SingNat (n :: Nat))
instance NFData (SingNat (n :: Nat)) where
  rnf :: SingNat n -> ()
rnf SZ = ()
  rnf (SS n :: SingNat n
n) = SingNat n -> ()
forall a. NFData a => a -> ()
rnf SingNat n
n

type instance Sing = SingNat

instance SingI 'Z where
  sing :: Sing 'Z
sing = Sing 'Z
SingNat 'Z
SZ
instance (SingI n, KnownPeano n) => SingI ('S n) where
  sing :: Sing ('S n)
sing = SingNat n -> SingNat ('S n)
forall (n :: Peano).
(SingI n, KnownPeano n) =>
SingNat n -> SingNat ('S n)
SS (SingI n => Sing n
forall k (a :: k). SingI a => Sing a
sing @n)

----------------------------------------------------------------------------
-- Lists
----------------------------------------------------------------------------

type family Length l :: Peano where
  Length l = RLength l

type family At (n :: Peano) s where
  At 'Z (x ': _) = x
  At ('S n) (_ ': xs) = At n xs
  At a '[] =
    TypeError
      ('Text "You try to access to non-existing element of the stack, n = " ':<>:
         'ShowType (FromPeano a))

type family Drop (n :: Peano) (s :: [k]) :: [k] where
  Drop  'Z s = s
  Drop ('S _) '[] = '[]
  Drop ('S n) (_ ': s) = Drop n s

type family Take (n :: Peano) (s :: [k]) :: [k] where
  Take  'Z _ = '[]
  Take _ '[] = '[]
  Take ('S n) (a ': s) = a ': Take n s

----------------------------------------------------------------------------
-- Morley-specific utils
----------------------------------------------------------------------------

-- Note that we could define type families to return 'Constraint' instead
-- of defining standalone constraint in form `c ~ 'True`, but apparently
-- such constraint would be weaker, e. g. there is an example when with
-- current approach there is no warning, but if we change the approach
-- to return 'Constraint' from type family then GHC complains about
-- non-exhaustive patterns (@gromak).
-- Also we use these `Bool` type families in more than one place, we generate
-- two constraints: one gives more information to GHC and another one produces
-- better error messages on failure.

-- | Comparison of type-level naturals, as a function.
--
-- It is as lazy on the list argument as possible - there is no
-- need to know the whole list if the natural argument is small enough.
-- This property is important if we want to be able to extract reusable
-- parts of code which are aware only of relevant part of stack.
type family IsLongerThan (l :: [k]) (a :: Peano) :: Bool where
  IsLongerThan (_ ': _) 'Z = 'True
  IsLongerThan (_ ': xs) ('S a) = IsLongerThan xs a
  IsLongerThan '[] _ = 'False

-- | Comparison of type-level naturals, as a constraint.
type LongerThan l a = IsLongerThan l a ~ 'True

-- | Similar to 'IsLongerThan', but returns 'True' when list length
-- equals to the passed number.
type family IsLongerOrSameLength (l :: [k]) (a :: Peano) :: Bool where
  IsLongerOrSameLength _ 'Z = 'True
  IsLongerOrSameLength (_ ': xs) ('S a) = IsLongerOrSameLength xs a
  IsLongerOrSameLength '[] ('S _) = 'False

-- | 'IsLongerOrSameLength' in form of constraint that gives most
-- information to GHC.
type LongerOrSameLength l a = IsLongerOrSameLength l a ~ 'True

{- | Evaluates list length.

This type family is a best-effort attempt to display neat error messages
when list is known only partially.

For instance, when called on @Int ': Int ': s@, the result will be
@OfLengthWithTail 2 s@ - compare with result of simple 'Length' -
@1 + 1 + Length s@.

For concrete types this will be identical to calling @FromPeano (Length l)@.
-}
type family OfLengthWithTail (acc :: GHC.Nat) (l :: [k]) :: GHC.Nat where
  OfLengthWithTail a '[] = a
  OfLengthWithTail a (_ ': xs) = OfLengthWithTail (a + 1) xs

type LengthWithTail l = OfLengthWithTail 0 l

-- | Comparison of type-level naturals, raises human-readable compile error
-- when does not hold.
--
-- Here we use the same approach as for 'RequireLongerOrSameLength', this
-- type family is internal.
type family RequireLongerThan' (l :: [k]) (a :: Nat) :: Constraint where
  RequireLongerThan' l a =
    If (IsLongerThan l a)
       (() :: Constraint)
       (TypeError
          ('Text "Stack element #" ':<>: 'ShowType (FromPeano a) ':<>:
           'Text " is not accessible" ':$$:
           'Text "Current stack has size of only " ':<>:
           'ShowType (LengthWithTail l) ':<>:
           'Text ":" ':$$: 'ShowType l
           ))

class (RequireLongerThan' l a, LongerThan l a) =>
  RequireLongerThan (l :: [k]) (a :: Peano)
instance (RequireLongerThan' l a, LongerThan l a) =>
  RequireLongerThan l a

-- | 'IsLongerOrSameLength' in form of constraint that produces
-- good error message. Should be used together with 'LongerThan'
-- because 'LongerThan' gives GHC more information.
-- We use it in combination, so that it gives enough information to
-- GHC and also producess good error messages.
type family RequireLongerOrSameLength' (l :: [k]) (a :: Peano) :: Constraint where
  RequireLongerOrSameLength' l a =
    If (IsLongerOrSameLength l a)
       (() :: Constraint)
       (TypeError
          ('Text "Expected stack with length >= " ':<>: 'ShowType (FromPeano a) ':$$:
           'Text "Current stack has size of only " ':<>:
           'ShowType (LengthWithTail l) ':<>:
           'Text ":" ':$$: 'ShowType l
           ))

-- | We can have
-- `RequireLongerOrSameLength = (RequireLongerOrSameLength' l a, LongerOrSameLength l a)`,
-- but apparently the printed error message can be caused by `LongerOrSameLength`
-- rather than `RequireLongerOrSameLength'`.
-- We do not know for sure how it all works, but we think that if we require constraint X before
-- Y (using multiple `=>`s) then X will always be evaluated first.
class (RequireLongerOrSameLength' l a, LongerOrSameLength l a) =>
  RequireLongerOrSameLength (l :: [k]) (a :: Peano)
instance (RequireLongerOrSameLength' l a, LongerOrSameLength l a) =>
  RequireLongerOrSameLength l a

instance MockableConstraint (RequireLongerOrSameLength l a) where
  provideConstraintUnsafe :: Dict (RequireLongerOrSameLength l a)
provideConstraintUnsafe = Dict (RequireLongerOrSameLength '[] 'Z)
-> Dict (RequireLongerOrSameLength l a)
forall a b. a -> b
unsafeCoerce (Dict (RequireLongerOrSameLength '[] 'Z)
 -> Dict (RequireLongerOrSameLength l a))
-> Dict (RequireLongerOrSameLength '[] 'Z)
-> Dict (RequireLongerOrSameLength l a)
forall a b. (a -> b) -> a -> b
$ RequireLongerOrSameLength '[] 'Z =>
Dict (RequireLongerOrSameLength '[] 'Z)
forall (a :: Constraint). a => Dict a
Dict @(RequireLongerOrSameLength '[] 'Z)

instance MockableConstraint (RequireLongerThan l a) where
  provideConstraintUnsafe :: Dict (RequireLongerThan l a)
provideConstraintUnsafe = Dict (RequireLongerThan '[()] 'Z) -> Dict (RequireLongerThan l a)
forall a b. a -> b
unsafeCoerce (Dict (RequireLongerThan '[()] 'Z) -> Dict (RequireLongerThan l a))
-> Dict (RequireLongerThan '[()] 'Z)
-> Dict (RequireLongerThan l a)
forall a b. (a -> b) -> a -> b
$ RequireLongerThan '[()] 'Z => Dict (RequireLongerThan '[()] 'Z)
forall (a :: Constraint). a => Dict a
Dict @(RequireLongerThan '[()] 'Z)

----------------------------------------------------------------------------
-- Length constraints 'Dict'ionaries
----------------------------------------------------------------------------

requireLongerThan
  :: Rec any stk
  -> Sing n
  -> Maybe (Dict (RequireLongerThan stk n))
requireLongerThan :: Rec any stk -> Sing n -> Maybe (Dict (RequireLongerThan stk n))
requireLongerThan RNil _           = Maybe (Dict (RequireLongerThan stk n))
forall a. Maybe a
Nothing
requireLongerThan (_ :& _xs :: Rec any rs
_xs) SZ    = Dict (RequireLongerThan stk n)
-> Maybe (Dict (RequireLongerThan stk n))
forall a. a -> Maybe a
Just Dict (RequireLongerThan stk n)
forall (a :: Constraint). a => Dict a
Dict
requireLongerThan (_ :& xs :: Rec any rs
xs) (SS n) = do
  Dict (RequireLongerThan rs n)
Dict <- Rec any rs -> Sing n -> Maybe (Dict (RequireLongerThan rs n))
forall k (any :: k -> *) (stk :: [k]) (n :: Peano).
Rec any stk -> Sing n -> Maybe (Dict (RequireLongerThan stk n))
requireLongerThan Rec any rs
xs Sing n
SingNat n
n
  Dict (RequireLongerThan stk n)
-> Maybe (Dict (RequireLongerThan stk n))
forall (m :: * -> *) a. Monad m => a -> m a
return Dict (RequireLongerThan stk n)
forall (a :: Constraint). a => Dict a
Dict

requireLongerOrSameLength
  :: Rec any stk
  -> Sing n
  -> Maybe (Dict (RequireLongerOrSameLength stk n))
requireLongerOrSameLength :: Rec any stk
-> Sing n -> Maybe (Dict (RequireLongerOrSameLength stk n))
requireLongerOrSameLength _ SZ             = Dict (RequireLongerOrSameLength stk n)
-> Maybe (Dict (RequireLongerOrSameLength stk n))
forall a. a -> Maybe a
Just Dict (RequireLongerOrSameLength stk n)
forall (a :: Constraint). a => Dict a
Dict
requireLongerOrSameLength RNil (SS _)      = Maybe (Dict (RequireLongerOrSameLength stk n))
forall a. Maybe a
Nothing
requireLongerOrSameLength (_ :& xs :: Rec any rs
xs) (SS n) = do
  Dict (RequireLongerOrSameLength rs n)
Dict <- Rec any rs
-> Sing n -> Maybe (Dict (RequireLongerOrSameLength rs n))
forall k (any :: k -> *) (stk :: [k]) (n :: Peano).
Rec any stk
-> Sing n -> Maybe (Dict (RequireLongerOrSameLength stk n))
requireLongerOrSameLength Rec any rs
xs Sing n
SingNat n
n
  Dict (RequireLongerOrSameLength stk n)
-> Maybe (Dict (RequireLongerOrSameLength stk n))
forall (m :: * -> *) a. Monad m => a -> m a
return Dict (RequireLongerOrSameLength stk n)
forall (a :: Constraint). a => Dict a
Dict