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

-- | General type utilities.
module Morley.Util.Type
  ( IsEq
  , type (==)
  , If
  , type (++)
  , IsElem
  , type (/)
  , type (//)
  , FailUnlessEqual
  , FailUnlessEqualElse
  , FailWhen
  , FailWhenElse
  , FailWhenElsePoly
  , FailUnless
  , FailUnlessElse
  , FailUnlessElsePoly
  , AllUnique
  , RequireAllUnique
  , ReifyList (..)
  , PatternMatch
  , PatternMatchL
  , KnownList (..)
  , KList (..)
  , RSplit
  , rsplit
  , Some1 (..)
  , recordToSomeList

  , onFirst
  , knownListFromSingI
  ) where

import Data.Constraint (Dict(..), (\\))
import Data.Eq.Singletons (DefaultEq)
import Data.Singletons (SingI(sing))
import Data.Type.Bool (If, Not, type (&&))
import Data.Type.Equality (type (==))
import Data.Vinyl.Core (Rec(..))
import Data.Vinyl.Functor qualified as Vinyl
import Data.Vinyl.Recursive (recordToList, rmap)
import Data.Vinyl.TypeLevel qualified as Vinyl
import GHC.TypeLits (ErrorMessage(..), Symbol, TypeError)
import Prelude.Singletons (SList(..))

-- $setup
-- >>> import GHC.TypeLits (TypeError, ErrorMessage (..))
-- >>> import Type.Errors (DelayError)
-- >>> import Morley.Util.Peano (Peano, pattern S, pattern Z)
-- >>> import Data.Constraint ((:-)(..), Dict(..))

-- | Equality constraint in form of a typeclass.
class a ~ b => IsEq a b
instance a ~ b => IsEq a b

infixr 5 ++

-- | Append for type-level lists. We use this synonym instead of using the one
-- from "Data.Vinyl.TypeLevel" directly because the latter is missing a fixity
-- declaration. See the
-- [vinyl pull request](https://github.com/VinylRecords/Vinyl/pull/165).

type (++) :: [k] -> [k] -> [k]
type as ++ bs = as Vinyl.++ bs

type family IsElem (a :: k) (l :: [k]) :: Bool where
  IsElem _ '[] = 'False
  IsElem a (a ': _) = 'True
  IsElem a (_ ': as) = IsElem a as

-- | Remove all occurences of the given element.
type family (l :: [k]) / (a :: k) where
  '[] / _ = '[]
  (a ': xs) / a = xs / a
  (b ': xs) / a = b ': (xs / a)

-- | Difference between two lists.
type family (l1 :: [k]) // (l2 :: [k]) :: [k] where
  l // '[] = l
  l // (x ': xs) = (l / x) // xs

{- | Constrain two types to be equal. If they are found not to be, produce
an error message. If they are shown to be equal, impose an additional
given constraint.

>>> :k! FailUnlessEqualElse Int Int ('Text "This should not result in a failure") ()
FailUnlessEqualElse Int Int ('Text "This should not result in a failure") () :: Constraint
= (() :: Constraint, Int ~ Int)

>>> :k! FailUnlessEqualElse Bool Int ('Text "This should result in a failure") ()
FailUnlessEqualElse Bool Int ('Text "This should result in a failure") () :: Constraint
= ((TypeError ...), Bool ~ Int)

the @~@ constraint might seem redundant, but, without it, GHC would reject

>>> :{
foo :: FailUnlessEqualElse a b ('Text "MyError") () => a -> b
foo = id
:}

GHC needs to \"see\" the type equality @~@ in order to actually \"learn\"
something from a type family's result.

We use 'DefaultEq' here, rather than 'Data.Type.Equality.==', so this will
work with type variables known to be equal, even if nothing is known about
them concretely. For example, the following will compile:

>>> :{
foo :: FailUnlessEqualElse a b ('Text "MyError") () => a -> b
foo x = x
--
bar :: a -> a
bar = foo
:}

If we used @(==)@, then bar would be rejected.

@(==)@ has its place (see the comments below it in "Data.Type.Equality")
but I don't think it has anything to offer us here. In particular, the equality
constraint we bundle up seems to win us all the reasoning power that @(==)@ would
provide and more. The following adaptation of the example in the
@Data.Type.Equality@ comments compiles:

>>> :{
foo :: FailUnlessEqualElse (Maybe a) (Maybe b) ('Text "yo") ()
       :- FailUnlessEqualElse a b ('Text "heya") ()
foo = Sub Dict
:}

In this example, the entire consequent follows from the equality constraint
in the antecedent; the `FailUnlessElsePoly` part of it is irrelevant, so we
don't need to be able to reason from it. If we /were/ reasoning solely using
@(==)@, @foo@ would be rejected because the error messages are different.
-}
type FailUnlessEqualElse :: forall k. k -> k -> ErrorMessage -> Constraint -> Constraint
type FailUnlessEqualElse a b err els =
  ( FailUnlessElsePoly (a `DefaultEq` b) err els
  , a ~ b
  )

-- | Constrain two types to be equal, and produce an error message if
-- they are found not to be.
--
-- >>> :k! FailUnlessEqual Int Int ('Text "This should not result in a failure")
-- FailUnlessEqual Int Int ('Text "This should not result in a failure") :: Constraint
-- = (() :: Constraint, Int ~ Int)
--
-- >>> :k! FailUnlessEqual Bool Int ('Text "This should result in a failure")
-- FailUnlessEqual Bool Int ('Text "This should result in a failure") :: Constraint
-- = ((TypeError ...), Bool ~ Int)
type FailUnlessEqual :: forall k. k -> k -> ErrorMessage -> Constraint
type FailUnlessEqual a b err = FailUnlessEqualElse a b err ()

-- | A version of 'FailWhenElsePoly' that imposes an equality constraint on its
-- 'Bool' argument.
type FailWhenElse :: Bool -> ErrorMessage -> Constraint -> Constraint
type FailWhenElse cond msg els = FailUnlessEqualElse cond 'False msg els

-- | Fail with given error if the condition does not hold. Otherwise, return the
-- third argument.
type FailWhenElsePoly :: forall k. Bool -> ErrorMessage -> k -> k
type family FailWhenElsePoly cond msg els where
  FailWhenElsePoly 'True msg _els = TypeError msg
  FailWhenElsePoly 'False _ els = els

-- | A version of 'FailUnlessElsePoly' that imposes an equality constraint on its
-- 'Bool' argument.
type FailUnlessElse :: Bool -> ErrorMessage -> Constraint -> Constraint
type FailUnlessElse cond msg els = FailUnlessEqualElse cond 'True msg els

-- | Fail with given error if the condition does not hold. Otherwise,
-- return the third argument.
type FailUnlessElsePoly :: forall k. Bool -> ErrorMessage -> k -> k
type family FailUnlessElsePoly b e k where
  FailUnlessElsePoly 'False msg _els = TypeError msg
  FailUnlessElsePoly 'True _ els = els

-- | Fail with given error if the condition does not hold.
type FailUnless :: Bool -> ErrorMessage -> Constraint
type FailUnless cond msg = FailUnlessElse cond msg ()

-- | Fail with given error if the condition holds.
type FailWhen :: Bool -> ErrorMessage -> Constraint
type FailWhen cond msg = FailWhenElse cond msg ()

type family AllUnique (l :: [k]) :: Bool where
  AllUnique '[] = 'True
  AllUnique (x : xs) = Not (IsElem x xs) && AllUnique xs

type RequireAllUnique desc l = RequireAllUnique' desc l l

type family RequireAllUnique' (desc :: Symbol) (l :: [k]) (origL :: [k]) :: Constraint where
  RequireAllUnique' _ '[] _ = ()
  RequireAllUnique' desc (x : xs) origL =
    FailWhenElse
       (IsElem x xs)
       ('Text "Duplicated " ':<>: 'Text desc ':<>: 'Text ":" ':$$:
        'ShowType x ':$$:
        'Text "Full list: " ':<>:
        'ShowType origL
       )
       (RequireAllUnique' desc xs origL)

-- | Make sure given type is evaluated.
-- This type family fits only for types of 'Type' kind.
type family PatternMatch (a :: Type) :: Constraint where
  PatternMatch Int = ((), ())
  PatternMatch _ = ()

type family PatternMatchL (l :: [k]) :: Constraint where
  PatternMatchL '[] = ((), ())
  PatternMatchL _ = ()

-- | Bring type-level list at term-level using given function
-- to demote its individual elements.
class ReifyList (c :: k -> Constraint) (l :: [k]) where
  reifyList :: (forall a. c a => Proxy a -> r) -> [r]

instance ReifyList c '[] where
  reifyList :: forall r. (forall (a :: k). c a => Proxy a -> r) -> [r]
reifyList forall (a :: k). c a => Proxy a -> r
_ = []

instance (c x, ReifyList c xs) => ReifyList c (x ': xs) where
  reifyList :: forall r. (forall (a :: a). c a => Proxy a -> r) -> [r]
reifyList forall (a :: a). c a => Proxy a -> r
reifyElem = Proxy x -> r
forall (a :: a). c a => Proxy a -> r
reifyElem (forall (t :: a). Proxy t
forall {k} (t :: k). Proxy t
Proxy @x) r -> [r] -> [r]
forall a. a -> [a] -> [a]
: forall k (c :: k -> Constraint) (l :: [k]) r.
ReifyList c l =>
(forall (a :: k). c a => Proxy a -> r) -> [r]
reifyList @_ @c @xs Proxy a -> r
forall (a :: a). c a => Proxy a -> r
reifyElem

-- | Similar to @SingI []@, but does not require individual elements to be also
-- instance of @SingI@.
class KnownList l where
  klist :: KList l
instance KnownList '[] where
  klist :: KList '[]
klist = KList '[]
forall k. KList '[]
KNil
instance KnownList xs => KnownList (x ': xs) where
  klist :: KList (x : xs)
klist = Proxy x -> Proxy xs -> KList (x : xs)
forall {k} (xs :: [k]) (x :: k).
KnownList xs =>
Proxy x -> Proxy xs -> KList (x : xs)
KCons Proxy x
forall {k} (t :: k). Proxy t
Proxy Proxy xs
forall {k} (t :: k). Proxy t
Proxy

-- | Given a type-level list that is 'SingI', construct evidence that it is also
-- a 'KnownList'. Note that 'KnownList' is weaker, hence this construction
-- is always possible.
knownListFromSingI :: forall xs. SingI xs => Dict (KnownList xs)
knownListFromSingI :: forall {k} (xs :: [k]). SingI xs => Dict (KnownList xs)
knownListFromSingI = SList xs -> Dict (KnownList xs)
forall {k} (ys :: [k]). SList ys -> Dict (KnownList ys)
go (forall (a :: [k]). SingI a => Sing a
forall {k} (a :: k). SingI a => Sing a
sing @xs)
  where
    go :: forall ys. SList ys -> Dict (KnownList ys)
    go :: forall {k} (ys :: [k]). SList ys -> Dict (KnownList ys)
go = \case
      SList ys
SNil -> Dict (KnownList ys)
forall (a :: Constraint). a => Dict a
Dict
      SCons Sing n1
_ Sing n2
ys -> Dict (KnownList ys)
KnownList n2 => Dict (KnownList ys)
forall (a :: Constraint). a => Dict a
Dict (KnownList n2 => Dict (KnownList ys))
-> Dict (KnownList n2) -> Dict (KnownList ys)
forall (c :: Constraint) e r. HasDict c e => (c => r) -> e -> r
\\ SList n2 -> Dict (KnownList n2)
forall {k} (ys :: [k]). SList ys -> Dict (KnownList ys)
go Sing n2
SList n2
ys

-- | @Data.List.Singletons.SList@ analogue for 'KnownList'.
data KList (l :: [k]) where
  KNil :: KList '[]
  KCons :: KnownList xs => Proxy x -> Proxy xs -> KList (x ': xs)

type RSplit l r = KnownList l

-- | Split a record into two pieces.
rsplit
  :: forall k (l :: [k]) (r :: [k]) f.
      (RSplit l r)
  => Rec f (l ++ r) -> (Rec f l, Rec f r)
rsplit :: forall k (l :: [k]) (r :: [k]) (f :: k -> *).
RSplit l r =>
Rec f (l ++ r) -> (Rec f l, Rec f r)
rsplit = case forall (l :: [k]). KnownList l => KList l
forall {k} (l :: [k]). KnownList l => KList l
klist @l of
  KList l
KNil -> (Rec f '[]
forall {u} (a :: u -> *). Rec a '[]
RNil, )
  KCons{} -> \(f r
x :& Rec f rs
r) ->
    let (Rec f xs
x1, Rec f r
r1) = Rec f (xs ++ r) -> (Rec f xs, Rec f r)
forall k (l :: [k]) (r :: [k]) (f :: k -> *).
RSplit l r =>
Rec f (l ++ r) -> (Rec f l, Rec f r)
rsplit Rec f rs
Rec f (xs ++ r)
r
    in (f r
x f r -> Rec f xs -> Rec f (r : xs)
forall {u} (a :: u -> *) (r :: u) (rs :: [u]).
a r -> Rec a rs -> Rec a (r : rs)
:& Rec f xs
x1, Rec f r
r1)

-- | A value of type parametrized with /some/ type parameter.
data Some1 (f :: k -> Type) =
  forall a. Some1 (f a)

deriving stock instance (forall a. Show (f a)) => Show (Some1 f)

recordToSomeList :: Rec f l -> [Some1 f]
recordToSomeList :: forall {k} (f :: k -> *) (l :: [k]). Rec f l -> [Some1 f]
recordToSomeList = Rec (Const (Some1 f)) l -> [Some1 f]
forall {u} a (rs :: [u]). Rec (Const a) rs -> [a]
recordToList (Rec (Const (Some1 f)) l -> [Some1 f])
-> (Rec f l -> Rec (Const (Some1 f)) l) -> Rec f l -> [Some1 f]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall (x :: k). f x -> Const (Some1 f) x)
-> Rec f l -> Rec (Const (Some1 f)) l
forall {u} (f :: u -> *) (g :: u -> *) (rs :: [u]).
(forall (x :: u). f x -> g x) -> Rec f rs -> Rec g rs
rmap (Some1 f -> Const (Some1 f) x
forall k a (b :: k). a -> Const a b
Vinyl.Const (Some1 f -> Const (Some1 f) x)
-> (f x -> Some1 f) -> f x -> Const (Some1 f) x
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f x -> Some1 f
forall k (f :: k -> *) (a :: k). f a -> Some1 f
Some1)

-- | Utility function to help transform the first argument of a binfunctor.
onFirst :: Bifunctor p => p a c -> (a -> b) -> p b c
onFirst :: forall (p :: * -> * -> *) a c b.
Bifunctor p =>
p a c -> (a -> b) -> p b c
onFirst = ((a -> b) -> p a c -> p b c) -> p a c -> (a -> b) -> p b c
forall a b c. (a -> b -> c) -> b -> a -> c
flip (a -> b) -> p a c -> p b c
forall a b c. (a -> b) -> p a c -> p b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first