kind inference with rows/variants

kinds.hs

#!/usr/bin/env cabal
{- cabal:
build-depends: base, mtl, containers, uniplate
ghc-options: -Wall
-}


-- | An example of a kind inference for data types using
-- unification-based constraint solving.
--
-- See the blog post:
-- <https://gilmi.me/blog/post/2023/09/30/kind-inference>

{-# Language GHC2021 #-}
{-# Language LambdaCase #-}

import Data.Data (Data)
import Data.List (intercalate)
import GHC.Generics (Generic)
import Data.Tuple (swap)
import Data.Maybe (listToMaybe)
import Data.Foldable (for_)
import Data.Traversable (for)
import Control.Monad (foldM)
import Control.Monad.State qualified as Mtl
import Control.Monad.Except qualified as Mtl
import Data.Map qualified as Map

-- * Tests

main :: IO ()
main = do
  putStrLn "========================"
  putStrLn "-- Expected successes --"
  for_ expectedSuccesses inferAndPrint

  putStrLn "========================"
  putStrLn "-- Expected failures  --"
  for_ expectedFailures inferAndPrint
  putStrLn "========================"

expectedSuccesses :: [[Datatype ()]]
expectedSuccesses =
  [ -- We start with a simple data type of kind 'Type'.
    [ Datatype () (MkTypeName "Person")
      []
      [ Variant "Person"
        [ TypeName () $ MkTypeName "String"
        , TypeName () $ MkTypeName "Int"
        ]
      ]
    ]

  , -- A data type with a type variable.
    [ Datatype () (MkTypeName "Option")
      [MkTypeVar "a"]
      [ Variant "Some" [TypeVar () $ MkTypeVar "a"]
      , Variant "None" []
      ]
    ]

  , -- A data type with two type variables.
    [ Datatype () (MkTypeName "Result")
      [MkTypeVar "a", MkTypeVar "b"]
      [ Variant "Ok" [TypeVar () $ MkTypeVar "a"]
      , Variant "Err" [TypeVar () $ MkTypeVar "b"]
      ]
    ]

  , -- A more complex, recursive data type with higher-kinded type variables.
    [ Datatype () (MkTypeName "Cofree")
      [MkTypeVar "f", MkTypeVar "a"]
      [ Variant "Cofree"
        [ TypeVar () $ MkTypeVar "a"
        , TypeApp ()
          (TypeVar () $ MkTypeVar "f")
          ( TypeApp ()
            ( TypeApp ()
              (TypeName () $ MkTypeName "Cofree")
              (TypeVar () $ MkTypeVar "f")
            )
            (TypeVar () $ MkTypeVar "a")
          )
        ]
      ]
    ]

  , -- A data type where the type variable does not appear in the variants.
    [ Datatype () (MkTypeName "Phantom")
      [MkTypeVar "a"]
      [ Variant "A" [TypeName () $ MkTypeName "Int"]
      , Variant "B" [TypeApp () (TypeName () $ MkTypeName "List") (TypeName () $ MkTypeName "Bool")]
      ]
    ]

  , -- Mutually recursive data types.
    [ Datatype () (MkTypeName "Select")
      [MkTypeVar "ann", MkTypeVar "lit"]
      [ Variant "Select"
        [ TypeVar () $ MkTypeVar "ann"
        , TypeApp ()
          ( TypeApp ()
            (TypeName () $ MkTypeName "Expr")
            (TypeVar () $ MkTypeVar "ann")
          )
          (TypeVar () $ MkTypeVar "lit")
        ]
      ]
    , Datatype () (MkTypeName "Expr")
      [MkTypeVar "ann", MkTypeVar "lit"]
      [ Variant "Subselect"
        [ TypeApp ()
          ( TypeApp ()
            (TypeName () $ MkTypeName "Select")
            (TypeVar () $ MkTypeVar "ann")
          )
          (TypeVar () $ MkTypeVar "lit")
        ]
      , Variant "Lit" [TypeApp () (TypeVar () $ MkTypeVar "lit") (TypeVar () $ MkTypeVar "ann")]
      ]
    ]

  , -- Empty data types with proxy types used more than once.
    [ Datatype () (MkTypeName "Person") [] []
    , Datatype () (MkTypeName "Company") [MkTypeVar "taxId"] [Variant "TaxId" [TypeVar () $ MkTypeVar "taxId"]]
    , Datatype () (MkTypeName "ObjectId")
      []
      [ Variant "Person"
        [ TypeApp ()
          (TypeName () $ MkTypeName "Id")
          (TypeName () $ MkTypeName "Person")
        ]
      , Variant "Company"
        [ TypeApp ()
          (TypeName () $ MkTypeName "Id")
          (TypeName () $ MkTypeName "Company")
        ]
      ]
    ]

  , -- Record.
    [ Datatype () (MkTypeName "Person")
      [MkTypeVar "p"]
      [ Variant "Person"
        [ TypeRec () (Just $ MkTypeVar "p")
          [ ("name", TypeName () (MkTypeName "String"))
          , ("age", TypeName () (MkTypeName "Int"))
          ]
        ]
      ]
    ]
  ]

expectedFailures :: [[Datatype ()]]
expectedFailures =
  [ -- Type variable 'a' is used as a 'Type' and a higher-kinded type.
    [ Datatype () (MkTypeName "Fail1")
      [MkTypeVar "a"]
      [ Variant "A" [TypeVar () $ MkTypeVar "a"]
      , Variant "B" [TypeApp () (TypeVar () $ MkTypeVar "a") (TypeName () $ MkTypeName "Bool")]
      ]
    ]

  , -- Data type 'Fail2' is higher-kinded but does not take type parameters
    -- in the variant 'A'.
    [ Datatype () (MkTypeName "Fail2")
      [MkTypeVar "a"]
      [ Variant "A" [TypeName () $ MkTypeName "Fail2"]
      ]
    ]

  , -- contains a type variable with two different kinds.
    [ Datatype () (MkTypeName "Fail3")
      [MkTypeVar "f", MkTypeVar "a"]
      [ Variant "A"
        [ TypeApp ()
          (TypeVar () $ MkTypeVar "f")
          ( TypeApp ()
            (TypeVar () $ MkTypeVar "a")
            (TypeVar () $ MkTypeVar "a")
          )
        ]
      ]
    ]

  , -- Data type usage is not fully saturated.
    [ Datatype () (MkTypeName "Tree")
      [MkTypeVar "a"]
      [ Variant "Node"
        [ TypeVar () $ MkTypeVar "a"
        , TypeName () $ MkTypeName "Tree"
        , TypeName () $ MkTypeName "Tree"
        ]
      ]
    ]
  ]

-- | Infer a group of data types and print the result.
inferAndPrint :: [Datatype ()] -> IO ()
inferAndPrint datatypes = do
  putStrLn "========================"
  putStrLn ""
  for_ datatypes $ putStrLn . ppDatatype
  putStrLn "------------------------"
  putStrLn ""
  let result = infer builtinKindEnv datatypes
  case result of
    Left err -> do
      putStrLn (ppError err)
    Right datatypes' ->
      for_ datatypes' $ putStrLn . ppDatatype'

-- | A environment of kinds for built-in types.
builtinKindEnv :: Map.Map TypeName Kind
builtinKindEnv = Map.fromList
  [ (MkTypeName "Int", Type)
  , (MkTypeName "Bool", Type)
  , (MkTypeName "String", Type)
  , (MkTypeName "List", KindFun Type Type)
  , ( MkTypeName "Id" -- Something like: `Id object = Id Int`.
    , KindScheme [MkKindVar "k1"] $ KindFun (KindVar (MkKindVar "k1")) Type
    )
  ]

-- * Data types

-- | The representation of a data type definition.
data Datatype a
  = Datatype
    { -- | A place to put kind annotation in.
      dtAnn :: a
    , -- | The name of the data type.
      dtName :: TypeName
    , -- | Type parameters.
      dtParameters :: [TypeVar]
    , -- | Alternative variants.
      dtVariants :: [Variant a]
    }
  deriving (Show, Eq, Data, Generic, Functor, Foldable, Traversable)

-- | A Variant of a data type definition.
data Variant a
  = Variant
    { -- | A type constructor.
      vTypeConstructor :: String
    , -- | A list of types.
      vTypes :: [Type a]
    }
  deriving (Show, Eq, Data, Generic, Functor, Foldable, Traversable)

-- | A name of known types.
newtype TypeName = MkTypeName { getTypeName :: String }
  deriving (Show, Eq, Ord, Data, Generic)

-- | A type variable.
newtype TypeVar = MkTypeVar { getTypeVar :: String }
  deriving (Show, Eq, Ord, Data, Generic)

-- | A representation of a type with a place for kind annotation.
data Type a
  = -- | A type variable.
    TypeVar a TypeVar
  | -- | A named type.
    TypeName a TypeName
  | -- | An application of two types, of the form `t1 t2`.
    TypeApp a (Type a) (Type a)
  | -- | A potentially polymorphic record.
    TypeRec a (Maybe TypeVar) [(String, Type a)]
  | -- | A closed variant.
    TypeVariant a [(String, Type a)]
  | -- | A lower-bound polymorphic variant.
    TypeVariantLB a [(String, Type a)] TypeVar
  | -- | An upper-bound polymorphic variant.
    TypeVariantUB a TypeVar [(String, Type a)]
  deriving (Show, Eq, Data, Generic, Functor, Foldable, Traversable)

-- | A representation of a kind.
data Kind
  = -- | For types like `Int`.
    Type
  | -- | For rows like `{ a :: Int, b :: String }`.
    Row
  | -- | For variants like `[ A : Int | B : String ]`.
    Variants
  | -- | For types like `Option`.
    KindFun Kind Kind
  | -- | For polymorphic kinds.
    KindVar KindVar
  | -- | For closing over polymorphic kinds.
    KindScheme [KindVar] Kind
  deriving (Show, Eq, Data, Generic)

-- | A kind variable.
newtype KindVar = MkKindVar { getKindVar :: String }
  deriving (Show, Eq, Ord, Data, Generic)

-- * Algorithm

-- ** Inference

-- | Infer the kind of a group of data types that should be solved together
--   (because they are mutually recursive).
infer :: Map.Map TypeName Kind -> [Datatype ()] -> Either Error [Datatype Kind]
infer kindEnv datatypes =
  -- initialize our `InferenceM` which is State + Except
  flip Mtl.evalState (initialState kindEnv) $ Mtl.runExceptT $ do
    -- Invent a kind variable for each data type
    for_ datatypes $ \(Datatype _ name _ _) -> do
      kindvar <- freshKindVar
      declareNamedType name kindvar
    -- Elaborate all of the data types
    datatypes' <- traverse elaborate datatypes
    -- Solve the constraints
    solveConstraints
    for datatypes' $ \(Datatype kindvar name vars variants) -> do
      -- Substitute the kind variable for a kind
      -- for the data type
      kind <- lookupKindVarInSubstitution kindvar
      -- ... and for all types
      variants' <- for variants $ traverse lookupKindVarInSubstitution
      -- generalize the data type's kind, and return.
      pure (Datatype (generalize kind) name vars variants')

-- ** Elaboration

-- | Invent kind variables for types we don't know and add constraints
--   on them according to their usage.
elaborate :: Datatype () -> InferenceM (Datatype KindVar)
elaborate (Datatype _ datatypeName vars variants) = do
  -- We go over each of the data type parameters and
  -- generate a fresh kind variable for them.
  varKinds <- for vars $ \var -> do
    kindvar <- freshKindVar
    declareTypeVar var kindvar
    pure kindvar

  -- We go over the variants, elaborate each field,
  -- and return the elaborated variants.
  variants' <- for variants $ \(Variant name fields) -> do
    Variant name <$>
      for fields
        ( \field -> do
          field' <- elaborateType field
          -- a constraint on fields: their kind must be `Type`.
          newEqualityConstraint (KindVar $ getAnn field') Type
          pure field'
        )

  -- We grab the kind variable of the data type
  -- so we can add a constraint on it.
  datatypeKindvar <- lookupNameKindVar datatypeName
  -- A type of the form `T a b c ... =` has the kind:
  -- `aKind -> bKind -> cKind -> ... -> Type`.
  -- We add that as a constraint.
  let kind = foldr KindFun Type $ map KindVar varKinds
  newEqualityConstraint (KindVar datatypeKindvar) kind

  -- We return the elaborated data type after annotating
  -- all types with kind variables and generating constraints.
  pure (Datatype datatypeKindvar datatypeName vars variants')

-- | Elaborate a type with a kind variable and add constraints
--   according to usage.
elaborateType :: Type () -> InferenceM (Type KindVar)
elaborateType = \case
  -- for type variables and type names,
  -- we lookup the kind variables we generated when we ran into
  -- the declaration of them.
  TypeVar () var ->
    fmap (\kindvar -> TypeVar kindvar var) (lookupVarKindVar var)

  TypeName () name ->
    fmap (\kindvar -> TypeName kindvar name) (lookupNameKindVar name)

  TypeRec () var types -> do
    types' <- traverse (traverse elaborateType) types
    typeRecKindvar <- freshKindVar
    newEqualityConstraint (KindVar typeRecKindvar) Type
    for_ var $ \typevar -> do
      kindvar <- lookupVarKindVar typevar
      newEqualityConstraint (KindVar kindvar) Row
    pure $ TypeRec typeRecKindvar var types'

  TypeVariant () types -> do
    types' <- traverse (traverse elaborateType) types
    typeVariantKindvar <- freshKindVar
    newEqualityConstraint (KindVar typeVariantKindvar) Type
    pure $ TypeVariant typeVariantKindvar types'

  TypeVariantLB () types var -> do
    types' <- traverse (traverse elaborateType) types
    typeVariantKindvar <- freshKindVar
    newEqualityConstraint (KindVar typeVariantKindvar) Type
    kindvar <- lookupVarKindVar var
    newEqualityConstraint (KindVar kindvar) Variants
    pure $ TypeVariantLB typeVariantKindvar types' var

  TypeVariantUB () var types -> do
    types' <- traverse (traverse elaborateType) types
    typeVariantKindvar <- freshKindVar
    newEqualityConstraint (KindVar typeVariantKindvar) Type
    kindvar <- lookupVarKindVar var
    newEqualityConstraint (KindVar kindvar) Variants
    pure $ TypeVariantUB typeVariantKindvar var types'

  -- for type application
  TypeApp () t1 t2 -> do
    -- we elaborate both types
    t1Kindvar <- elaborateType t1
    t2Kindvar <- elaborateType t2
    -- then we generate a kind variable for the type application
    typeAppKindvar <- freshKindVar
    -- then we constrain the type application kind variable
    -- it should unify with `t2Kind -> typeAppKind`.
    newEqualityConstraint
      (KindVar $ getAnn t1Kindvar)
      (KindFun (KindVar $ getAnn t2Kindvar) (KindVar typeAppKindvar))

    pure (TypeApp typeAppKindvar t1Kindvar t2Kindvar)

-- ** Constraint solving and generating a substitution.

-- | Solve constraints according to logic.
--   this process is iterative. We continue fetching
--   the next constraint and try to solve it.
--
--   Each step can either reduce or increase the number of constraints,
--   and we are done when there are no more constraints to solve,
--   or if we ran into a constraint that cannot be solved.
solveConstraints :: InferenceM ()
solveConstraints = do
  -- Pop the next constraint we should solve.
  constraint <- do
    c <- listToMaybe . constraints <$> Mtl.get
    Mtl.modify $ \s -> s { constraints = drop 1 $ constraints s }
    pure c

  case constraint of
    -- If we have two 'Type's, the unify. We can skip to the next constraint.
    Just (Equality Type Type) -> solveConstraints
    -- We have an equality between two kind functions.
    -- We add two new equality constraints matching the two firsts
    -- with the two seconds.
    Just (Equality (KindFun k1 k2) (KindFun k3 k4)) -> do
      Mtl.modify $ \s ->
        s { constraints = Equality k1 k3 : Equality k2 k4 : constraints s }
      solveConstraints

    Just (Equality Row Row) -> solveConstraints
    -- When we run into a kind scheme, we instantiate it
    -- (we look at the kind and replace all closed kind variables
    -- with fresh kind variables), and add an equality constraint
    -- between the other kind and the instantiated kind.
    Just (Equality (KindScheme vars kind) k) -> do
      kind' <- instantiate kind vars
      Mtl.modify $ \s ->
        s { constraints = Equality k kind' : constraints s }
      solveConstraints
    -- Same as the previous scenario.
    Just (Equality k (KindScheme vars kind)) -> do
      kind' <- instantiate kind vars
      Mtl.modify $ \s ->
        s { constraints = Equality kind' k : constraints s }
      solveConstraints
    -- If we run into a kind variable on one of the sides,
    -- we replace all instances of it with the other kind and continue.
    Just (Equality (KindVar var) k) -> do
      replaceInState var k
      solveConstraints
    -- The same as the previous scenario, but the kind var is on the other side.
    Just (Equality k (KindVar var)) -> do
      replaceInState var k
      solveConstraints
    -- If we have an equality constraint between a 'Type' and
    -- a 'KindFun', we cannot unify the two, and unification fails.
    Just (Equality k1 k2) -> Mtl.throwError (UnificationFailed k1 k2)
    -- If there are no more constraints, we are done. Good job!
    Nothing -> pure ()

-- | Instantiate a kind.
--   We look at the kind and replace all closed kind variables
--   with fresh kind variables.
instantiate :: Kind -> [KindVar] -> InferenceM Kind
instantiate = foldM replaceKindVarWithFreshKindVar

-- | Replace a kind variable with a fresh variable in the kind.
replaceKindVarWithFreshKindVar :: Kind -> KindVar -> InferenceM Kind
replaceKindVarWithFreshKindVar kind var = do
  kindvar <- freshKindVar
  -- Uniplate.transformBi lets us perform reflection and
  -- apply a function to all instances of a certain type
  -- in a value. Think of it like `fmap`, but for any type.
  --
  -- It is a bit slow though, so it's worth replacing it with
  -- hand rolled recursion or a functor, but its convenient.
  pure $ flip transformKindVars kind $ \case
    KindVar kv | kv == var -> KindVar kindvar
    x -> x

-- | Replace every instance of 'KindVar var' in our state with 'kind'.
--   And add it to the substitution.
replaceInState :: KindVar -> Kind -> InferenceM ()
replaceInState var kind = do
  occursCheck var kind
  s <- Mtl.get
  let
    -- Uniplate.transformBi lets us perform reflection and
    -- apply a function to all instances of a certain type
    -- in a value. Think of it like `fmap`, but for any type.
    --
    -- Note that we are changing all instances of `Kind` of the form
    -- `KindVar v | v == var` in all of `State`! This includes both the
    -- `substitution` and the remaining `constraints`,
    --
    -- It is a bit slow though, so it's worth replacing it with
    -- hand rolled recursion or a functor, but its convenient.
    transformKV = \case
      KindVar v | v == var -> kind
      x -> x
    transformC (Equality k1 k2) =
      Equality (transformKindVars transformKV k1) (transformKindVars transformKV k2)
    s' = s
      { constraints = map transformC (constraints s)
      , substitution =
        Map.insert var kind $ Map.map (transformKindVars transformKV) (substitution s)
      }
  Mtl.put s'

-- | We check that the kind variable does not appear in the kind
--   and throw an error if it does.
occursCheck :: KindVar -> Kind -> InferenceM ()
occursCheck var kind =
  if KindVar var == kind || null [ () | kv <- getKindVars kind, kv == var ]
    then pure ()
    else do
      -- We try to find the type of the kind variable by doing reverse lookup,
      -- but this might not succeed before the kind variable might be generated
      -- during constraint solving.
      -- We might be able to find the type if we look at the substitution as well,
      -- but for now lets leave it at this "best effort" attempt.
      reverseEnv <- map swap . Map.toList . env <$> Mtl.get
      let typ = either (TypeName ()) (TypeVar ()) <$> lookup var reverseEnv
      Mtl.throwError (OccursCheckFailed typ var kind)

-- ** Generalization

-- | Close over kind variables we did not solve.
generalize :: Kind -> Kind
generalize kind = KindScheme (getKindVars kind) kind

-----------------------------------
-- * Utilities for working State

-- ** Types

-- | We combine the capabilities of Except and State
--   For our kind inference code.
type InferenceM a = Mtl.ExceptT Error (Mtl.State State) a

-- | The errors that can be thrown in the process.
data Error
  = UnboundVar TypeVar
  | UnboundName TypeName
  | UnificationFailed Kind Kind
  | OccursCheckFailed (Maybe (Type ())) KindVar Kind
  deriving (Show)

-- | The state we keep during an inference cycle
data State = State
  { -- | Mapping from named types or type variables to kind variables.
    -- When we declare a new data type or a type variable, we'll add it here.
    -- When run into a type variable or a type name during elaboration,
    -- we search its matching kind here.
    env :: Map.Map (Either TypeName TypeVar) KindVar
  , -- | Mapping from existing named types to their kinds.
    -- Kinds for types that are supplied before the inference process can be found here.
    kindEnv :: Map.Map TypeName Kind
  , -- | Used for generating fresh kind variables.
    counter :: Int
  , -- | When we learn information about kinds during elaboration, we'll add it here.
    constraints :: [Constraint]
  , -- | The constraint solving process will generate this mapping from
    -- the kind variables we collected to the kind they should represent.
    -- If we don't find the kind variable in the substitution, that means
    -- it is a free variable we should close over.
    substitution :: Map.Map KindVar Kind
  }
  deriving (Show, Eq, Data, Generic)

-- | A constraint on kinds.
data Constraint
  = Equality Kind Kind
    -- ^ The two kinds should unify.
    -- If one of the kinds is a kind scheme, we will instantiate it, and
    -- add an equality constraint of the other kind with the instantiated kind.
  deriving (Show, Eq, Data, Generic)

-- ** Utilities

-- | The state at the start of the process.
initialState :: Map.Map TypeName Kind -> State
initialState kindEnv =
  State mempty kindEnv 0 mempty mempty

-- | Generate a fresh kind variables.
freshKindVar :: InferenceM KindVar
freshKindVar = do
  s <- Mtl.get
  let kindvar = MkKindVar ("k" <> show (counter s))
  Mtl.put s { counter = 1 + counter s }
  pure kindvar

-- | Insert declared type names into the environment.
declareNamedType :: TypeName -> KindVar -> InferenceM ()
declareNamedType name kindvar =
  Mtl.modify $ \s ->
    s { env = Map.insert (Left name) kindvar (env s) }

-- | Insert declared type variables into the environment.
declareTypeVar :: TypeVar -> KindVar -> InferenceM ()
declareTypeVar var kindvar =
  Mtl.modify $ \s ->
    s { env = Map.insert (Right var) kindvar (env s) }

-- | Add a new equality constraint to the state.
newEqualityConstraint :: Kind -> Kind -> InferenceM ()
newEqualityConstraint k1 k2 =
  Mtl.modify $ \s ->
    s { constraints = Equality k1 k2 : constraints s }

-- | Find the kind variable of a named type in the environment.
lookupNameKindVar :: TypeName -> InferenceM KindVar
lookupNameKindVar name = do
  state <- Mtl.get
  -- We first look the named type in the supplied kind env.
  case Map.lookup name (kindEnv state) of
    -- If we find it, we generate a new kind variable for it
    -- and constraint it to be this type, so that each use has it own
    -- type variable (later used for instantiation).
    Just kind -> do
      kindvar <- freshKindVar
      newEqualityConstraint (KindVar kindvar) kind
      pure kindvar
    -- If it's not a supplied type, it means we are actively inferring it,
    -- and we need to use the same kind variable for all uses.
    -- We'll look it up in our environment of declared types.
    Nothing ->
      maybe
        -- If we still can't find it, we error.
        (Mtl.throwError $ UnboundName name)
        pure
        . Map.lookup (Left name)
        $ env state

-- | Find the kind variable of a type variable in the environment.
lookupVarKindVar :: TypeVar -> InferenceM KindVar
lookupVarKindVar var =
  maybe
    (Mtl.throwError $ UnboundVar var)
    pure
    . Map.lookup (Right var)
    . env =<< Mtl.get

-- | Look up what the kind of a kind variable is in the substitution
--   produced by constraint solving.
--   If there was no constraint on the kind variable, it won't appear
--   in the substitution, which means it can stay a kind variable which
--   we will close over later.
lookupKindVarInSubstitution :: KindVar -> InferenceM Kind
lookupKindVarInSubstitution kindvar =
  maybe (KindVar kindvar) id . Map.lookup kindvar . substitution <$> Mtl.get

getKindVars :: Kind -> [KindVar]
getKindVars = \case
  Type -> []
  KindVar v -> [v]
  KindFun k1 k2 ->
    getKindVars k1 <> getKindVars k2
  KindScheme vars _ ->
    vars
  Row -> []
  Variants -> []

transformKindVars :: (Kind -> Kind) -> Kind -> Kind
transformKindVars f = \case
  Type -> Type
  Row -> Row
  Variants -> Variants
  KindVar v -> f (KindVar v)
  KindFun k1 k2 ->
    KindFun (transformKindVars f k1) (transformKindVars f k2)
  KindScheme vars kind -> KindScheme vars (transformKindVars f kind)


-- | Get the annotation of a type.
getAnn :: Type a -> a
getAnn = \case
  TypeVar a _ -> a
  TypeRec a _ _ -> a
  TypeName a _ -> a
  TypeApp a _ _ -> a
  TypeVariant a _ -> a
  TypeVariantLB a _ _ -> a
  TypeVariantUB a _ _ -> a

-- * Prettyprinting

ppDatatype :: Datatype a -> String
ppDatatype (Datatype _ name vars variants) =
  let
    parameters
      | null vars = ""
      | otherwise = " " <> unwords (map getTypeVar vars)
  in
    getTypeName name <> parameters <> " =\n" <> unlines (map ppVariant variants)

ppDatatype' :: Datatype Kind -> String
ppDatatype' (Datatype kind name vars variants) =
  getTypeName name <> " : " <> ppKind False kind <> "\n" <>
  getTypeName name <> " " <> unwords (map getTypeVar vars) <> " =\n" <> unlines (map ppVariant variants)

ppVariant :: Variant a -> String
ppVariant (Variant con types) =
  "  | " <> con <> " " <> unwords (map (ppType True) types)

ppType :: Bool -> Type a -> String
ppType parens = \case
  TypeVar _ var -> getTypeVar var
  TypeRec _ var types ->
    let
      var' = case var of
        Nothing -> ""
        Just v -> " | " <> getTypeVar v
    in
      "{ " <> intercalate ", " (map (\(k, v) -> k <> ": " <> ppType False v) types) <> var' <> " }"
  TypeVariant _ types ->
      "[ " <> intercalate " | " (map (\(k, v) -> k <> ": " <> ppType False v) types) <> " ]"
  TypeVariantLB _ types _ ->
      "[> " <> intercalate " | " (map (\(k, v) -> k <> ": " <> ppType False v) types) <> " ]"
  TypeVariantUB _ _ types ->
      "[< " <> intercalate " | " (map (\(k, v) -> k <> ": " <> ppType False v) types) <> " ]"
  TypeName _ name -> getTypeName name
  TypeApp _ t1 t2 ->
    let result = ppType False t1 <> " " <> ppType True t2
    in if parens then "(" <> result <> ")" else result

ppKind :: Bool -> Kind -> String
ppKind parens = \case
  Type -> "Type"
  Row -> "Row"
  Variants -> "Variants"
  KindVar v -> getKindVar v
  KindFun k1 k2 ->
    let str = ppKind True k1 <> " -> " <> ppKind False k2
    in if parens then "(" <> str <> ")" else str
  KindScheme vars kind
    | null vars -> ppKind parens kind
    | otherwise ->
      let str = "forall " <> unwords (map getKindVar vars) <> ". " <> ppKind False kind
      in if parens then "(" <> str <> ")" else str

ppError :: Error -> String
ppError = \case
  UnboundVar var ->
    "Unbound type variable: '" <> getTypeVar var <> "'.\n"
  UnboundName name ->
    "Unbound type variable: '" <> getTypeName name <> "'.\n"
  UnificationFailed k1 k2 ->
    unlines
      [ "Unification failed between the following kinds:"
      , "  * " <> ppKind False k1
      , "  * " <> ppKind False k2
      ]
  OccursCheckFailed mtype kindvar kind ->
    unlines
      [ "Occurs check failed" <>
        ( case mtype of
          Nothing -> ""
          Just typ -> " for type '" <> ppType False typ <> "'"
        ) <> ":"
      , "  * " <> getKindVar kindvar
      , "  * " <> ppKind False kind
      ]