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

{-# OPTIONS_HADDOCK not-home #-}

module Morley.Michelson.Typed.Scope.Internal.ForbidT
  ( module Morley.Michelson.Typed.Scope.Internal.ForbidT
  ) where

import Data.Type.Bool (type (||))
import GHC.TypeLits (ErrorMessage(..), Symbol, TypeError)
import Type.Errors (WhenStuck)

import Morley.Michelson.Typed.T (T(..))
import Morley.Util.Sing
import Morley.Util.Type (FailWhen)

{- $setup
>>> import Morley.Michelson.Typed (T(..), SingT(..))
>>> import Morley.Util.Peano (pattern S, pattern Z)
>>> import Data.Type.Equality ((:~:)(..))
-}

-- | Type-level symbol for type predicates used in 'ContainsT'
data TPredicateSym
  = PSOp -- ^ Contains @operation@
  | PSContract -- ^ Contains @contract 'a@
  | PSTicket -- ^ Contains @ticket 'a@
  | PSBigMap -- ^ Contains @big_map 'k 'v@
  | PSNestedBigMaps -- ^ Contains @big_map 'k (big_map 'k 'v)@
  | PSSaplingState -- ^ Contains @sapling_state 'n@
  | PSNonComparable -- ^ Contains non-comparable types

genSingletonsType ''TPredicateSym

type TPredicateSymTypeName :: TPredicateSym -> Symbol
type family TPredicateSymTypeName p where
 TPredicateSymTypeName 'PSOp            = "`operation`"
 TPredicateSymTypeName 'PSContract      = "`contract`"
 TPredicateSymTypeName 'PSTicket        = "`ticket`"
 TPredicateSymTypeName 'PSBigMap        = "`big_map`"
 TPredicateSymTypeName 'PSNestedBigMaps = "nested `big_map`s"
 TPredicateSymTypeName 'PSSaplingState  = "`sapling_state`"
 TPredicateSymTypeName 'PSNonComparable = "non-comparable types"

type ContainsT :: TPredicateSym -> T -> Bool
type family ContainsT p t where
  ContainsT 'PSOp 'TOperation = 'True
  ContainsT 'PSContract ('TContract _) = 'True
  ContainsT 'PSTicket ('TTicket _) = 'True
  ContainsT 'PSBigMap ('TBigMap _ _) = 'True
  ContainsT 'PSSaplingState ('TSaplingState _) = 'True
  ContainsT 'PSNestedBigMaps ('TBigMap _ v) = ContainsT 'PSBigMap v
  -- non-comparable
  ContainsT 'PSNonComparable 'TBls12381Fr = 'True
  ContainsT 'PSNonComparable 'TBls12381G1 = 'True
  ContainsT 'PSNonComparable 'TBls12381G2 = 'True
  ContainsT 'PSNonComparable ('TList _) = 'True
  ContainsT 'PSNonComparable ('TSet _) = 'True
  ContainsT 'PSNonComparable 'TOperation = 'True
  ContainsT 'PSNonComparable ('TContract _) = 'True
  ContainsT 'PSNonComparable ('TTicket _) = 'True
  ContainsT 'PSNonComparable ('TLambda _ _) = 'True
  ContainsT 'PSNonComparable ('TMap _ _) = 'True
  ContainsT 'PSNonComparable ('TBigMap _ _) = 'True
  ContainsT 'PSNonComparable 'TChest = 'True
  ContainsT 'PSNonComparable 'TChestKey = 'True
  ContainsT 'PSNonComparable ('TSaplingState _) = 'True
  ContainsT 'PSNonComparable ('TSaplingTransaction _) = 'True
  -- recursion
  ContainsT p ('TOption t) = ContainsT p t
  ContainsT p ('TList t) = ContainsT p t
  ContainsT p ('TSet t) = ContainsT p t
  ContainsT p ('TTicket t) = ContainsT p t
  ContainsT p ('TPair a b) = ContainsT p a || ContainsT p b
  ContainsT p ('TOr a b) = ContainsT p a || ContainsT p b
  ContainsT p ('TBigMap k v) = ContainsT p k || ContainsT p v
  ContainsT p ('TMap k v) = ContainsT p k || ContainsT p v
  -- rest
  ContainsT _ _ = 'False

type ForbidTErrorMsg :: TPredicateSym -> T -> ErrorMessage
type family ForbidTErrorMsg p t where
  ForbidTErrorMsg 'PSNestedBigMaps t =
    'Text "Nested `big_map`s found in" ':$$: 'ShowType t ':$$: 'Text "are not allowed"
  ForbidTErrorMsg 'PSNonComparable t =
    'Text "Non-comparable type"
    ':$$: 'ShowType t
    ':$$: 'Text "is not allowed in this scope"
  ForbidTErrorMsg p t =
    'Text "Type " ':<>: 'Text (TPredicateSymTypeName p) ':<>: 'Text " found in"
    ':$$: 'ShowType t
    ':$$: 'Text "is not allowed in this scope"

{- | Constraint for classes forbidding type presence based on predicate defined
by 'TPredicateSym'.

Not just a type alias in order to be able to partially apply it (e.g. in
'Each').

Reports errors when a type does not satisfy predicate:

>>> () :: ForbidT PSOp TOperation => ()
...
... Type `operation` found in
... 'TOperation
... is not allowed in this scope
...
>>> () :: ForbidT PSContract (TContract TUnit) => ()
...
... Type `contract` found in
... 'TContract 'TUnit
... is not allowed in this scope
...
>>> () :: ForbidT PSTicket (TTicket TUnit) => ()
...
... Type `ticket` found in
... 'TTicket 'TUnit
... is not allowed in this scope
...
>>> () :: ForbidT PSBigMap (TBigMap TUnit TUnit) => ()
...
... Type `big_map` found in
... 'TBigMap 'TUnit 'TUnit
... is not allowed in this scope
...
>>> () :: ForbidT PSSaplingState (TSaplingState Z) => ()
...
... Type `sapling_state` found in
... 'TSaplingState 'Z
... is not allowed in this scope
...
>>> () :: ForbidT PSNestedBigMaps (TBigMap TUnit (TBigMap TUnit TUnit)) => ()
...
... Nested `big_map`s found in
... 'TBigMap 'TUnit ('TBigMap 'TUnit 'TUnit)
... are not allowed
...

When the type is ambiguous or polymorphic, suggests adding the corresponding
constraint:

>>> (const () :: ForbidOp t => f t -> ()) undefined
...
... Can't check if type
... t0
... contains `operation`. Perhaps you need to add
... ForbidT 'PSOp t0
... constraint? You can also try adding a type annotation.
...

This constraint implies @ContainsT ~ False@:

>>> :{
foo :: ForbidT p t => ContainsT p t :~: False
foo = Refl
:}
-}
type ForbidT :: TPredicateSym -> T -> Constraint
class ForbidManyT (ForbidT p t) '[p] t => ForbidT p t
instance ForbidManyT (ForbidT p t) '[p] t => ForbidT p t

type ForbidOp = ForbidT 'PSOp                       -- ^ Convenience synonym
type ForbidContract = ForbidT 'PSContract           -- ^ Convenience synonym
type ForbidTicket = ForbidT 'PSTicket               -- ^ Convenience synonym
type ForbidSaplingState = ForbidT 'PSSaplingState   -- ^ Convenience synonym
type ForbidBigMap = ForbidT 'PSBigMap               -- ^ Convenience synonym
type ForbidNestedBigMaps = ForbidT 'PSNestedBigMaps -- ^ Convenience synonym
type ForbidNonComparable = ForbidT 'PSNonComparable -- ^ Convenience synonym

{- | Class abstracting multiple 'ForbidT' constraints.

This version accepts a custom error message displayed if 'ContainsT' gets stuck.
You likely want to use 'ForbidManyT', which sets this message.

It's a class and not a recursive type family because recursive type families
take noticeably longer to compile.
-}
type ForbidManyT' :: Constraint -> [TPredicateSym] -> [TPredicateSym] -> T -> Constraint
class ForbidManySuperC ps t => ForbidManyT' c ps' ps t where
  type ForbidManySuperC ps t :: Constraint

instance ForbidManyT' c ps' '[] t where
  type ForbidManySuperC '[] _ = ()
instance
  ( ForbidManyT' c ps' ps t
  , DelayedContainsTCheck c ps' p t (ContainsT p t)
  ) => ForbidManyT' c ps' (p : ps) t where
    type ForbidManySuperC (p : ps) t =
      -- NB: WhenStuck isn't in the superclass constraints by design, as
      -- FailWhen unifies it with False.
      ( ContainsT p t ~ 'False
      , ForbidManySuperC ps t
      )

-- | This is a horrible hack designed to delay GHC's test for equality of
-- ContainsT ... with False. Unfortunately, due to the use of incoherence, more
-- type signatures will be required downstream. But the impact seems somewhat
-- minimal in practice.
class a ~ 'False => DelayedContainsTCheck c ps' p t a
instance DelayedContainsTCheck c ps' p t 'False
instance {-# incoherent #-}
  ( WhenStuck t (TypeError (ForbidManyTStuckErr c ps' t))
  , FailWhen a (ForbidTErrorMsg p t)
  ) => DelayedContainsTCheck c ps' p t a

{- | Abstracts multiple 'ForbidT' constraints. Accepts a constraint that will be
suggested in an error message in case 'ContainsT' gets stuck. This is used with
scope constraints in "Morley.Michelson.Typed.Scope.Internal.Scopes".
-}
type ForbidManyT c ps = ForbidManyT' c ps ps

-- | Given a list of 'TPredicateSym', pretty-print a list of corresponding type
-- names. Uses 'TPredicateSymTypeName' for a single type, separates the last two
-- types with @or@, others with @,@.
--
-- Doesn't do the oxford comma, because it's rather tricky to implement, and
-- compilation types aren't great already.
type TPredicateSymTypeNames :: [TPredicateSym] -> ErrorMessage
type family TPredicateSymTypeNames ps where
  TPredicateSymTypeNames '[] = 'Text ""
  TPredicateSymTypeNames '[p] = 'Text (TPredicateSymTypeName p)
  TPredicateSymTypeNames '[p, q] =
     'Text (TPredicateSymTypeName p) ':<>: 'Text " or " ':<>: 'Text (TPredicateSymTypeName q)
  TPredicateSymTypeNames (p : ps) =
    'Text (TPredicateSymTypeName p) ':<>: 'Text ", " ':<>: TPredicateSymTypeNames ps

type ForbidManyTStuckErr :: Constraint -> [TPredicateSym] -> T -> ErrorMessage
type ForbidManyTStuckErr c ps t =
  'Text "Can't check if type" ':$$:
  'ShowType t ':$$:
  'Text "contains " ':<>: TPredicateSymTypeNames ps ':<>:
  'Text ". Perhaps you need to add" ':$$:
  'ShowType c ':$$:
  'Text "constraint? You can also try adding a type annotation."