{-
(c) The University of Glasgow 2006
(c) The GRASP/AQUA Project, Glasgow University, 1992-1999
-}

{-# LANGUAGE CPP #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE TypeFamilies #-}

module TcTypeable(mkTypeableBinds) where

#include "HsVersions.h"

import GhcPrelude

import BasicTypes ( Boxity(..), neverInlinePragma, SourceText(..) )
import TcBinds( addTypecheckedBinds )
import IfaceEnv( newGlobalBinder )
import TyCoRep( Type(..), TyLit(..) )
import TcEnv
import TcEvidence ( mkWpTyApps )
import TcRnMonad
import HscTypes ( lookupId )
import PrelNames
import TysPrim ( primTyCons )
import TysWiredIn ( tupleTyCon, sumTyCon, runtimeRepTyCon
                  , vecCountTyCon, vecElemTyCon
                  , nilDataCon, consDataCon )
import Name
import Id
import Type
import TyCon
import DataCon
import Module
import HsSyn
import DynFlags
import Bag
import Var ( VarBndr(..) )
import CoreMap
import Constants
import Fingerprint(Fingerprint(..), fingerprintString, fingerprintFingerprints)
import Outputable
import FastString ( FastString, mkFastString, fsLit )

import Control.Monad.Trans.State
import Control.Monad.Trans.Class (lift)
import Data.Maybe ( isJust )
import Data.Word( Word64 )

{- Note [Grand plan for Typeable]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The overall plan is this:

1. Generate a binding for each module p:M
   (done in TcTypeable by mkModIdBindings)
       M.$trModule :: GHC.Types.Module
       M.$trModule = Module "p" "M"
   ("tr" is short for "type representation"; see GHC.Types)

   We might want to add the filename too.
   This can be used for the lightweight stack-tracing stuff too

   Record the Name M.$trModule in the tcg_tr_module field of TcGblEnv

2. Generate a binding for every data type declaration T in module M,
       M.$tcT :: GHC.Types.TyCon
       M.$tcT = TyCon ...fingerprint info...
                      $trModule
                      "T"
                      0#
                      kind_rep

   Here 0# is the number of arguments expected by the tycon to fully determine
   its kind. kind_rep is a value of type GHC.Types.KindRep, which gives a
   recipe for computing the kind of an instantiation of the tycon (see
   Note [Representing TyCon kinds: KindRep] later in this file for details).

   We define (in TyCon)

        type TyConRepName = Name

   to use for these M.$tcT "tycon rep names". Note that these must be
   treated as "never exported" names by Backpack (see
   Note [Handling never-exported TyThings under Backpack]). Consequently
   they get slightly special treatment in RnModIface.rnIfaceDecl.

3. Record the TyConRepName in T's TyCon, including for promoted
   data and type constructors, and kinds like * and #.

   The TyConRepName is not an "implicit Id".  It's more like a record
   selector: the TyCon knows its name but you have to go to the
   interface file to find its type, value, etc

4. Solve Typeable constraints.  This is done by a custom Typeable solver,
   currently in TcInteract, that use M.$tcT so solve (Typeable T).

There are many wrinkles:

* The timing of when we produce this bindings is rather important: they must be
  defined after the rest of the module has been typechecked since we need to be
  able to lookup Module and TyCon in the type environment and we may be
  currently compiling GHC.Types (where they are defined).

* GHC.Prim doesn't have any associated object code, so we need to put the
  representations for types defined in this module elsewhere. We chose this
  place to be GHC.Types. TcTypeable.mkPrimTypeableBinds is responsible for
  injecting the bindings for the GHC.Prim representions when compiling
  GHC.Types.

* TyCon.tyConRepModOcc is responsible for determining where to find
  the representation binding for a given type. This is where we handle
  the special case for GHC.Prim.

* To save space and reduce dependencies, we need use quite low-level
  representations for TyCon and Module.  See GHC.Types
  Note [Runtime representation of modules and tycons]

* The KindReps can unfortunately get quite large. Moreover, the simplifier will
  float out various pieces of them, resulting in numerous top-level bindings.
  Consequently we mark the KindRep bindings as noinline, ensuring that the
  float-outs don't make it into the interface file. This is important since
  there is generally little benefit to inlining KindReps and they would
  otherwise strongly affect compiler performance.

* In general there are lots of things of kind *, * -> *, and * -> * -> *. To
  reduce the number of bindings we need to produce, we generate their KindReps
  once in GHC.Types. These are referred to as "built-in" KindReps below.

* Even though KindReps aren't inlined, this scheme still has more of an effect on
  compilation time than I'd like. This is especially true in the case of
  families of type constructors (e.g. tuples and unboxed sums). The problem is
  particularly bad in the case of sums, since each arity-N tycon brings with it
  N promoted datacons, each with a KindRep whose size also scales with N.
  Consequently we currently simply don't allow sums to be Typeable.

  In general we might consider moving some or all of this generation logic back
  to the solver since the performance hit we take in doing this at
  type-definition time is non-trivial and Typeable isn't very widely used. This
  is discussed in #13261.

-}

-- | Generate the Typeable bindings for a module. This is the only
-- entry-point of this module and is invoked by the typechecker driver in
-- 'tcRnSrcDecls'.
--
-- See Note [Grand plan for Typeable] in TcTypeable.
mkTypeableBinds :: TcM TcGblEnv
mkTypeableBinds
  = do { -- Create a binding for $trModule.
         -- Do this before processing any data type declarations,
         -- which need tcg_tr_module to be initialised
       ; tcg_env <- mkModIdBindings
         -- Now we can generate the TyCon representations...
         -- First we handle the primitive TyCons if we are compiling GHC.Types
       ; (tcg_env, prim_todos) <- setGblEnv tcg_env mkPrimTypeableTodos

         -- Then we produce bindings for the user-defined types in this module.
       ; setGblEnv tcg_env $
    do { mod <- getModule
       ; let tycons = filter needs_typeable_binds (tcg_tcs tcg_env)
             mod_id = case tcg_tr_module tcg_env of  -- Should be set by now
                        Just mod_id -> mod_id
                        Nothing     -> pprPanic "tcMkTypeableBinds" (ppr tycons)
       ; traceTc "mkTypeableBinds" (ppr tycons)
       ; this_mod_todos <- todoForTyCons mod mod_id tycons
       ; mkTypeRepTodoBinds (this_mod_todos : prim_todos)
       } }
  where
    needs_typeable_binds tc
      | tc `elem` [runtimeRepTyCon, vecCountTyCon, vecElemTyCon]
      = False
      | otherwise =
          isAlgTyCon tc
       || isDataFamilyTyCon tc
       || isClassTyCon tc

{- *********************************************************************
*                                                                      *
            Building top-level binding for $trModule
*                                                                      *
********************************************************************* -}

mkModIdBindings :: TcM TcGblEnv
mkModIdBindings
  = do { mod <- getModule
       ; loc <- getSrcSpanM
       ; mod_nm        <- newGlobalBinder mod (mkVarOcc "$trModule") loc
       ; trModuleTyCon <- tcLookupTyCon trModuleTyConName
       ; let mod_id = mkExportedVanillaId mod_nm (mkTyConApp trModuleTyCon [])
       ; mod_bind      <- mkVarBind mod_id <$> mkModIdRHS mod

       ; tcg_env <- tcExtendGlobalValEnv [mod_id] getGblEnv
       ; return (tcg_env { tcg_tr_module = Just mod_id }
                 `addTypecheckedBinds` [unitBag mod_bind]) }

mkModIdRHS :: Module -> TcM (LHsExpr GhcTc)
mkModIdRHS mod
  = do { trModuleDataCon <- tcLookupDataCon trModuleDataConName
       ; trNameLit <- mkTrNameLit
       ; return $ nlHsDataCon trModuleDataCon
                  `nlHsApp` trNameLit (unitIdFS (moduleUnitId mod))
                  `nlHsApp` trNameLit (moduleNameFS (moduleName mod))
       }

{- *********************************************************************
*                                                                      *
                Building type-representation bindings
*                                                                      *
********************************************************************* -}

-- | Information we need about a 'TyCon' to generate its representation. We
-- carry the 'Id' in order to share it between the generation of the @TyCon@ and
-- @KindRep@ bindings.
data TypeableTyCon
    = TypeableTyCon
      { tycon        :: !TyCon
      , tycon_rep_id :: !Id
      }

-- | A group of 'TyCon's in need of type-rep bindings.
data TypeRepTodo
    = TypeRepTodo
      { mod_rep_expr    :: LHsExpr GhcTc    -- ^ Module's typerep binding
      , pkg_fingerprint :: !Fingerprint     -- ^ Package name fingerprint
      , mod_fingerprint :: !Fingerprint     -- ^ Module name fingerprint
      , todo_tycons     :: [TypeableTyCon]
        -- ^ The 'TyCon's in need of bindings kinds
      }
    | ExportedKindRepsTodo [(Kind, Id)]
      -- ^ Build exported 'KindRep' bindings for the given set of kinds.

todoForTyCons :: Module -> Id -> [TyCon] -> TcM TypeRepTodo
todoForTyCons mod mod_id tycons = do
    trTyConTy <- mkTyConTy <$> tcLookupTyCon trTyConTyConName
    let mk_rep_id :: TyConRepName -> Id
        mk_rep_id rep_name = mkExportedVanillaId rep_name trTyConTy

    let typeable_tycons :: [TypeableTyCon]
        typeable_tycons =
            [ TypeableTyCon { tycon = tc''
                            , tycon_rep_id = mk_rep_id rep_name
                            }
            | tc     <- tycons
            , tc'    <- tc : tyConATs tc
              -- We need type representations for any associated types
            , let promoted = map promoteDataCon (tyConDataCons tc')
            , tc''   <- tc' : promoted
              -- Don't make bindings for data-family instance tycons.
              -- Do, however, make them for their promoted datacon (see #13915).
            , not $ isFamInstTyCon tc''
            , Just rep_name <- pure $ tyConRepName_maybe tc''
            , typeIsTypeable $ dropForAlls $ tyConKind tc''
            ]
    return TypeRepTodo { mod_rep_expr    = nlHsVar mod_id
                       , pkg_fingerprint = pkg_fpr
                       , mod_fingerprint = mod_fpr
                       , todo_tycons     = typeable_tycons
                       }
  where
    mod_fpr = fingerprintString $ moduleNameString $ moduleName mod
    pkg_fpr = fingerprintString $ unitIdString $ moduleUnitId mod

todoForExportedKindReps :: [(Kind, Name)] -> TcM TypeRepTodo
todoForExportedKindReps kinds = do
    trKindRepTy <- mkTyConTy <$> tcLookupTyCon kindRepTyConName
    let mkId (k, name) = (k, mkExportedVanillaId name trKindRepTy)
    return $ ExportedKindRepsTodo $ map mkId kinds

-- | Generate TyCon bindings for a set of type constructors
mkTypeRepTodoBinds :: [TypeRepTodo] -> TcM TcGblEnv
mkTypeRepTodoBinds [] = getGblEnv
mkTypeRepTodoBinds todos
  = do { stuff <- collect_stuff

         -- First extend the type environment with all of the bindings
         -- which we are going to produce since we may need to refer to them
         -- while generating kind representations (namely, when we want to
         -- represent a TyConApp in a kind, we must be able to look up the
         -- TyCon associated with the applied type constructor).
       ; let produced_bndrs :: [Id]
             produced_bndrs = [ tycon_rep_id
                              | todo@(TypeRepTodo{}) <- todos
                              , TypeableTyCon {..} <- todo_tycons todo
                              ] ++
                              [ rep_id
                              | ExportedKindRepsTodo kinds <- todos
                              , (_, rep_id) <- kinds
                              ]
       ; gbl_env <- tcExtendGlobalValEnv produced_bndrs getGblEnv

       ; let mk_binds :: TypeRepTodo -> KindRepM [LHsBinds GhcTc]
             mk_binds todo@(TypeRepTodo {}) =
                 mapM (mkTyConRepBinds stuff todo) (todo_tycons todo)
             mk_binds (ExportedKindRepsTodo kinds) =
                 mkExportedKindReps stuff kinds >> return []

       ; (gbl_env, binds) <- setGblEnv gbl_env
                             $ runKindRepM (mapM mk_binds todos)
       ; return $ gbl_env `addTypecheckedBinds` concat binds }

-- | Generate bindings for the type representation of a wired-in 'TyCon's
-- defined by the virtual "GHC.Prim" module. This is where we inject the
-- representation bindings for these primitive types into "GHC.Types"
--
-- See Note [Grand plan for Typeable] in this module.
mkPrimTypeableTodos :: TcM (TcGblEnv, [TypeRepTodo])
mkPrimTypeableTodos
  = do { mod <- getModule
       ; if mod == gHC_TYPES
           then do { -- Build Module binding for GHC.Prim
                     trModuleTyCon <- tcLookupTyCon trModuleTyConName
                   ; let ghc_prim_module_id =
                             mkExportedVanillaId trGhcPrimModuleName
                                                 (mkTyConTy trModuleTyCon)

                   ; ghc_prim_module_bind <- mkVarBind ghc_prim_module_id
                                             <$> mkModIdRHS gHC_PRIM

                     -- Extend our environment with above
                   ; gbl_env <- tcExtendGlobalValEnv [ghc_prim_module_id]
                                                     getGblEnv
                   ; let gbl_env' = gbl_env `addTypecheckedBinds`
                                    [unitBag ghc_prim_module_bind]

                     -- Build TypeRepTodos for built-in KindReps
                   ; todo1 <- todoForExportedKindReps builtInKindReps
                     -- Build TypeRepTodos for types in GHC.Prim
                   ; todo2 <- todoForTyCons gHC_PRIM ghc_prim_module_id
                                            ghcPrimTypeableTyCons
                   ; return ( gbl_env' , [todo1, todo2])
                   }
           else do gbl_env <- getGblEnv
                   return (gbl_env, [])
       }

-- | This is the list of primitive 'TyCon's for which we must generate bindings
-- in "GHC.Types". This should include all types defined in "GHC.Prim".
--
-- The majority of the types we need here are contained in 'primTyCons'.
-- However, not all of them: in particular unboxed tuples are absent since we
-- don't want to include them in the original name cache. See
-- Note [Built-in syntax and the OrigNameCache] in IfaceEnv for more.
ghcPrimTypeableTyCons :: [TyCon]
ghcPrimTypeableTyCons = concat
    [ [ runtimeRepTyCon, vecCountTyCon, vecElemTyCon, funTyCon ]
    , map (tupleTyCon Unboxed) [0..mAX_TUPLE_SIZE]
    , map sumTyCon [2..mAX_SUM_SIZE]
    , primTyCons
    ]

data TypeableStuff
    = Stuff { dflags         :: DynFlags
            , trTyConDataCon :: DataCon         -- ^ of @TyCon@
            , trNameLit      :: FastString -> LHsExpr GhcTc
                                                -- ^ To construct @TrName@s
              -- The various TyCon and DataCons of KindRep
            , kindRepTyCon           :: TyCon
            , kindRepTyConAppDataCon :: DataCon
            , kindRepVarDataCon      :: DataCon
            , kindRepAppDataCon      :: DataCon
            , kindRepFunDataCon      :: DataCon
            , kindRepTYPEDataCon     :: DataCon
            , kindRepTypeLitSDataCon :: DataCon
            , typeLitSymbolDataCon   :: DataCon
            , typeLitNatDataCon      :: DataCon
            }

-- | Collect various tidbits which we'll need to generate TyCon representations.
collect_stuff :: TcM TypeableStuff
collect_stuff = do
    dflags <- getDynFlags
    trTyConDataCon         <- tcLookupDataCon trTyConDataConName
    kindRepTyCon           <- tcLookupTyCon   kindRepTyConName
    kindRepTyConAppDataCon <- tcLookupDataCon kindRepTyConAppDataConName
    kindRepVarDataCon      <- tcLookupDataCon kindRepVarDataConName
    kindRepAppDataCon      <- tcLookupDataCon kindRepAppDataConName
    kindRepFunDataCon      <- tcLookupDataCon kindRepFunDataConName
    kindRepTYPEDataCon     <- tcLookupDataCon kindRepTYPEDataConName
    kindRepTypeLitSDataCon <- tcLookupDataCon kindRepTypeLitSDataConName
    typeLitSymbolDataCon   <- tcLookupDataCon typeLitSymbolDataConName
    typeLitNatDataCon      <- tcLookupDataCon typeLitNatDataConName
    trNameLit              <- mkTrNameLit
    return Stuff {..}

-- | Lookup the necessary pieces to construct the @trNameLit@. We do this so we
-- can save the work of repeating lookups when constructing many TyCon
-- representations.
mkTrNameLit :: TcM (FastString -> LHsExpr GhcTc)
mkTrNameLit = do
    trNameSDataCon <- tcLookupDataCon trNameSDataConName
    let trNameLit :: FastString -> LHsExpr GhcTc
        trNameLit fs = nlHsPar $ nlHsDataCon trNameSDataCon
                       `nlHsApp` nlHsLit (mkHsStringPrimLit fs)
    return trNameLit

-- | Make Typeable bindings for the given 'TyCon'.
mkTyConRepBinds :: TypeableStuff -> TypeRepTodo
                -> TypeableTyCon -> KindRepM (LHsBinds GhcTc)
mkTyConRepBinds stuff todo (TypeableTyCon {..})
  = do -- Make a KindRep
       let (bndrs, kind) = splitForAllVarBndrs (tyConKind tycon)
       liftTc $ traceTc "mkTyConKindRepBinds"
                        (ppr tycon $$ ppr (tyConKind tycon) $$ ppr kind)
       let ctx = mkDeBruijnContext (map binderVar bndrs)
       kind_rep <- getKindRep stuff ctx kind

       -- Make the TyCon binding
       let tycon_rep_rhs = mkTyConRepTyConRHS stuff todo tycon kind_rep
           tycon_rep_bind = mkVarBind tycon_rep_id tycon_rep_rhs
       return $ unitBag tycon_rep_bind

-- | Here is where we define the set of Typeable types. These exclude type
-- families and polytypes.
tyConIsTypeable :: TyCon -> Bool
tyConIsTypeable tc =
       isJust (tyConRepName_maybe tc)
    && typeIsTypeable (dropForAlls $ tyConKind tc)
      -- Ensure that the kind of the TyCon, with its initial foralls removed,
      -- is representable (e.g. has no higher-rank polymorphism or type
      -- synonyms).

-- | Is a particular 'Type' representable by @Typeable@? Here we look for
-- polytypes and types containing casts (which may be, for instance, a type
-- family).
typeIsTypeable :: Type -> Bool
-- We handle types of the form (TYPE rep) specifically to avoid
-- looping on (tyConIsTypeable RuntimeRep)
typeIsTypeable ty
  | Just ty' <- coreView ty         = typeIsTypeable ty'
typeIsTypeable ty
  | isJust (kindRep_maybe ty)       = True
typeIsTypeable (TyVarTy _)          = True
typeIsTypeable (AppTy a b)          = typeIsTypeable a && typeIsTypeable b
typeIsTypeable (FunTy _ a b)        = typeIsTypeable a && typeIsTypeable b
typeIsTypeable (TyConApp tc args)   = tyConIsTypeable tc
                                   && all typeIsTypeable args
typeIsTypeable (ForAllTy{})         = False
typeIsTypeable (LitTy _)            = True
typeIsTypeable (CastTy{})           = False
typeIsTypeable (CoercionTy{})       = False

-- | Maps kinds to 'KindRep' bindings. This binding may either be defined in
-- some other module (in which case the @Maybe (LHsExpr Id@ will be 'Nothing')
-- or a binding which we generated in the current module (in which case it will
-- be 'Just' the RHS of the binding).
type KindRepEnv = TypeMap (Id, Maybe (LHsExpr GhcTc))

-- | A monad within which we will generate 'KindRep's. Here we keep an
-- environment containing 'KindRep's which we've already generated so we can
-- re-use them opportunistically.
newtype KindRepM a = KindRepM { unKindRepM :: StateT KindRepEnv TcRn a }
                   deriving (Functor, Applicative, Monad)

liftTc :: TcRn a -> KindRepM a
liftTc = KindRepM . lift

-- | We generate @KindRep@s for a few common kinds in @GHC.Types@ so that they
-- can be reused across modules.
builtInKindReps :: [(Kind, Name)]
builtInKindReps =
    [ (star, starKindRepName)
    , (mkVisFunTy star star, starArrStarKindRepName)
    , (mkVisFunTys [star, star] star, starArrStarArrStarKindRepName)
    ]
  where
    star = liftedTypeKind

initialKindRepEnv :: TcRn KindRepEnv
initialKindRepEnv = foldlM add_kind_rep emptyTypeMap builtInKindReps
  where
    add_kind_rep acc (k,n) = do
        id <- tcLookupId n
        return $! extendTypeMap acc k (id, Nothing)

-- | Performed while compiling "GHC.Types" to generate the built-in 'KindRep's.
mkExportedKindReps :: TypeableStuff
                   -> [(Kind, Id)]  -- ^ the kinds to generate bindings for
                   -> KindRepM ()
mkExportedKindReps stuff = mapM_ kindrep_binding
  where
    empty_scope = mkDeBruijnContext []

    kindrep_binding :: (Kind, Id) -> KindRepM ()
    kindrep_binding (kind, rep_bndr) = do
        -- We build the binding manually here instead of using mkKindRepRhs
        -- since the latter would find the built-in 'KindRep's in the
        -- 'KindRepEnv' (by virtue of being in 'initialKindRepEnv').
        rhs <- mkKindRepRhs stuff empty_scope kind
        addKindRepBind empty_scope kind rep_bndr rhs

addKindRepBind :: CmEnv -> Kind -> Id -> LHsExpr GhcTc -> KindRepM ()
addKindRepBind in_scope k bndr rhs =
    KindRepM $ modify' $
    \env -> extendTypeMapWithScope env in_scope k (bndr, Just rhs)

-- | Run a 'KindRepM' and add the produced 'KindRep's to the typechecking
-- environment.
runKindRepM :: KindRepM a -> TcRn (TcGblEnv, a)
runKindRepM (KindRepM action) = do
    kindRepEnv <- initialKindRepEnv
    (res, reps_env) <- runStateT action kindRepEnv
    let rep_binds = foldTypeMap to_bind_pair [] reps_env
        to_bind_pair (bndr, Just rhs) rest = (bndr, rhs) : rest
        to_bind_pair (_, Nothing) rest = rest
    tcg_env <- tcExtendGlobalValEnv (map fst rep_binds) getGblEnv
    let binds = map (uncurry mkVarBind) rep_binds
        tcg_env' = tcg_env `addTypecheckedBinds` [listToBag binds]
    return (tcg_env', res)

-- | Produce or find a 'KindRep' for the given kind.
getKindRep :: TypeableStuff -> CmEnv  -- ^ in-scope kind variables
           -> Kind   -- ^ the kind we want a 'KindRep' for
           -> KindRepM (LHsExpr GhcTc)
getKindRep stuff@(Stuff {..}) in_scope = go
  where
    go :: Kind -> KindRepM (LHsExpr GhcTc)
    go = KindRepM . StateT . go'

    go' :: Kind -> KindRepEnv -> TcRn (LHsExpr GhcTc, KindRepEnv)
    go' k env
        -- Look through type synonyms
      | Just k' <- tcView k = go' k' env

        -- We've already generated the needed KindRep
      | Just (id, _) <- lookupTypeMapWithScope env in_scope k
      = return (nlHsVar id, env)

        -- We need to construct a new KindRep binding
      | otherwise
      = do -- Place a NOINLINE pragma on KindReps since they tend to be quite
           -- large and bloat interface files.
           rep_bndr <- (`setInlinePragma` neverInlinePragma)
                   <$> newSysLocalId (fsLit "$krep") (mkTyConTy kindRepTyCon)

           -- do we need to tie a knot here?
           flip runStateT env $ unKindRepM $ do
               rhs <- mkKindRepRhs stuff in_scope k
               addKindRepBind in_scope k rep_bndr rhs
               return $ nlHsVar rep_bndr

-- | Construct the right-hand-side of the 'KindRep' for the given 'Kind' and
-- in-scope kind variable set.
mkKindRepRhs :: TypeableStuff
             -> CmEnv       -- ^ in-scope kind variables
             -> Kind        -- ^ the kind we want a 'KindRep' for
             -> KindRepM (LHsExpr GhcTc) -- ^ RHS expression
mkKindRepRhs stuff@(Stuff {..}) in_scope = new_kind_rep
  where
    new_kind_rep k
        -- We handle (TYPE LiftedRep) etc separately to make it
        -- clear to consumers (e.g. serializers) that there is
        -- a loop here (as TYPE :: RuntimeRep -> TYPE 'LiftedRep)
      | not (tcIsConstraintKind k)
              -- Typeable respects the Constraint/Type distinction
              -- so do not follow the special case here
      , Just arg <- kindRep_maybe k
      , Just (tc, []) <- splitTyConApp_maybe arg
      , Just dc <- isPromotedDataCon_maybe tc
      = return $ nlHsDataCon kindRepTYPEDataCon `nlHsApp` nlHsDataCon dc

    new_kind_rep (TyVarTy v)
      | Just idx <- lookupCME in_scope v
      = return $ nlHsDataCon kindRepVarDataCon
                 `nlHsApp` nlHsIntLit (fromIntegral idx)
      | otherwise
      = pprPanic "mkTyConKindRepBinds.go(tyvar)" (ppr v)

    new_kind_rep (AppTy t1 t2)
      = do rep1 <- getKindRep stuff in_scope t1
           rep2 <- getKindRep stuff in_scope t2
           return $ nlHsDataCon kindRepAppDataCon
                    `nlHsApp` rep1 `nlHsApp` rep2

    new_kind_rep k@(TyConApp tc tys)
      | Just rep_name <- tyConRepName_maybe tc
      = do rep_id <- liftTc $ lookupId rep_name
           tys' <- mapM (getKindRep stuff in_scope) tys
           return $ nlHsDataCon kindRepTyConAppDataCon
                    `nlHsApp` nlHsVar rep_id
                    `nlHsApp` mkList (mkTyConTy kindRepTyCon) tys'
      | otherwise
      = pprPanic "mkTyConKindRepBinds(TyConApp)" (ppr tc $$ ppr k)

    new_kind_rep (ForAllTy (Bndr var _) ty)
      = pprPanic "mkTyConKindRepBinds(ForAllTy)" (ppr var $$ ppr ty)

    new_kind_rep (FunTy _ t1 t2)
      = do rep1 <- getKindRep stuff in_scope t1
           rep2 <- getKindRep stuff in_scope t2
           return $ nlHsDataCon kindRepFunDataCon
                    `nlHsApp` rep1 `nlHsApp` rep2

    new_kind_rep (LitTy (NumTyLit n))
      = return $ nlHsDataCon kindRepTypeLitSDataCon
                 `nlHsApp` nlHsDataCon typeLitNatDataCon
                 `nlHsApp` nlHsLit (mkHsStringPrimLit $ mkFastString $ show n)

    new_kind_rep (LitTy (StrTyLit s))
      = return $ nlHsDataCon kindRepTypeLitSDataCon
                 `nlHsApp` nlHsDataCon typeLitSymbolDataCon
                 `nlHsApp` nlHsLit (mkHsStringPrimLit $ mkFastString $ show s)

    new_kind_rep (CastTy ty co)
      = pprPanic "mkTyConKindRepBinds.go(cast)" (ppr ty $$ ppr co)

    new_kind_rep (CoercionTy co)
      = pprPanic "mkTyConKindRepBinds.go(coercion)" (ppr co)

-- | Produce the right-hand-side of a @TyCon@ representation.
mkTyConRepTyConRHS :: TypeableStuff -> TypeRepTodo
                   -> TyCon      -- ^ the 'TyCon' we are producing a binding for
                   -> LHsExpr GhcTc -- ^ its 'KindRep'
                   -> LHsExpr GhcTc
mkTyConRepTyConRHS (Stuff {..}) todo tycon kind_rep
  =           nlHsDataCon trTyConDataCon
    `nlHsApp` nlHsLit (word64 dflags high)
    `nlHsApp` nlHsLit (word64 dflags low)
    `nlHsApp` mod_rep_expr todo
    `nlHsApp` trNameLit (mkFastString tycon_str)
    `nlHsApp` nlHsLit (int n_kind_vars)
    `nlHsApp` kind_rep
  where
    n_kind_vars = length $ filter isNamedTyConBinder (tyConBinders tycon)
    tycon_str = add_tick (occNameString (getOccName tycon))
    add_tick s | isPromotedDataCon tycon = '\'' : s
               | otherwise               = s

    -- This must match the computation done in
    -- Data.Typeable.Internal.mkTyConFingerprint.
    Fingerprint high low = fingerprintFingerprints [ pkg_fingerprint todo
                                                   , mod_fingerprint todo
                                                   , fingerprintString tycon_str
                                                   ]

    int :: Int -> HsLit GhcTc
    int n = HsIntPrim (SourceText $ show n) (toInteger n)

word64 :: DynFlags -> Word64 -> HsLit GhcTc
word64 dflags n
  | wORD_SIZE dflags == 4 = HsWord64Prim NoSourceText (toInteger n)
  | otherwise             = HsWordPrim   NoSourceText (toInteger n)

{-
Note [Representing TyCon kinds: KindRep]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
One of the operations supported by Typeable is typeRepKind,

    typeRepKind :: TypeRep (a :: k) -> TypeRep k

Implementing this is a bit tricky for poly-kinded types like

    data Proxy (a :: k) :: Type
    -- Proxy :: forall k. k -> Type

The TypeRep encoding of `Proxy Type Int` looks like this:

    $tcProxy :: GHC.Types.TyCon
    $trInt   :: TypeRep Int
    TrType   :: TypeRep Type

    $trProxyType :: TypeRep (Proxy Type :: Type -> Type)
    $trProxyType = TrTyCon $tcProxy
                           [TrType]  -- kind variable instantiation
                           (tyConKind $tcProxy [TrType]) -- The TypeRep of
                                                         -- Type -> Type

    $trProxy :: TypeRep (Proxy Type Int)
    $trProxy = TrApp $trProxyType $trInt TrType

    $tkProxy :: GHC.Types.KindRep
    $tkProxy = KindRepFun (KindRepVar 0)
                          (KindRepTyConApp (KindRepTYPE LiftedRep) [])

Note how $trProxyType cannot use 'TrApp', because TypeRep cannot represent
polymorphic types.  So instead

 * $trProxyType uses 'TrTyCon' to apply Proxy to (the representations)
   of all its kind arguments. We can't represent a tycon that is
   applied to only some of its kind arguments.

 * In $tcProxy, the GHC.Types.TyCon structure for Proxy, we store a
   GHC.Types.KindRep, which represents the polymorphic kind of Proxy
       Proxy :: forall k. k->Type

 * A KindRep is just a recipe that we can instantiate with the
   argument kinds, using Data.Typeable.Internal.tyConKind and
   store in the relevant 'TypeRep' constructor.

   Data.Typeable.Internal.typeRepKind looks up the stored kinds.

 * In a KindRep, the kind variables are represented by 0-indexed
   de Bruijn numbers:

    type KindBndr = Int   -- de Bruijn index

    data KindRep = KindRepTyConApp TyCon [KindRep]
                 | KindRepVar !KindBndr
                 | KindRepApp KindRep KindRep
                 | KindRepFun KindRep KindRep
                 ...
-}

mkList :: Type -> [LHsExpr GhcTc] -> LHsExpr GhcTc
mkList ty = foldr consApp (nilExpr ty)
  where
    cons = consExpr ty
    consApp :: LHsExpr GhcTc -> LHsExpr GhcTc -> LHsExpr GhcTc
    consApp x xs = cons `nlHsApp` x `nlHsApp` xs

    nilExpr :: Type -> LHsExpr GhcTc
    nilExpr ty = mkLHsWrap (mkWpTyApps [ty]) (nlHsDataCon nilDataCon)

    consExpr :: Type -> LHsExpr GhcTc
    consExpr ty = mkLHsWrap (mkWpTyApps [ty]) (nlHsDataCon consDataCon)


{-# LANGUAGE LambdaCase #-}

module TcTypeNats
  ( typeNatTyCons
  , typeNatCoAxiomRules
  , BuiltInSynFamily(..)

    -- If you define a new built-in type family, make sure to export its TyCon
    -- from here as well.
    -- See Note [Adding built-in type families]
  , typeNatAddTyCon
  , typeNatMulTyCon
  , typeNatExpTyCon
  , typeNatLeqTyCon
  , typeNatSubTyCon
  , typeNatDivTyCon
  , typeNatModTyCon
  , typeNatLogTyCon
  , typeNatCmpTyCon
  , typeSymbolCmpTyCon
  , typeSymbolAppendTyCon
  ) where

import GhcPrelude

import Type
import Pair
import TcType     ( TcType, tcEqType )
import TyCon      ( TyCon, FamTyConFlav(..), mkFamilyTyCon
                  , Injectivity(..) )
import Coercion   ( Role(..) )
import TcRnTypes  ( Xi )
import CoAxiom    ( CoAxiomRule(..), BuiltInSynFamily(..), TypeEqn )
import Name       ( Name, BuiltInSyntax(..) )
import TysWiredIn
import TysPrim    ( mkTemplateAnonTyConBinders )
import PrelNames  ( gHC_TYPELITS
                  , gHC_TYPENATS
                  , typeNatAddTyFamNameKey
                  , typeNatMulTyFamNameKey
                  , typeNatExpTyFamNameKey
                  , typeNatLeqTyFamNameKey
                  , typeNatSubTyFamNameKey
                  , typeNatDivTyFamNameKey
                  , typeNatModTyFamNameKey
                  , typeNatLogTyFamNameKey
                  , typeNatCmpTyFamNameKey
                  , typeSymbolCmpTyFamNameKey
                  , typeSymbolAppendFamNameKey
                  )
import FastString ( FastString
                  , fsLit, nilFS, nullFS, unpackFS, mkFastString, appendFS
                  )
import qualified Data.Map as Map
import Data.Maybe ( isJust )
import Control.Monad ( guard )
import Data.List  ( isPrefixOf, isSuffixOf )

{-
Note [Type-level literals]
~~~~~~~~~~~~~~~~~~~~~~~~~~
There are currently two forms of type-level literals: natural numbers, and
symbols (even though this module is named TcTypeNats, it covers both).

Type-level literals are supported by CoAxiomRules (conditional axioms), which
power the built-in type families (see Note [Adding built-in type families]).
Currently, all built-in type families are for the express purpose of supporting
type-level literals.

See also the Wiki page:

    https://ghc.haskell.org/trac/ghc/wiki/TypeNats

Note [Adding built-in type families]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
There are a few steps to adding a built-in type family:

* Adding a unique for the type family TyCon

  These go in PrelNames. It will likely be of the form
  @myTyFamNameKey = mkPreludeTyConUnique xyz@, where @xyz@ is a number that
  has not been chosen before in PrelNames. There are several examples already
  in PrelNamesâ€”see, for instance, typeNatAddTyFamNameKey.

* Adding the type family TyCon itself

  This goes in TcTypeNats. There are plenty of examples of how to define
  theseâ€”see, for instance, typeNatAddTyCon.

  Once your TyCon has been defined, be sure to:

  - Export it from TcTypeNats. (Not doing so caused #14632.)
  - Include it in the typeNatTyCons list, defined in TcTypeNats.

* Exposing associated type family axioms

  When defining the type family TyCon, you will need to define an axiom for
  the type family in general (see, for instance, axAddDef), and perhaps other
  auxiliary axioms for special cases of the type family (see, for instance,
  axAdd0L and axAdd0R).

  After you have defined all of these axioms, be sure to include them in the
  typeNatCoAxiomRules list, defined in TcTypeNats.
  (Not doing so caused #14934.)

* Define the type family somewhere

  Finally, you will need to define the type family somewhere, likely in @base@.
  Currently, all of the built-in type families are defined in GHC.TypeLits or
  GHC.TypeNats, so those are likely candidates.

  Since the behavior of your built-in type family is specified in TcTypeNats,
  you should give an open type family definition with no instances, like so:

    type family MyTypeFam (m :: Nat) (n :: Nat) :: Nat

  Changing the argument and result kinds as appropriate.

* Update the relevant test cases

  The GHC test suite will likely need to be updated after you add your built-in
  type family. For instance:

  - The T9181 test prints the :browse contents of GHC.TypeLits, so if you added
    a test there, the expected output of T9181 will need to change.
  - The TcTypeNatSimple and TcTypeSymbolSimple tests have compile-time unit
    tests, as well as TcTypeNatSimpleRun and TcTypeSymbolSimpleRun, which have
    runtime unit tests. Consider adding further unit tests to those if your
    built-in type family deals with Nats or Symbols, respectively.
-}

{-------------------------------------------------------------------------------
Built-in type constructors for functions on type-level nats
-}

-- The list of built-in type family TyCons that GHC uses.
-- If you define a built-in type family, make sure to add it to this list.
-- See Note [Adding built-in type families]
typeNatTyCons :: [TyCon]
typeNatTyCons =
  [ typeNatAddTyCon
  , typeNatMulTyCon
  , typeNatExpTyCon
  , typeNatLeqTyCon
  , typeNatSubTyCon
  , typeNatDivTyCon
  , typeNatModTyCon
  , typeNatLogTyCon
  , typeNatCmpTyCon
  , typeSymbolCmpTyCon
  , typeSymbolAppendTyCon
  ]

typeNatAddTyCon :: TyCon
typeNatAddTyCon = mkTypeNatFunTyCon2 name
  BuiltInSynFamily
    { sfMatchFam      = matchFamAdd
    , sfInteractTop   = interactTopAdd
    , sfInteractInert = interactInertAdd
    }
  where
  name = mkWiredInTyConName UserSyntax gHC_TYPENATS (fsLit "+")
            typeNatAddTyFamNameKey typeNatAddTyCon

typeNatSubTyCon :: TyCon
typeNatSubTyCon = mkTypeNatFunTyCon2 name
  BuiltInSynFamily
    { sfMatchFam      = matchFamSub
    , sfInteractTop   = interactTopSub
    , sfInteractInert = interactInertSub
    }
  where
  name = mkWiredInTyConName UserSyntax gHC_TYPENATS (fsLit "-")
            typeNatSubTyFamNameKey typeNatSubTyCon

typeNatMulTyCon :: TyCon
typeNatMulTyCon = mkTypeNatFunTyCon2 name
  BuiltInSynFamily
    { sfMatchFam      = matchFamMul
    , sfInteractTop   = interactTopMul
    , sfInteractInert = interactInertMul
    }
  where
  name = mkWiredInTyConName UserSyntax gHC_TYPENATS (fsLit "*")
            typeNatMulTyFamNameKey typeNatMulTyCon

typeNatDivTyCon :: TyCon
typeNatDivTyCon = mkTypeNatFunTyCon2 name
  BuiltInSynFamily
    { sfMatchFam      = matchFamDiv
    , sfInteractTop   = interactTopDiv
    , sfInteractInert = interactInertDiv
    }
  where
  name = mkWiredInTyConName UserSyntax gHC_TYPENATS (fsLit "Div")
            typeNatDivTyFamNameKey typeNatDivTyCon

typeNatModTyCon :: TyCon
typeNatModTyCon = mkTypeNatFunTyCon2 name
  BuiltInSynFamily
    { sfMatchFam      = matchFamMod
    , sfInteractTop   = interactTopMod
    , sfInteractInert = interactInertMod
    }
  where
  name = mkWiredInTyConName UserSyntax gHC_TYPENATS (fsLit "Mod")
            typeNatModTyFamNameKey typeNatModTyCon

typeNatExpTyCon :: TyCon
typeNatExpTyCon = mkTypeNatFunTyCon2 name
  BuiltInSynFamily
    { sfMatchFam      = matchFamExp
    , sfInteractTop   = interactTopExp
    , sfInteractInert = interactInertExp
    }
  where
  name = mkWiredInTyConName UserSyntax gHC_TYPENATS (fsLit "^")
                typeNatExpTyFamNameKey typeNatExpTyCon

typeNatLogTyCon :: TyCon
typeNatLogTyCon = mkTypeNatFunTyCon1 name
  BuiltInSynFamily
    { sfMatchFam      = matchFamLog
    , sfInteractTop   = interactTopLog
    , sfInteractInert = interactInertLog
    }
  where
  name = mkWiredInTyConName UserSyntax gHC_TYPENATS (fsLit "Log2")
            typeNatLogTyFamNameKey typeNatLogTyCon

typeNatLeqTyCon :: TyCon
typeNatLeqTyCon =
  mkFamilyTyCon name
    (mkTemplateAnonTyConBinders [ typeNatKind, typeNatKind ])
    boolTy
    Nothing
    (BuiltInSynFamTyCon ops)
    Nothing
    NotInjective

  where
  name = mkWiredInTyConName UserSyntax gHC_TYPENATS (fsLit "<=?")
                typeNatLeqTyFamNameKey typeNatLeqTyCon
  ops = BuiltInSynFamily
    { sfMatchFam      = matchFamLeq
    , sfInteractTop   = interactTopLeq
    , sfInteractInert = interactInertLeq
    }

typeNatCmpTyCon :: TyCon
typeNatCmpTyCon =
  mkFamilyTyCon name
    (mkTemplateAnonTyConBinders [ typeNatKind, typeNatKind ])
    orderingKind
    Nothing
    (BuiltInSynFamTyCon ops)
    Nothing
    NotInjective

  where
  name = mkWiredInTyConName UserSyntax gHC_TYPENATS (fsLit "CmpNat")
                typeNatCmpTyFamNameKey typeNatCmpTyCon
  ops = BuiltInSynFamily
    { sfMatchFam      = matchFamCmpNat
    , sfInteractTop   = interactTopCmpNat
    , sfInteractInert = \_ _ _ _ -> []
    }

typeSymbolCmpTyCon :: TyCon
typeSymbolCmpTyCon =
  mkFamilyTyCon name
    (mkTemplateAnonTyConBinders [ typeSymbolKind, typeSymbolKind ])
    orderingKind
    Nothing
    (BuiltInSynFamTyCon ops)
    Nothing
    NotInjective

  where
  name = mkWiredInTyConName UserSyntax gHC_TYPELITS (fsLit "CmpSymbol")
                typeSymbolCmpTyFamNameKey typeSymbolCmpTyCon
  ops = BuiltInSynFamily
    { sfMatchFam      = matchFamCmpSymbol
    , sfInteractTop   = interactTopCmpSymbol
    , sfInteractInert = \_ _ _ _ -> []
    }

typeSymbolAppendTyCon :: TyCon
typeSymbolAppendTyCon = mkTypeSymbolFunTyCon2 name
  BuiltInSynFamily
    { sfMatchFam      = matchFamAppendSymbol
    , sfInteractTop   = interactTopAppendSymbol
    , sfInteractInert = interactInertAppendSymbol
    }
  where
  name = mkWiredInTyConName UserSyntax gHC_TYPELITS (fsLit "AppendSymbol")
                typeSymbolAppendFamNameKey typeSymbolAppendTyCon

-- Make a unary built-in constructor of kind: Nat -> Nat
mkTypeNatFunTyCon1 :: Name -> BuiltInSynFamily -> TyCon
mkTypeNatFunTyCon1 op tcb =
  mkFamilyTyCon op
    (mkTemplateAnonTyConBinders [ typeNatKind ])
    typeNatKind
    Nothing
    (BuiltInSynFamTyCon tcb)
    Nothing
    NotInjective

-- Make a binary built-in constructor of kind: Nat -> Nat -> Nat
mkTypeNatFunTyCon2 :: Name -> BuiltInSynFamily -> TyCon
mkTypeNatFunTyCon2 op tcb =
  mkFamilyTyCon op
    (mkTemplateAnonTyConBinders [ typeNatKind, typeNatKind ])
    typeNatKind
    Nothing
    (BuiltInSynFamTyCon tcb)
    Nothing
    NotInjective

-- Make a binary built-in constructor of kind: Symbol -> Symbol -> Symbol
mkTypeSymbolFunTyCon2 :: Name -> BuiltInSynFamily -> TyCon
mkTypeSymbolFunTyCon2 op tcb =
  mkFamilyTyCon op
    (mkTemplateAnonTyConBinders [ typeSymbolKind, typeSymbolKind ])
    typeSymbolKind
    Nothing
    (BuiltInSynFamTyCon tcb)
    Nothing
    NotInjective

{-------------------------------------------------------------------------------
Built-in rules axioms
-------------------------------------------------------------------------------}

-- If you add additional rules, please remember to add them to
-- `typeNatCoAxiomRules` also.
-- See Note [Adding built-in type families]
axAddDef
  , axMulDef
  , axExpDef
  , axLeqDef
  , axCmpNatDef
  , axCmpSymbolDef
  , axAppendSymbolDef
  , axAdd0L
  , axAdd0R
  , axMul0L
  , axMul0R
  , axMul1L
  , axMul1R
  , axExp1L
  , axExp0R
  , axExp1R
  , axLeqRefl
  , axCmpNatRefl
  , axCmpSymbolRefl
  , axLeq0L
  , axSubDef
  , axSub0R
  , axAppendSymbol0R
  , axAppendSymbol0L
  , axDivDef
  , axDiv1
  , axModDef
  , axMod1
  , axLogDef
  :: CoAxiomRule

axAddDef = mkBinAxiom "AddDef" typeNatAddTyCon $
              \x y -> Just $ num (x + y)

axMulDef = mkBinAxiom "MulDef" typeNatMulTyCon $
              \x y -> Just $ num (x * y)

axExpDef = mkBinAxiom "ExpDef" typeNatExpTyCon $
              \x y -> Just $ num (x ^ y)

axLeqDef = mkBinAxiom "LeqDef" typeNatLeqTyCon $
              \x y -> Just $ bool (x <= y)

axCmpNatDef   = mkBinAxiom "CmpNatDef" typeNatCmpTyCon
              $ \x y -> Just $ ordering (compare x y)

axCmpSymbolDef =
  CoAxiomRule
    { coaxrName      = fsLit "CmpSymbolDef"
    , coaxrAsmpRoles = [Nominal, Nominal]
    , coaxrRole      = Nominal
    , coaxrProves    = \cs ->
        do [Pair s1 s2, Pair t1 t2] <- return cs
           s2' <- isStrLitTy s2
           t2' <- isStrLitTy t2
           return (mkTyConApp typeSymbolCmpTyCon [s1,t1] ===
                   ordering (compare s2' t2')) }

axAppendSymbolDef = CoAxiomRule
    { coaxrName      = fsLit "AppendSymbolDef"
    , coaxrAsmpRoles = [Nominal, Nominal]
    , coaxrRole      = Nominal
    , coaxrProves    = \cs ->
        do [Pair s1 s2, Pair t1 t2] <- return cs
           s2' <- isStrLitTy s2
           t2' <- isStrLitTy t2
           let z = mkStrLitTy (appendFS s2' t2')
           return (mkTyConApp typeSymbolAppendTyCon [s1, t1] === z)
    }

axSubDef = mkBinAxiom "SubDef" typeNatSubTyCon $
              \x y -> fmap num (minus x y)

axDivDef = mkBinAxiom "DivDef" typeNatDivTyCon $
              \x y -> do guard (y /= 0)
                         return (num (div x y))

axModDef = mkBinAxiom "ModDef" typeNatModTyCon $
              \x y -> do guard (y /= 0)
                         return (num (mod x y))

axLogDef = mkUnAxiom "LogDef" typeNatLogTyCon $
              \x -> do (a,_) <- genLog x 2
                       return (num a)

axAdd0L     = mkAxiom1 "Add0L"    $ \(Pair s t) -> (num 0 .+. s) === t
axAdd0R     = mkAxiom1 "Add0R"    $ \(Pair s t) -> (s .+. num 0) === t
axSub0R     = mkAxiom1 "Sub0R"    $ \(Pair s t) -> (s .-. num 0) === t
axMul0L     = mkAxiom1 "Mul0L"    $ \(Pair s _) -> (num 0 .*. s) === num 0
axMul0R     = mkAxiom1 "Mul0R"    $ \(Pair s _) -> (s .*. num 0) === num 0
axMul1L     = mkAxiom1 "Mul1L"    $ \(Pair s t) -> (num 1 .*. s) === t
axMul1R     = mkAxiom1 "Mul1R"    $ \(Pair s t) -> (s .*. num 1) === t
axDiv1      = mkAxiom1 "Div1"     $ \(Pair s t) -> (tDiv s (num 1) === t)
axMod1      = mkAxiom1 "Mod1"     $ \(Pair s _) -> (tMod s (num 1) === num 0)
                                    -- XXX: Shouldn't we check that _ is 0?
axExp1L     = mkAxiom1 "Exp1L"    $ \(Pair s _) -> (num 1 .^. s) === num 1
axExp0R     = mkAxiom1 "Exp0R"    $ \(Pair s _) -> (s .^. num 0) === num 1
axExp1R     = mkAxiom1 "Exp1R"    $ \(Pair s t) -> (s .^. num 1) === t
axLeqRefl   = mkAxiom1 "LeqRefl"  $ \(Pair s _) -> (s <== s) === bool True
axCmpNatRefl    = mkAxiom1 "CmpNatRefl"
                $ \(Pair s _) -> (cmpNat s s) === ordering EQ
axCmpSymbolRefl = mkAxiom1 "CmpSymbolRefl"
                $ \(Pair s _) -> (cmpSymbol s s) === ordering EQ
axLeq0L     = mkAxiom1 "Leq0L"    $ \(Pair s _) -> (num 0 <== s) === bool True
axAppendSymbol0R  = mkAxiom1 "Concat0R"
            $ \(Pair s t) -> (mkStrLitTy nilFS `appendSymbol` s) === t
axAppendSymbol0L  = mkAxiom1 "Concat0L"
            $ \(Pair s t) -> (s `appendSymbol` mkStrLitTy nilFS) === t

-- The list of built-in type family axioms that GHC uses.
-- If you define new axioms, make sure to include them in this list.
-- See Note [Adding built-in type families]
typeNatCoAxiomRules :: Map.Map FastString CoAxiomRule
typeNatCoAxiomRules = Map.fromList $ map (\x -> (coaxrName x, x))
  [ axAddDef
  , axMulDef
  , axExpDef
  , axLeqDef
  , axCmpNatDef
  , axCmpSymbolDef
  , axAppendSymbolDef
  , axAdd0L
  , axAdd0R
  , axMul0L
  , axMul0R
  , axMul1L
  , axMul1R
  , axExp1L
  , axExp0R
  , axExp1R
  , axLeqRefl
  , axCmpNatRefl
  , axCmpSymbolRefl
  , axLeq0L
  , axSubDef
  , axSub0R
  , axAppendSymbol0R
  , axAppendSymbol0L
  , axDivDef
  , axDiv1
  , axModDef
  , axMod1
  , axLogDef
  ]

{-------------------------------------------------------------------------------
Various utilities for making axioms and types
-------------------------------------------------------------------------------}

(.+.) :: Type -> Type -> Type
s .+. t = mkTyConApp typeNatAddTyCon [s,t]

(.-.) :: Type -> Type -> Type
s .-. t = mkTyConApp typeNatSubTyCon [s,t]

(.*.) :: Type -> Type -> Type
s .*. t = mkTyConApp typeNatMulTyCon [s,t]

tDiv :: Type -> Type -> Type
tDiv s t = mkTyConApp typeNatDivTyCon [s,t]

tMod :: Type -> Type -> Type
tMod s t = mkTyConApp typeNatModTyCon [s,t]

(.^.) :: Type -> Type -> Type
s .^. t = mkTyConApp typeNatExpTyCon [s,t]

(<==) :: Type -> Type -> Type
s <== t = mkTyConApp typeNatLeqTyCon [s,t]

cmpNat :: Type -> Type -> Type
cmpNat s t = mkTyConApp typeNatCmpTyCon [s,t]

cmpSymbol :: Type -> Type -> Type
cmpSymbol s t = mkTyConApp typeSymbolCmpTyCon [s,t]

appendSymbol :: Type -> Type -> Type
appendSymbol s t = mkTyConApp typeSymbolAppendTyCon [s, t]

(===) :: Type -> Type -> Pair Type
x === y = Pair x y

num :: Integer -> Type
num = mkNumLitTy

bool :: Bool -> Type
bool b = if b then mkTyConApp promotedTrueDataCon []
              else mkTyConApp promotedFalseDataCon []

isBoolLitTy :: Type -> Maybe Bool
isBoolLitTy tc =
  do (tc,[]) <- splitTyConApp_maybe tc
     case () of
       _ | tc == promotedFalseDataCon -> return False
         | tc == promotedTrueDataCon  -> return True
         | otherwise                   -> Nothing

orderingKind :: Kind
orderingKind = mkTyConApp orderingTyCon []

ordering :: Ordering -> Type
ordering o =
  case o of
    LT -> mkTyConApp promotedLTDataCon []
    EQ -> mkTyConApp promotedEQDataCon []
    GT -> mkTyConApp promotedGTDataCon []

isOrderingLitTy :: Type -> Maybe Ordering
isOrderingLitTy tc =
  do (tc1,[]) <- splitTyConApp_maybe tc
     case () of
       _ | tc1 == promotedLTDataCon -> return LT
         | tc1 == promotedEQDataCon -> return EQ
         | tc1 == promotedGTDataCon -> return GT
         | otherwise                -> Nothing

known :: (Integer -> Bool) -> TcType -> Bool
known p x = case isNumLitTy x of
              Just a  -> p a
              Nothing -> False

mkUnAxiom :: String -> TyCon -> (Integer -> Maybe Type) -> CoAxiomRule
mkUnAxiom str tc f =
  CoAxiomRule
    { coaxrName      = fsLit str
    , coaxrAsmpRoles = [Nominal]
    , coaxrRole      = Nominal
    , coaxrProves    = \cs ->
        do [Pair s1 s2] <- return cs
           s2' <- isNumLitTy s2
           z   <- f s2'
           return (mkTyConApp tc [s1] === z)
    }

-- For the definitional axioms
mkBinAxiom :: String -> TyCon ->
              (Integer -> Integer -> Maybe Type) -> CoAxiomRule
mkBinAxiom str tc f =
  CoAxiomRule
    { coaxrName      = fsLit str
    , coaxrAsmpRoles = [Nominal, Nominal]
    , coaxrRole      = Nominal
    , coaxrProves    = \cs ->
        do [Pair s1 s2, Pair t1 t2] <- return cs
           s2' <- isNumLitTy s2
           t2' <- isNumLitTy t2
           z   <- f s2' t2'
           return (mkTyConApp tc [s1,t1] === z)
    }

mkAxiom1 :: String -> (TypeEqn -> TypeEqn) -> CoAxiomRule
mkAxiom1 str f =
  CoAxiomRule
    { coaxrName      = fsLit str
    , coaxrAsmpRoles = [Nominal]
    , coaxrRole      = Nominal
    , coaxrProves    = \case [eqn] -> Just (f eqn)
                             _     -> Nothing
    }

{-------------------------------------------------------------------------------
Evaluation
-------------------------------------------------------------------------------}

matchFamAdd :: [Type] -> Maybe (CoAxiomRule, [Type], Type)
matchFamAdd [s,t]
  | Just 0 <- mbX = Just (axAdd0L, [t], t)
  | Just 0 <- mbY = Just (axAdd0R, [s], s)
  | Just x <- mbX, Just y <- mbY =
    Just (axAddDef, [s,t], num (x + y))
  where mbX = isNumLitTy s
        mbY = isNumLitTy t
matchFamAdd _ = Nothing

matchFamSub :: [Type] -> Maybe (CoAxiomRule, [Type], Type)
matchFamSub [s,t]
  | Just 0 <- mbY = Just (axSub0R, [s], s)
  | Just x <- mbX, Just y <- mbY, Just z <- minus x y =
    Just (axSubDef, [s,t], num z)
  where mbX = isNumLitTy s
        mbY = isNumLitTy t
matchFamSub _ = Nothing

matchFamMul :: [Type] -> Maybe (CoAxiomRule, [Type], Type)
matchFamMul [s,t]
  | Just 0 <- mbX = Just (axMul0L, [t], num 0)
  | Just 0 <- mbY = Just (axMul0R, [s], num 0)
  | Just 1 <- mbX = Just (axMul1L, [t], t)
  | Just 1 <- mbY = Just (axMul1R, [s], s)
  | Just x <- mbX, Just y <- mbY =
    Just (axMulDef, [s,t], num (x * y))
  where mbX = isNumLitTy s
        mbY = isNumLitTy t
matchFamMul _ = Nothing

matchFamDiv :: [Type] -> Maybe (CoAxiomRule, [Type], Type)
matchFamDiv [s,t]
  | Just 1 <- mbY = Just (axDiv1, [s], s)
  | Just x <- mbX, Just y <- mbY, y /= 0 = Just (axDivDef, [s,t], num (div x y))
  where mbX = isNumLitTy s
        mbY = isNumLitTy t
matchFamDiv _ = Nothing

matchFamMod :: [Type] -> Maybe (CoAxiomRule, [Type], Type)
matchFamMod [s,t]
  | Just 1 <- mbY = Just (axMod1, [s], num 0)
  | Just x <- mbX, Just y <- mbY, y /= 0 = Just (axModDef, [s,t], num (mod x y))
  where mbX = isNumLitTy s
        mbY = isNumLitTy t
matchFamMod _ = Nothing

matchFamExp :: [Type] -> Maybe (CoAxiomRule, [Type], Type)
matchFamExp [s,t]
  | Just 0 <- mbY = Just (axExp0R, [s], num 1)
  | Just 1 <- mbX = Just (axExp1L, [t], num 1)
  | Just 1 <- mbY = Just (axExp1R, [s], s)
  | Just x <- mbX, Just y <- mbY =
    Just (axExpDef, [s,t], num (x ^ y))
  where mbX = isNumLitTy s
        mbY = isNumLitTy t
matchFamExp _ = Nothing

matchFamLog :: [Type] -> Maybe (CoAxiomRule, [Type], Type)
matchFamLog [s]
  | Just x <- mbX, Just (n,_) <- genLog x 2 = Just (axLogDef, [s], num n)
  where mbX = isNumLitTy s
matchFamLog _ = Nothing

matchFamLeq :: [Type] -> Maybe (CoAxiomRule, [Type], Type)
matchFamLeq [s,t]
  | Just 0 <- mbX = Just (axLeq0L, [t], bool True)
  | Just x <- mbX, Just y <- mbY =
    Just (axLeqDef, [s,t], bool (x <= y))
  | tcEqType s t  = Just (axLeqRefl, [s], bool True)
  where mbX = isNumLitTy s
        mbY = isNumLitTy t
matchFamLeq _ = Nothing

matchFamCmpNat :: [Type] -> Maybe (CoAxiomRule, [Type], Type)
matchFamCmpNat [s,t]
  | Just x <- mbX, Just y <- mbY =
    Just (axCmpNatDef, [s,t], ordering (compare x y))
  | tcEqType s t = Just (axCmpNatRefl, [s], ordering EQ)
  where mbX = isNumLitTy s
        mbY = isNumLitTy t
matchFamCmpNat _ = Nothing

matchFamCmpSymbol :: [Type] -> Maybe (CoAxiomRule, [Type], Type)
matchFamCmpSymbol [s,t]
  | Just x <- mbX, Just y <- mbY =
    Just (axCmpSymbolDef, [s,t], ordering (compare x y))
  | tcEqType s t = Just (axCmpSymbolRefl, [s], ordering EQ)
  where mbX = isStrLitTy s
        mbY = isStrLitTy t
matchFamCmpSymbol _ = Nothing

matchFamAppendSymbol :: [Type] -> Maybe (CoAxiomRule, [Type], Type)
matchFamAppendSymbol [s,t]
  | Just x <- mbX, nullFS x = Just (axAppendSymbol0R, [t], t)
  | Just y <- mbY, nullFS y = Just (axAppendSymbol0L, [s], s)
  | Just x <- mbX, Just y <- mbY =
    Just (axAppendSymbolDef, [s,t], mkStrLitTy (appendFS x y))
  where
  mbX = isStrLitTy s
  mbY = isStrLitTy t
matchFamAppendSymbol _ = Nothing

{-------------------------------------------------------------------------------
Interact with axioms
-------------------------------------------------------------------------------}

interactTopAdd :: [Xi] -> Xi -> [Pair Type]
interactTopAdd [s,t] r
  | Just 0 <- mbZ = [ s === num 0, t === num 0 ]                          -- (s + t ~ 0) => (s ~ 0, t ~ 0)
  | Just x <- mbX, Just z <- mbZ, Just y <- minus z x = [t === num y]     -- (5 + t ~ 8) => (t ~ 3)
  | Just y <- mbY, Just z <- mbZ, Just x <- minus z y = [s === num x]     -- (s + 5 ~ 8) => (s ~ 3)
  where
  mbX = isNumLitTy s
  mbY = isNumLitTy t
  mbZ = isNumLitTy r
interactTopAdd _ _ = []

{-
Note [Weakened interaction rule for subtraction]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

A simpler interaction here might be:

  `s - t ~ r` --> `t + r ~ s`

This would enable us to reuse all the code for addition.
Unfortunately, this works a little too well at the moment.
Consider the following example:

    0 - 5 ~ r --> 5 + r ~ 0 --> (5 = 0, r = 0)

This (correctly) spots that the constraint cannot be solved.

However, this may be a problem if the constraint did not
need to be solved in the first place!  Consider the following example:

f :: Proxy (If (5 <=? 0) (0 - 5) (5 - 0)) -> Proxy 5
f = id

Currently, GHC is strict while evaluating functions, so this does not
work, because even though the `If` should evaluate to `5 - 0`, we
also evaluate the "then" branch which generates the constraint `0 - 5 ~ r`,
which fails.

So, for the time being, we only add an improvement when the RHS is a constant,
which happens to work OK for the moment, although clearly we need to do
something more general.
-}
interactTopSub :: [Xi] -> Xi -> [Pair Type]
interactTopSub [s,t] r
  | Just z <- mbZ = [ s === (num z .+. t) ]         -- (s - t ~ 5) => (5 + t ~ s)
  where
  mbZ = isNumLitTy r
interactTopSub _ _ = []

interactTopMul :: [Xi] -> Xi -> [Pair Type]
interactTopMul [s,t] r
  | Just 1 <- mbZ = [ s === num 1, t === num 1 ]                        -- (s * t ~ 1)  => (s ~ 1, t ~ 1)
  | Just x <- mbX, Just z <- mbZ, Just y <- divide z x = [t === num y]  -- (3 * t ~ 15) => (t ~ 5)
  | Just y <- mbY, Just z <- mbZ, Just x <- divide z y = [s === num x]  -- (s * 3 ~ 15) => (s ~ 5)
  where
  mbX = isNumLitTy s
  mbY = isNumLitTy t
  mbZ = isNumLitTy r
interactTopMul _ _ = []

interactTopDiv :: [Xi] -> Xi -> [Pair Type]
interactTopDiv _ _ = []   -- I can't think of anything...

interactTopMod :: [Xi] -> Xi -> [Pair Type]
interactTopMod _ _ = []   -- I can't think of anything...

interactTopExp :: [Xi] -> Xi -> [Pair Type]
interactTopExp [s,t] r
  | Just 0 <- mbZ = [ s === num 0 ]                                       -- (s ^ t ~ 0) => (s ~ 0)
  | Just x <- mbX, Just z <- mbZ, Just y <- logExact  z x = [t === num y] -- (2 ^ t ~ 8) => (t ~ 3)
  | Just y <- mbY, Just z <- mbZ, Just x <- rootExact z y = [s === num x] -- (s ^ 2 ~ 9) => (s ~ 3)
  where
  mbX = isNumLitTy s
  mbY = isNumLitTy t
  mbZ = isNumLitTy r
interactTopExp _ _ = []

interactTopLog :: [Xi] -> Xi -> [Pair Type]
interactTopLog _ _ = []   -- I can't think of anything...

interactTopLeq :: [Xi] -> Xi -> [Pair Type]
interactTopLeq [s,t] r
  | Just 0 <- mbY, Just True <- mbZ = [ s === num 0 ]                     -- (s <= 0) => (s ~ 0)
  where
  mbY = isNumLitTy t
  mbZ = isBoolLitTy r
interactTopLeq _ _ = []

interactTopCmpNat :: [Xi] -> Xi -> [Pair Type]
interactTopCmpNat [s,t] r
  | Just EQ <- isOrderingLitTy r = [ s === t ]
interactTopCmpNat _ _ = []

interactTopCmpSymbol :: [Xi] -> Xi -> [Pair Type]
interactTopCmpSymbol [s,t] r
  | Just EQ <- isOrderingLitTy r = [ s === t ]
interactTopCmpSymbol _ _ = []

interactTopAppendSymbol :: [Xi] -> Xi -> [Pair Type]
interactTopAppendSymbol [s,t] r
  -- (AppendSymbol a b ~ "") => (a ~ "", b ~ "")
  | Just z <- mbZ, nullFS z =
    [s === mkStrLitTy nilFS, t === mkStrLitTy nilFS ]

  -- (AppendSymbol "foo" b ~ "foobar") => (b ~ "bar")
  | Just x <- fmap unpackFS mbX, Just z <- fmap unpackFS mbZ, x `isPrefixOf` z =
    [ t === mkStrLitTy (mkFastString $ drop (length x) z) ]

  -- (AppendSymbol f "bar" ~ "foobar") => (f ~ "foo")
  | Just y <- fmap unpackFS mbY, Just z <- fmap unpackFS mbZ, y `isSuffixOf` z =
    [ t === mkStrLitTy (mkFastString $ take (length z - length y) z) ]

  where
  mbX = isStrLitTy s
  mbY = isStrLitTy t
  mbZ = isStrLitTy r

interactTopAppendSymbol _ _ = []

{-------------------------------------------------------------------------------
Interaction with inerts
-------------------------------------------------------------------------------}

interactInertAdd :: [Xi] -> Xi -> [Xi] -> Xi -> [Pair Type]
interactInertAdd [x1,y1] z1 [x2,y2] z2
  | sameZ && tcEqType x1 x2         = [ y1 === y2 ]
  | sameZ && tcEqType y1 y2         = [ x1 === x2 ]
  where sameZ = tcEqType z1 z2
interactInertAdd _ _ _ _ = []

interactInertSub :: [Xi] -> Xi -> [Xi] -> Xi -> [Pair Type]
interactInertSub [x1,y1] z1 [x2,y2] z2
  | sameZ && tcEqType x1 x2         = [ y1 === y2 ]
  | sameZ && tcEqType y1 y2         = [ x1 === x2 ]
  where sameZ = tcEqType z1 z2
interactInertSub _ _ _ _ = []

interactInertMul :: [Xi] -> Xi -> [Xi] -> Xi -> [Pair Type]
interactInertMul [x1,y1] z1 [x2,y2] z2
  | sameZ && known (/= 0) x1 && tcEqType x1 x2 = [ y1 === y2 ]
  | sameZ && known (/= 0) y1 && tcEqType y1 y2 = [ x1 === x2 ]
  where sameZ   = tcEqType z1 z2

interactInertMul _ _ _ _ = []

interactInertDiv :: [Xi] -> Xi -> [Xi] -> Xi -> [Pair Type]
interactInertDiv _ _ _ _ = []

interactInertMod :: [Xi] -> Xi -> [Xi] -> Xi -> [Pair Type]
interactInertMod _ _ _ _ = []

interactInertExp :: [Xi] -> Xi -> [Xi] -> Xi -> [Pair Type]
interactInertExp [x1,y1] z1 [x2,y2] z2
  | sameZ && known (> 1) x1 && tcEqType x1 x2 = [ y1 === y2 ]
  | sameZ && known (> 0) y1 && tcEqType y1 y2 = [ x1 === x2 ]
  where sameZ = tcEqType z1 z2

interactInertExp _ _ _ _ = []

interactInertLog :: [Xi] -> Xi -> [Xi] -> Xi -> [Pair Type]
interactInertLog _ _ _ _ = []

interactInertLeq :: [Xi] -> Xi -> [Xi] -> Xi -> [Pair Type]
interactInertLeq [x1,y1] z1 [x2,y2] z2
  | bothTrue && tcEqType x1 y2 && tcEqType y1 x2 = [ x1 === y1 ]
  | bothTrue && tcEqType y1 x2                 = [ (x1 <== y2) === bool True ]
  | bothTrue && tcEqType y2 x1                 = [ (x2 <== y1) === bool True ]
  where bothTrue = isJust $ do True <- isBoolLitTy z1
                               True <- isBoolLitTy z2
                               return ()

interactInertLeq _ _ _ _ = []

interactInertAppendSymbol :: [Xi] -> Xi -> [Xi] -> Xi -> [Pair Type]
interactInertAppendSymbol [x1,y1] z1 [x2,y2] z2
  | sameZ && tcEqType x1 x2         = [ y1 === y2 ]
  | sameZ && tcEqType y1 y2         = [ x1 === x2 ]
  where sameZ = tcEqType z1 z2
interactInertAppendSymbol _ _ _ _ = []

{- -----------------------------------------------------------------------------
These inverse functions are used for simplifying propositions using
concrete natural numbers.
----------------------------------------------------------------------------- -}

-- | Subtract two natural numbers.
minus :: Integer -> Integer -> Maybe Integer
minus x y = if x >= y then Just (x - y) else Nothing

-- | Compute the exact logarithm of a natural number.
-- The logarithm base is the second argument.
logExact :: Integer -> Integer -> Maybe Integer
logExact x y = do (z,True) <- genLog x y
                  return z

-- | Divide two natural numbers.
divide :: Integer -> Integer -> Maybe Integer
divide _ 0  = Nothing
divide x y  = case divMod x y of
                (a,0) -> Just a
                _     -> Nothing

-- | Compute the exact root of a natural number.
-- The second argument specifies which root we are computing.
rootExact :: Integer -> Integer -> Maybe Integer
rootExact x y = do (z,True) <- genRoot x y
                   return z

{- | Compute the n-th root of a natural number, rounded down to
the closest natural number.  The boolean indicates if the result
is exact (i.e., True means no rounding was done, False means rounded down).
The second argument specifies which root we are computing. -}
genRoot :: Integer -> Integer -> Maybe (Integer, Bool)
genRoot _  0    = Nothing
genRoot x0 1    = Just (x0, True)
genRoot x0 root = Just (search 0 (x0+1))
  where
  search from to = let x = from + div (to - from) 2
                       a = x ^ root
                   in case compare a x0 of
                        EQ              -> (x, True)
                        LT | x /= from  -> search x to
                           | otherwise  -> (from, False)
                        GT | x /= to    -> search from x
                           | otherwise  -> (from, False)

{- | Compute the logarithm of a number in the given base, rounded down to the
closest integer.  The boolean indicates if we the result is exact
(i.e., True means no rounding happened, False means we rounded down).
The logarithm base is the second argument. -}
genLog :: Integer -> Integer -> Maybe (Integer, Bool)
genLog x 0    = if x == 1 then Just (0, True) else Nothing
genLog _ 1    = Nothing
genLog 0 _    = Nothing
genLog x base = Just (exactLoop 0 x)
  where
  exactLoop s i
    | i == 1     = (s,True)
    | i < base   = (s,False)
    | otherwise  =
        let s1 = s + 1
        in s1 `seq` case divMod i base of
                      (j,r)
                        | r == 0    -> exactLoop s1 j
                        | otherwise -> (underLoop s1 j, False)

  underLoop s i
    | i < base  = s
    | otherwise = let s1 = s + 1 in s1 `seq` underLoop s1 (div i base)