content: {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE PatternSynonyms #-}
  {-# LANGUAGE TemplateHaskell #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE UndecidableInstances #-}
  {-# OPTIONS_GHC -fno-warn-orphans #-}
  
  -- |
  -- SPDX-License-Identifier: BSD-3-Clause
  --
  -- Types for the Swarm programming language and related utilities.
  module Swarm.Language.Types (
    -- * Basic definitions
  
    -- ** Base types and type constructors
    BaseTy (..),
    baseTyName,
    TyCon (..),
    Arity (..),
    tcArity,
    Var,
  
    -- ** Type structure functor
    TypeF (..),
  
    -- ** Recursive types
    Nat (..),
    natToInt,
    unfoldRec,
    SubstRec (..),
  
    -- * @Type@
    Type,
    tyVars,
    pattern TyConApp,
    pattern TyBase,
    pattern TyVar,
    pattern TyVoid,
    pattern TyUnit,
    pattern TyInt,
    pattern TyText,
    pattern TyDir,
    pattern TyBool,
    pattern TyActor,
    pattern TyKey,
    pattern TyType,
    pattern (:+:),
    pattern (:*:),
    pattern (:->:),
    pattern TyRcd,
    pattern TyCmd,
    pattern TyDelay,
    pattern TyUser,
    pattern TyRec,
  
    -- * @UType@
    IntVar (..),
    UType,
    pattern UTyConApp,
    pattern UTyBase,
    pattern UTyVar,
    pattern UTyVar',
    pattern UTyVoid,
    pattern UTyUnit,
    pattern UTyInt,
    pattern UTyText,
    pattern UTyDir,
    pattern UTyBool,
    pattern UTyActor,
    pattern UTyKey,
    pattern UTyType,
    pattern UTySum,
    pattern UTyProd,
    pattern UTyFun,
    pattern UTyRcd,
    pattern UTyCmd,
    pattern UTyDelay,
    pattern UTyUser,
    pattern UTyRec,
  
    -- ** Utilities
    ucata,
    mkVarName,
    fuvs,
    hasAnyUVars,
    isTopLevelConstructor,
    UnchainableFun (..),
  
    -- * Polytypes
    ImplicitQuantification (..),
    Poly,
    mkPoly,
    mkQPoly,
    mkTrivPoly,
    unPoly,
    ptVars,
    ptBody,
    Polytype,
    RawPolytype,
    pattern PolyUnit,
    UPolytype,
    quantify,
    absQuantify,
    forgetQ,
  
    -- * Contexts
    TCtx,
    UCtx,
    TVCtx,
  
    -- * User type definitions
    TydefInfo (..),
    tydefType,
    tydefArity,
    substTydef,
    ExpandTydefErr (..),
    expandTydef,
    expandTydefs,
    elimTydef,
    TDCtx,
    emptyTDCtx,
    lookupTD,
    lookupTDR,
    addBindingTD,
    withBindingTD,
    resolveUserTy,
  
    -- * WHNF
    whnfType,
  
    -- * The 'WithU' class
    WithU (..),
  ) where
  
  import Control.Algebra (Has, run)
  import Control.Carrier.Reader (runReader)
  import Control.Carrier.Throw.Either (runThrow)
  import Control.Effect.Reader (Reader, ask, local)
  import Control.Effect.Throw (Throw, throwError)
  import Control.Lens (Plated (..), makeLenses, rewriteM, view)
  import Control.Monad.Free
  import Data.Aeson (FromJSON (..), FromJSON1 (..), ToJSON (..), ToJSON1 (..), genericLiftParseJSON, genericLiftToJSON, genericParseJSON, genericToJSON)
  import Data.Data (Data)
  import Data.Data.Lens (uniplate)
  import Data.Eq.Deriving (deriveEq1)
  import Data.Fix
  import Data.Foldable (fold)
  import Data.Functor.Classes (Eq1)
  import Data.Hashable (Hashable (..))
  import Data.Hashable.Lifted (Hashable1)
  import Data.Kind qualified
  import Data.List.NonEmpty ((<|))
  import Data.List.NonEmpty qualified as NE
  import Data.Map.Strict (Map)
  import Data.Map.Strict qualified as M
  import Data.MonoidMap (MonoidMap)
  import Data.MonoidMap qualified as MM
  import Data.MonoidMap.JSON ()
  import Data.Ord.Deriving (deriveOrd1)
  import Data.Semigroup (Sum (..))
  import Data.Set (Set)
  import Data.Set qualified as S
  import Data.String (IsString (..))
  import Data.Text (Text)
  import Data.Text qualified as T
  import GHC.Generics (Generic, Generic1)
  import Prettyprinter (align, braces, brackets, concatWith, flatAlt, hsep, pretty, punctuate, softline, (<+>))
  import Swarm.Language.Context (Ctx)
  import Swarm.Language.Context qualified as Ctx
  import Swarm.Language.TDVar (TDVar, mkTDVar, mkTDVar', tdVarName)
  import Swarm.Language.Var (Var)
  import Swarm.Pretty (PrettyPrec (..), pparens, pparens', ppr, prettyBinding)
  import Swarm.Util (showT, unsnocNE)
  import Swarm.Util.JSON (optionsMinimize, optionsUnwrapUnary)
  import Text.Show.Deriving (deriveShow1)
  import Witch (into)
  
  ------------------------------------------------------------
  -- Base types
  ------------------------------------------------------------
  
  -- | Base types.
  data BaseTy
    = -- | The void type, with no inhabitants.
      BVoid
    | -- | The unit type, with a single inhabitant.
      BUnit
    | -- | Signed, arbitrary-size integers.
      BInt
    | -- | Unicode strings.
      BText
    | -- | Directions.
      BDir
    | -- | Booleans.
      BBool
    | -- | "Actors", i.e. anything that can do stuff. Internally, these
      --   are all just "robots", but we give them a more generic
      --   in-game name because they could represent other things like
      --   aliens, animals, seeds, ...
      BActor
    | -- | Keys, i.e. things that can be pressed on the keyboard
      BKey
    | -- | The type of types
      BType
    deriving (Eq, Ord, Show, Bounded, Enum, Data, Generic, Hashable, FromJSON, ToJSON)
  
  baseTyName :: BaseTy -> Text
  baseTyName = into @Text . drop 1 . show
  
  instance PrettyPrec BaseTy where
    prettyPrec _ = pretty . drop 1 . show
  
  ------------------------------------------------------------
  -- Type constructors
  ------------------------------------------------------------
  
  -- | Type constructors.
  data TyCon
    = -- | Base types are (nullary) type constructors.
      TCBase BaseTy
    | -- | Command types.
      TCCmd
    | -- | Delayed computations.
      TCDelay
    | -- | Sum types.
      TCSum
    | -- | Product types.
      TCProd
    | -- | Function types.
      TCFun
    | -- | User-defined type constructor.
      TCUser TDVar
    deriving (Eq, Ord, Show, Data, Generic, Hashable)
  
  instance ToJSON TyCon where
    toJSON = genericToJSON optionsMinimize
  
  instance FromJSON TyCon where
    parseJSON = genericParseJSON optionsMinimize
  
  instance PrettyPrec TyCon where
    prettyPrec _ = \case
      TCBase b -> ppr b
      TCCmd -> "Cmd"
      TCDelay -> "Delay"
      TCSum -> "Sum"
      TCProd -> "Prod"
      TCFun -> "Fun"
      TCUser t -> ppr t
  
  -- | The arity of a type, /i.e./ the number of type parameters it
  --   expects.
  newtype Arity = Arity {getArity :: Int}
    deriving (Eq, Ord, Show, Generic, Data, Hashable)
  
  instance ToJSON Arity where
    toJSON = genericToJSON optionsUnwrapUnary
  
  instance FromJSON Arity where
    parseJSON = genericParseJSON optionsUnwrapUnary
  
  instance PrettyPrec Arity where
    prettyPrec _ (Arity a) = pretty a
  
  ------------------------------------------------------------
  -- Types
  ------------------------------------------------------------
  
  -- | Peano naturals, for use as de Bruijn indices in recursive types.
  data Nat where
    NZ :: Nat
    NS :: Nat -> Nat
    deriving (Eq, Ord, Show, Data, Generic, Hashable, FromJSON, ToJSON)
  
  natToInt :: Nat -> Int
  natToInt NZ = 0
  natToInt (NS n) = 1 + natToInt n
  
  -- | A "structure functor" encoding the shape of type expressions.
  --   Actual types are then represented by taking a fixed point of this
  --   functor.  We represent types in this way, via a "two-level type",
  --   so that we can easily use generic recursion schemes to implement
  --   things like substitution.
  data TypeF t
    = -- | A type constructor applied to some type arguments. For now,
      --   all type constructor applications are required to be fully
      --   saturated (higher kinds are not supported), so we just
      --   directly store a list of all arguments (as opposed to
      --   iterating binary application).
      TyConF TyCon [t]
    | -- | A type variable.  The first Var represents the original name,
      --   and should be ignored except for use in e.g. error messages.
      --   The second Var is the real name of the variable; it may be the
      --   same as the first, or it may be different e.g. if it is a
      --   freshly generated skolem variable.
      TyVarF Var Var
    | -- | Record type.
      TyRcdF (Map Var t)
    | -- | A recursive type variable bound by an enclosing Rec,
      --   represented by a de Bruijn index.
      TyRecVarF Nat
    | -- | Recursive type.  The variable is just a suggestion for a name to use
      --   when pretty-printing; the actual bound variables are represented
      --   via de Bruijn indices.
      TyRecF Var t
    deriving (Show, Eq, Ord, Functor, Foldable, Traversable, Generic, Generic1, Data, Hashable)
  
  deriveEq1 ''TypeF
  deriveOrd1 ''TypeF
  deriveShow1 ''TypeF
  
  instance Hashable1 TypeF -- needs the Eq1 instance
  
  instance ToJSON1 TypeF where
    liftToJSON = genericLiftToJSON optionsMinimize
  
  instance FromJSON1 TypeF where
    liftParseJSON = genericLiftParseJSON optionsMinimize
  
  -- | @Type@ is now defined as the fixed point of 'TypeF'.  It would be
  --   annoying to manually apply and match against 'Fix' constructors
  --   everywhere, so we provide pattern synonyms that allow us to work
  --   with 'Type' as if it were defined in a directly recursive way.
  type Type = Fix TypeF
  
  instance Plated Type where
    plate = uniplate
  
  newtype IntVar = IntVar Int
    deriving (Show, Data, Eq, Ord, Generic, Hashable)
  
  instance PrettyPrec IntVar where
    prettyPrec _ = ppr . mkVarName "u"
  
  -- | 'UType's are like 'Type's, but also contain unification
  --   variables.  'UType' is defined via 'Free', which is also a kind
  --   of fixed point (in fact, @Free TypeF@ is the /free monad/ over
  --   'TypeF').
  --
  --   Just as with 'Type', we provide a bunch of pattern synonyms for
  --   working with 'UType' as if it were defined directly.
  type UType = Free TypeF IntVar
  
  -- orphan instance
  instance (Eq1 f, Hashable x, Hashable (f (Free f x))) => Hashable (Free f x)
  
  -- | A generic /fold/ for things defined via 'Free' (including, in
  --   particular, 'UType').
  ucata :: Functor t => (v -> a) -> (t a -> a) -> Free t v -> a
  ucata f _ (Pure v) = f v
  ucata f g (Free t) = g (fmap (ucata f g) t)
  
  -- | A quick-and-dirty method for turning an 'IntVar' (used internally
  --   as a unification variable) into a unique variable name, by
  --   appending a number to the given name.
  mkVarName :: Text -> IntVar -> Var
  mkVarName nm (IntVar v) = T.append nm (showT v)
  
  -- | Get all the free unification variables in a 'UType'.
  fuvs :: UType -> Set IntVar
  fuvs = ucata S.singleton fold
  
  -- | Check whether a type contains any unification variables at all.
  hasAnyUVars :: UType -> Bool
  hasAnyUVars = ucata (const True) or
  
  -- | Check whether a type consists of a top-level type constructor
  --   immediately applied to unification variables.
  isTopLevelConstructor :: UType -> Maybe (TypeF ())
  isTopLevelConstructor = \case
    Free (TyRcdF m) | all isPure m -> Just (TyRcdF M.empty)
    UTyConApp c ts | all isPure ts -> Just (TyConF c [])
    _ -> Nothing
  
  isPure :: Free f a -> Bool
  isPure (Pure {}) = True
  isPure _ = False
  
  -- | For convenience, so we can write /e.g./ @"a"@ instead of @TyVar "a"@.
  instance IsString Type where
    fromString = TyVar . into @Text
  
  instance IsString UType where
    fromString = UTyVar . into @Text
  
  --------------------------------------------------
  -- Recursive type utilities
  
  -- | @unfoldRec x t@ unfolds the recursive type @rec x. t@ one step,
  --   to @t [(rec x. t) / x]@.
  unfoldRec :: SubstRec t => Var -> t -> t
  unfoldRec x ty = substRec (TyRecF x ty) ty NZ
  
  -- | Class of type-like things where we can substitute for a bound de
  --   Bruijn variable.
  class SubstRec t where
    -- | @substRec s t n@ substitutes @s@ for the bound de Bruijn variable
    --   @n@ everywhere in @t@.
    substRec :: TypeF t -> t -> Nat -> t
  
  instance SubstRec (Free TypeF v) where
    substRec s = ucata (\i _ -> Pure i) $ \f i -> case f of
      TyRecVarF j
        | i == j -> Free s
        | otherwise -> Free (TyRecVarF j)
      TyRecF x g -> Free (TyRecF x (g (NS i)))
      _ -> Free (fmap ($ i) f)
  
  instance SubstRec Type where
    substRec s = foldFix $ \f i -> case f of
      TyRecVarF j
        | i == j -> Fix s
        | otherwise -> Fix (TyRecVarF j)
      TyRecF x g -> Fix (TyRecF x (g (NS i)))
      _ -> Fix (fmap ($ i) f)
  
  --------------------------------------------------
  -- Pretty-printing machinery for types
  
  -- | Split a function type chain, so that we can pretty print
  --   the type parameters aligned on each line when they don't fit.
  class UnchainableFun t where
    unchainFun :: t -> NE.NonEmpty t
  
  instance UnchainableFun Type where
    unchainFun (a :->: ty) = a <| unchainFun ty
    unchainFun ty = pure ty
  
  instance UnchainableFun (Free TypeF ty) where
    unchainFun (Free (TyConF TCFun [ty1, ty2])) = ty1 <| unchainFun ty2
    unchainFun ty = pure ty
  
  instance (UnchainableFun t, PrettyPrec t, SubstRec t) => PrettyPrec (TypeF t) where
    prettyPrec p = \case
      TyVarF v _ -> ppr v
      TyRcdF m -> brackets $ hsep (punctuate "," (map prettyBinding (M.assocs m)))
      -- Special cases for type constructors with special syntax.
      -- Always use parentheses around sum and product types, see #1625
      TyConF TCSum [ty1, ty2] ->
        pparens (p > 0) $
          prettyPrec 2 ty1 <+> "+" <+> prettyPrec 2 ty2
      TyConF TCProd [ty1, ty2] ->
        pparens (p > 0) $
          prettyPrec 2 ty1 <+> "*" <+> prettyPrec 2 ty2
      TyConF TCDelay [ty] -> braces $ ppr ty
      TyConF TCFun [ty1, ty2] ->
        let (iniF, lastF) = unsnocNE $ ty1 <| unchainFun ty2
            funs = (prettyPrec 2 <$> iniF) <> [prettyPrec 1 lastF]
            inLine l r = l <+> "->" <+> r
            multiLine l r = l <+> "->" <> softline <> r
         in pparens' (p > 1) . align $
              flatAlt (concatWith multiLine funs) (concatWith inLine funs)
      TyRecF x ty ->
        pparens (p > 0) $
          "rec" <+> ppr x <> "." <+> prettyPrec 0 (substRec (TyVarF x x) ty NZ)
      -- This case shouldn't be possible, since TyRecVar should only occur inside a TyRec,
      -- and pretty-printing the TyRec (above) will substitute a variable name for
      -- any bound TyRecVars before recursing.
      TyRecVarF i -> pretty (show (natToInt i))
      -- Fallthrough cases for type constructor application.  Handles base
      -- types, Cmd, user-defined types, or ill-kinded things like 'Int
      -- Bool'.
      TyConF c [] -> ppr c
      TyConF c tys -> pparens (p > 9) $ ppr c <+> hsep (map (prettyPrec 10) tys)
  
  ------------------------------------------------------------
  -- Generic folding over type representations
  ------------------------------------------------------------
  
  -- | Type class for various type representations (e.g. 'Type',
  --   'UType') and generic operations we can perform on them.  This
  --   helps us avoid code duplication in some cases, by writing a
  --   single generic function which works for any 'Typical' instance.
  class Typical t where
    -- | Fold a type into a summary value of some type @a@, given an
    --   "empty" value of type @a@ (for use at e.g. Pure nodes in a
    --   UType)
    foldT :: a -> (TypeF a -> a) -> t -> a
  
    -- | Refold a type into another type.  This is a special case of
    --   'foldT' when we want to produce another type, since then we can
    --   do something more sensible with 'Pure' nodes in a 'UType'
    --   (i.e. preserve them instead of replacing them with a default
    --   value).
    refoldT :: (TypeF t -> t) -> t -> t
  
    -- | Unroll one level of structure.
    unrollT :: t -> Maybe (TypeF t)
  
    -- | Roll up one level of structure.
    rollT :: TypeF t -> t
  
    -- | It should be possible to convert 'Type' to any type-ish thing.
    fromType :: Type -> t
  
  instance Typical Type where
    foldT _ = foldFix
    refoldT = foldFix
    unrollT (Fix t) = Just t
    rollT = Fix
    fromType = id
  
  instance Typical UType where
    foldT e = ucata (const e)
    refoldT = ucata Pure
    unrollT (Free t) = Just t
    unrollT _ = Nothing
    rollT = Free
    fromType = toU
  
  -- | Get all the type variables (/not/ unification variables)
  --   contained in a 'Type' or 'UType'.
  tyVars :: Typical t => t -> Set Var
  tyVars = foldT S.empty (\case TyVarF _ x -> S.singleton x; f -> fold f)
  
  ------------------------------------------------------------
  -- Polytypes
  ------------------------------------------------------------
  
  data ImplicitQuantification = Unquantified | Quantified
    deriving (Eq, Ord, Read, Show)
  
  -- | A @Poly q t@ is a universally quantified @t@.  The variables in
  --   the list are bound inside the @t@.  For example, the type @forall
  --   a. a -> a@ would be represented as @Forall ["a"] (TyFun "a"
  --   "a")@.
  --
  --   The type index @q@ is a phantom type index indicating whether the
  --   type has been implicitly quantified.  Immediately after a
  --   polytype is parsed it is 'Unquantified' and unsafe to use.  For
  --   example, the type @a -> a@ would parse literally as @Forall []
  --   (TyFun "a" "a") :: Poly Unquantified Type@, where the type
  --   variable @a@ is not in the list of bound variables.  Later, after
  --   running through 'quantify', it would become a @Poly Quantified
  --   Type@, either @Forall ["a"] (TyFun "a" "a")@ if the type variable
  --   is implicitly quantified, or unchanged if the type variable @a@
  --   was already in scope.
  --
  --   The @Poly@ constructor intentionally unexported, so that the
  --   only way to create a @Poly Quantified@ is through the 'quantify'
  --   function.
  data Poly (q :: ImplicitQuantification) t = Forall {_ptVars :: [Var], ptBody :: t}
    deriving (Show, Eq, Functor, Foldable, Traversable, Data, Generic, FromJSON, ToJSON, Hashable)
  
  -- | Create a raw, unquantified @Poly@ value.
  mkPoly :: [Var] -> t -> Poly 'Unquantified t
  mkPoly = Forall
  
  -- | Create a polytype while performing implicit quantification.
  mkQPoly :: Typical t => t -> Poly 'Quantified t
  mkQPoly = absQuantify . Forall []
  
  -- | Create a trivial "polytype" with no bound variables.  This is
  --   somewhat unsafe --- only use this if you are sure that the polytype
  --   you want has no type variables.
  mkTrivPoly :: t -> Poly q t
  mkTrivPoly = Forall []
  
  -- | Project out the variables and body of a 'Poly' value.  It's only
  --   possible to project from a 'Poly Quantified' since the list of
  --   variables might be meaningless for a type that has not had
  --   implicit quantification applied.
  unPoly :: Poly 'Quantified t -> ([Var], t)
  unPoly (Forall xs t) = (xs, t)
  
  -- | Project out the variables of a 'Poly Quantified' value.
  ptVars :: Poly 'Quantified t -> [Var]
  ptVars (Forall xs _) = xs
  
  -- | Forget that a polytype has been properly quantified.
  forgetQ :: Poly 'Quantified t -> Poly 'Unquantified t
  forgetQ (Forall xs t) = Forall xs t
  
  -- | A regular polytype (without unification variables).  A @Polytype@
  --   (as opposed to a @RawPolytype@) is guaranteed to have implicit
  --   quantification properly applied, so that all type variables bound
  --   by the forall are explicitly listed.
  type Polytype = Poly 'Quantified Type
  
  -- | A raw polytype (without unification variables), which corresponds
  --   exactly to the way a polytype was written in source code.  In
  --   particular there may be type variables which are used in the type
  --   but not listed explicitly, which are to be implicitly quantified.
  type RawPolytype = Poly 'Unquantified Type
  
  instance PrettyPrec (Poly q Type) where
    prettyPrec _ (Forall [] t) = ppr t
    prettyPrec _ (Forall xs t) = hsep ("∀" : map ppr xs) <> "." <+> ppr t
  
  -- | A polytype with unification variables.
  type UPolytype = Poly 'Quantified UType
  
  instance PrettyPrec (Poly q UType) where
    prettyPrec _ (Forall [] t) = ppr t
    prettyPrec _ (Forall xs t) = hsep ("∀" : map ppr xs) <> "." <+> ppr t
  
  ------------------------------------------------------------
  -- WithU
  ------------------------------------------------------------
  
  -- | In several cases we have two versions of something: a "normal"
  --   version, and a @U@ version with unification variables in it
  --   (/e.g./ 'Type' vs 'UType', 'Polytype' vs 'UPolytype', 'TCtx' vs
  --   'UCtx'). This class abstracts over the process of converting back
  --   and forth between them.
  --
  --   In particular, @'WithU' t@ represents the fact that the type @t@
  --   also has a @U@ counterpart, with a way to convert back and forth.
  --   Note, however, that converting back may be "unsafe" in the sense
  --   that it requires an extra burden of proof to guarantee that it is
  --   used only on inputs that are safe.
  class WithU t where
    -- | The associated "@U@-version" of the type @t@.
    type U t :: Data.Kind.Type
  
    -- | Convert from @t@ to its associated "@U@-version".  This
    --   direction is always safe (we simply have no unification
    --   variables even though the type allows it).
    toU :: t -> U t
  
    -- | Convert from the associated "@U@-version" back to @t@.
    --   Generally, this direction requires somehow knowing that there
    --   are no longer any unification variables in the value being
    --   converted.
    fromU :: U t -> Maybe t
  
  -- | 'Type' is an instance of 'WithU', with associated type 'UType'.
  instance WithU Type where
    type U Type = UType
    toU = foldFix Free
    fromU = ucata (const Nothing) (fmap wrapFix . sequence)
  
  -- | A 'WithU' instance can be lifted through any functor (including,
  --   in particular, 'Ctx' and 'Poly').
  instance (WithU t, Traversable f) => WithU (f t) where
    type U (f t) = f (U t)
    toU = fmap toU
    fromU = traverse fromU
  
  ------------------------------------------------------------
  -- Pattern synonyms
  ------------------------------------------------------------
  
  --------------------------------------------------
  -- Type
  
  pattern TyConApp :: TyCon -> [Type] -> Type
  pattern TyConApp c tys = Fix (TyConF c tys)
  
  pattern TyBase :: BaseTy -> Type
  pattern TyBase b = TyConApp (TCBase b) []
  
  pattern TyVar :: Var -> Type
  pattern TyVar x <- Fix (TyVarF _ x)
    where
      TyVar x = Fix (TyVarF x x)
  
  pattern TyVoid :: Type
  pattern TyVoid = TyBase BVoid
  
  pattern TyUnit :: Type
  pattern TyUnit = TyBase BUnit
  
  pattern TyInt :: Type
  pattern TyInt = TyBase BInt
  
  pattern TyText :: Type
  pattern TyText = TyBase BText
  
  pattern TyDir :: Type
  pattern TyDir = TyBase BDir
  
  pattern TyBool :: Type
  pattern TyBool = TyBase BBool
  
  pattern TyActor :: Type
  pattern TyActor = TyBase BActor
  
  pattern TyKey :: Type
  pattern TyKey = TyBase BKey
  
  pattern TyType :: Type
  pattern TyType = TyBase BType
  
  infixr 5 :+:
  
  pattern (:+:) :: Type -> Type -> Type
  pattern ty1 :+: ty2 = TyConApp TCSum [ty1, ty2]
  
  infixr 6 :*:
  
  pattern (:*:) :: Type -> Type -> Type
  pattern ty1 :*: ty2 = TyConApp TCProd [ty1, ty2]
  
  infixr 1 :->:
  
  pattern (:->:) :: Type -> Type -> Type
  pattern ty1 :->: ty2 = TyConApp TCFun [ty1, ty2]
  
  pattern TyRcd :: Map Var Type -> Type
  pattern TyRcd m = Fix (TyRcdF m)
  
  pattern TyCmd :: Type -> Type
  pattern TyCmd ty = TyConApp TCCmd [ty]
  
  pattern TyDelay :: Type -> Type
  pattern TyDelay ty = TyConApp TCDelay [ty]
  
  pattern TyUser :: TDVar -> [Type] -> Type
  pattern TyUser v tys = TyConApp (TCUser v) tys
  
  pattern TyRec :: Var -> Type -> Type
  pattern TyRec x ty = Fix (TyRecF x ty)
  
  --------------------------------------------------
  -- UType
  
  pattern UTyConApp :: TyCon -> [UType] -> UType
  pattern UTyConApp c tys = Free (TyConF c tys)
  
  pattern UTyBase :: BaseTy -> UType
  pattern UTyBase b = UTyConApp (TCBase b) []
  
  pattern UTyVar :: Var -> UType
  pattern UTyVar x <- Free (TyVarF _ x)
    where
      UTyVar x = Free (TyVarF x x)
  
  pattern UTyVar' :: Var -> Var -> UType
  pattern UTyVar' x y = Free (TyVarF x y)
  
  pattern UTyVoid :: UType
  pattern UTyVoid = UTyBase BVoid
  
  pattern UTyUnit :: UType
  pattern UTyUnit = UTyBase BUnit
  
  pattern UTyInt :: UType
  pattern UTyInt = UTyBase BInt
  
  pattern UTyText :: UType
  pattern UTyText = UTyBase BText
  
  pattern UTyDir :: UType
  pattern UTyDir = UTyBase BDir
  
  pattern UTyBool :: UType
  pattern UTyBool = UTyBase BBool
  
  pattern UTyActor :: UType
  pattern UTyActor = UTyBase BActor
  
  pattern UTyKey :: UType
  pattern UTyKey = UTyBase BKey
  
  pattern UTyType :: UType
  pattern UTyType = UTyBase BType
  
  pattern UTySum :: UType -> UType -> UType
  pattern UTySum ty1 ty2 = UTyConApp TCSum [ty1, ty2]
  
  pattern UTyProd :: UType -> UType -> UType
  pattern UTyProd ty1 ty2 = UTyConApp TCProd [ty1, ty2]
  
  pattern UTyFun :: UType -> UType -> UType
  pattern UTyFun ty1 ty2 = UTyConApp TCFun [ty1, ty2]
  
  pattern UTyRcd :: Map Var UType -> UType
  pattern UTyRcd m = Free (TyRcdF m)
  
  pattern UTyCmd :: UType -> UType
  pattern UTyCmd ty = UTyConApp TCCmd [ty]
  
  pattern UTyDelay :: UType -> UType
  pattern UTyDelay ty = UTyConApp TCDelay [ty]
  
  pattern UTyUser :: TDVar -> [UType] -> UType
  pattern UTyUser v tys = UTyConApp (TCUser v) tys
  
  pattern UTyRec :: Var -> UType -> UType
  pattern UTyRec x ty = Free (TyRecF x ty)
  
  pattern PolyUnit :: Polytype
  pattern PolyUnit = Forall [] (TyCmd TyUnit)
  
  ------------------------------------------------------------
  -- Contexts
  ------------------------------------------------------------
  
  -- | A @TCtx@ is a mapping from variables to polytypes.  We generally
  --   get one of these at the end of the type inference process.
  type TCtx = Ctx Var Polytype
  
  -- | A @UCtx@ is a mapping from variables to polytypes with
  --   unification variables.  We generally have one of these while we
  --   are in the midst of the type inference process.
  type UCtx = Ctx Var UPolytype
  
  -- | A @TVCtx@ tracks which type variables are in scope, and what
  --   skolem variables were assigned to them.
  type TVCtx = Ctx Var UType
  
  ------------------------------------------------------------
  -- Implicit quantification of polytypes
  ------------------------------------------------------------
  
  -- | Implicitly quantify any otherwise unbound type variables.
  quantify ::
    (Has (Reader TVCtx) sig m, Typical ty) =>
    Poly 'Unquantified ty ->
    m (Poly 'Quantified ty)
  quantify (Forall xs ty) = do
    inScope <- ask @TVCtx
    -- Look at all variables which occur in the type but are not
    -- explicitly bound by the forall, and and are not bound in the
    -- context.  Such variables must be implicitly quantified.
    let implicit =
          (tyVars ty `S.difference` S.fromList xs)
            `S.difference` M.keysSet (Ctx.unCtx inScope)
    pure $ Forall (xs ++ S.toList implicit) ty
  
  -- | Absolute implicit quantification, i.e. assume there are no type
  --   variables in scope.
  absQuantify :: Typical t => Poly 'Unquantified t -> Poly 'Quantified t
  absQuantify = run . runReader @TVCtx Ctx.empty . quantify
  
  ------------------------------------------------------------
  -- Type definitions
  ------------------------------------------------------------
  
  data TydefInfo = TydefInfo
    { _tydefType :: Polytype
    , _tydefArity :: Arity
    }
    deriving (Eq, Show, Generic, Data, FromJSON, ToJSON, Hashable)
  
  makeLenses ''TydefInfo
  
  -- | A @TDCtx@ is a mapping from user-defined type constructor names
  --   to their definitions and arities/kinds.  It also stores the
  --   latest version of each name (for any names with more than one
  --   version), so we can tell when a type definition has been
  --   shadowed.
  data TDCtx = TDCtx
    { getTDCtx :: Ctx TDVar TydefInfo
    , getTDVersions :: MonoidMap Text (Sum Int)
    }
    deriving (Eq, Generic, Show, ToJSON)
  
  -- Need to write manual Hashable instance since MonoidMap
  -- does not have instances of its own.
  
  instance Hashable TDCtx where
    hashWithSalt s (TDCtx ctx versions) =
      s
        `hashWithSalt` ctx
        `hashWithSalt` (getSum <$> MM.toMap versions)
  
  -- | The empty type definition context.
  emptyTDCtx :: TDCtx
  emptyTDCtx = TDCtx Ctx.empty MM.empty
  
  -- | Look up a variable in the type definition context.
  lookupTD :: TDVar -> TDCtx -> Maybe TydefInfo
  lookupTD x = Ctx.lookup x . getTDCtx
  
  -- | Look up a variable in an ambient type definition context.
  lookupTDR :: Has (Reader TDCtx) sig m => TDVar -> m (Maybe TydefInfo)
  lookupTDR x = lookupTD x <$> ask
  
  -- | Add a binding of a variable name to a type definition, giving it
  --   an appropriate version number if it shadows other variables with
  --   the same name.
  addBindingTD :: Text -> TydefInfo -> TDCtx -> TDCtx
  addBindingTD x info (TDCtx tdCtx tdVersions) =
    case Ctx.lookup (mkTDVar x) tdCtx of
      Nothing -> TDCtx (Ctx.addBinding (mkTDVar x) info tdCtx) tdVersions
      Just _ ->
        let newVersion = 1 + MM.get x tdVersions
         in TDCtx (Ctx.addBinding (mkTDVar' (getSum newVersion) x) info tdCtx) (MM.set x newVersion tdVersions)
  
  -- | Locally extend the ambient type definition context with an
  --   additional binding, via 'addBindingTD'.
  withBindingTD :: Has (Reader TDCtx) sig m => Text -> TydefInfo -> m a -> m a
  withBindingTD x info = local (addBindingTD x info)
  
  -- | Given a parsed variable representing a user-defined type, figure
  --   out which version is currently in scope and set the version
  --   number of the variable appropriately.
  resolveUserTy :: Has (Reader TDCtx) sig m => TDVar -> m TDVar
  resolveUserTy v = do
    tdCtx <- ask
    let x = tdVarName v
    let ver = getSum (MM.get x (getTDVersions tdCtx))
    pure $ mkTDVar' ver x
  
  newtype ExpandTydefErr = UnexpandedUserType {getUnexpanded :: TDVar}
    deriving (Eq, Show)
  
  -- | Expand an application "T ty1 ty2 ... tyn" by looking up the
  --   definition of T and substituting ty1 .. tyn for its arguments.
  --
  --   Note that this has already been kind-checked so we know the
  --   number of arguments must match up, and user types must be
  --   defined; we don't worry about what happens if the lists have
  --   different lengths since in theory that can never happen.
  --   However, if T does not exist, we throw an error containing its
  --   name.
  expandTydef ::
    ( Has (Reader TDCtx) sig m
    , Has (Throw ExpandTydefErr) sig m
    , Typical t
    ) =>
    TDVar ->
    [t] ->
    m t
  expandTydef userTyCon tys = do
    mtydefInfo <- lookupTDR userTyCon
    case mtydefInfo of
      Nothing -> throwError (UnexpandedUserType userTyCon)
      Just tydefInfo -> pure $ substTydef tydefInfo tys
  
  -- | Expand *all* applications of user-defined type constructors
  --   everywhere in a type.
  expandTydefs ::
    (Has (Reader TDCtx) sig m, Has (Throw ExpandTydefErr) sig m, Typical t, Plated t) =>
    t ->
    m t
  expandTydefs = rewriteM expand
   where
    -- expand :: t -> m (Maybe t)
    expand t = case unrollT t of
      Just (TyConF (TCUser u) tys) -> Just <$> expandTydef u tys
      _ -> pure Nothing
  
  -- | Given the definition of a type alias, substitute the given
  --   arguments into its body and return the resulting type.
  substTydef :: forall t. Typical t => TydefInfo -> [t] -> t
  substTydef (TydefInfo (Forall as ty) _) tys = refoldT @t substF (fromType ty)
   where
    argMap = M.fromList $ zip as tys
  
    substF tyF@(TyVarF _ x) = case M.lookup x argMap of
      Nothing -> rollT tyF
      Just ty' -> ty'
    substF tyF = rollT tyF
  
  -- | Replace a type alias with its definition everywhere it occurs
  --   inside a type.  Typically this is done when reporting the type of
  --   a term containing a local tydef: since the tydef is local we
  --   can't use it in the reported type.
  elimTydef :: forall t. Typical t => TDVar -> TydefInfo -> t -> t
  elimTydef x tdInfo = refoldT substF
   where
    substF = \case
      TyConF (TCUser u) tys | u == x -> substTydef tdInfo tys
      tyF -> rollT tyF
  
  ------------------------------------------------------------
  -- Arity
  ------------------------------------------------------------
  
  -- | The arity, /i.e./ number of type arguments, of each type
  --   constructor.  Eventually, if we generalize to higher-kinded
  --   types, we'll need to upgrade this to return a full-fledged kind
  --   instead of just an arity.
  tcArity :: TDCtx -> TyCon -> Maybe Arity
  tcArity tydefs =
    fmap Arity . \case
      TCBase _ -> Just 0
      TCCmd -> Just 1
      TCDelay -> Just 1
      TCSum -> Just 2
      TCProd -> Just 2
      TCFun -> Just 2
      TCUser t -> getArity . view tydefArity <$> lookupTD t tydefs
  
  ------------------------------------------------------------
  -- Reducing types to WHNF
  ------------------------------------------------------------
  
  -- | Reduce a type to weak head normal form, i.e. keep unfolding type
  --   aliases and recursive types just until the top-level constructor
  --   of the type is neither @rec@ nor an application of a defined type
  --   alias.
  whnfType :: TDCtx -> Type -> Type
  whnfType tdCtx = run . runReader tdCtx . go
   where
    go :: Has (Reader TDCtx) sig m => Type -> m Type
    go = \case
      TyUser u tys -> do
        res <- runThrow @ExpandTydefErr (expandTydef u tys)
        case res of
          Left _ -> pure $ TyUser u tys
          Right expTy -> go expTy
      TyRec x ty -> go (unfoldRec x ty)
      ty -> pure ty
  