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

{-# OPTIONS_GHC -Wno-orphans #-}
{-# OPTIONS_GHC -Wno-unused-top-binds #-}
-- This is needed because `singletons-th` generates redundant constraints.
-- GHC may not throw the warning always, see #791.
{-# OPTIONS_GHC -Wno-redundant-constraints #-}
{-# LANGUAGE UndecidableSuperClasses #-}
{-# LANGUAGE DeriveLift #-}

-- | 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 Morley.Util.Peano
  ( -- * General
    Peano
  , pattern S
  , pattern Z
  , ToPeano
  , FromPeano
  , SingNat (SZ, SS)

  , peanoSing
  , peanoSing'
  , withPeanoSingI
  , withSomePeano

  -- * Utility type synonyms
  , IsoNatPeano
  , SingIPeano

  -- * Peano Arithmetic
  , type (>)
  , type (>=)
  , peanoSingDecrement
  , peanoSingAdd
  , Decrement
  , AddPeano
  , SubPeano
  , MinPeano
  , MaxPeano

  -- * Lists
  , Length
  , At
  , Drop
  , LazyTake
  , Take
  , Head
  , Tail

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

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

  -- * Length constraints 'Dict'ionaries
  , isGreaterThan
  , isGreaterEqualThan

  -- * Inductive proofs
  , additivity
  , associativity
  , minIdempotency
  , commutativity
  , transitivity

  -- * Helpers
  , toNatural
  , someSingNat
  ) where

import Data.Constraint (Dict(..), (\\))
import Data.Data (Data(..))
import Data.Singletons (Sing, SingI(..), SomeSing(..))
import Data.Type.Equality (type (:~:)(..))
import Data.Vinyl (Rec(..))
import Data.Vinyl.Core (Head, Tail)
import Data.Vinyl.TypeLevel (Nat(..), RLength)
import GHC.TypeLits (ErrorMessage(..), TypeError)
import GHC.TypeNats (type (+), type (-))
import GHC.TypeNats qualified as GHC
import Language.Haskell.TH.Syntax (Lift)
import Unsafe.Coerce (unsafeCoerce)

import Morley.Util.Sing (genSingletonsType)
import Morley.Util.StubbedProof
import Morley.Util.Type (FailUnless)

-- This is very obviously a false positive.
{-# ANN module ("HLint: ignore Use 'natVal' from Universum" :: Text) #-}
{-# ANN module ("HLint: ignore Use 'someNatVal' from Universum" :: Text) #-}

----------------------------------------------------------------------------
-- 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

deriving stock instance Eq Nat
deriving stock instance Show Nat
deriving stock instance Generic Nat
deriving stock instance Data Nat
deriving anyclass instance NFData Nat

$(genSingletonsType ''Nat)

deriving stock instance Show (SingNat (n :: Nat))
deriving stock instance Eq (SingNat (n :: Nat))
deriving stock instance Lift (SingNat (n :: Nat))
instance NFData (SingNat (n :: Nat)) where
  rnf SZ = ()
  rnf (SS n) = rnf n


-- | A constraint asserting that GHC's @Nat@ @n@ and @Peano@ @p@ are the same (up to an
-- isomorphism)
type IsoNatPeano (n :: GHC.Nat) (p :: Peano) = (n ~ FromPeano p, ToPeano n ~ p)

-- | A synonym for @SingI (ToPeano n)@. Essentially requires that we can construct a 'Peano'
-- singleton for a given 'Nat'
type SingIPeano (n :: GHC.Nat) = SingI (ToPeano n)

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

-- | Get the peano singleton for a given type-level nat literal.
--
-- >>> peanoSing @2
-- SS (SS SZ)
peanoSing :: forall (n :: GHC.Nat). SingIPeano n => SingNat (ToPeano n)
peanoSing = sing @(ToPeano n)

-- | Same as 'peanoSing', but only requires 'KnownNat' instance.
--
-- Witnesses half the equivalence between @KnownNat n@ and @SingI (ToPeano n)@
peanoSing' :: forall (n :: GHC.Nat). KnownNat n => SingNat (ToPeano n)
peanoSing' = go (natVal @n Proxy)
  where
  go :: forall m. Natural -> SingNat m
  go = \case
    0 -> SZ \\ (unsafeCoerce Refl :: m :~: 'Z)
    n -> SS (go @(Decrement m) (n - 1)) \\ (unsafeCoerce Refl :: m :~: 'S (Decrement m))

-- | Run a computation requiring @SingI (ToPeano n)@ in a context which only has
-- @KnownNat n@. Mostly useful when used with 'SomeNat'
withPeanoSingI :: forall (n :: GHC.Nat) r. KnownNat n => (SingIPeano n => r) -> r
withPeanoSingI act = act \\ go @(ToPeano n) (natVal @n Proxy)
  where
  go :: forall (m :: Peano). Natural -> Dict (SingI m)
  go = \case
    0 -> Dict \\ (unsafeCoerce Refl :: m :~: 'Z)
    n -> Dict \\ go @(Decrement m) (n - 1) \\ (unsafeCoerce Refl :: m :~: 'S (Decrement m))

-- | Lift a given term-level 'Natural' to the type level for a given computation. The computation is
-- expected to accept a 'Proxy' for the sake of convenience: it's easier to get at the type-level
-- natural with @ScopedTypeVariables@ when pattern-matching on the proxy, e.g.
--
-- > (x :: Natural) `withSomePeano` \(_ :: Proxy n) -> doSomeNatComputation @n
withSomePeano :: Natural -> (forall n. (KnownNat n, SingIPeano n) => Proxy n -> r) -> r
withSomePeano n f = case someNatVal n of
  SomeNat (pn :: Proxy n) -> withPeanoSingI @n $ f pn


----------------------------------------------------------------------------
-- Peano Arithmetic
----------------------------------------------------------------------------

type family Decrement (a :: Peano) :: Peano where
  Decrement 'Z = TypeError ('Text "Expected n > 0")
  Decrement ('S n) = n

-- | Utility to 'Decrement' a Peano 'Sing'leton.
--
-- Useful when dealing with the constraint
peanoSingDecrement :: Sing n -> Maybe (Sing (Decrement n))
peanoSingDecrement = \case
  SZ -> Nothing
  SS n -> pure n

-- | 'Peano' naturals comparisson
type family (>) (x :: Peano) (y :: Peano) :: Bool where
  'Z   > _    = 'False
  'S _ > 'Z   = 'True
  'S x > 'S y = x > y

-- | 'Peano' naturals comparisson
type family (>=) (x :: Peano) (y :: Peano) :: Bool where
  _      >= 'Z     = 'True
  'Z     >= _      = 'False
  ('S x) >= ('S y) = x >= y

-- | 'Peano' naturals addition
type family AddPeano (n :: Peano) (m :: Peano) :: Peano where
  AddPeano 'Z x = x
  AddPeano ('S x) y = 'S (AddPeano x y)

-- | 'Peano' naturals subtraction
type family SubPeano (n :: Peano) (m :: Peano) :: Peano where
  SubPeano 'Z ('S m) = TypeError ('Text "Subtracting " ':<>: 'ShowType (FromPeano ('S m))
    ':<>: 'Text " from zero")
  SubPeano n 'Z = n
  SubPeano ('S n) ('S m) = SubPeano n m

-- | Out of two 'Peano' naturals, return the smaller one
type family MinPeano (n :: Peano) (m :: Peano) :: Peano where
  MinPeano _ 'Z = 'Z
  MinPeano 'Z _ = 'Z
  MinPeano ('S n) ('S m) = 'S (MinPeano n m)

-- | Out of two 'Peano' naturals, return the larger one
type family MaxPeano (n :: Peano) (m :: Peano) :: Peano where
  MaxPeano n 'Z = n
  MaxPeano 'Z m = m
  MaxPeano ('S n) ('S m) = 'S (MaxPeano n m)

-- | Singleton addition
peanoSingAdd :: SingNat n -> SingNat m -> SingNat (AddPeano n m)
peanoSingAdd (SS n) (SS (m :: SingNat m')) = SS (SS (peanoSingAdd n m)) \\ associativity @m' n
peanoSingAdd SZ m = m
peanoSingAdd n SZ = n \\ commutativity n SZ

----------------------------------------------------------------------------
-- 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 tried to access a 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

-- | A type-level version of 'Data.List.take'. Note that in many cases,
-- using 'LazyTake' instead will improve type inference.
type family Take (n :: Peano) (s :: [k]) :: [k] where
  Take  'Z _ = '[]
  Take _ '[] = '[]
  Take ('S n) (a ': s) = a ': Take n s

-- | A "lazier" version of 'Take'. @LazyTake n xs@ will always produce
-- a list of exactly @n@ elements. If @xs@ has less than @n@ elements, then
-- some of the elements of the result will be stuck. Similarly, if some tail
-- of @xs@ is stuck or ambiguous, then elements of the result past that point
-- will be stuck.
--
-- @
-- LazyTake ('ToPeano' 4) '[1,2,3,4,5] ~ '[1,2,3,4]
-- LazyTake (ToPeano 4) (1 ': 2 ': 3 ': 4 ': Any) ~ '[1,2,3,4]
-- LazyTake (ToPeano 4) '[1,2] ~
--   '[1, 2, Head '[], Head (Tail '[])]
-- @
--
-- @LazyTake@ is often better than @Take@ for driving type inference.
-- For example, given the constraint
--
-- @
-- xs ~ Take (ToPeano 3) xs ++ q ': Drop (ToPeano 4) xs
-- @
--
-- GHC can't infer anything about the shape of @xs@. However, the constraint
--
-- @
-- xs ~ LazyTake (ToPeano 3) xs ++ q ': Drop (ToPeano 4) xs
-- @
--
-- will allow GHC to infer
--
-- @
-- xs ~ x1 ': x2 ': x3 ': q ': Drop (ToPeano 4) xs
-- @
type LazyTake :: Peano -> [k] -> [k]
type family LazyTake n xs where
  LazyTake 'Z _ = '[]
  LazyTake ('S n) xs = Head xs ': LazyTake n (Tail xs)

----------------------------------------------------------------------------
-- 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 =
    FailUnless
      (IsLongerThan l a)
      ('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 =
    FailUnless
      (IsLongerOrSameLength l a)
      ('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

----------------------------------------------------------------------------
-- 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

----------------------------------------------------------------------------
-- Arith constraints 'Dict'ionaries
----------------------------------------------------------------------------

isGreaterThan
  :: Sing a -> Sing b
  -> Maybe (Dict ((a > b) ~ 'True))
isGreaterThan SZ _          = Nothing
isGreaterThan (SS _) SZ     = pure Dict
isGreaterThan (SS a) (SS b) = isGreaterThan a b

isGreaterEqualThan
  :: Sing a -> Sing b
  -> Maybe (Dict ((a >= b) ~ 'True))
isGreaterEqualThan _ SZ          = pure Dict
isGreaterEqualThan SZ _          = Nothing
isGreaterEqualThan (SS a) (SS b) = isGreaterEqualThan a b

----------------------------------------------------------------------------
-- Inductive proofs
----------------------------------------------------------------------------

-- | Proof that for naturals, @k + (m + 1) = n@ entails @n > k@
additivity :: forall k m n. AddPeano k ('S m) ~ n
           => SingNat k -> n > k :~: 'True
additivity k = stubProof $ case k of
  SZ -> Refl
  _ -> Refl \\ lemma k \\ lemma2 @_ @m k \\ associativity @m k
  where

    lemma2 :: forall k' m'. SingNat k' -> 'S (AddPeano k' m') > k' :~: 'True
    lemma2 SZ = Refl
    lemma2 (SS k') = Refl \\ lemma2 @_ @m' k'

    lemma :: SingNat k' -> 'S k' > k' :~: 'True
    lemma SZ = Refl
    lemma (SS n') = Refl \\ lemma n'

-- | Proof that for naturals, @x + (y + 1) = (x + y) + 1@
associativity :: forall y x. SingNat x -> AddPeano x ('S y) :~: 'S (AddPeano x y)
associativity x = stubProof $ case x of
  SZ -> Refl
  SS x' -> Refl \\ associativity @y x'

-- | Proof that @x + y = y + x@
commutativity :: SingNat x -> SingNat y -> AddPeano x y :~: AddPeano y x
commutativity x y = stubProof $ case x of
  SZ -> case y of
    SZ -> Refl
    SS y' -> Refl \\ commutativity x y'
  SS (x' :: SingNat x') -> Refl \\ associativity @x' y \\ commutativity x' y

-- | Proof that for naturals, @min(n, n) = n@
minIdempotency :: SingNat n -> MinPeano n n :~: n
minIdempotency n = stubProof $ case n of
  SZ -> Refl
  SS n' -> Refl \\ minIdempotency n'

-- | Proof that for naturals, @x >= y > z@ implies @x > z@
transitivity :: (x >= y ~ 'True, y > z ~ 'True)
             => SingNat x -> SingNat y -> SingNat z -> x > z :~: 'True
transitivity x y z = stubProof $ case (x, y, z) of
  (SS _, SS _, SZ) -> Refl
  (SS x', SS y', SS z') -> Refl \\ transitivity x' y' z'
  (SZ, y', _) -> case y' of

----------------------------------------------------------------------------
-- Helpers
----------------------------------------------------------------------------

toNatural :: Peano -> Natural
toNatural Z = 0
toNatural (S x) = 1 + toNatural x

someSingNat :: Natural -> SomeSing Peano
someSingNat 0 = SomeSing SZ
someSingNat n = case someSingNat (n - 1) of
  SomeSing sn -> SomeSing (SS sn)