-- 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 _ = 0
instance KnownPeano a => KnownPeano ('S a) where
  peanoVal _ = peanoVal' @a + 1

peanoVal' :: forall n. KnownPeano n => Natural
peanoVal' = peanoVal (Proxy @n)

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

instance KnownPeano a => MockableConstraint (KnownPeano a) where
  provideConstraintUnsafe = 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 SZ = ()
  rnf (SS n) = rnf n

type instance Sing = SingNat

instance SingI 'Z where
  sing = SZ
instance (SingI n, KnownPeano n) => SingI ('S n) where
  sing = SS (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 = unsafeCoerce $ Dict @(RequireLongerOrSameLength '[] 'Z)

instance MockableConstraint (RequireLongerThan l a) where
  provideConstraintUnsafe = unsafeCoerce $ Dict @(RequireLongerThan '[()] 'Z)

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

requireLongerThan
  :: Rec any stk
  -> Sing n
  -> Maybe (Dict (RequireLongerThan stk n))
requireLongerThan RNil _           = Nothing
requireLongerThan (_ :& _xs) SZ    = Just Dict
requireLongerThan (_ :& xs) (SS n) = do
  Dict <- requireLongerThan xs n
  return Dict

requireLongerOrSameLength
  :: Rec any stk
  -> Sing n
  -> Maybe (Dict (RequireLongerOrSameLength stk n))
requireLongerOrSameLength _ SZ             = Just Dict
requireLongerOrSameLength RNil (SS _)      = Nothing
requireLongerOrSameLength (_ :& xs) (SS n) = do
  Dict <- requireLongerOrSameLength xs n
  return Dict