{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeApplications #-}

-- | Constraint generation for declarations.
--
-- Processes top-level data declarations and value bindings from a module.
module Aihc.Tc.Generate.Decl
  ( tcModule,
    tcModuleScc,
    moduleBindings,
    moduleInstances,
    moduleClasses,
    TcBindingResult (..),
  )
where

import Aihc.Parser.Syntax
  ( Annotation,
    BangType (..),
    BinderHead (..),
    CallConv (..),
    CaseAlt (..),
    ClassDecl (..),
    ClassDeclItem (..),
    DataConDecl (..),
    DataDecl (..),
    DataFamilyDecl (..),
    DataFamilyInst (..),
    Decl (..),
    Expr (..),
    Extension,
    FieldDecl (..),
    ForallTelescope (..),
    ForeignDecl (..),
    ForeignDirection (..),
    GadtBody (..),
    InstanceDecl (..),
    InstanceDeclItem (..),
    Match (..),
    MatchHeadForm (..),
    Module (..),
    Name (..),
    NameType (..),
    NewtypeDecl (..),
    Pattern (..),
    Pragma (..),
    PragmaType (..),
    PragmaUnpackKind (..),
    Rhs (..),
    SourceSpan (..),
    TupleFlavor (..),
    TyVarBinder,
    Type (..),
    TypeSynDecl (..),
    UnqualifiedName (..),
    ValueDecl (..),
    applyExtensionSetting,
    applyImpliedExtensions,
    binderHeadName,
    binderHeadParams,
    fromAnnotation,
    gadtBodyResultType,
    instanceHeadName,
    instanceHeadTypes,
    mkAnnotation,
    moduleName,
    nameQualifier,
    nameText,
    peelClassDeclItemAnn,
    peelDeclAnn,
    unqualifiedNameAnns,
  )
import Aihc.Resolve (PackageId (..), ResolutionAnnotation (..), ResolvedName (..))
import Aihc.Tc.Annotations
  ( TcAnnotation (..),
    TcClassAnnotation (..),
    TcClassMethodAnnotation (..),
    TcDerivingAnnotation (..),
    TcDictBinderAnnotation (..),
    TcForeignAbiType (..),
    TcForeignEffect (..),
    TcForeignImportAnnotation (..),
    TcForeignMarshal (..),
    TcInstanceAnnotation (..),
    TcInstanceMethodAnnotation (..),
    annotateDecl,
  )
import Aihc.Tc.Constraint
import Aihc.Tc.Deriving (annotateAttachedDerivingTc, annotateStandaloneDerivingTc)
import Aihc.Tc.Deriving.Context (derivingPlanInstanceInfo, finalizeDerivingModulesTc)
import Aihc.Tc.Env (ClassInfo (..), DataConFieldInfo (..), DataConFieldUnpack (..), DataConInfo (..), DataConSourceForm (..), DataFamilyInstanceInfo (..), DataTypeInfo (..), InstanceInfo (..), TyConFlavor (..), TyConInfo (..), TypeSynonymInfo (..), dataFamilyAxiomName, dataFamilyRepresentationName)
import Aihc.Tc.Error (TcErrorKind (..))
import Aihc.Tc.Evidence (EvTerm (..))
import Aihc.Tc.Finalize (finalizeModuleTc)
import Aihc.Tc.Generalize (generalizeAndCommitIgnoring, predMetaVars)
import Aihc.Tc.Generate.Bind (inferRhsWithLocals)
import Aihc.Tc.Generate.Expr (inferExpr)
import Aihc.Tc.Generate.Pattern
import Aihc.Tc.Instantiate qualified
import Aihc.Tc.Kind (ParamInfo (..), TvKindEnv, checkRuntimeType, checkSurfaceType, classPredicateArgKinds, defaultKindMetas, freeTypeVars, freshKindMeta, makeParamEnv, makeParamEnvWith, sigToScheme, standaloneKindSigToScheme, surfacePredToPred, tyConKindFromParams, tyConKindFromParamsWith, unifyKinds)
import Aihc.Tc.Monad
import Aihc.Tc.Solve (SolveResult (..), solveConstraints, solveWithImpls)
import Aihc.Tc.Solve.Dict (DictResult (..), solveDictWithGivens)
import Aihc.Tc.Solve.InertSet (InertSet (..))
import Aihc.Tc.Types
import Aihc.Tc.Zonk (defaultPredKinds, defaultTyConKindScheme, defaultTyVarKinds, defaultTypeKinds, defaultTypeSchemeKinds, zonkType)
import Control.Monad (foldM, forM_, unless, when, zipWithM, zipWithM_, (>=>))
import Control.Monad.Trans.Class (lift)
import Control.Monad.Trans.State.Strict (get, modify')
import Data.Graph (SCC (..), stronglyConnComp)
import Data.List (find, mapAccumL, nub, nubBy, partition, (\\))
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Maybe (catMaybes, fromMaybe, listToMaybe, mapMaybe, maybeToList)
import Data.Set qualified as Set
import Data.Text (Text)
import Data.Text qualified as T

-- | Merge concrete source spans embedded in a list of annotations.
sourceSpanFromAnns :: [Annotation] -> SourceSpan
sourceSpanFromAnns anns =
  case mapMaybe (fromAnnotation @SourceSpan) anns of
    [] -> NoSourceSpan
    s : _ -> s

peelDeclSpan :: SourceSpan -> Decl -> SourceSpan
peelDeclSpan ambient (DeclAnn ann inner) =
  peelDeclSpan (fromMaybe ambient (fromAnnotation @SourceSpan ann)) inner
peelDeclSpan ambient _ = ambient

-- | Result of type-checking a single binding.
data TcBindingResult = TcBindingResult
  { -- | Canonical binder identity. Symbolic binders are stored without
    -- prefix-position parentheses, e.g. @++@ rather than @(++)@.
    tbName :: !Text,
    -- | Human-facing rendering for diagnostics and golden output.
    tbDisplayName :: !Text,
    tbType :: !TcType
  }
  deriving (Show, Read)

data UserSig = UserSig
  { userSigName :: !Text,
    userSigType :: !Type,
    userSigSpan :: !SourceSpan
  }
  deriving (Show)

data CheckedSig = CheckedSig
  { checkedSigName :: !Text,
    checkedSigScheme :: !TypeScheme,
    checkedSigSpan :: !SourceSpan
  }
  deriving (Show)

moduleBindings :: Module -> [TcBindingResult]
moduleBindings modu =
  concatMap declBindings (moduleDecls modu)

-- | Recover instance-environment entries from finalized module annotations.
moduleInstances :: Module -> [InstanceInfo]
moduleInstances modu =
  map (setInstanceOrigin (resolvedModuleOrigin modu)) (concatMap declInstances (moduleDecls modu))

setInstanceOrigin :: (Text, Text) -> InstanceInfo -> InstanceInfo
setInstanceOrigin origin instanceInfo = instanceInfo {iiDictOrigin = Just origin}

resolvedModuleOrigin :: Module -> (Text, Text)
resolvedModuleOrigin resolvedModule =
  fromMaybe ("", fromMaybe "Main" (moduleName resolvedModule)) $ do
    resolved <- listToMaybe (mapMaybe definitionResolution (moduleDecls resolvedModule))
    case resolutionTarget resolved of
      ResolvedTopLevel packageId name ->
        pure (packageIdText packageId, fromMaybe (fromMaybe "Main" (moduleName resolvedModule)) (nameQualifier name))
      _ -> Nothing

definitionResolution :: Decl -> Maybe ResolutionAnnotation
definitionResolution declaration =
  case peelDeclAnn declaration of
    DeclValue (FunctionBind name _) -> nameResolution name
    DeclValue (PatternBind _ pattern' _) -> patternResolution pattern'
    DeclData dataDeclaration -> nameResolution (binderHeadName (dataDeclHead dataDeclaration))
    DeclNewtype newtypeDeclaration -> nameResolution (binderHeadName (newtypeDeclHead newtypeDeclaration))
    DeclClass classDeclaration -> nameResolution (binderHeadName (classDeclHead classDeclaration))
    _ -> Nothing

patternResolution :: Pattern -> Maybe ResolutionAnnotation
patternResolution pattern' =
  case pattern' of
    PVar name -> nameResolution name
    PAnn _ inner -> patternResolution inner
    PParen inner -> patternResolution inner
    PStrict inner -> patternResolution inner
    PIrrefutable inner -> patternResolution inner
    PAs name _ -> nameResolution name
    PTypeSig inner _ -> patternResolution inner
    _ -> Nothing

nameResolution :: UnqualifiedName -> Maybe ResolutionAnnotation
nameResolution = listToMaybe . mapMaybe fromAnnotation . unqualifiedNameAnns

-- | Recover class-environment entries from finalized module annotations.
moduleClasses :: Module -> [ClassInfo]
moduleClasses modu = map setOrigin (concatMap declClasses (moduleDecls modu))
  where
    setOrigin classInfo = classInfo {ciOrigin = Just (resolvedModuleOrigin modu)}

declClasses :: Decl -> [ClassInfo]
declClasses decl =
  case decl of
    DeclAnn ann inner ->
      annotationClasses ann inner <> declClasses inner
    _ -> []

annotationClasses :: Annotation -> Decl -> [ClassInfo]
annotationClasses ann decl =
  case (fromAnnotation ann, peelDeclAnn decl) of
    (Just classAnn, DeclClass classDecl) ->
      [ ClassInfo
          { ciName = unqualifiedNameText (binderHeadName (classDeclHead classDecl)),
            ciTyCon = tcClassTyCon classAnn,
            ciOrigin = Nothing,
            ciTyVars = tcClassTyVars classAnn,
            ciSuperClassTypes = map tcDictBinderType (tcClassSuperClasses classAnn),
            ciMethods =
              [ (tcClassMethodName method, typeToScheme (tcClassMethodType method))
              | method <- tcClassMethods classAnn
              ],
            ciDefaultMethods = tcClassDefaultMethods classAnn,
            ciDefaultSignatures =
              [ (methodName, typeToScheme signature)
              | (methodName, signature) <- tcClassDefaultSignatures classAnn
              ]
          }
      ]
    _ -> []

typeToScheme :: TcType -> TypeScheme
typeToScheme ty =
  let (tyVars, body) = peelForAlls ty
   in case body of
        TcQualTy predicates result -> ForAll tyVars predicates result
        result -> ForAll tyVars [] result

declInstances :: Decl -> [InstanceInfo]
declInstances decl =
  case decl of
    DeclAnn ann inner ->
      annotationInstances ann inner <> declInstances inner
    _ -> []

annotationInstances :: Annotation -> Decl -> [InstanceInfo]
annotationInstances ann decl =
  explicitInstance <> derivedInstances
  where
    explicitInstance =
      case (fromAnnotation @TcInstanceAnnotation ann, peelDeclAnn decl) of
        (Just instAnn, DeclInstance instanceDecl)
          | Just className <- instanceHeadName (instanceDeclHead instanceDecl) ->
              [ InstanceInfo
                  { iiClassName = nameText className,
                    iiDictName = tcInstanceDictName instAnn,
                    iiDictOrigin = Nothing,
                    iiDictType = tcInstanceDictType instAnn,
                    iiTyVars = tcInstanceTyVars instAnn,
                    iiContext = map dictBinderPred (tcInstanceContextDicts instAnn),
                    iiHead = tcInstanceHeadTypes instAnn
                  }
              ]
        _ -> []
    derivedInstances =
      case fromAnnotation @TcDerivingAnnotation ann of
        Just derivingAnnotation -> mapMaybe derivingPlanInstanceInfo (tcDerivingPlans derivingAnnotation)
        Nothing -> []

dictBinderPred :: TcDictBinderAnnotation -> Pred
dictBinderPred dictBinder =
  case constraintTypeToPred (tcDictBinderType dictBinder) of
    Just predicate -> predicate
    Nothing -> error "invalid checked dictionary binder type"

declBindings :: Decl -> [TcBindingResult]
declBindings decl =
  case decl of
    DeclAnn ann inner ->
      annotationBindings ann inner <> declBindings inner
    DeclData dataDecl ->
      concatMap dataConBindings (dataDeclConstructors dataDecl)
        <> concatMap recordSelectorBindings (dataDeclConstructors dataDecl)
    DeclNewtype newtypeDecl ->
      maybe [] (\constructor -> dataConBindings constructor <> recordSelectorBindings constructor) (newtypeDeclConstructor newtypeDecl)
    DeclDataFamilyInst familyInst ->
      concatMap dataConBindings (dataFamilyInstConstructors familyInst)
    _ -> []

annotationBindings :: Annotation -> Decl -> [TcBindingResult]
annotationBindings ann decl =
  tcAnnotationBindings ann decl
    <> classAnnotationBindings ann decl
    <> instanceAnnotationBindings ann
    <> derivingAnnotationBindings ann

tcAnnotationBindings :: Annotation -> Decl -> [TcBindingResult]
tcAnnotationBindings ann decl =
  case fromAnnotation ann of
    Nothing -> []
    Just tcAnn ->
      case decl of
        DeclValue valueDecl ->
          [ TcBindingResult name displayName (tcAnnType tcAnn)
          | (name, displayName) <- valueDeclBindingNames valueDecl
          ]
        DeclData dataDecl ->
          let name = unqualifiedNameText (binderHeadName (dataDeclHead dataDecl))
           in [TcBindingResult name name (tcAnnType tcAnn)]
        DeclNewtype newtypeDecl ->
          let name = unqualifiedNameText (binderHeadName (newtypeDeclHead newtypeDecl))
           in [TcBindingResult name name (tcAnnType tcAnn)]
        DeclDataFamilyDecl familyDecl ->
          let name = unqualifiedNameText (binderHeadName (dataFamilyDeclHead familyDecl))
           in [TcBindingResult name name (tcAnnType tcAnn)]
        DeclForeign foreignDecl ->
          let name = unqualifiedNameText (foreignName foreignDecl)
              displayName = renderBinderName (foreignName foreignDecl)
           in [TcBindingResult name displayName (tcAnnType tcAnn)]
        _ -> []

classAnnotationBindings :: Annotation -> Decl -> [TcBindingResult]
classAnnotationBindings ann decl =
  case (fromAnnotation ann, decl) of
    (Just classAnn, DeclClass {}) ->
      [ TcBindingResult (tcClassMethodName method) (tcClassMethodName method) (tcClassMethodType method)
      | method <- tcClassMethods classAnn
      ]
        <> [ TcBindingResult (defaultMethodName (tcClassMethodName method)) (defaultMethodName (tcClassMethodName method)) (classDefaultWorkerType classAnn method)
           | method <- tcClassMethods classAnn,
             tcClassMethodName method `elem` tcClassDefaultMethods classAnn
           ]
    _ -> []

classDefaultWorkerType :: TcClassAnnotation -> TcClassMethodAnnotation -> TcType
classDefaultWorkerType classAnnotation method =
  case lookup (tcClassMethodName method) (tcClassDefaultSignatures classAnnotation) of
    Just signature
      | Just classPredicate <- constraintTypeToPred (tcClassMethodDictType method) ->
          let (tyVars, body) = peelForAlls signature
              qualifiedBody =
                case body of
                  TcQualTy predicates result -> TcQualTy (classPredicate : predicates) result
                  result -> TcQualTy [classPredicate] result
           in foldr TcForAllTy qualifiedBody tyVars
    _ -> tcClassMethodType method

instanceAnnotationBindings :: Annotation -> [TcBindingResult]
instanceAnnotationBindings ann =
  case fromAnnotation ann of
    Just instAnn ->
      [TcBindingResult (tcInstanceDictName instAnn) (tcInstanceDictName instAnn) (tcInstanceDictType instAnn)]
    Nothing -> []

derivingAnnotationBindings :: Annotation -> [TcBindingResult]
derivingAnnotationBindings ann =
  case fromAnnotation ann of
    Just derivingAnnotation ->
      [ TcBindingResult (iiDictName instanceInfo) (iiDictName instanceInfo) (iiDictType instanceInfo)
      | plan <- tcDerivingPlans derivingAnnotation,
        Just instanceInfo <- [derivingPlanInstanceInfo plan]
      ]
    Nothing -> []

dataConBindings :: DataConDecl -> [TcBindingResult]
dataConBindings dataConDecl =
  case dataConDecl of
    DataConAnn ann inner ->
      case fromAnnotation ann of
        Just tcAnn ->
          [ TcBindingResult name displayName (tcAnnType tcAnn)
          | (name, displayName) <- dataConBindingNames inner
          ]
        Nothing -> dataConBindings inner
    _ -> []

recordSelectorBindings :: DataConDecl -> [TcBindingResult]
recordSelectorBindings declaration =
  case declaration of
    DataConAnn ann inner ->
      case fromAnnotation ann of
        Just tcAnn -> selectorBindingsFromConstructorType (tcAnnType tcAnn) inner
        Nothing -> recordSelectorBindings inner
    _ -> []
  where
    selectorBindingsFromConstructorType constructorType inner =
      let (typeVariables, qualifiedConstructor) = peelForAlls constructorType
          (predicates, body) =
            case qualifiedConstructor of
              TcQualTy context result -> (context, result)
              result -> ([], result)
          (fieldTypes, resultType) = splitFunctionType body
          (_, sourceFields, _) = dataConSourceLayout inner
       in [ TcBindingResult label label (foldr TcForAllTy (qualify predicates (TcFunTy resultType fieldType)) typeVariables)
          | ((maybeLabel, _), fieldType) <- zip sourceFields fieldTypes,
            Just label <- [maybeLabel]
          ]
    qualify [] ty = ty
    qualify predicates ty = TcQualTy predicates ty

valueDeclBindingNames :: ValueDecl -> [(Text, Text)]
valueDeclBindingNames valueDecl =
  case valueDecl of
    FunctionBind binder _ -> [binderBindingName binder]
    PatternBind _ pat _ -> patternBindingNames pat

patternBindingNames :: Pattern -> [(Text, Text)]
patternBindingNames pat =
  case pat of
    PVar name -> [binderBindingName name]
    PAnn _ inner -> patternBindingNames inner
    PParen inner -> patternBindingNames inner
    PAs name inner -> binderBindingName name : patternBindingNames inner
    PStrict inner -> patternBindingNames inner
    PIrrefutable inner -> patternBindingNames inner
    PInfix lhs _ rhs -> patternBindingNames lhs <> patternBindingNames rhs
    PCon _ _ pats -> concatMap patternBindingNames pats
    _ -> []

dataConBindingNames :: DataConDecl -> [(Text, Text)]
dataConBindingNames dataConDecl =
  case dataConDecl of
    DataConAnn _ inner -> dataConBindingNames inner
    PrefixCon _ _ name _ -> [dataConBindingName name]
    InfixCon _ _ _ name _ -> [dataConBindingName name]
    RecordCon _ _ name _ -> [dataConBindingName name]
    GadtCon _ _ names _ -> map dataConBindingName names
    TupleCon _ _ flavor fields ->
      let name = tupleConText flavor (length fields)
       in [(name, name)]
    UnboxedSumCon _ _ pos arity _ ->
      let name = unboxedSumConText pos arity
       in [(name, name)]
    ListCon {} -> [("[]", "[]")]

binderBindingName :: UnqualifiedName -> (Text, Text)
binderBindingName name =
  (unqualifiedNameText name, renderBinderName name)

dataConBindingName :: UnqualifiedName -> (Text, Text)
dataConBindingName name =
  let raw = unqualifiedNameText name
   in (raw, raw)

-- | Type-check a module, returning the same syntax tree annotated with the
-- inferred interface. Call 'moduleBindings' when a flat compatibility view is
-- needed by older callers.
tcModule :: Module -> TcM Module
tcModule m = do
  modules <- tcModuleScc [m]
  case modules of
    [result] -> pure result
    _ -> pure m

-- | Type-check one strongly connected module component. Data declarations and
-- explicit signatures are registered for the whole component before any
-- value body is checked, allowing a module to refer back to a signed binding
-- in another member of the same import cycle.
tcModuleScc :: [Module] -> TcM [Module]
tcModuleScc modules = do
  -- Phase 1: register type constructor headers before expanding synonym
  -- bodies, then register value-level declarations against those expanded
  -- types. This permits forward references from synonyms to data types while
  -- making aliases available in constructor fields and class methods.
  let declarations = concatMap moduleDecls modules
      standaloneKindSignatures = collectStandaloneKindSignatures declarations
  standaloneKindSchemes <- traverse (standaloneKindSigToScheme >=> defaultTypeSchemeKinds) standaloneKindSignatures
  mapM_ (registerTypeDeclHeader standaloneKindSchemes) declarations
  mapM_ registerTypeSynonymBody declarations
  mapM_ (\modu -> mapM_ (registerStructuralDecl (resolvedModuleOrigin modu)) (moduleDecls modu)) modules
  -- Deriving strategy and context inference depends only on registered type,
  -- class, and explicit-instance information. Finalize the entire SCC as one
  -- batch before checking signatures and bodies so sibling derived instances
  -- are mutually visible and ordinary values can use them.
  defaultGlobalKindMetas
  derivingAnnotated <- mapM annotateModuleDerivingTc modules
  derivingFinalized <- finalizeDerivingModulesTc (map resolvedModuleOrigin modules) derivingAnnotated
  -- Phase 2: collect type signatures and convert them to schemes.
  rawSigs <- mapM (collectUserSigs . moduleDecls) derivingFinalized
  schemes <- mapM (traverse checkUserSig) rawSigs
  mapM_ (uncurry registerCheckedSig) (concatMap Map.toList schemes)
  pending <- zipWithM tcModuleBody schemes derivingFinalized
  -- No module interface in the SCC may retain state-local kind metavariables.
  defaultGlobalKindMetas
  annotated <- mapM annotatePendingModule pending
  mapM finalizeModuleTc annotated

registerCheckedSig :: TcTermKey -> CheckedSig -> TcM ()
registerCheckedSig key sig = do
  extendTermEnvPermanent (checkedSigName sig) binder
  extendTermKeyEnvPermanent key binder
  where
    binder = TcIdBinder (checkedSigScheme sig) Closed

data PendingModule = PendingModule
  { pendingSyntax :: !Module,
    pendingValueResults :: ![TcBindingResult]
  }

tcModuleBody :: Map TcTermKey CheckedSig -> Module -> TcM PendingModule
tcModuleBody schemes m = do
  -- Phase 3: group and type-check value bindings using signatures.
  let sourceGroups = zip [0 :: Int ..] (groupValueDecls (moduleDecls m))
  grouped <- sortDeclGroups sourceGroups
  groupResults <- mapM (tcDeclGroup schemes) grouped
  let valueResults = concatMap tcGroupBindingResults groupResults
      checkedGroups =
        Map.fromList
          [ (tcGroupId result, decls)
          | result <- groupResults,
            Just decls <- [tcGroupAnnotatedDecls result]
          ]
      valueAnnotatedModule =
        m {moduleDecls = concatMap (renderCheckedGroup checkedGroups) sourceGroups}
  -- Phase 4: type-check instance method bodies. They are not top-level
  -- value bindings, but their occurrences still need the same instantiation
  -- and evidence records as ordinary expressions.
  classDecls <- mapM tcClassDeclBodies (moduleDecls valueAnnotatedModule)
  instanceDecls <- mapM tcInstanceDeclBodies classDecls
  let pendingModule = valueAnnotatedModule {moduleDecls = instanceDecls}
  -- Phase 5: reject source top-level values whose finalized types are
  -- unlifted. Generated declarations without source spans are permitted so
  -- downstream passes can introduce internal unlifted bindings.
  checkTopLevelUnliftedBindings sourceGroups groupResults
  pure (PendingModule pendingModule valueResults)

annotatePendingModule :: PendingModule -> TcM Module
annotatePendingModule pending = do
  -- Only bindings that checked without errors are eligible for value
  -- annotations. Failed bindings remain in the recovery environment, but
  -- they must not be rendered as successful inferred types.
  annotateModuleTc (Set.fromList (map tbName (pendingValueResults pending))) (pendingSyntax pending)

defaultGlobalKindMetas :: TcM ()
defaultGlobalKindMetas = do
  state <- lift get
  terms <- traverse defaultBinderKinds (tcsGlobalTerms state)
  tyCons <- traverse defaultTyConInfoKinds (tcsGlobalTyCons state)
  dataTypes <- traverse defaultDataTypeKinds (tcsDataTypes state)
  classes <- traverse defaultClassKinds (tcsClasses state)
  instances <- mapM defaultInstanceKinds (tcsInstances state)
  dataFamilyInstances <- mapM defaultDataFamilyInstanceKinds (tcsDataFamilyInstances state)
  lift $
    modify' $ \current ->
      current
        { tcsGlobalTerms = terms,
          tcsGlobalTyCons = tyCons,
          tcsDataTypes = dataTypes,
          tcsClasses = classes,
          tcsInstances = instances,
          tcsDataFamilyInstances = dataFamilyInstances
        }
  where
    defaultBinderKinds binder =
      case binder of
        TcIdBinder scheme closedness -> TcIdBinder <$> defaultTypeSchemeKinds scheme <*> pure closedness
        TcMonoIdBinder ty -> TcMonoIdBinder <$> defaultTypeKinds ty
    defaultTyConInfoKinds info = do
      kindScheme <- defaultTyConKindScheme (tyConKindScheme (tciTyCon info))
      synonym <- traverse defaultTypeSynonymKinds (tciTypeSynonym info)
      pure
        info
          { tciTyCon = setTyConKindScheme kindScheme (tciTyCon info),
            tciTypeSynonym = synonym
          }
    defaultTypeSynonymKinds synonym =
      TypeSynonymInfo
        <$> mapM defaultTyVarKinds (tsiParams synonym)
        <*> traverse defaultTypeKinds (tsiBody synonym)
    defaultDataTypeKinds info = do
      tyVars <- mapM defaultTyVarKinds (dtiTyVars info)
      resultKind <- defaultKindMetas (dtiResultKind info)
      constructors <- mapM defaultDataConKinds (dtiConstructors info)
      kindScheme <- defaultTyConKindScheme (tyConKindScheme (dtiTyCon info))
      pure
        info
          { dtiTyCon = setTyConKindScheme kindScheme (dtiTyCon info),
            dtiTyVars = tyVars,
            dtiResultKind = resultKind,
            dtiConstructors = constructors
          }
    defaultDataConKinds info = do
      universalTyVars <- mapM defaultTyVarKinds (dciUnivTyVars info)
      existentialTyVars <- mapM defaultTyVarKinds (dciExTyVars info)
      predicates <- mapM defaultPredKinds (dciTheta info)
      fields <- mapM defaultDataConFieldKinds (dciFields info)
      resultType <- defaultTypeKinds (dciResTy info)
      pure
        info
          { dciUnivTyVars = universalTyVars,
            dciExTyVars = existentialTyVars,
            dciTheta = predicates,
            dciFields = fields,
            dciResTy = resultType
          }
    defaultDataConFieldKinds field = do
      fieldType' <- defaultTypeKinds (dcfiType field)
      pure field {dcfiType = fieldType'}
    defaultClassKinds info = do
      classKindScheme <- defaultTyConKindScheme (tyConKindScheme (ciTyCon info))
      tyVars <- mapM defaultTyVarKinds (ciTyVars info)
      superClassTypes <- mapM defaultTypeKinds (ciSuperClassTypes info)
      methods <- mapM (traverse defaultTypeSchemeKinds) (ciMethods info)
      defaultSignatures <- mapM (traverse defaultTypeSchemeKinds) (ciDefaultSignatures info)
      pure
        info
          { ciTyCon = setTyConKindScheme classKindScheme (ciTyCon info),
            ciTyVars = tyVars,
            ciSuperClassTypes = superClassTypes,
            ciMethods = methods,
            ciDefaultSignatures = defaultSignatures
          }
    defaultInstanceKinds info =
      InstanceInfo
        (iiClassName info)
        (iiDictName info)
        (iiDictOrigin info)
        <$> defaultTypeKinds (iiDictType info)
        <*> mapM defaultTyVarKinds (iiTyVars info)
        <*> mapM defaultPredKinds (iiContext info)
        <*> mapM defaultTypeKinds (iiHead info)
    defaultDataFamilyInstanceKinds info = do
      familyType <- defaultTypeKinds (dfiiFamilyType info)
      tyVars <- mapM defaultTyVarKinds (dfiiTyVars info)
      let representationTyCon = dfiiRepresentationTyCon info
      representationKindScheme <- defaultTyConKindScheme (tyConKindScheme representationTyCon)
      pure
        info
          { dfiiFamilyType = familyType,
            dfiiTyVars = tyVars,
            dfiiRepresentationTyCon = setTyConKindScheme representationKindScheme representationTyCon
          }

data TcDeclGroupResult = TcDeclGroupResult
  { tcGroupId :: !Int,
    tcGroupBindingResults :: ![TcBindingResult],
    tcGroupAnnotatedDecls :: !(Maybe [Decl])
  }

checkTopLevelUnliftedBindings :: [(Int, DeclGroup)] -> [TcDeclGroupResult] -> TcM ()
checkTopLevelUnliftedBindings sourceGroups results =
  forM_ results $ \result ->
    case Map.lookup (tcGroupId result) groupsById of
      Just group
        | sourceSpan <- declGroupSourceSpan group,
          sourceSpan /= NoSourceSpan ->
            forM_ (tcGroupBindingResults result) $ \binding ->
              when (isUnliftedType (tbType binding)) $
                emitError sourceSpan (TopLevelUnliftedBinding (tbDisplayName binding) (tbType binding))
      _ -> pure ()
  where
    groupsById = Map.fromList sourceGroups

declGroupSourceSpan :: DeclGroup -> SourceSpan
declGroupSourceSpan group =
  case group of
    SingleDecl decl -> peelDeclSpan NoSourceSpan decl
    MergedFunctionBind sourceSpan _ _ _ -> sourceSpan

renderCheckedGroup :: Map Int [Decl] -> (Int, DeclGroup) -> [Decl]
renderCheckedGroup checkedGroups (groupId, group) =
  fromMaybe (renderDeclGroup group) (Map.lookup groupId checkedGroups)

annotateModuleTc :: Set.Set Text -> Module -> TcM Module
annotateModuleTc checkedValueNames m = do
  let classMethods = collectClassMethodNames (moduleDecls m)
  decls <- mapM (annotateDeclTc classMethods checkedValueNames) (moduleDecls m)
  pure (m {moduleDecls = decls})

annotateModuleDerivingTc :: Module -> TcM Module
annotateModuleDerivingTc modu = do
  declarations <- mapM (annotateDeclDerivingTc extensions) (moduleDecls modu)
  pure modu {moduleDecls = declarations}
  where
    extensions = moduleEnabledExtensions modu

annotateDeclDerivingTc :: [Extension] -> Decl -> TcM Decl
annotateDeclDerivingTc extensions decl =
  case decl of
    DeclAnn annotation inner -> DeclAnn annotation <$> annotateDeclDerivingTc extensions inner
    DeclData dataDecl ->
      annotateAttachedDerivingTc extensions DataTyCon (dataDeclHead dataDecl) (dataDeclDeriving dataDecl) decl
    DeclNewtype newtypeDecl ->
      annotateAttachedDerivingTc extensions NewtypeTyCon (newtypeDeclHead newtypeDecl) (newtypeDeclDeriving newtypeDecl) decl
    DeclStandaloneDeriving derivingDecl -> annotateStandaloneDerivingTc extensions derivingDecl
    _ -> pure decl

moduleEnabledExtensions :: Module -> [Extension]
moduleEnabledExtensions modu =
  applyImpliedExtensions $
    foldr applyExtensionSetting [] (moduleLanguagePragmas modu)

annotateDeclTc :: Map Text [Text] -> Set.Set Text -> Decl -> TcM Decl
annotateDeclTc classMethods checkedValueNames decl =
  case decl of
    DeclAnn ann inner -> DeclAnn ann <$> annotateDeclTc classMethods checkedValueNames inner
    DeclValue valueDecl
      | valueDeclWasChecked checkedValueNames valueDecl -> do
          (ty, valueDecl') <- annotateValueDeclTc valueDecl
          pure (annotateDeclAt (valueDeclSpan valueDecl) (TcAnnotation ty [] [] [] []) (DeclValue valueDecl'))
      | otherwise -> pure decl
    DeclData dataDecl -> annotateDataDeclTc dataDecl
    DeclNewtype newtypeDecl -> annotateNewtypeDeclTc newtypeDecl
    DeclTypeSyn typeSynDecl -> annotateTypeSynDeclTc typeSynDecl
    DeclDataFamilyDecl familyDecl -> annotateDataFamilyDeclTc familyDecl
    DeclDataFamilyInst familyInst -> annotateDataFamilyInstTc familyInst
    DeclForeign foreignDecl
      | isForeignImport foreignDecl -> annotateForeignDeclTc foreignDecl
    DeclClass classDecl -> annotateClassDeclTc classDecl
    DeclInstance instanceDecl -> annotateInstanceDeclTc instanceDecl
    DeclStandaloneDeriving {} -> pure decl
    _ -> pure decl

valueDeclWasChecked :: Set.Set Text -> ValueDecl -> Bool
valueDeclWasChecked checkedValueNames valueDecl =
  any (`Set.member` checkedValueNames) (valueDeclBinderNames valueDecl)

valueDeclBinderNames :: ValueDecl -> [Text]
valueDeclBinderNames valueDecl =
  case valueDecl of
    FunctionBind binder _ -> [unqualifiedNameText binder]
    PatternBind _ pat _ -> patternBinders pat

annotateClassDeclTc :: ClassDecl -> TcM Decl
annotateClassDeclTc classDecl = do
  let className = unqualifiedNameText (binderHeadName (classDeclHead classDecl))
  classInfo <- lookupClass className
  case classInfo of
    Nothing -> missingTypeInfo ("class " <> T.unpack className)
    Just info -> do
      methods <- zipWithM annotateClassMethod [0 :: Int ..] (classDeclMethodNames classDecl)
      items <- mapM annotateClassDefaultItem (classDeclItems classDecl)
      pure
        ( DeclAnn
            ( mkAnnotation
                TcClassAnnotation
                  { tcClassTyCon = ciTyCon info,
                    tcClassTyVars = ciTyVars info,
                    tcClassSuperClasses = map constraintTypeDictBinder (ciSuperClassTypes info),
                    tcClassMethods = methods,
                    tcClassDefaultMethods = ciDefaultMethods info,
                    tcClassDefaultSignatures =
                      [(methodName, schemeToType signature) | (methodName, signature) <- ciDefaultSignatures info]
                  }
            )
            (DeclClass (classDecl {classDeclItems = items}))
        )

annotateClassDefaultItem :: ClassDeclItem -> TcM ClassDeclItem
annotateClassDefaultItem item =
  case item of
    ClassItemAnn ann inner -> ClassItemAnn ann <$> annotateClassDefaultItem inner
    ClassItemDefault valueDecl ->
      case valueDeclBinderName valueDecl of
        Just (_, methodName) -> do
          methodTy <- bindingType (defaultMethodName methodName)
          pure (ClassItemAnn (mkAnnotation (TcInstanceMethodAnnotation methodName methodTy)) item)
        Nothing -> pure item
    _ -> pure item

annotateClassMethod :: Int -> Text -> TcM TcClassMethodAnnotation
annotateClassMethod index methodName = do
  methodTy <- bindingType methodName
  let (tvs, _) = peelForAlls methodTy
  dictTy <- selectorDictTypeTc methodName methodTy
  pure
    TcClassMethodAnnotation
      { tcClassMethodName = methodName,
        tcClassMethodType = methodTy,
        tcClassMethodTyVars = tvs,
        tcClassMethodDictType = dictTy,
        tcClassMethodIndex = index
      }

annotateDataDeclTc :: DataDecl -> TcM Decl
annotateDataDeclTc dataDecl = do
  let tyName = unqualifiedNameText (binderHeadName (dataDeclHead dataDecl))
  ty <- tyConBindingType tyName
  constructors <- mapM annotateDataConDeclTc (dataDeclConstructors dataDecl)
  let annotatedHead = annotateBinderHeadName (TcAnnotation ty [] [] [] []) (dataDeclHead dataDecl)
  pure (DeclData (dataDecl {dataDeclHead = annotatedHead, dataDeclConstructors = constructors}))

annotateNewtypeDeclTc :: NewtypeDecl -> TcM Decl
annotateNewtypeDeclTc newtypeDecl = do
  let tyName = unqualifiedNameText (binderHeadName (newtypeDeclHead newtypeDecl))
  ty <- tyConBindingType tyName
  constructor <- mapM annotateDataConDeclTc (newtypeDeclConstructor newtypeDecl)
  let annotatedHead = annotateBinderHeadName (TcAnnotation ty [] [] [] []) (newtypeDeclHead newtypeDecl)
  pure (DeclNewtype (newtypeDecl {newtypeDeclHead = annotatedHead, newtypeDeclConstructor = constructor}))

annotateTypeSynDeclTc :: TypeSynDecl -> TcM Decl
annotateTypeSynDeclTc typeSynDecl = do
  let tyName = unqualifiedNameText (binderHeadName (typeSynHead typeSynDecl))
  ty <- tyConBindingType tyName
  let annotatedHead = annotateBinderHeadName (TcAnnotation ty [] [] [] []) (typeSynHead typeSynDecl)
  pure (DeclTypeSyn (typeSynDecl {typeSynHead = annotatedHead}))

annotateDataFamilyDeclTc :: DataFamilyDecl -> TcM Decl
annotateDataFamilyDeclTc familyDecl = do
  let familyName = unqualifiedNameText (binderHeadName (dataFamilyDeclHead familyDecl))
  ty <- tyConBindingType familyName
  let annotatedHead = annotateBinderHeadName (TcAnnotation ty [] [] [] []) (dataFamilyDeclHead familyDecl)
  pure (DeclDataFamilyDecl (familyDecl {dataFamilyDeclHead = annotatedHead}))

annotateDataFamilyInstTc :: DataFamilyInst -> TcM Decl
annotateDataFamilyInstTc familyInst = do
  constructors <- mapM annotateRegisteredDataConDeclTc (dataFamilyInstConstructors familyInst)
  let annotated = DeclDataFamilyInst (familyInst {dataFamilyInstConstructors = constructors})
      constructorNames = concatMap (map fst . dataConBindingNames) constructors
  familyInstances <- getDataFamilyInstances
  case constructorNames of
    firstConstructor : _ ->
      case find (elem firstConstructor . dfiiConstructorNames) familyInstances of
        Just familyInstance -> pure (DeclAnn (mkAnnotation familyInstance) annotated)
        Nothing -> pure annotated
    [] -> pure annotated

annotateRegisteredDataConDeclTc :: DataConDecl -> TcM DataConDecl
annotateRegisteredDataConDeclTc dataConDecl =
  case dataConBindingNames dataConDecl of
    [] -> pure dataConDecl
    (name, _) : _ -> do
      maybeBinder <- lookupTerm name
      case maybeBinder of
        Just (TcIdBinder scheme _) -> annotateWithType (schemeToType scheme)
        Just (TcMonoIdBinder ty) -> annotateWithType ty
        Nothing -> pure dataConDecl
  where
    annotateWithType ty = do
      zonkedTy <- zonkType ty
      pure (DataConAnn (mkAnnotation (TcAnnotation zonkedTy [] [] [] [])) dataConDecl)

annotateBinderHeadName :: TcAnnotation -> BinderHead UnqualifiedName -> BinderHead UnqualifiedName
annotateBinderHeadName tcAnn head' =
  case head' of
    PrefixBinderHead name params ->
      PrefixBinderHead (annotateUnqualifiedName tcAnn name) params
    InfixBinderHead lhs name rhs params ->
      InfixBinderHead lhs (annotateUnqualifiedName tcAnn name) rhs params

annotateUnqualifiedName :: TcAnnotation -> UnqualifiedName -> UnqualifiedName
annotateUnqualifiedName tcAnn name =
  name {unqualifiedNameAnns = unqualifiedNameAnns name <> [mkAnnotation tcAnn]}

annotateDataConDeclTc :: DataConDecl -> TcM DataConDecl
annotateDataConDeclTc dataConDecl = do
  case dataConBindingNames dataConDecl of
    [] -> pure dataConDecl
    (name, _) : _ -> do
      ty <- dataConBindingType name
      selectors <- annotateRecordSelectorNames dataConDecl
      pure (DataConAnn (mkAnnotation (TcAnnotation ty [] [] [] [])) selectors)

annotateRecordSelectorNames :: DataConDecl -> TcM DataConDecl
annotateRecordSelectorNames declaration =
  case declaration of
    RecordCon forallVars context constructor fields ->
      RecordCon forallVars context constructor <$> mapM annotateField fields
    GadtCon forallBinders context constructors (GadtRecordBody fields result) ->
      GadtCon forallBinders context constructors . (`GadtRecordBody` result) <$> mapM annotateField fields
    _ -> pure declaration
  where
    annotateField field = do
      names <- mapM annotateSelectorName (fieldNames field)
      pure field {fieldNames = names}
    annotateSelectorName name = do
      ty <- bindingType (unqualifiedNameText name)
      pure (annotateUnqualifiedName (TcAnnotation ty [] [] [] []) name)

dataConBindingType :: Text -> TcM TcType
dataConBindingType name = do
  mBinder <- lookupTerm name
  case mBinder of
    Just (TcIdBinder scheme _) -> zonkType (schemeToType scheme)
    Just (TcMonoIdBinder ty) -> zonkType ty
    Nothing -> missingTypeInfo ("data constructor " <> T.unpack name)

annotateForeignDeclTc :: ForeignDecl -> TcM Decl
annotateForeignDeclTc foreignDecl = do
  ty <- bindingType (unqualifiedNameText (foreignName foreignDecl))
  let sourceSpan = unqualifiedNameSpan (foreignName foreignDecl)
      annotated = annotateDeclAt sourceSpan (TcAnnotation ty [] [] [] []) (DeclForeign foreignDecl)
  case foreignCallConv foreignDecl of
    CCall -> do
      plan <- checkForeignImportType sourceSpan ty
      pure (DeclAnn (mkAnnotation plan) annotated)
    _ -> pure annotated

checkForeignImportType :: SourceSpan -> TcType -> TcM TcForeignImportAnnotation
checkForeignImportType sourceSpan ty = do
  let (argumentTypes, resultType) = splitFunctionType ty
      (effect, valueResultType) =
        case resultType of
          TcTyCon (TyCon "IO" 1) [ioResult] -> (TcForeignRealWorld, ioResult)
          _ -> (TcForeignPure, resultType)
  arguments <- mapM (checkForeignValueType sourceSpan) argumentTypes
  result <- checkForeignValueType sourceSpan valueResultType
  pure
    TcForeignImportAnnotation
      { tcForeignArguments = arguments,
        tcForeignResult = result,
        tcForeignEffect = effect
      }

splitFunctionType :: TcType -> ([TcType], TcType)
splitFunctionType ty =
  case ty of
    TcFunTy argument result ->
      let (arguments, finalResult) = splitFunctionType result
       in (argument : arguments, finalResult)
    _ -> ([], ty)

checkForeignValueType :: SourceSpan -> TcType -> TcM TcForeignMarshal
checkForeignValueType sourceSpan ty =
  case ty of
    TcTyCon (TyCon "Int" 0) [] -> pure (intMarshal ty ["I#"])
    TcTyCon (TyCon "Int#" 0) [] -> pure (intMarshal ty [])
    TcTyCon (TyCon "CInt" 0) [] -> pure (int32Marshal ty ["CInt", "I32#"])
    TcTyCon (TyCon "Int32" 0) [] -> pure (int32Marshal ty ["I32#"])
    TcTyCon (TyCon "Int32#" 0) [] -> pure (int32Marshal ty [])
    TcTyCon (TyCon "Word64" 0) [] -> pure (word64Marshal ty ["W64#"])
    TcTyCon (TyCon "Word64#" 0) [] -> pure (word64Marshal ty [])
    TcTyCon (TyCon "Addr#" 0) [] -> pure (addrMarshal ty [])
    TcTyCon (TyCon "Ptr" 1) [_] -> pure (addrMarshal ty ["Ptr"])
    _ -> do
      emitError sourceSpan (OtherError ("unsupported foreign import value type: " <> show ty))
      pure (int32Marshal ty [])
  where
    intMarshal sourceType constructors =
      TcForeignMarshal
        { tcForeignSourceType = sourceType,
          tcForeignPrimitiveType = TcTyCon (TyCon "Int#" 0) [],
          tcForeignConstructors = constructors,
          tcForeignAbiType = TcForeignInt
        }
    int32Marshal sourceType constructors =
      TcForeignMarshal
        { tcForeignSourceType = sourceType,
          tcForeignPrimitiveType = TcTyCon (TyCon "Int32#" 0) [],
          tcForeignConstructors = constructors,
          tcForeignAbiType = TcForeignInt32
        }
    word64Marshal sourceType constructors =
      TcForeignMarshal
        { tcForeignSourceType = sourceType,
          tcForeignPrimitiveType = TcTyCon (TyCon "Word64#" 0) [],
          tcForeignConstructors = constructors,
          tcForeignAbiType = TcForeignWord64
        }
    addrMarshal sourceType constructors =
      TcForeignMarshal
        { tcForeignSourceType = sourceType,
          tcForeignPrimitiveType = TcTyCon (TyCon "Addr#" 0) [],
          tcForeignConstructors = constructors,
          tcForeignAbiType = TcForeignAddr
        }

annotateDeclAt :: SourceSpan -> TcAnnotation -> Decl -> Decl
annotateDeclAt NoSourceSpan tcAnn decl =
  annotateDecl tcAnn decl
annotateDeclAt sp tcAnn decl =
  DeclAnn (mkAnnotation sp) (annotateDecl tcAnn decl)

valueDeclSpan :: ValueDecl -> SourceSpan
valueDeclSpan valueDecl =
  case valueDecl of
    FunctionBind name _ -> unqualifiedNameSpan name
    PatternBind _ pat _ -> patternSpan pat

unqualifiedNameSpan :: UnqualifiedName -> SourceSpan
unqualifiedNameSpan =
  sourceSpanFromAnns . unqualifiedNameAnns

tyConBindingType :: Text -> TcM TcType
tyConBindingType name = do
  mInfo <- lookupTyCon name
  case mInfo of
    Just info -> kindToTcType <$> defaultKindMetas (tyConKind (tciTyCon info))
    Nothing -> missingTypeInfo ("type constructor " <> T.unpack name)

annotateValueDeclTc :: ValueDecl -> TcM (TcType, ValueDecl)
annotateValueDeclTc valueDecl =
  case valueDecl of
    FunctionBind name matches -> do
      bindingTy <- bindingType (unqualifiedNameText name)
      pure (bindingTy, FunctionBind name matches)
    PatternBind anns pat rhs ->
      case patternBinderName pat of
        Just (_, name) -> do
          bindingTy <- bindingType name
          pure (bindingTy, PatternBind anns pat rhs)
        Nothing -> do
          ty <- missingTypeInfo ("top-level pattern binding " <> show pat)
          pure (ty, valueDecl)

annotateInstanceDeclTc :: InstanceDecl -> TcM Decl
annotateInstanceDeclTc instanceDecl =
  case (instanceHeadName (instanceDeclHead instanceDecl), instanceHeadTypes (instanceDeclHead instanceDecl)) of
    (_, []) -> pure (DeclInstance instanceDecl)
    (Nothing, _) -> pure (DeclInstance instanceDecl)
    (Just className, headArgTypes) -> do
      (rawTvIds, tvEnv) <- makeInstanceTyVarEnv instanceDecl headArgTypes
      let classNameText = nameText className
      rawHeadTys <- checkInstanceHeadTypes classNameText tvEnv headArgTypes
      rawContext <- mapM (surfacePredToPred tvEnv) (instanceDeclContext instanceDecl)
      tvIds <- mapM defaultTyVarKinds rawTvIds
      headTys <- mapM defaultTypeKinds rawHeadTys
      context <- mapM defaultPredKinds rawContext
      let dictName = instanceDictName classNameText headTys
      classInfo <- lookupClass classNameText
      info <- maybe (missingTypeInfo ("class " <> T.unpack classNameText)) pure classInfo
      let classSubstitution =
            Map.fromList [(tvUnique tyVar, ty) | (tyVar, ty) <- zip (ciTyVars info) headTys]
          superClassTypes = map (substType classSubstitution) (ciSuperClassTypes info)
          defaults = ciDefaultMethods info
      superClasses <- mapM constraintTypePred superClassTypes
      superClassEvidence <- mapM (solveInstanceSuperClass classNameText context) superClasses
      let contextDicts = map predDictBinder context
          dictTy = foldr TcForAllTy (TcQualTy context (predType (ClassPred (ciTyCon info) headTys))) tvIds
          methodOrder = map fst (ciMethods info)
          instAnn =
            TcInstanceAnnotation
              { tcInstanceDictName = dictName,
                tcInstanceDictType = dictTy,
                tcInstanceClassTyCon = ciTyCon info,
                tcInstanceTyVars = tvIds,
                tcInstanceHeadTypes = headTys,
                tcInstanceClassTyVars = ciTyVars info,
                tcInstanceClassOrigin = ciOrigin info,
                tcInstanceClassSuperClasses = map constraintTypeDictBinder (ciSuperClassTypes info),
                tcInstanceContextDicts = contextDicts,
                tcInstanceSuperClasses = zip (map predDictBinder superClasses) superClassEvidence,
                tcInstanceMethodOrder = methodOrder,
                tcInstanceDefaultMethods = defaults
              }
      items <- mapM (annotateInstanceItemTc headTys) (instanceDeclItems instanceDecl)
      pure (DeclAnn (mkAnnotation instAnn) (DeclInstance (instanceDecl {instanceDeclItems = items})))

solveInstanceSuperClass :: Text -> [Pred] -> Pred -> TcM EvTerm
solveInstanceSuperClass className givens predicate = do
  evidenceVariable <- freshEvVar
  let constraint = mkWantedCt predicate evidenceVariable (InstOrigin className) NoSourceSpan
  result <- solveDictWithGivens givens constraint
  case result of
    DictSolved -> do
      evidence <- lookupEvidence evidenceVariable
      case evidence of
        Just term -> pure term
        Nothing -> missingTypeInfo ("superclass evidence for " <> T.unpack className)
    DictStuck stuck -> do
      emitError (ctLoc stuck) (UnsolvedWanted (ctPred stuck) (ctOrigin stuck))
      pure (EvVarTerm evidenceVariable)

annotateInstanceItemTc :: [TcType] -> InstanceDeclItem -> TcM InstanceDeclItem
annotateInstanceItemTc headTys item =
  case item of
    InstanceItemAnn ann inner -> InstanceItemAnn ann <$> annotateInstanceItemTc headTys inner
    InstanceItemBind (FunctionBind name matches) -> do
      let methodName = unqualifiedNameText name
      methodTy <- methodExpectedType headTys (unqualifiedNameText name)
      pure (InstanceItemAnn (mkAnnotation (TcInstanceMethodAnnotation methodName methodTy)) (InstanceItemBind (FunctionBind name matches)))
    InstanceItemBind (PatternBind anns pat rhs) ->
      case patternBinderName pat of
        Just (_, methodName) -> do
          methodTy <- methodExpectedType headTys methodName
          pure (InstanceItemAnn (mkAnnotation (TcInstanceMethodAnnotation methodName methodTy)) (InstanceItemBind (PatternBind anns pat rhs)))
        Nothing -> pure item
    _ -> pure item

tcClassDeclBodies :: Decl -> TcM Decl
tcClassDeclBodies (DeclAnn ann inner) =
  DeclAnn ann <$> tcClassDeclBodies inner
tcClassDeclBodies (DeclClass classDecl) = do
  items <- mapM tcClassDefaultBody (classDeclItems classDecl)
  pure (DeclClass (classDecl {classDeclItems = items}))
tcClassDeclBodies decl = pure decl

tcClassDefaultBody :: ClassDeclItem -> TcM ClassDeclItem
tcClassDefaultBody item =
  case item of
    ClassItemAnn ann inner -> ClassItemAnn ann <$> tcClassDefaultBody inner
    ClassItemDefault valueDecl -> do
      checked <- tcClassDefaultValue valueDecl
      pure (ClassItemDefault checked)
    _ -> pure item

tcClassDefaultValue :: ValueDecl -> TcM ValueDecl
tcClassDefaultValue valueDecl =
  case valueDeclBinderName valueDecl of
    Nothing -> pure valueDecl
    Just (_, methodName) -> do
      binder <- lookupTerm (defaultMethodName methodName)
      case binder of
        Just (TcIdBinder (ForAll _ givens methodTy) _) ->
          case valueDecl of
            FunctionBind name matches -> do
              let (argumentTypes, resultType) = splitFunTy methodTy (matchArity matches)
              results <- mapM (tcMatchEquation Nothing argumentTypes resultType) matches
              solveInstanceBodyConstraints givens [(constraints, implications) | (_, constraints, implications) <- results]
              pure (FunctionBind name [match | (match, _, _) <- results])
            PatternBind annotations pattern' rhs -> do
              results <- mapM (tcMatchEquation Nothing [] methodTy) [zeroArgMatch (patternSpan pattern') rhs]
              solveInstanceBodyConstraints givens [(constraints, implications) | (_, constraints, implications) <- results]
              case results of
                [(match, _, _)] -> pure (PatternBind annotations pattern' (matchRhs match))
                _ -> pure valueDecl
        _ -> missingTypeInfo ("class default method " <> T.unpack methodName)

tcInstanceDeclBodies :: Decl -> TcM Decl
tcInstanceDeclBodies (DeclAnn ann inner) =
  DeclAnn ann <$> tcInstanceDeclBodies inner
tcInstanceDeclBodies (DeclInstance instanceDecl) =
  case (instanceHeadName (instanceDeclHead instanceDecl), instanceHeadTypes (instanceDeclHead instanceDecl)) of
    (_, []) -> pure (DeclInstance instanceDecl)
    (Nothing, _) -> pure (DeclInstance instanceDecl)
    (Just className, headArgTypes) -> do
      (_, tvEnv) <- makeInstanceTyVarEnv instanceDecl headArgTypes
      rawHeadTys <- checkInstanceHeadTypes (nameText className) tvEnv headArgTypes
      rawGivens <- mapM (surfacePredToPred tvEnv) (instanceDeclContext instanceDecl)
      headTys <- mapM defaultTypeKinds rawHeadTys
      givens <- mapM defaultPredKinds rawGivens
      items <- mapM (tcInstanceItemBody givens headTys) (instanceDeclItems instanceDecl)
      pure (DeclInstance (instanceDecl {instanceDeclItems = items}))
tcInstanceDeclBodies decl =
  pure decl

tcInstanceItemBody :: [Pred] -> [TcType] -> InstanceDeclItem -> TcM InstanceDeclItem
tcInstanceItemBody givens headTys item =
  case item of
    InstanceItemAnn ann inner ->
      InstanceItemAnn ann <$> tcInstanceItemBody givens headTys inner
    InstanceItemBind (FunctionBind name matches) -> do
      methodScheme <- methodExpectedScheme headTys (unqualifiedNameText name)
      (methodGivens, methodTy) <- skolemizeQualified methodScheme
      let (argTys, resTy) = splitFunTy methodTy (matchArity matches)
      results <- mapM (tcMatchEquation Nothing argTys resTy) matches
      solveInstanceBodyConstraints (givens <> methodGivens) [(cts, impls) | (_match, cts, impls) <- results]
      pure (InstanceItemBind (FunctionBind name [match | (match, _cts, _impls) <- results]))
    InstanceItemBind (PatternBind _ pat rhs) ->
      case patternBinderName pat of
        Just (_, methodName) -> do
          methodScheme <- methodExpectedScheme headTys methodName
          (methodGivens, methodTy) <- skolemizeQualified methodScheme
          results <- mapM (tcMatchEquation Nothing [] methodTy) [zeroArgMatch (patternSpan pat) rhs]
          solveInstanceBodyConstraints (givens <> methodGivens) [(cts, impls) | (_match, cts, impls) <- results]
          case results of
            [(match, _cts, _impls)] ->
              pure (replaceInstancePatternBindRhs (matchRhs match) item)
            _ -> pure item
        Nothing -> pure item
    _ -> pure item

matchArity :: [Match] -> Int
matchArity (match : _) = length (matchPats match)
matchArity [] = 0

solveInstanceBodyConstraints :: [Pred] -> [([Ct], [Implication])] -> TcM ()
solveInstanceBodyConstraints givens results = do
  let (ctsList, implsList) = unzip results
      cts = concat ctsList
      impls = concat implsList
  solveBodyConstraintsWithGivens givens cts impls

solveBodyConstraintsWithGivens :: [Pred] -> [Ct] -> [Implication] -> TcM ()
solveBodyConstraintsWithGivens givens cts impls = do
  implications <- mapM addOuterGivens impls
  _ <- solveWithImpls cts implications
  mapM_ solveClassCt cts
  where
    addOuterGivens implication = do
      outerGivens <- mapM givenConstraint givens
      pure (implication {implGivenCts = outerGivens <> implGivenCts implication})
    givenConstraint predicate = do
      evidence <- freshEvVar
      let origin = InstOrigin "class body"
      pure ((mkWantedCt predicate evidence origin NoSourceSpan) {ctFlavor = Given})
    solveClassCt ct@Ct {ctPred = ClassPred {}} = do
      _ <- solveDictWithGivens givens ct
      pure ()
    solveClassCt _ = pure ()

bindingType :: Text -> TcM TcType
bindingType name = do
  mBinder <- lookupTerm name
  case mBinder of
    Just binder -> pure (binderType binder)
    Nothing -> missingTypeInfo ("binding " <> T.unpack name)

binderType :: TcBinder -> TcType
binderType (TcIdBinder scheme _) = schemeToType scheme
binderType (TcMonoIdBinder ty) = ty

methodExpectedType :: [TcType] -> Text -> TcM TcType
methodExpectedType headTys methodName = schemeToType <$> methodExpectedScheme headTys methodName

methodExpectedScheme :: [TcType] -> Text -> TcM TypeScheme
methodExpectedScheme headTys methodName = do
  mBinder <- lookupTerm methodName
  case mBinder of
    Just (TcIdBinder (ForAll tyVars predicates body) _) ->
      case splitClassReceiver predicates headTys of
        Just (subst, methodPredicates) ->
          pure
            ( ForAll
                (filter (\tyVar -> not (Map.member (tvUnique tyVar) subst)) tyVars)
                (map (substPred subst) methodPredicates)
                (substType subst body)
            )
        Nothing -> missingTypeInfo ("class method receiver for " <> T.unpack methodName)
    Just (TcMonoIdBinder ty) -> pure (ForAll [] [] ty)
    Nothing -> missingTypeInfo ("class method " <> T.unpack methodName)

splitClassReceiver :: [Pred] -> [TcType] -> Maybe (Map Unique TcType, [Pred])
splitClassReceiver [] _ = Nothing
splitClassReceiver (predicate : predicates) headTys =
  case predicate of
    ClassPred _ classArgs -> (,predicates) <$> matchTypes classArgs headTys
    _ -> do
      (subst, rest) <- splitClassReceiver predicates headTys
      pure (subst, predicate : rest)

missingTypeInfo :: String -> TcM a
missingTypeInfo msg =
  abortTc ("internal type annotation error: missing " <> msg)

selectorDictTypeTc :: Text -> TcType -> TcM TcType
selectorDictTypeTc methodName methodTy =
  case snd (peelForAlls methodTy) of
    TcQualTy (pred' : _) _ -> pure (predType pred')
    _ -> missingTypeInfo ("class dictionary type for method selector " <> T.unpack methodName)

peelForAlls :: TcType -> ([TyVarId], TcType)
peelForAlls (TcForAllTy tv body) =
  let (tvs, inner) = peelForAlls body
   in (tv : tvs, inner)
peelForAlls ty = ([], ty)

predDictBinder :: Pred -> TcDictBinderAnnotation
predDictBinder pred' =
  case pred' of
    ClassPred classTyCon args ->
      TcDictBinderAnnotation (tyConName classTyCon) args (predType pred')
    EqPred {} ->
      TcDictBinderAnnotation "<constraint>" [] (predType pred')

constraintTypeDictBinder :: TcType -> TcDictBinderAnnotation
constraintTypeDictBinder ty =
  case constraintTypeToPred ty of
    Just (ClassPred classTyCon args) -> TcDictBinderAnnotation (tyConName classTyCon) args ty
    _ -> TcDictBinderAnnotation "<constraint>" [] ty

constraintTypePred :: TcType -> TcM Pred
constraintTypePred ty =
  case constraintTypeToPred ty of
    Just predicate -> pure predicate
    Nothing -> missingTypeInfo ("class predicate for constraint " <> show ty)

constraintTypeToPred :: TcType -> Maybe Pred
constraintTypeToPred ty =
  case collectTypeApplications ty of
    (TcTyCon (TyCon "~" 2) [], [left, right]) -> Just (EqPred left right)
    (TcTyCon tyCon headArgs, arguments) ->
      Just (ClassPred tyCon (headArgs <> arguments))
    _ -> Nothing

collectTypeApplications :: TcType -> (TcType, [TcType])
collectTypeApplications ty =
  case ty of
    TcAppTy function argument ->
      let (headType, arguments) = collectTypeApplications function
       in (headType, arguments <> [argument])
    _ -> (ty, [])

collectClassMethodNames :: [Decl] -> Map Text [Text]
collectClassMethodNames = Map.fromList . mapMaybe collect
  where
    collect decl =
      case peelDeclAnn decl of
        DeclClass classDecl ->
          Just (unqualifiedNameText (binderHeadName (classDeclHead classDecl)), classDeclMethodNames classDecl)
        _ -> Nothing

classDeclMethodNames :: ClassDecl -> [Text]
classDeclMethodNames classDecl = concatMap classItemMethodNames (classDeclItems classDecl)

classItemMethodNames :: ClassDeclItem -> [Text]
classItemMethodNames item =
  case peelClassDeclItemAnn item of
    ClassItemTypeSig names _ -> map unqualifiedNameText names
    _ -> []

classDeclDefaultMethodNames :: ClassDecl -> [Text]
classDeclDefaultMethodNames classDecl = mapMaybe classItemDefaultMethodName (classDeclItems classDecl)

classItemDefaultMethodName :: ClassDeclItem -> Maybe Text
classItemDefaultMethodName item =
  case peelClassDeclItemAnn item of
    ClassItemDefault valueDecl -> snd <$> valueDeclBinderName valueDecl
    _ -> Nothing

valueDeclBinderName :: ValueDecl -> Maybe (Text, Text)
valueDeclBinderName valueDecl =
  case valueDecl of
    FunctionBind name _ -> Just (binderBindingName name)
    PatternBind _ pat _ -> patternBinderName pat

defaultMethodName :: Text -> Text
defaultMethodName methodName = "$dm" <> methodName

matchTypes :: [TcType] -> [TcType] -> Maybe (Map Unique TcType)
matchTypes patterns targets
  | length patterns /= length targets = Nothing
  | otherwise = foldM matchOne Map.empty (zip patterns targets)

matchOne :: Map Unique TcType -> (TcType, TcType) -> Maybe (Map Unique TcType)
matchOne subst (TcTyVar tv, target) =
  case Map.lookup (tvUnique tv) subst of
    Nothing -> Just (Map.insert (tvUnique tv) target subst)
    Just existing
      | existing == target -> Just subst
      | otherwise -> Nothing
matchOne subst (TcTyCon tc args, TcTyCon targetTc targetArgs)
  | tc == targetTc,
    length args == length targetArgs =
      foldM matchOne subst (zip args targetArgs)
matchOne subst (TcFunTy a b, TcFunTy targetA targetB) =
  matchOne subst (a, targetA) >>= \subst' -> matchOne subst' (b, targetB)
matchOne subst (TcAppTy f a, TcAppTy targetF targetA) =
  matchOne subst (f, targetF) >>= \subst' -> matchOne subst' (a, targetA)
matchOne subst (patternTy, targetTy)
  | patternTy == targetTy = Just subst
  | otherwise = Nothing

substPred :: Map Unique TcType -> Pred -> Pred
substPred subst (ClassPred className args) = ClassPred className (map (substType subst) args)
substPred subst (EqPred left right) = EqPred (substType subst left) (substType subst right)

substType :: Map Unique TcType -> TcType -> TcType
substType = Aihc.Tc.Instantiate.applySubst

-- | Collect type signatures from a list of declarations.
collectUserSigs :: [Decl] -> TcM (Map TcTermKey UserSig)
collectUserSigs decls = do
  signatures <- concat <$> mapM (extractSig NoSourceSpan) decls
  foldM insertSignature Map.empty signatures
  where
    insertSignature collected (key, signature)
      | Map.member key collected = abortTc ("duplicate source signature key: " <> show key)
      | otherwise = pure (Map.insert key signature collected)
    extractSig ambient (DeclTypeSig names ty) =
      mapM
        ( \n -> do
            key <- resolvedUnqualifiedTermKey n
            let name = unqualifiedNameText n
                sigSp = ambient `orSourceSpan` unqualifiedNameSpan n `orSourceSpan` typeSpan ty
            pure (key, UserSig name ty sigSp)
        )
        names
    extractSig ambient (DeclForeign foreignDecl)
      | isForeignImport foreignDecl =
          do
            key <- resolvedUnqualifiedTermKey (foreignName foreignDecl)
            let name = unqualifiedNameText (foreignName foreignDecl)
                sigSp = ambient `orSourceSpan` unqualifiedNameSpan (foreignName foreignDecl) `orSourceSpan` typeSpan (foreignType foreignDecl)
            pure [(key, UserSig name (foreignType foreignDecl) sigSp)]
    extractSig ambient (DeclAnn ann inner) =
      extractSig (fromMaybe ambient (fromAnnotation @SourceSpan ann)) inner
    extractSig _ _ = pure []

checkUserSig :: UserSig -> TcM CheckedSig
checkUserSig userSig = do
  scheme <- sigToScheme (userSigType userSig)
  pure
    CheckedSig
      { checkedSigName = userSigName userSig,
        checkedSigScheme = scheme,
        checkedSigSpan = userSigSpan userSig
      }

splitContext :: Type -> ([Type], Type)
splitContext (TAnn _ inner) = splitContext inner
splitContext (TContext preds inner) = (preds, inner)
splitContext ty = ([], ty)

makeInstanceTyVarEnv :: InstanceDecl -> [Type] -> TcM ([TyVarId], TvKindEnv)
makeInstanceTyVarEnv instanceDecl headArgTypes = do
  explicitParams <- makeParamEnv (instanceDeclForall instanceDecl)
  let explicitNames = map paramName explicitParams
      freeVars = nub (explicitNames <> concatMap freeTypeVars (instanceDeclContext instanceDecl <> headArgTypes))
      implicitNames = freeVars \\ explicitNames
      explicitEnv = Map.fromList [(paramName param, (paramTyVar param, paramKind param)) | param <- explicitParams]
  rawImplicitTyVars <- mapM freshSkolemTv implicitNames
  implicitKinds <- mapM (const freshKindMeta) implicitNames
  let implicitTyVars = zipWith setTyVarKind implicitKinds rawImplicitTyVars
      implicitEnv = Map.fromList (zip implicitNames (zip implicitTyVars implicitKinds))
  pure (map paramTyVar explicitParams <> implicitTyVars, explicitEnv <> implicitEnv)

checkInstanceHeadTypes :: Text -> TvKindEnv -> [Type] -> TcM [TcType]
checkInstanceHeadTypes className tvEnv headArgTypes = do
  argKinds <- classPredicateArgKinds className (length headArgTypes)
  zipWithM (checkSurfaceType tvEnv) headArgTypes argKinds

-- | Instantiate a type scheme with fresh skolems for type-checking while
-- preserving the scheme predicates as scoped givens for the checked body.
-- Unlike regular instantiation (which uses metas), this produces rigid
-- type variables that cannot be unified during constraint solving.
skolemizeQualified :: TypeScheme -> TcM ([Pred], TcType)
skolemizeQualified (ForAll tvs preds body) = do
  subst <- Map.fromList <$> mapM mkSubst tvs
  pure (map (substPred subst) preds, substType subst body)
  where
    mkSubst tv = do
      sk <- freshSkolemTv (tvName tv)
      pure (tvUnique tv, TcTyVar sk)

-- | Split a function type into argument types and result type.
splitFunTy :: TcType -> Int -> ([TcType], TcType)
splitFunTy ty 0 = ([], ty)
splitFunTy (TcFunTy a rest) n =
  let (args, res) = splitFunTy rest (n - 1)
   in (a : args, res)
splitFunTy ty _ = ([], ty)

-- | A group of declarations that should be typechecked together.
-- Multiple FunctionBind equations for the same name are merged.
data DeclGroup
  = SingleDecl Decl
  | MergedFunctionBind SourceSpan UnqualifiedName [Decl] [Match]

data DeclGraphKey
  = DeclGraphBinder !TcTermKey
  | DeclGraphSynthetic !Int
  deriving (Eq, Ord, Show)

-- | Group consecutive FunctionBind declarations with the same name.
groupValueDecls :: [Decl] -> [DeclGroup]
groupValueDecls [] = []
groupValueDecls (d : ds) = case extractFunctionBind d of
  Just (sp, name, matches) ->
    let (sameNameDecls, rest) = span (hasSameName name) ds
        groupDecls = d : sameNameDecls
        allMatches = matches ++ concatMap (maybe [] (\(_, _, ms) -> ms) . extractFunctionBind) sameNameDecls
     in MergedFunctionBind sp name groupDecls allMatches : groupValueDecls rest
  Nothing -> SingleDecl d : groupValueDecls ds

-- | Extract function bind info from a declaration.
extractFunctionBind :: Decl -> Maybe (SourceSpan, UnqualifiedName, [Match])
extractFunctionBind decl =
  case peelDeclAnn decl of
    DeclValue (FunctionBind name matches) ->
      let sp = peelDeclSpan NoSourceSpan decl
       in Just (sp, name, matches)
    _ -> Nothing

-- | Check if a declaration is a FunctionBind with the given name.
hasSameName :: UnqualifiedName -> Decl -> Bool
hasSameName name d = case extractFunctionBind d of
  Just (_, n, _) -> unqualifiedNameText n == unqualifiedNameText name
  Nothing -> False

-- | Sort top-level groups so acyclic forward references are checked after
-- their dependencies have been generalized into the global environment.
sortDeclGroups :: [(Int, DeclGroup)] -> TcM [(Int, DeclGroup)]
sortDeclGroups groups = do
  allBinders <- Set.unions <$> mapM (fmap Set.fromList . declGroupBinderKeys . snd) groups
  nodes <- mapM (declGraphNode allBinders) groups
  pure (concatMap flattenScc (stronglyConnComp nodes))
  where
    flattenScc (AcyclicSCC group) = [group]
    flattenScc (CyclicSCC cyclicGroups) = cyclicGroups

declGraphNode :: Set.Set TcTermKey -> (Int, DeclGroup) -> TcM ((Int, DeclGroup), DeclGraphKey, [DeclGraphKey])
declGraphNode allBinders numberedGroup@(groupId, group) = do
  key <- groupKey groupId group
  freeVars <- freeVarsGroup group
  let deps = map DeclGraphBinder (Set.toList (Set.intersection allBinders freeVars))
  pure (numberedGroup, key, deps)

groupKey :: Int -> DeclGroup -> TcM DeclGraphKey
groupKey ix group = do
  keys <- declGroupBinderKeys group
  case keys of
    key : _ -> pure (DeclGraphBinder key)
    [] -> pure (DeclGraphSynthetic ix)

declGroupBinderKeys :: DeclGroup -> TcM [TcTermKey]
declGroupBinderKeys group =
  case group of
    MergedFunctionBind _sp binder _decls _matches -> (: []) <$> resolvedUnqualifiedTermKey binder
    SingleDecl decl ->
      case peelDeclAnn decl of
        DeclValue (FunctionBind binder _) -> (: []) <$> resolvedUnqualifiedTermKey binder
        DeclValue (PatternBind _ pat _) -> maybe (pure []) (fmap (: []) . resolvedUnqualifiedTermKey) (patternBinderSyntaxName pat)
        _ -> pure []

freeVarsGroup :: DeclGroup -> TcM (Set.Set TcTermKey)
freeVarsGroup group =
  case group of
    MergedFunctionBind _sp binder _decls matches -> do
      vars <- Set.unions <$> mapM freeVarsMatch matches
      binderKey <- resolvedUnqualifiedTermKey binder
      pure (Set.delete binderKey vars)
    SingleDecl decl -> freeVarsDecl decl

renderDeclGroup :: DeclGroup -> [Decl]
renderDeclGroup group =
  case group of
    SingleDecl decl -> [decl]
    MergedFunctionBind _ _ decls _ -> decls

replaceFunctionDeclMatches :: [Match] -> [Decl] -> [Decl]
replaceFunctionDeclMatches matches decls =
  snd (mapAccumL replace matches decls)
  where
    replace remaining decl =
      let count = functionDeclMatchCount decl
          (here, rest) = splitAt count remaining
       in (rest, replaceDeclFunctionMatches here decl)

functionDeclMatchCount :: Decl -> Int
functionDeclMatchCount decl =
  case peelDeclAnn decl of
    DeclValue (FunctionBind _ matches) -> length matches
    _ -> 0

replaceDeclFunctionMatches :: [Match] -> Decl -> Decl
replaceDeclFunctionMatches matches decl =
  case decl of
    DeclAnn ann inner -> DeclAnn ann (replaceDeclFunctionMatches matches inner)
    DeclValue (FunctionBind name _) -> DeclValue (FunctionBind name matches)
    _ -> decl

replacePatternBindRhs :: Rhs Expr -> Decl -> Decl
replacePatternBindRhs rhs decl =
  case decl of
    DeclAnn ann inner -> DeclAnn ann (replacePatternBindRhs rhs inner)
    DeclValue (PatternBind mult pat _) -> DeclValue (PatternBind mult pat rhs)
    _ -> decl

replaceInstancePatternBindRhs :: Rhs Expr -> InstanceDeclItem -> InstanceDeclItem
replaceInstancePatternBindRhs rhs item =
  case item of
    InstanceItemAnn ann inner -> InstanceItemAnn ann (replaceInstancePatternBindRhs rhs inner)
    InstanceItemBind (PatternBind mult pat _) -> InstanceItemBind (PatternBind mult pat rhs)
    _ -> item

freeVarsDecl :: Decl -> TcM (Set.Set TcTermKey)
freeVarsDecl decl =
  case peelDeclAnn decl of
    DeclValue (FunctionBind binder matches) -> do
      vars <- Set.unions <$> mapM freeVarsMatch matches
      binderKey <- resolvedUnqualifiedTermKey binder
      pure (Set.delete binderKey vars)
    DeclValue (PatternBind _ pat rhs) -> do
      vars <- freeVarsRhs rhs
      binders <- patternBinderKeys pat
      pure (Set.difference vars binders)
    _ -> pure Set.empty

freeVarsMatch :: Match -> TcM (Set.Set TcTermKey)
freeVarsMatch match = do
  vars <- freeVarsRhs (matchRhs match)
  binders <- Set.unions <$> mapM patternBinderKeys (matchPats match)
  pure (Set.difference vars binders)

freeVarsRhs :: Rhs Expr -> TcM (Set.Set TcTermKey)
freeVarsRhs rhs =
  case rhs of
    UnguardedRhs _ expr maybeDecls -> do
      exprVars <- freeVarsExpr expr
      declVars <- maybe (pure Set.empty) freeVarsDecls maybeDecls
      pure (exprVars <> declVars)
    GuardedRhss _ _ maybeDecls ->
      maybe (pure Set.empty) freeVarsDecls maybeDecls

freeVarsDecls :: [Decl] -> TcM (Set.Set TcTermKey)
freeVarsDecls decls =
  Set.unions <$> mapM freeVarsDecl decls

freeVarsExpr :: Expr -> TcM (Set.Set TcTermKey)
freeVarsExpr expr =
  case expr of
    EVar name -> Set.singleton <$> resolvedTermKey name
    EAnn _ inner -> freeVarsExpr inner
    EIf cond trueBranch falseBranch -> Set.unions <$> mapM freeVarsExpr [cond, trueBranch, falseBranch]
    ELambdaPats pats body -> do
      bodyVars <- freeVarsExpr body
      binders <- Set.unions <$> mapM patternBinderKeys pats
      pure (Set.difference bodyVars binders)
    EInfix lhs op rhs -> do
      lhsVars <- freeVarsExpr lhs
      rhsVars <- freeVarsExpr rhs
      opKey <- resolvedTermKey op
      pure (Set.insert opKey (lhsVars <> rhsVars))
    ENegate inner -> freeVarsExpr inner
    ESectionL inner op -> do
      innerVars <- freeVarsExpr inner
      opKey <- resolvedTermKey op
      pure (Set.insert opKey innerVars)
    ESectionR op inner -> do
      innerVars <- freeVarsExpr inner
      opKey <- resolvedTermKey op
      pure (Set.insert opKey innerVars)
    ELetDecls decls body -> do
      declVars <- freeVarsDecls decls
      bodyVars <- freeVarsExpr body
      localBinders <- declBinderKeys decls
      pure (Set.difference (declVars <> bodyVars) localBinders)
    ECase scrut alts -> do
      scrutVars <- freeVarsExpr scrut
      altVars <- Set.unions <$> mapM freeVarsAlt alts
      pure (scrutVars <> altVars)
    ETypeSig inner _ -> freeVarsExpr inner
    EParen inner -> freeVarsExpr inner
    EList items -> Set.unions <$> mapM freeVarsExpr items
    ETuple _ items -> Set.unions <$> mapM (maybe (pure Set.empty) freeVarsExpr) items
    EApp fun arg -> do
      funVars <- freeVarsExpr fun
      argVars <- freeVarsExpr arg
      pure (funVars <> argVars)
    _ -> pure Set.empty

freeVarsAlt :: CaseAlt Expr -> TcM (Set.Set TcTermKey)
freeVarsAlt (CaseAlt _ pat rhs) = do
  vars <- freeVarsRhs rhs
  binders <- patternBinderKeys pat
  pure (Set.difference vars binders)

declBinderKeys :: [Decl] -> TcM (Set.Set TcTermKey)
declBinderKeys decls =
  Set.unions <$> mapM declBinderKeySet decls

declBinderKeySet :: Decl -> TcM (Set.Set TcTermKey)
declBinderKeySet decl =
  case peelDeclAnn decl of
    DeclValue (FunctionBind binder _) -> Set.singleton <$> resolvedUnqualifiedTermKey binder
    DeclValue (PatternBind _ pat _) -> patternBinderKeys pat
    _ -> pure Set.empty

patternBinderKeys :: Pattern -> TcM (Set.Set TcTermKey)
patternBinderKeys pat =
  case pat of
    PVar name -> Set.singleton <$> resolvedUnqualifiedTermKey name
    PAnn _ inner -> patternBinderKeys inner
    PParen inner -> patternBinderKeys inner
    PAs name inner -> do
      key <- resolvedUnqualifiedTermKey name
      Set.insert key <$> patternBinderKeys inner
    PStrict inner -> patternBinderKeys inner
    PIrrefutable inner -> patternBinderKeys inner
    PCon _ _ pats -> Set.unions <$> mapM patternBinderKeys pats
    PInfix lhs _ rhs -> do
      lhsKeys <- patternBinderKeys lhs
      rhsKeys <- patternBinderKeys rhs
      pure (lhsKeys <> rhsKeys)
    _ -> pure Set.empty

patternBinders :: Pattern -> [Text]
patternBinders pat =
  case pat of
    PVar name -> [unqualifiedNameText name]
    PAnn _ inner -> patternBinders inner
    PParen inner -> patternBinders inner
    PAs name inner -> unqualifiedNameText name : patternBinders inner
    PStrict inner -> patternBinders inner
    PIrrefutable inner -> patternBinders inner
    PCon _ _ pats -> concatMap patternBinders pats
    PInfix lhs _ rhs -> patternBinders lhs ++ patternBinders rhs
    _ -> []

-- | Type-check a declaration group.
tcDeclGroup :: Map TcTermKey CheckedSig -> (Int, DeclGroup) -> TcM TcDeclGroupResult
tcDeclGroup sigs (groupId, group) =
  case group of
    SingleDecl d -> tcSingleDeclGroup sigs groupId d
    MergedFunctionBind _sp binder decls matches -> tcMergedFunctionGroup sigs groupId binder decls matches

tcSingleDeclGroup :: Map TcTermKey CheckedSig -> Int -> Decl -> TcM TcDeclGroupResult
tcSingleDeclGroup sigs groupId d =
  case peelDeclAnn d of
    DeclValue (PatternBind _ pat rhs) ->
      case patternBinderSyntaxName pat of
        Just binder -> do
          key <- resolvedUnqualifiedTermKey binder
          let displayName = renderBinderName binder
              name = unqualifiedNameText binder
          (maybeMatches, bindings) <-
            case Map.lookup key sigs of
              Just sig ->
                tcFunctionWithSig displayName name sig [zeroArgMatch (patternSpan pat `orSourceSpan` peelDeclSpan NoSourceSpan d) rhs]
              Nothing ->
                tcFunctionInfer key displayName name [zeroArgMatch (patternSpan pat) rhs]
          let annotatedDecls = fmap (\case [match] -> [replacePatternBindRhs (matchRhs match) d]; _ -> [d]) maybeMatches
          pure (TcDeclGroupResult groupId bindings annotatedDecls)
        Nothing -> do
          ((rhs', ty), failed) <- withErrorTracking (tcRhs rhs)
          if failed
            then pure (TcDeclGroupResult groupId [] Nothing)
            else do
              zonkedTy <- zonkType ty
              let decl' = replacePatternBindRhs rhs' d
              pure (TcDeclGroupResult groupId [TcBindingResult "<pattern>" "<pattern>" zonkedTy] (Just [decl']))
    _ -> do
      bindings <- tcDecl d
      pure (TcDeclGroupResult groupId bindings Nothing)

tcMergedFunctionGroup :: Map TcTermKey CheckedSig -> Int -> UnqualifiedName -> [Decl] -> [Match] -> TcM TcDeclGroupResult
tcMergedFunctionGroup sigs groupId binder decls matches = do
  let name = unqualifiedNameText binder
      displayName = renderBinderName binder
  key <- resolvedUnqualifiedTermKey binder
  (maybeMatches, bindings) <- case Map.lookup key sigs of
    Just sig -> do
      -- Use the declared type signature for checking.
      tcFunctionWithSig displayName name sig matches
    Nothing -> do
      -- No signature: infer the type.
      tcFunctionInfer key displayName name matches
  let annotatedDecls = fmap (`replaceFunctionDeclMatches` decls) maybeMatches
  pure (TcDeclGroupResult groupId bindings annotatedDecls)

-- | Type-check a function with a known type signature.
-- The signature's type variables are opened as rigid skolems so that
-- the body is checked against them. GADT patterns generate implication
-- constraints using the signature's skolems as given equalities.
tcFunctionWithSig :: Text -> Text -> CheckedSig -> [Match] -> TcM (Maybe [Match], [TcBindingResult])
tcFunctionWithSig displayName name sig matches = do
  let scheme = checkedSigScheme sig
  (matches', failed) <-
    withErrorTracking $ do
      -- Open the scheme with skolems (not metas) for checking.
      (sigPreds, sigTy) <- skolemizeQualified scheme
      let nArgs = case matches of
            (m : _) -> length (matchPats m)
            [] -> 0
          (argTys, resTy) = splitFunTy sigTy nArgs
      -- Check each equation against the signature types.
      results <- mapM (tcMatchEquation (Just (TypeSignatureOrigin (checkedSigName sig) (checkedSigSpan sig))) argTys resTy) matches
      let (_matches', ctsList, implsList) = unzip3 results
          allCts = concat ctsList
          allImpls = concat implsList
      solveBodyConstraintsWithGivens sigPreds allCts allImpls
      rejectEscapingExistentials sigTy allImpls
      pure _matches'
  if failed
    then pure (Nothing, [])
    else do
      -- Report the declared scheme as the binding's type.
      let declaredTy = schemeToType scheme
      zonkedTy <- zonkType declaredTy
      pure (Just matches', [TcBindingResult name displayName zonkedTy])

-- | Type-check a function without a type signature (infer).
tcFunctionInfer :: TcTermKey -> Text -> Text -> [Match] -> TcM (Maybe [Match], [TcBindingResult])
tcFunctionInfer key displayName name matches = do
  placeholderTy <- freshMetaTv
  ((matches', ty, residualPreds), failed) <-
    withErrorTracking $ do
      extendTermEnvPermanent name (TcMonoIdBinder placeholderTy)
      extendTermKeyEnvPermanent key (TcMonoIdBinder placeholderTy)
      (matches', ty, cts', impls') <- tcMatches matches
      solveResult <- solveWithImpls cts' impls'
      rejectEscapingExistentials ty impls'
      residualPreds <- generalizableResidualPreds solveResult
      pure (matches', ty, residualPreds)
  if failed
    then pure (Nothing, [])
    else do
      scheme <- generalizeAndCommitIgnoring (Set.fromList [unqualifiedTermKey name, key]) ty residualPreds
      let schemeTy = schemeToType scheme
      zonkedTy <- zonkType schemeTy
      finalizeInferredTermEnvPermanent name key placeholderTy scheme
      pure (Just matches', [TcBindingResult name displayName zonkedTy])

generalizableResidualPreds :: SolveResult -> TcM [Pred]
generalizableResidualPreds solveResult = do
  residualCts <- mapM zonkCtPred (srResidual solveResult <> inertDicts (srInerts solveResult))
  let uniqueResidualCts = nubBy sameCtPred residualCts
      (polymorphicCts, concreteCts) = partition (predicateCanGeneralize . ctPred) uniqueResidualCts
  -- Every occurrence still needs evidence, even when equal predicates share
  -- one constraint in the generalized type.
  forM_ residualCts $ \ct ->
    when (predicateCanGeneralize (ctPred ct)) $
      bindEvidence (ctEvVar ct) (EvGiven (ctPred ct))
  -- A fully concrete residual cannot be discharged by a caller-supplied
  -- dictionary, so reject it at the originating expression.
  forM_ concreteCts $ \ct ->
    emitError (ctLoc ct) (UnsolvedWanted (ctPred ct) (ctOrigin ct))
  pure (map ctPred polymorphicCts)
  where
    zonkCtPred ct = do
      pred' <- zonkPred (ctPred ct)
      pure (ct {ctPred = pred'})

    sameCtPred left right = ctPred left == ctPred right

predicateCanGeneralize :: Pred -> Bool
predicateCanGeneralize =
  not . null . predMetaVars

rejectEscapingExistentials :: TcType -> [Implication] -> TcM ()
rejectEscapingExistentials outerType implications = do
  zonkedOuterType <- zonkType outerType
  let skolems = concatMap implSkols implications
      escaping = filter (`typeMentionsTyVar` zonkedOuterType) skolems
  unless (null escaping) $
    emitError
      NoSourceSpan
      ( OtherError
          ( "existential type variable escapes its pattern-match branch: "
              <> T.unpack (T.intercalate ", " (map tvName escaping))
          )
      )

typeMentionsTyVar :: TyVarId -> TcType -> Bool
typeMentionsTyVar target ty =
  case ty of
    TcTyVar tyVar -> tyVar == target || kindMentionsUnique (tvUnique target) (tvKind tyVar)
    TcMetaTv {} -> False
    TcTyCon _ arguments -> any (typeMentionsTyVar target) arguments
    TcFunTy argument result -> typeMentionsTyVar target argument || typeMentionsTyVar target result
    TcForAllTy tyVar body -> tyVar /= target && typeMentionsTyVar target body
    TcQualTy predicates body -> any (predicateMentionsTyVar target) predicates || typeMentionsTyVar target body
    TcAppTy function argument -> typeMentionsTyVar target function || typeMentionsTyVar target argument
    TcBuiltinTyCon _ _ arguments -> any (typeMentionsTyVar target) arguments

kindMentionsUnique :: Unique -> Kind -> Bool
kindMentionsUnique target kind =
  case kind of
    KTYPE runtimeRep -> runtimeRepMentionsUnique target runtimeRep
    KFun argument result -> kindMentionsUnique target argument || kindMentionsUnique target result
    KMeta unique -> unique == target
    _ -> False

runtimeRepMentionsUnique :: Unique -> RuntimeRep -> Bool
runtimeRepMentionsUnique target runtimeRep =
  case runtimeRep of
    RuntimeRepVar unique -> unique == target
    RuntimeRepMeta unique -> unique == target
    TupleRep fields -> any (runtimeRepMentionsUnique target) fields
    SumRep fields -> any (runtimeRepMentionsUnique target) fields
    _ -> False

predicateMentionsTyVar :: TyVarId -> Pred -> Bool
predicateMentionsTyVar target predicate =
  case predicate of
    ClassPred _ arguments -> any (typeMentionsTyVar target) arguments
    EqPred left right -> typeMentionsTyVar target left || typeMentionsTyVar target right

zonkPred :: Pred -> TcM Pred
zonkPred pred' =
  case pred' of
    ClassPred className args -> ClassPred className <$> mapM zonkType args
    EqPred left right -> EqPred <$> zonkType left <*> zonkType right

collectStandaloneKindSignatures :: [Decl] -> Map TcTypeKey Type
collectStandaloneKindSignatures = Map.fromList . mapMaybe collect
  where
    collect declaration =
      case declaration of
        DeclAnn _ inner -> collect inner
        DeclStandaloneKindSig name kind -> (,kind) <$> resolvedTypeKey name
        _ -> Nothing

resolvedTypeKey :: UnqualifiedName -> Maybe TcTypeKey
resolvedTypeKey name = do
  ResolutionAnnotation {resolutionTarget = ResolvedTopLevel packageId resolvedName} <- nameResolution name
  moduleName' <- nameQualifier resolvedName
  pure (packageId, moduleName', nameText resolvedName)

registerTypeDeclHeader :: Map TcTypeKey TypeScheme -> Decl -> TcM [TcBindingResult]
registerTypeDeclHeader kindSchemes (DeclData dataDecl) =
  registerDataDeclHeader (resolvedTypeKey (binderHeadName (dataDeclHead dataDecl)) >>= (`Map.lookup` kindSchemes)) dataDecl
registerTypeDeclHeader kindSchemes (DeclNewtype newtypeDecl) =
  registerNewtypeDeclHeader (resolvedTypeKey (binderHeadName (newtypeDeclHead newtypeDecl)) >>= (`Map.lookup` kindSchemes)) newtypeDecl
registerTypeDeclHeader kindSchemes (DeclDataFamilyDecl familyDecl) =
  registerDataFamilyDeclHeader (resolvedTypeKey (binderHeadName (dataFamilyDeclHead familyDecl)) >>= (`Map.lookup` kindSchemes)) familyDecl
registerTypeDeclHeader kindSchemes (DeclTypeSyn typeSynDecl) =
  registerTypeSynonymHeader (resolvedTypeKey (binderHeadName (typeSynHead typeSynDecl)) >>= (`Map.lookup` kindSchemes)) typeSynDecl
registerTypeDeclHeader kindSchemes (DeclAnn _ inner) = registerTypeDeclHeader kindSchemes inner
registerTypeDeclHeader _ _ = pure []

registerStructuralDecl :: (Text, Text) -> Decl -> TcM [TcBindingResult]
registerStructuralDecl _ (DeclData dataDecl) = registerDataConstructors dataDecl
registerStructuralDecl _ (DeclNewtype newtypeDecl) = registerNewtypeConstructor newtypeDecl
registerStructuralDecl _ (DeclDataFamilyInst familyInst) = registerDataFamilyInstance familyInst
registerStructuralDecl origin (DeclClass classDecl) = registerClassDecl origin classDecl
registerStructuralDecl origin (DeclInstance instanceDecl) = registerInstanceDecl origin instanceDecl
registerStructuralDecl origin (DeclAnn _ inner) = registerStructuralDecl origin inner
registerStructuralDecl _ _ = pure []

isForeignImport :: ForeignDecl -> Bool
isForeignImport foreignDecl =
  foreignDirection foreignDecl == ForeignImport

registerClassDecl :: (Text, Text) -> ClassDecl -> TcM [TcBindingResult]
registerClassDecl origin classDecl = do
  let classBinder = binderHeadName (classDeclHead classDecl)
      className = unqualifiedNameText classBinder
      params = binderHeadParams (classDeclHead classDecl)
  paramInfos <- makeParamEnv params
  let paramTyVars = map paramTyVar paramInfos
      paramKinds = map paramKind paramInfos
      paramTvEnv = Map.fromList [(paramName param, (paramTyVar param, paramKind param)) | param <- paramInfos]
  superClassTypes <- mapM (\ty -> checkSurfaceType paramTvEnv ty KConstraint) (fromMaybe [] (classDeclContext classDecl))
  classKind <- defaultKindMetas (foldr KFun KConstraint paramKinds)
  classTyCon <- mkDeclaredTyCon classBinder className (length params) classKind
  let classPred = ClassPred classTyCon (map TcTyVar paramTyVars)
  extendTyConEnvPermanent
    className
    TyConInfo
      { tciName = className,
        tciArity = length params,
        tciTyCon = classTyCon,
        tciFlavor = ClassTyCon,
        tciTypeSynonym = Nothing
      }
  methodResults <- concat <$> mapM (registerClassItem classPred paramTvEnv paramTyVars) (classDeclItems classDecl)
  methods <- mapM registeredMethod (classDeclMethodNames classDecl)
  defaultSignatures <- catMaybes <$> mapM (registerClassDefaultSignature paramTvEnv paramTyVars) (classDeclItems classDecl)
  let defaults = classDeclDefaultMethodNames classDecl
  defaultResults <- mapM (registerDefaultMethod defaults defaultSignatures) methods
  addClass
    ClassInfo
      { ciName = className,
        ciTyCon = classTyCon,
        ciOrigin = Just origin,
        ciTyVars = paramTyVars,
        ciSuperClassTypes = superClassTypes,
        ciMethods = methods,
        ciDefaultMethods = defaults,
        ciDefaultSignatures = defaultSignatures
      }
  pure (methodResults <> catMaybes defaultResults)
  where
    registeredMethod methodName = do
      binder <- lookupTerm methodName
      case binder of
        Just (TcIdBinder scheme _) -> pure (methodName, scheme)
        _ -> missingTypeInfo ("class method " <> T.unpack methodName)

    registerDefaultMethod defaults defaultSignatures (methodName, scheme)
      | methodName `elem` defaults = do
          let workerName = defaultMethodName methodName
              workerScheme = maybe scheme (defaultWorkerScheme scheme) (lookup methodName defaultSignatures)
              workerType = schemeToType workerScheme
          extendTermEnvPermanent workerName (TcIdBinder workerScheme Closed)
          pure (Just (TcBindingResult workerName workerName workerType))
      | otherwise = pure Nothing

    defaultWorkerScheme ordinaryScheme (ForAll tyVars predicates body) =
      case ordinaryScheme of
        ForAll _ (classPredicate : _) _ -> ForAll tyVars (classPredicate : predicates) body
        _ -> ForAll tyVars predicates body

registerClassItem :: Pred -> TvKindEnv -> [TyVarId] -> ClassDeclItem -> TcM [TcBindingResult]
registerClassItem classPred classTvEnv classTyVars item =
  case peelClassDeclItemAnn item of
    ClassItemTypeSig names ty -> do
      let (context, body) = splitContext ty
          classVarNames = Map.keys classTvEnv
          freeVars = freeTypeVars ty \\ classVarNames
      extraTyVars <- mapM freshSkolemTv freeVars
      extraKinds <- mapM (const freshKindMeta) freeVars
      let tvEnv = classTvEnv <> Map.fromList (zip freeVars (zip extraTyVars extraKinds))
      methodBody <- checkSurfaceType tvEnv body KType
      contextPreds <- mapM (surfacePredToPred tvEnv) context
      let preds = classPred : contextPreds
          scheme = ForAll (classTyVars <> extraTyVars) preds methodBody
          declaredTy = schemeToType scheme
      mapM
        ( \methodName -> do
            let name = unqualifiedNameText methodName
                displayName = renderBinderName methodName
            extendResolvedTermEnvPermanent methodName (TcIdBinder scheme Closed)
            zonkedTy <- zonkType declaredTy
            pure (TcBindingResult name displayName zonkedTy)
        )
        names
    _ -> pure []

registerClassDefaultSignature :: TvKindEnv -> [TyVarId] -> ClassDeclItem -> TcM (Maybe (Text, TypeScheme))
registerClassDefaultSignature classTvEnv classTyVars item =
  case peelClassDeclItemAnn item of
    ClassItemDefaultSig methodName ty -> do
      let (context, body) = splitContext ty
          classVarNames = Map.keys classTvEnv
          freeVars = freeTypeVars ty \\ classVarNames
      extraTyVars <- mapM freshSkolemTv freeVars
      extraKinds <- mapM (const freshKindMeta) freeVars
      let tvEnv = classTvEnv <> Map.fromList (zip freeVars (zip extraTyVars extraKinds))
      methodBody <- checkSurfaceType tvEnv body KType
      contextPreds <- mapM (surfacePredToPred tvEnv) context
      pure
        ( Just
            ( unqualifiedNameText methodName,
              ForAll (classTyVars <> extraTyVars) contextPreds methodBody
            )
        )
    _ -> pure Nothing

registerInstanceDecl :: (Text, Text) -> InstanceDecl -> TcM [TcBindingResult]
registerInstanceDecl origin instanceDecl =
  case instanceHeadName (instanceDeclHead instanceDecl) of
    Nothing -> pure []
    Just className -> do
      let headArgs = instanceHeadTypes (instanceDeclHead instanceDecl)
      (tvIds, tvEnv) <- makeInstanceTyVarEnv instanceDecl headArgs
      let classNameText = nameText className
      headTys <- checkInstanceHeadTypes classNameText tvEnv headArgs
      let dictName = instanceDictName classNameText headTys
      context <- mapM (surfacePredToPred tvEnv) (instanceDeclContext instanceDecl)
      classInfo <- lookupClass classNameText >>= maybe (missingTypeInfo ("class " <> T.unpack classNameText)) pure
      let dictTy = foldr TcForAllTy (TcQualTy context (predType (ClassPred (ciTyCon classInfo) headTys))) tvIds
      addInstance
        InstanceInfo
          { iiClassName = classNameText,
            iiDictName = dictName,
            iiDictOrigin = Just origin,
            iiDictType = dictTy,
            iiTyVars = tvIds,
            iiContext = context,
            iiHead = headTys
          }
      pure [TcBindingResult dictName dictName dictTy]

predType :: Pred -> TcType
predType (ClassPred classTyCon args) = TcTyCon classTyCon args
predType (EqPred left right) = TcTyCon (TyCon "~" 2) [left, right]

instanceDictName :: Text -> [TcType] -> Text
instanceDictName className tys = "$f" <> className <> T.concat (map typeSuffix tys)

typeSuffix :: TcType -> Text
typeSuffix ty =
  case ty of
    TcTyVar tv -> tvName tv
    TcTyCon tc [] -> tyConName tc
    TcTyCon (TyCon "[]" _) [_] -> "List"
    TcTyCon tc args -> tyConName tc <> T.concat (map typeSuffix args)
    _ -> "T"

registerDataFamilyDeclHeader :: Maybe TypeScheme -> DataFamilyDecl -> TcM [TcBindingResult]
registerDataFamilyDeclHeader maybeKindScheme familyDecl = do
  let familyBinder = binderHeadName (dataFamilyDeclHead familyDecl)
      familyName = unqualifiedNameText familyBinder
      params = binderHeadParams (dataFamilyDeclHead familyDecl)
      arity = length params
  (_, paramInfos) <- typeDeclParamInfos maybeKindScheme params
  inferredKind <- tyConKindFromParams paramInfos (dataFamilyDeclKind familyDecl)
  familyTyCon <-
    case maybeKindScheme of
      Just kindScheme -> declaredTyConWithKindScheme familyBinder familyName arity kindScheme
      Nothing -> mkDeclaredTyCon familyBinder familyName arity inferredKind
  let declaredKind = maybe inferredKind (const (tyConKind familyTyCon)) maybeKindScheme
  extendTyConEnvPermanent
    familyName
    TyConInfo
      { tciName = familyName,
        tciArity = arity,
        tciTyCon = familyTyCon,
        tciFlavor = DataFamilyTyCon,
        tciTypeSynonym = Nothing
      }
  zonkedKind <- defaultKindMetas declaredKind
  pure [TcBindingResult familyName familyName (kindToTcType zonkedKind)]

registerDataFamilyInstance :: DataFamilyInst -> TcM [TcBindingResult]
registerDataFamilyInstance familyInst = do
  paramInfos <- dataFamilyInstanceParams familyInst
  let tvEnv =
        Map.fromList
          [ (paramName param, (paramTyVar param, paramKind param))
          | param <- paramInfos
          ]
      constructorNames = concatMap (map fst . dataConBindingNames) (dataFamilyInstConstructors familyInst)
  familyType <- checkSurfaceType tvEnv (dataFamilyInstHead familyInst) KType
  case (familyType, constructorNames) of
    (_, []) -> do
      emitError NoSourceSpan (OtherError "data-family instances without constructors are not supported")
      pure []
    (TcTyCon familyTyCon _, firstConstructor : _) -> do
      maybeFamilyInfo <- lookupTyConByIdentity familyTyCon
      case maybeFamilyInfo of
        Just familyInfo
          | tciFlavor familyInfo == DataFamilyTyCon -> do
              representationKind <- tyConKindFromParams paramInfos (dataFamilyInstKind familyInst)
              let familyName = tciName familyInfo
                  representationName = dataFamilyRepresentationName familyName firstConstructor
                  representationTyCon = mkTyCon representationName (length paramInfos) representationKind
                  axiomName = dataFamilyAxiomName familyName firstConstructor
                  instanceInfo =
                    DataFamilyInstanceInfo
                      { dfiiFamilyName = familyName,
                        dfiiFamilyType = familyType,
                        dfiiTyVars = map paramTyVar paramInfos,
                        dfiiRepresentationTyCon = representationTyCon,
                        dfiiAxiomName = axiomName,
                        dfiiConstructorNames = constructorNames,
                        dfiiIsNewtype = dataFamilyInstIsNewtype familyInst
                      }
              addDataFamilyInstance instanceInfo
              mapM (registerDataConWithResult paramInfos familyType) (dataFamilyInstConstructors familyInst)
        _ -> do
          emitError NoSourceSpan (OtherError ("data-family instance head does not name a data family: " <> T.unpack (tyConName familyTyCon)))
          pure []
    _ -> do
      emitError NoSourceSpan (OtherError ("invalid data-family instance head: " <> show familyType))
      pure []

dataFamilyInstanceParams :: DataFamilyInst -> TcM [ParamInfo]
dataFamilyInstanceParams familyInst = do
  explicitParams <- makeParamEnv (dataFamilyInstForall familyInst)
  let explicitNames = map paramName explicitParams
      implicitNames = freeTypeVars (dataFamilyInstHead familyInst) \\ explicitNames
  implicitParams <- mapM makeImplicitParam implicitNames
  pure (explicitParams <> implicitParams)
  where
    makeImplicitParam name = do
      rawTyVar <- freshSkolemTv name
      kind <- freshKindMeta
      pure
        ParamInfo
          { paramName = name,
            paramTyVar = setTyVarKind kind rawTyVar,
            paramKind = kind
          }

-- | Register a data declaration's type constructor and data constructors.
--
-- For @data Bool = True | False@, this produces:
--   - @Bool :: *@
--   - @True :: Bool@
--   - @False :: Bool@
typeDeclParamInfos :: Maybe TypeScheme -> [TyVarBinder] -> TcM ([ParamInfo], [ParamInfo])
typeDeclParamInfos maybeKindScheme params =
  case maybeKindScheme of
    Nothing -> ([],) <$> makeParamEnv params
    Just scheme@(ForAll kindTyVars _ _) -> do
      let kindParams = map kindParam kindTyVars
          kindEnv = Map.fromList [(paramName param, (paramTyVar param, paramKind param)) | param <- kindParams]
          expectedKinds = takeVisibleArgumentKinds (length params) (kindFromTypeScheme scheme)
      paramInfos <- makeParamEnvWith kindEnv params
      if length expectedKinds == length params
        then do
          zipWithM_ (unifyKinds . paramKind) paramInfos expectedKinds
          let checkedParams = zipWith setParamKind expectedKinds paramInfos
          pure (kindParams, checkedParams)
        else do
          emitError NoSourceSpan (OtherError "standalone kind signature arity does not match its type declaration")
          pure (kindParams, paramInfos)
  where
    kindParam tyVar = ParamInfo (tvName tyVar) tyVar (tvKind tyVar)
    setParamKind kind param =
      param
        { paramTyVar = setTyVarKind kind (paramTyVar param),
          paramKind = kind
        }

dataDeclParamInfos :: Maybe TypeScheme -> DataDecl -> TcM ([ParamInfo], [ParamInfo])
dataDeclParamInfos maybeKindScheme declaration =
  typeDeclParamInfos maybeKindScheme (binderHeadParams (dataDeclHead declaration))

quantifiedKindScheme :: TyCon -> Maybe TypeScheme
quantifiedKindScheme tyCon =
  case tyConKindScheme tyCon of
    ForAll [] _ _ -> Nothing
    scheme -> Just scheme

takeVisibleArgumentKinds :: Int -> Kind -> [Kind]
takeVisibleArgumentKinds = go
  where
    go remaining (KFun argument result)
      | remaining > 0 = argument : go (remaining - 1) result
    go _ _ = []

registerDataDeclHeader :: Maybe TypeScheme -> DataDecl -> TcM [TcBindingResult]
registerDataDeclHeader maybeKindScheme dd = do
  let tyBinder = binderHeadName (dataDeclHead dd)
      tyName = unqualifiedNameText tyBinder
      params = binderHeadParams (dataDeclHead dd)
      arity = length params
  (kindParams, paramInfos) <- dataDeclParamInfos maybeKindScheme dd
  let kindEnv = Map.fromList [(paramName param, (paramTyVar param, paramKind param)) | param <- kindParams]
  inferredKind <- tyConKindFromParamsWith kindEnv paramInfos (dataDeclKind dd)
  tc <-
    case maybeKindScheme of
      Just kindScheme -> declaredTyConWithKindScheme tyBinder tyName arity kindScheme
      Nothing -> dataDeclTyCon tyBinder tyName arity inferredKind
  let declaredKind = maybe inferredKind (const (tyConKind tc)) maybeKindScheme
  extendTyConEnvPermanent
    tyName
    TyConInfo
      { tciName = tyName,
        tciArity = arity,
        tciTyCon = tc,
        tciFlavor = DataTyCon,
        tciTypeSynonym = Nothing
      }
  zonkedKind <- defaultKindMetas declaredKind
  let tyConResult = TcBindingResult tyName tyName (kindToTcType zonkedKind)
  pure [tyConResult]

registerDataConstructors :: DataDecl -> TcM [TcBindingResult]
registerDataConstructors dataDecl = do
  let tyBinder = binderHeadName (dataDeclHead dataDecl)
      tyName = unqualifiedNameText tyBinder
  maybeInfo <- lookupDeclaredTyCon tyBinder
  case maybeInfo of
    Nothing -> missingTypeInfo ("data type " <> T.unpack tyName)
    Just info -> do
      (kindParams, paramInfos) <- dataDeclParamInfos (quantifiedKindScheme (tciTyCon info)) dataDecl
      bindings <- mapM (registerDataCon (tciTyCon info) kindParams paramInfos) (dataDeclConstructors dataDecl)
      constructors <- concat <$> mapM (checkedDataConInfos (tciTyCon info)) (dataDeclConstructors dataDecl)
      selectorBindings <- registerRecordSelectors constructors
      let tyVars = map paramTyVar paramInfos
          resultKind = typeKind (TcTyCon (tciTyCon info) (map TcTyVar tyVars))
      addDataType
        DataTypeInfo
          { dtiName = tyName,
            dtiTyCon = tciTyCon info,
            dtiTyVars = tyVars,
            dtiResultKind = resultKind,
            dtiFlavor = DataTyCon,
            dtiConstructors = constructors
          }
      pure (bindings <> selectorBindings)

-- | Register a newtype declaration's type constructor and representation
-- constructor.  Newtype erasure/coercion semantics are handled elsewhere; at
-- this stage the type checker only needs the source-level names and types.
registerNewtypeDeclHeader :: Maybe TypeScheme -> NewtypeDecl -> TcM [TcBindingResult]
registerNewtypeDeclHeader maybeKindScheme nd = do
  let tyBinder = binderHeadName (newtypeDeclHead nd)
      tyName = unqualifiedNameText tyBinder
      params = binderHeadParams (newtypeDeclHead nd)
      arity = length params
  (_, paramInfos) <- typeDeclParamInfos maybeKindScheme params
  inferredKind <- tyConKindFromParams paramInfos (newtypeDeclKind nd)
  tc <-
    case maybeKindScheme of
      Just kindScheme -> declaredTyConWithKindScheme tyBinder tyName arity kindScheme
      Nothing -> mkDeclaredTyCon tyBinder tyName arity inferredKind
  let declaredKind = maybe inferredKind (const (tyConKind tc)) maybeKindScheme
  extendTyConEnvPermanent
    tyName
    TyConInfo
      { tciName = tyName,
        tciArity = arity,
        tciTyCon = tc,
        tciFlavor = NewtypeTyCon,
        tciTypeSynonym = Nothing
      }
  zonkedKind <- defaultKindMetas declaredKind
  let tyConResult = TcBindingResult tyName tyName (kindToTcType zonkedKind)
  pure [tyConResult]

registerNewtypeConstructor :: NewtypeDecl -> TcM [TcBindingResult]
registerNewtypeConstructor newtypeDecl = do
  let tyBinder = binderHeadName (newtypeDeclHead newtypeDecl)
      tyName = unqualifiedNameText tyBinder
  maybeInfo <- lookupDeclaredTyCon tyBinder
  case maybeInfo of
    Nothing -> missingTypeInfo ("newtype " <> T.unpack tyName)
    Just info -> do
      (kindParams, paramInfos) <- typeDeclParamInfos (quantifiedKindScheme (tciTyCon info)) (binderHeadParams (newtypeDeclHead newtypeDecl))
      constructor <- mapM (registerDataCon (tciTyCon info) kindParams paramInfos) (newtypeDeclConstructor newtypeDecl)
      constructors <- maybe (pure []) (checkedDataConInfos (tciTyCon info)) (newtypeDeclConstructor newtypeDecl)
      selectorBindings <- registerRecordSelectors constructors
      let tyVars = map paramTyVar paramInfos
          resultKind = typeKind (TcTyCon (tciTyCon info) (map TcTyVar tyVars))
      addDataType
        DataTypeInfo
          { dtiName = tyName,
            dtiTyCon = tciTyCon info,
            dtiTyVars = tyVars,
            dtiResultKind = resultKind,
            dtiFlavor = NewtypeTyCon,
            dtiConstructors = constructors
          }
      pure (maybeToList constructor <> selectorBindings)

registerRecordSelectors :: [DataConInfo] -> TcM [TcBindingResult]
registerRecordSelectors constructors =
  mapM registerSelector (Map.toList selectors)
  where
    selectors =
      Map.fromListWith
        (++)
        [ (label, [(constructor, field)])
        | constructor <- constructors,
          field <- dciFields constructor,
          Just label <- [dcfiLabel field]
        ]
    registerSelector (label, (constructor, field) : _) = do
      let scheme =
            ForAll
              (dciUnivTyVars constructor)
              (dciTheta constructor)
              (TcFunTy (dciResTy constructor) (dcfiType field))
      extendTermEnvPermanent label (TcIdBinder scheme Closed)
      zonkedType <- zonkType (schemeToType scheme)
      pure (TcBindingResult label label zonkedType)
    registerSelector (label, []) =
      abortTc ("record selector has no fields: " <> T.unpack label)

registerTypeSynonymHeader :: Maybe TypeScheme -> TypeSynDecl -> TcM [TcBindingResult]
registerTypeSynonymHeader maybeKindScheme typeSynDecl = do
  let tyBinder = binderHeadName (typeSynHead typeSynDecl)
      tyName = unqualifiedNameText tyBinder
      params = binderHeadParams (typeSynHead typeSynDecl)
      arity = length params
  (_, paramInfos) <- typeDeclParamInfos maybeKindScheme params
  inferredResultKind <- freshKindMeta
  let inferredKind = foldr (KFun . paramKind) inferredResultKind paramInfos
  tyCon <-
    case maybeKindScheme of
      Just kindScheme -> declaredTyConWithKindScheme tyBinder tyName arity kindScheme
      Nothing -> mkDeclaredTyCon tyBinder tyName arity inferredKind
  let declaredKind = maybe inferredKind (const (tyConKind tyCon)) maybeKindScheme
  let synonym = TypeSynonymInfo (map paramTyVar paramInfos) Nothing
  extendTyConEnvPermanent
    tyName
    TyConInfo
      { tciName = tyName,
        tciArity = arity,
        tciTyCon = tyCon,
        tciFlavor = SynonymTyCon,
        tciTypeSynonym = Just synonym
      }
  pure [TcBindingResult tyName tyName (kindToTcType declaredKind)]

registerTypeSynonymBody :: Decl -> TcM ()
registerTypeSynonymBody (DeclAnn _ inner) = registerTypeSynonymBody inner
registerTypeSynonymBody (DeclTypeSyn typeSynDecl) = do
  let tyBinder = binderHeadName (typeSynHead typeSynDecl)
      tyName = unqualifiedNameText tyBinder
  maybeInfo <- lookupDeclaredTyCon tyBinder
  case maybeInfo of
    Just info
      | Just synonym <- tciTypeSynonym info -> do
          let params = tsiParams synonym
              tvEnv = Map.fromList [(tvName param, (param, tvKind param)) | param <- params]
              resultKind = typeResultKind (length params) (tyConKind (tciTyCon info))
          body <- checkSurfaceType tvEnv (typeSynBody typeSynDecl) resultKind
          replaceTyConEnvPermanent (info {tciTypeSynonym = Just (synonym {tsiBody = Just body})})
    _ -> missingTypeInfo ("type synonym " <> T.unpack tyName)
registerTypeSynonymBody _ = pure ()

typeResultKind :: Int -> Kind -> Kind
typeResultKind remaining kind
  | remaining <= 0 = kind
typeResultKind remaining (KFun _ result) = typeResultKind (remaining - 1) result
typeResultKind _ kind = kind

dataDeclTyCon :: UnqualifiedName -> Text -> Int -> Kind -> TcM TyCon
dataDeclTyCon binder "List" 1 kind = mkDeclaredTyCon binder "[]" 1 kind
dataDeclTyCon binder name arity kind = mkDeclaredTyCon binder name arity kind

mkDeclaredTyCon :: UnqualifiedName -> Text -> Int -> Kind -> TcM TyCon
mkDeclaredTyCon binder name arity kind =
  case nameResolution binder of
    Just ResolutionAnnotation {resolutionTarget = ResolvedTopLevel packageId resolvedName}
      | Just definingModuleName <- nameQualifier resolvedName ->
          pure (mkTyConWithOrigin packageId definingModuleName name arity kind)
    _ -> abortTc ("type declaration has no package or module identity: " <> T.unpack name)

declaredTyConWithKindScheme :: UnqualifiedName -> Text -> Int -> TypeScheme -> TcM TyCon
declaredTyConWithKindScheme binder name arity kindScheme =
  case nameResolution binder of
    Just ResolutionAnnotation {resolutionTarget = ResolvedTopLevel packageId resolvedName}
      | Just definingModuleName <- nameQualifier resolvedName ->
          pure (mkTyConWithOriginScheme packageId definingModuleName name arity kindScheme)
    _ -> abortTc ("type declaration has no package or module identity: " <> T.unpack name)

-- | Register a single data constructor as a polymorphic binding.
-- Returns the binding result for the constructor.
registerDataCon :: TyCon -> [ParamInfo] -> [ParamInfo] -> DataConDecl -> TcM TcBindingResult
registerDataCon tc kindParams paramInfos =
  registerDataConWithResult (kindParams <> paramInfos) (TcTyCon tc (map (TcTyVar . paramTyVar) paramInfos))

registerDataConWithResult :: [ParamInfo] -> TcType -> DataConDecl -> TcM TcBindingResult
registerDataConWithResult paramInfos resTy con = case con of
  DataConAnn _ inner -> registerDataConWithResult paramInfos resTy inner
  PrefixCon forallVars context conName args ->
    registerH98DataCon forallVars context (Just conName) (unqualifiedNameText conName) (map bangType args)
  InfixCon forallVars context lhs conName rhs ->
    registerH98DataCon forallVars context (Just conName) (unqualifiedNameText conName) (map bangType [lhs, rhs])
  RecordCon forallVars context conName fields ->
    registerH98DataCon forallVars context (Just conName) (unqualifiedNameText conName) (map bangType (recordBangFields fields))
  TupleCon forallVars context flavor fields ->
    registerH98DataCon forallVars context Nothing (tupleConText flavor (length fields)) (map bangType fields)
  UnboxedSumCon forallVars context pos arity field ->
    registerH98DataCon forallVars context Nothing (unboxedSumConText pos arity) [bangType field]
  ListCon forallVars context ->
    registerH98DataCon forallVars context Nothing "[]" []
  GadtCon forallBinders context names body -> do
    constructorParams <- makeParamEnv (concatMap forallTelescopeBinders forallBinders)
    let constructorEnv =
          Map.fromList
            [ (paramName param, (paramTyVar param, paramKind param))
            | param <- constructorParams
            ]
            <> paramEnv
        constructorTyVars = map paramTyVar constructorParams
    let resultSurfTy = gadtBodyResultType body
        argSurfTys = gadtBodyArgTypes body
    gadtResTy <- checkSurfaceType constructorEnv resultSurfTy KType
    gadtArgTys <- mapM (checkRuntimeType constructorEnv) argSurfTys
    predicates <- mapM (surfacePredToPred constructorEnv) context
    let conTy = foldr TcFunTy gadtResTy gadtArgTys
        candidateTyVars = paramVarIds <> constructorTyVars
        quantifiedTyVars = filter (\tyVar -> typeMentionsTyVar tyVar conTy || any (predicateMentionsTyVar tyVar) predicates) candidateTyVars
        gadtScheme = ForAll quantifiedTyVars predicates conTy
    mapM_
      ( \n -> do
          let nm = unqualifiedNameText n
          extendResolvedTermEnvPermanent n (TcIdBinder gadtScheme Closed)
          markGadtCon nm
      )
      names
    case names of
      (n : _) -> do
        zonkedTy <- zonkType conTy
        let name = unqualifiedNameText n
         in pure (TcBindingResult name name zonkedTy)
      [] -> pure (TcBindingResult "<gadt>" "<gadt>" gadtResTy)
  where
    paramEnv =
      Map.fromList
        [ (paramName param, (paramTyVar param, paramKind param))
        | param <- paramInfos
        ]
    paramVarIds = map paramTyVar paramInfos
    registerH98DataCon forallVars context maybeName name fieldTypes = do
      constructorParams <- makeParamEnv forallVars
      let constructorEnv =
            Map.fromList
              [ (paramName param, (paramTyVar param, paramKind param))
              | param <- constructorParams
              ]
              <> paramEnv
          constructorTyVars = map paramTyVar constructorParams
      argTys <- mapM (checkRuntimeType constructorEnv) fieldTypes
      predicates <- mapM (surfacePredToPred constructorEnv) context
      let conTy = foldr TcFunTy resTy argTys
          scheme = ForAll (paramVarIds <> constructorTyVars) predicates conTy
      case maybeName of
        Just sourceName -> extendResolvedTermEnvPermanent sourceName (TcIdBinder scheme Closed)
        Nothing ->
          case resTy of
            TcTyCon resultTyCon _ -> extendTyConTermEnvPermanent resultTyCon name (TcIdBinder scheme Closed)
            _ -> extendTermEnvPermanent name (TcIdBinder scheme Closed)
      zonkedTy <- zonkType (schemeToType scheme)
      pure (TcBindingResult name name zonkedTy)

tupleConText :: TupleFlavor -> Int -> Text
tupleConText flavor arity =
  case flavor of
    Boxed -> "(" <> commas arity <> ")"
    Unboxed -> "(#" <> commas arity <> "#)"

unboxedSumConText :: Int -> Int -> Text
unboxedSumConText pos arity = "(#" <> bars (pos - 1) <> "_" <> bars (arity - pos) <> "#)"

commas :: Int -> Text
commas n
  | n <= 1 = ""
  | otherwise = mconcat (replicate (n - 1) ",")

bars :: Int -> Text
bars n
  | n <= 0 = ""
  | otherwise = mconcat (replicate n "|")

-- | Extract argument types from a GadtBody.
gadtBodyArgTypes :: GadtBody -> [Type]
gadtBodyArgTypes (GadtPrefixBody argsWithKinds _) = map (bangType . fst) argsWithKinds
gadtBodyArgTypes (GadtRecordBody fields _) = map bangType (recordBangFields fields)

recordBangFields :: [FieldDecl] -> [BangType]
recordBangFields = concatMap $ \field -> replicate (length (fieldNames field)) (fieldType field)

checkedDataConInfos :: TyCon -> DataConDecl -> TcM [DataConInfo]
checkedDataConInfos tyCon declaration = do
  let (sourceForm, sourceFields, constructorNames) = dataConSourceLayout declaration
      origin = (tyConPackageId tyCon, tyConModuleName tyCon)
  mapM (checkedDataConInfo origin sourceForm sourceFields) constructorNames

checkedDataConInfo :: (PackageId, Text) -> DataConSourceForm -> [(Maybe Text, BangType)] -> Text -> TcM DataConInfo
checkedDataConInfo origin sourceForm sourceFields constructorName = do
  maybeBinder <- lookupTerm constructorName
  case maybeBinder of
    Just (TcIdBinder (ForAll tyVars predicates constructorType) _) -> do
      let (argumentTypes, resultType) = splitFunctionType constructorType
      if length sourceFields /= length argumentTypes
        then abortTc ("constructor metadata arity disagrees with checked type for " <> T.unpack constructorName)
        else do
          let (universalTyVars, existentialTyVars) = partition (`typeMentionsTyVar` resultType) tyVars
          pure
            DataConInfo
              { dciName = constructorName,
                dciOrigin = origin,
                dciUnivTyVars = universalTyVars,
                dciExTyVars = existentialTyVars,
                dciTheta = predicates,
                dciFields = zipWith checkedFieldInfo sourceFields argumentTypes,
                dciResTy = resultType,
                dciSourceForm = sourceForm
              }
    Just TcMonoIdBinder {} ->
      abortTc ("data constructor has a monomorphic binder: " <> T.unpack constructorName)
    Nothing ->
      missingTypeInfo ("data constructor " <> T.unpack constructorName)

checkedFieldInfo :: (Maybe Text, BangType) -> TcType -> DataConFieldInfo
checkedFieldInfo (label, bang) fieldType' =
  DataConFieldInfo
    { dcfiLabel = label,
      dcfiType = fieldType',
      dcfiStrict = bangStrict bang,
      dcfiLazy = bangLazy bang,
      dcfiUnpack = fieldUnpack bang
    }

fieldUnpack :: BangType -> DataConFieldUnpack
fieldUnpack bang =
  case [unpack | Pragma (PragmaUnpack unpack) _ <- bangPragmas bang] of
    UnpackPragma : _ -> UnpackField
    NoUnpackPragma : _ -> NoUnpackField
    [] -> NoFieldUnpack

dataConSourceLayout :: DataConDecl -> (DataConSourceForm, [(Maybe Text, BangType)], [Text])
dataConSourceLayout declaration =
  case declaration of
    DataConAnn _ inner -> dataConSourceLayout inner
    PrefixCon _ _ constructor fields ->
      (PrefixDataCon, map (Nothing,) fields, [unqualifiedNameText constructor])
    InfixCon _ _ left constructor right ->
      (InfixDataCon, map (Nothing,) [left, right], [unqualifiedNameText constructor])
    RecordCon _ _ constructor fields ->
      (RecordDataCon, recordSourceFields fields, [unqualifiedNameText constructor])
    TupleCon _ _ flavor fields ->
      (PrefixDataCon, map (Nothing,) fields, [tupleConText flavor (length fields)])
    UnboxedSumCon _ _ position arity field ->
      (PrefixDataCon, [(Nothing, field)], [unboxedSumConText position arity])
    ListCon {} ->
      (PrefixDataCon, [], ["[]"])
    GadtCon _ _ constructors body ->
      let names = map unqualifiedNameText constructors
       in case body of
            GadtPrefixBody fields _ -> (PrefixDataCon, map ((Nothing,) . fst) fields, names)
            GadtRecordBody fields _ -> (RecordDataCon, recordSourceFields fields, names)

recordSourceFields :: [FieldDecl] -> [(Maybe Text, BangType)]
recordSourceFields = concatMap $ \field ->
  [(Just (unqualifiedNameText label), fieldType field) | label <- fieldNames field]

-- | Type-check a declaration, returning binding results for value bindings.
tcDecl :: Decl -> TcM [TcBindingResult]
tcDecl (DeclValue vd) = tcValueDecl vd
tcDecl (DeclAnn _ inner) = tcDecl inner
tcDecl _ = pure []

-- | Type-check a value declaration.
tcValueDecl :: ValueDecl -> TcM [TcBindingResult]
tcValueDecl (FunctionBind binder matches) = do
  let name = unqualifiedNameText binder
      displayName = renderBinderName binder
  key <- resolvedUnqualifiedTermKey binder
  snd <$> tcFunctionInfer key displayName name matches
tcValueDecl (PatternBind _ pat rhs) = case patternBinderName pat of
  -- Bare variable pattern (e.g. @x = 5@, @(.>.) = (++)@): type-check as a
  -- zero-argument function so that the binding gets generalized and registered
  -- in the environment.
  Just (displayName, name) -> do
    case patternBinderSyntaxName pat of
      Just binder -> do
        key <- resolvedUnqualifiedTermKey binder
        snd <$> tcFunctionInfer key displayName name [zeroArgMatch (patternSpan pat) rhs]
      Nothing -> abortTc "a named pattern binding does not have binder syntax"
  -- Non-trivial pattern binding: infer the RHS type without generalization.
  Nothing -> do
    (_rhs', ty) <- tcRhs rhs
    zonkedTy <- zonkType ty
    pure [TcBindingResult "<pattern>" "<pattern>" zonkedTy]

-- | Extract the binder name from a pattern binding's LHS, if it is a bare
-- variable pattern.  Returns @(displayName, envName)@ for simple variable
-- patterns (possibly wrapped in parens or annotations), 'Nothing' for
-- non-trivial patterns like tuples or constructors.
patternBinderSyntaxName :: Pattern -> Maybe UnqualifiedName
patternBinderSyntaxName (PVar n) = Just n
patternBinderSyntaxName (PParen inner) = patternBinderSyntaxName inner
patternBinderSyntaxName (PAnn _ inner) = patternBinderSyntaxName inner
patternBinderSyntaxName _ = Nothing

patternBinderName :: Pattern -> Maybe (Text, Text)
patternBinderName pat =
  (\n -> (renderBinderName n, unqualifiedNameText n)) <$> patternBinderSyntaxName pat

zeroArgMatch :: SourceSpan -> Rhs Expr -> Match
zeroArgMatch sp rhs =
  Match
    { matchAnns = sourceSpanAnn sp,
      matchHeadForm = MatchHeadPrefix,
      matchPats = [],
      matchRhs = rhs
    }

sourceSpanAnn :: SourceSpan -> [Annotation]
sourceSpanAnn NoSourceSpan = []
sourceSpanAnn sp = [mkAnnotation sp]

orSourceSpan :: SourceSpan -> SourceSpan -> SourceSpan
orSourceSpan NoSourceSpan fallback = fallback
orSourceSpan sp _ = sp

patternSpan :: Pattern -> SourceSpan
patternSpan pat =
  case pat of
    PAnn ann inner ->
      fromMaybe (patternSpan inner) (fromAnnotation ann)
    PVar name -> sourceSpanFromAnns (unqualifiedNameAnns name)
    PParen inner -> patternSpan inner
    PAs name _ -> sourceSpanFromAnns (unqualifiedNameAnns name)
    PStrict inner -> patternSpan inner
    PIrrefutable inner -> patternSpan inner
    _ -> NoSourceSpan

typeSpan :: Type -> SourceSpan
typeSpan ty =
  case ty of
    TAnn ann inner ->
      fromMaybe (typeSpan inner) (fromAnnotation @SourceSpan ann)
    TParen inner -> typeSpan inner
    TForall _ inner -> typeSpan inner
    TContext _ inner -> typeSpan inner
    TKindSig inner _ -> typeSpan inner
    _ -> NoSourceSpan

rhsExprSpan :: Rhs Expr -> SourceSpan
rhsExprSpan rhs =
  case rhs of
    UnguardedRhs anns expr _ -> exprSpan expr `orSourceSpan` sourceSpanFromAnns anns
    GuardedRhss anns _ _ -> sourceSpanFromAnns anns

exprSpan :: Expr -> SourceSpan
exprSpan expr =
  case expr of
    EAnn ann inner ->
      fromMaybe (exprSpan inner) (fromAnnotation @SourceSpan ann)
    EParen inner -> exprSpan inner
    ETypeSig inner _ -> exprSpan inner
    _ -> NoSourceSpan

-- | Convert a type scheme to a displayable type.
schemeToType :: TypeScheme -> TcType
schemeToType (ForAll [] [] ty) = ty
schemeToType (ForAll tvs [] ty) = foldr TcForAllTy ty tvs
schemeToType (ForAll [] preds ty) = TcQualTy preds ty
schemeToType (ForAll tvs preds ty) = foldr TcForAllTy (TcQualTy preds ty) tvs

-- | Type-check a list of matches (equations for a function binding).
--
-- All equations must have the same number of patterns and produce
-- a consistent function type. We infer the type from each equation
-- and unify them.
tcMatches :: [Match] -> TcM ([Match], TcType, [Ct], [Implication])
tcMatches [] = do
  ty <- freshMetaTv
  pure ([], ty, [], [])
tcMatches matches@(m0 : _) = do
  let nArgs = length (matchPats m0)
  if nArgs == 0
    then do
      -- No patterns: just infer the RHS of the first match.
      (rhs0, ty0, cts0) <- inferRhsExpr (matchRhs m0)
      restResults <- mapM (unifyMatchRhs ty0) (drop 1 matches)
      let firstMatch = m0 {matchRhs = rhs0}
          restMatches = map fst restResults
          restCts = concatMap snd restResults
      pure (firstMatch : restMatches, ty0, cts0 ++ restCts, [])
    else do
      -- Create fresh meta-variables for the argument types and result type.
      argTys <- mapM (const freshMetaTv) [1 .. nArgs]
      resTy <- freshMetaTv
      -- Process each equation.
      results <- mapM (tcMatchEquation Nothing argTys resTy) matches
      let (matches', ctsList, implsList) = unzip3 results
          allCts = concat ctsList
          allImpls = concat implsList
          funTy = foldr TcFunTy resTy argTys
      pure (matches', funTy, allCts, allImpls)

-- | Type-check a single match equation against expected arg/result types.
-- Returns flat wanted constraints and implication constraints.
tcMatchEquation :: Maybe TypeOrigin -> [TcType] -> TcType -> Match -> TcM (Match, [Ct], [Implication])
tcMatchEquation expectedOrigin argTys resTy match = do
  let pats = matchPats match
      sp = sourceSpanFromAnns (matchAnns match)
  patCheck <- checkPatternsWithGivens sp (zip pats argTys)
  -- Infer the RHS under the extended environment.
  (rhs', rhsTy, rhsCts) <- withPatternBindings (pcBindings patCheck) (inferRhsExpr (matchRhs match))
  -- RHS type must match the expected result type.
  ev <- freshEvVar
  let rhsSp = rhsExprSpan (matchRhs match) `orSourceSpan` sp
      resCt =
        mkWantedEqCt
          TypeTrace
            { typeTraceType = rhsTy,
              typeTraceRole = ActualType,
              typeTraceOrigin = ExpressionTypeOrigin rhsSp
            }
          TypeTrace
            { typeTraceType = resTy,
              typeTraceRole = ExpectedType,
              typeTraceOrigin = fromMaybe (ConstraintTypeOrigin (AppOrigin sp)) expectedOrigin
            }
          ev
          (AppOrigin rhsSp)
          rhsSp
  let pats' = map (annotatePatternBindings (pcBindings patCheck)) (pcPatterns patCheck)
      givenCts = pcGivenCts patCheck
      bodyWanteds = pcWantedCts patCheck ++ rhsCts ++ [resCt]
  if null givenCts && null (pcSkolems patCheck)
    then -- No constructor-local type variables or givens: keep flat wanteds.
      pure (match {matchPats = pats', matchRhs = rhs'}, bodyWanteds, [])
    else do
      -- GADT givens: wrap body wanteds in an implication.
      level <- getTcLevel
      let impl =
            Implication
              { implSkols = pcSkolems patCheck,
                implGivenEvs = map ctEvVar givenCts,
                implGivenCts = givenCts,
                implWantedCts = bodyWanteds,
                implTcLevel = level,
                implInfo = AppOrigin sp
              }
      pure (match {matchPats = pats', matchRhs = rhs'}, [], [impl])

-- | Unify an additional match equation's RHS with the expected type.
unifyMatchRhs :: TcType -> Match -> TcM (Match, [Ct])
unifyMatchRhs expectedTy match = do
  (rhs', rhsTy, rhsCts) <- inferRhsExpr (matchRhs match)
  ev <- freshEvVar
  let sp = sourceSpanFromAnns (matchAnns match)
      rhsSp = rhsExprSpan (matchRhs match) `orSourceSpan` sp
      eqCt =
        mkWantedEqCt
          TypeTrace
            { typeTraceType = rhsTy,
              typeTraceRole = ActualType,
              typeTraceOrigin = ExpressionTypeOrigin rhsSp
            }
          TypeTrace
            { typeTraceType = expectedTy,
              typeTraceRole = ExpectedType,
              typeTraceOrigin = ConstraintTypeOrigin (AppOrigin sp)
            }
          ev
          (AppOrigin rhsSp)
          rhsSp
  pure (match {matchRhs = rhs'}, rhsCts ++ [eqCt])

-- | Infer the type of a right-hand side expression.
inferRhsExpr :: Rhs Expr -> TcM (Rhs Expr, TcType, [Ct])
inferRhsExpr = inferRhsWithLocals inferExpr

-- | Type-check a right-hand side (solving constraints immediately).
tcRhs :: Rhs Expr -> TcM (Rhs Expr, TcType)
tcRhs rhs = do
  (rhs', ty, cts) <- inferRhsWithLocals inferExpr rhs
  _ <- solveConstraints cts
  pure (rhs', ty)

-- | Render an unqualified name for display.
-- Operators (NameVarSym, NameConSym) are wrapped in parentheses.
renderBinderName :: UnqualifiedName -> Text
renderBinderName uname =
  case unqualifiedNameType uname of
    NameVarSym -> "(" <> unqualifiedNameText uname <> ")"
    NameConSym -> "(" <> unqualifiedNameText uname <> ")"
    _ -> unqualifiedNameText uname
