------------------------------------------------------------------------------
-- | Author:  Michael Hanus
--   Version: November 2025
--
-- This module contains operations to add type information to the body
-- of each function defined in a FlatCurry program.
-- The type information (defined as `TypeInfo`) consists of
-- a list of all typed variables visible in the expression and the result type
-- of the expression.
------------------------------------------------------------------------------

module FlatCurry.AddTypes (
  -- * Annotated programs and expressions
  AProg(..), AFuncDecl(..), ARule(..), AExpr(..), ABranchExpr(..),
  annOfAExpr, mapAExpr, fromAProg, fromAFuncDecl, fromAExpr,
  -- * Type annotations
  TypeInfo, typeInfo, tiTypedVars, tiType,
  -- * Transforming type-annotated programs
  annProg2Prog, annFunc2Func,
  -- * Operations to add type annotations
  addTypesInProg, addTypesInProgWithImports, addTypesInFunc, addTypesInBody
  ) where

import Data.List ( init, isPrefixOf, last, maximum )

import Control.Monad.Trans.State
import qualified Data.Map as Map
import FlatCurry.Files   ( readFlatCurryInt )
import FlatCurry.Goodies
import FlatCurry.Types

------------------------------------------------------------------------------
--- Annotated FlatCurry program.
data AProg a = AProg String [String] [TypeDecl] [AFuncDecl a] [OpDecl]
 deriving (Eq, Show)

--- Annotated function declaration.
data AFuncDecl a = AFunc QName Arity Visibility TypeExpr (ARule a)
 deriving (Eq, Show)

--- Annotated function rule where only the right-hand side contains annotations.
data ARule a
  = ARule     [VarIndex] (AExpr a)
  | AExternal String
 deriving (Eq, Show)

-- | Annotated FlatCurry expressions.
-- In contrast to the type `AExpr` defined in `FlatCurry.Annotated.Types`,
-- this type does not contain annotations for variables in let and free
-- declarations (since they already contain type information) and patterns.
data AExpr a
  = AVar   a VarIndex
  | ALit   a Literal
  | AComb  a CombType QName [AExpr a]
  | ALet   a [(VarIndex, TypeExpr, AExpr a)] (AExpr a)
  | AFree  a [(VarIndex, TypeExpr)] (AExpr a)
  | AOr    a (AExpr a) (AExpr a)
  | ACase  a CaseType (AExpr a) [ABranchExpr a]
  | ATyped a (AExpr a) TypeExpr
 deriving (Eq, Show)

-- | Type of annotated FlatCurry case branches.
data ABranchExpr a = ABranch Pattern (AExpr a)
 deriving (Eq, Show)

-- | Gets annotation of annotated FlatCurry expression.
annOfAExpr :: AExpr a -> a
annOfAExpr (AVar   a _)     = a
annOfAExpr (ALit   a _)     = a
annOfAExpr (AComb  a _ _ _) = a
annOfAExpr (ALet   a _ _)   = a
annOfAExpr (AFree  a _ _)    = a
annOfAExpr (AOr    a _ _)   = a
annOfAExpr (ACase  a _ _ _) = a
annOfAExpr (ATyped a _ _)   = a

-- | Transforms all annotations of an expression. by applying a function
--   on them.
mapAExpr :: (a -> b) -> AExpr a -> AExpr b
mapAExpr f (AVar   a v) = AVar (f a) v
mapAExpr f (ALit   a l) = ALit (f a) l
mapAExpr f (AComb  a ct qn es) = AComb (f a) ct qn (map (mapAExpr f) es)
mapAExpr f (ALet   a bs e) = ALet (f a)
                                  (map (\ (v,vt,be) -> (v,vt,mapAExpr f be)) bs)
                                  (mapAExpr f e)
mapAExpr f (AFree  a vts e) = AFree (f a) vts (mapAExpr f e)
mapAExpr f (AOr    a e1 e2) = AOr (f a) (mapAExpr f e1) (mapAExpr f e2)
mapAExpr f (ACase  a ct ce brs) =
  ACase (f a) ct (mapAExpr f ce)
        (map (\ (ABranch p be) -> ABranch p (mapAExpr f be)) brs)
mapAExpr f (ATyped a e te) = ATyped (f a) (mapAExpr f e) te

------------------------------------------------------------------------------

-- | Strip annotations from the body of all annotated functions in a program.
fromAProg :: AProg _ -> Prog
fromAProg (AProg mn imps tds fds ops) =
  Prog mn imps tds (map fromAFuncDecl fds) ops

-- | Strip annotations from the body of an annotated function definition.
fromAFuncDecl :: AFuncDecl _ -> FuncDecl
fromAFuncDecl (AFunc qn ar vis te rl) = Func qn ar vis te (fromAR rl)
 where fromAR (ARule vs ae) = Rule vs (fromAExpr ae)
       fromAR (AExternal s) = External s

-- | Strip annotations from an annotated FlatCurry expression.
fromAExpr :: AExpr _ -> Expr
fromAExpr (AVar   _ v) = Var v
fromAExpr (ALit   _ l) = Lit l
fromAExpr (AComb  _ ct qn es) = Comb ct qn (map fromAExpr es)
fromAExpr (ALet   _ bs e) = Let (map (\ (v,vt,be) -> (v,vt,fromAExpr be)) bs)
                                (fromAExpr e)
fromAExpr (AFree  _ vts e) = Free vts (fromAExpr e)
fromAExpr (AOr    _ e1 e2) = Or (fromAExpr e1) (fromAExpr e2)
fromAExpr (ACase  _ ct ce brs) =
  Case ct (fromAExpr ce)
       (map (\ (ABranch p be) -> Branch p (fromAExpr be)) brs)
fromAExpr (ATyped _ e te) = Typed (fromAExpr e) te

------------------------------------------------------------------------------
-- | The type annotation of an expression consists of a list of
--   typed variables (visible in the current expression) and
--   the result type of the expression.
data TypeInfo = TypeInfo [(VarIndex,TypeExpr)] TypeExpr
 deriving (Eq, Show)

-- | Construct a type annotation.
typeInfo :: [(VarIndex,TypeExpr)] -> TypeExpr -> TypeInfo
typeInfo = TypeInfo

-- | The typed variables comöponent of a type annotation.
tiTypedVars :: TypeInfo -> [(VarIndex,TypeExpr)]
tiTypedVars (TypeInfo tvs _) = tvs

-- | The type comöponent of a type annotation.
tiType :: TypeInfo -> TypeExpr
tiType (TypeInfo _ te) = te

-- Applies a type substitution to a type annotation.
applyTS2TypeInfo :: TSub -> TypeInfo -> TypeInfo
applyTS2TypeInfo ts (TypeInfo tvars texp) =
  TypeInfo (map (\ (v,tv) -> (v, applyTS ts tv)) tvars) (applyTS ts texp)

------------------------------------------------------------------------------
-- | Transforms a type-annotated program into a FlatCurry program
--   by applying the transformation (first argument) to each
--   type-annotated body which processes and eliminates type annotations.
annProg2Prog ::
     (AExpr TypeInfo -> Expr) -- ^ transformation on annotated expressions
  -> AProg TypeInfo           -- ^ type-annotated program to be transformed
  -> Prog                      -- ^ resulting FlatCurry program
annProg2Prog fromae (AProg mname imps tdecls fdecls ops) =
  Prog mname imps tdecls (map (annFunc2Func fromae) fdecls) ops

-- | Transforms a type-annotated function declaration into a FlatCurry
--   function declaration by applying the transformation (first argument) to
--   each type-annotated body which processes and eliminates type annotations.
annFunc2Func ::
     (AExpr TypeInfo -> Expr) -- ^ transformation on annotated expressions
  -> AFuncDecl TypeInfo       -- ^ type-annotated function to be transformed
  -> FuncDecl                  -- ^ resulting FlatCurry function
annFunc2Func fromae (AFunc qn ar vis te rule) =
  Func qn ar vis te
       (case rule of ARule vs rhs -> Rule vs (fromae rhs)
                     AExternal s  -> External s)

------------------------------------------------------------------------------
-- | Transforms a program by adding type annotations to the each function's
--   body after reading the interfaces of all imported modules.
addTypesInProg ::
     Prog                -- ^ program to be transformed
  -> IO (AProg TypeInfo) -- ^ transformed program
addTypesInProg prog = do
  impints <- mapM readFlatCurryInt (progImports prog)
  return $ addTypesInProgWithImports impints prog

-- | Transforms a program by adding type annotations to the each function's
--   body.
addTypesInProgWithImports ::
     [Prog]          -- ^ interfaces/programs of all imported modules
  -> Prog            -- ^ program to be transformed
  -> AProg TypeInfo  -- ^ transformed program
addTypesInProgWithImports impprogs prog@(Prog mname imps tdecls fdecls ops) =
  AProg mname imps tdecls
    (map (addTypesInFunc (typeOfQN ctypes) (typeOfQN ftypes))
         (--filter (\fd -> not ("_" `isPrefixOf` snd (funcName fd)))
                 fdecls))
    ops
 where
  ctypes = Map.fromList (concatMap consTypesOfProg (prog:impprogs))
  ftypes = Map.fromList (concatMap funcTypesOfProg (prog:impprogs))

  typeOfQN qnts qn =
    maybe (error $ "Type of entity " ++ show qn ++ " not found!")
          id
          (Map.lookup qn qnts)

-- Extract all constructor types defined in a program.
consTypesOfProg :: Prog -> [(QName,TypeExpr)]
consTypesOfProg prog =
  concatMap consTypes (progTypes prog)
 where
  consTypes (Type tn _ tvs cdecls) = map consTypeOfCDecl cdecls
   where
    consTypeOfCDecl (Cons cn _ _ texps) =
      (cn, stripForall (foldr FuncType (resultType tn tvs) texps))
  consTypes (TypeSyn _ _ _ _) = []
  consTypes (TypeNew tn _ tvs (NewCons cn _ texp)) =
    [(cn, stripForall (FuncType texp (resultType tn tvs)))]

  resultType tn tvks = TCons tn (map (TVar . fst) tvks)

-- Extract all function types defined in a program.
funcTypesOfProg :: Prog -> [(QName,TypeExpr)]
funcTypesOfProg prog =
  map (\fd -> (funcName fd, stripForall (funcType fd))) (progFuncs prog)

------------------------------------------------------------------------------
-- The state used to generate type information in an expression
-- contains a list of typed variables (visible in the current expression),
-- the currently largest type variable index, and the current type substitution.
data TIState = TIState
  { tiVars    :: [(VarIndex,TypeExpr)]
  , tiMaxTVar :: Int
  , tiTSub    :: TSub
  }
 deriving Show

-- Initial state where the type variables of the given type expressio
-- will not be used to generate fresh variables.
initTIState :: TypeExpr -> TIState
initTIState te = TIState [] (maximum (0 : allTVarsInTExp te)) idTSub

-- The type of the state monad for annotating expressions with type information.
type TransState a = State TIState a

-- Get the type of a variable.
getVarType :: VarIndex -> TransState TypeExpr
getVarType v = do
  ti <- get
  maybe (error $ "Variable " ++ show v ++ " not found in " ++ show ti)
        return
        (lookup v (tiVars ti))

-- Get current list of type variables from state.
getTypedVars :: TransState [(VarIndex,TypeExpr)]
getTypedVars = do
  ti <- get
  return (tiVars ti)

-- Sets the current list of type variables in the state.
setTypedVars :: [(VarIndex,TypeExpr)] -> TransState ()
setTypedVars tvars = do
  ti <- get
  put $ ti { tiVars = tvars }

-- Add list of type variables to state (and adjust free tvar index).
addTypedVars :: [(VarIndex,TypeExpr)] -> TransState ()
addTypedVars tvs = do
  ti <- get
  put $ ti { tiVars = tvs ++ tiVars ti
           , tiMaxTVar = maximum (tiMaxTVar ti :
                                  concatMap (allTVarsInTExp . snd) tvs)
           }

-- Get a fresh type variable index.
getFreshTVar :: TransState TVarIndex
getFreshTVar = do
  ti <- get
  let newtvi = tiMaxTVar ti + 1
  put ti { tiMaxTVar = newtvi }
  return newtvi

-- Get a fresh variant of a type expression according to the current state
-- (and adjust free tvar index).
getFreshTExp :: TypeExpr -> TransState TypeExpr
getFreshTExp texpr = do
  ti <- get
  let ftexpr = freshTE (tiMaxTVar ti + 1) texpr
  put $ ti { tiMaxTVar = maximum (tiMaxTVar ti : allTVarsInTExp ftexpr) }
  return ftexpr

freshTE :: Int -> TypeExpr -> TypeExpr
freshTE mtv texp = case texp of
  TVar v             -> TVar (v + mtv)
  FuncType t1 t2     -> FuncType (freshTE mtv t1) (freshTE mtv t2)
  TCons tc tes       -> TCons tc (map (freshTE mtv) tes)
  ForallType tvs fte -> ForallType (map (\(v,k) -> (v+mtv,k)) tvs)
                                   (freshTE mtv fte)

-- Apply the given type substitution to the current state.
addTSub :: TSub -> TransState ()
addTSub tsub
  | tsub == []
  = return ()
  | otherwise
  = do ti <- get
       put $ ti { tiTSub = compTSub tsub (tiTSub ti)
                , tiVars = map (\ (v,tv) -> (v, applyTS tsub tv)) (tiVars ti)
                }

-- Apply the current type substitution to the given type expression.
applyCurrentTSub :: TypeExpr -> TransState TypeExpr
applyCurrentTSub te = do
  ti <- get
  return (applyTS (tiTSub ti) te)

-- Returns an annotated expression where the result type is given and
-- the type variables are taken from the current state.
-- The hole to insert the annotation is defined by the abstraction given
-- as the second argument.
returnAExpr :: TypeExpr -> (TypeInfo -> AExpr TypeInfo)
            -> TransState (AExpr TypeInfo)
returnAExpr rtype presult = do
  ti <- get
  return (presult (TypeInfo (tiVars ti) rtype))

-- Gets all type variables occurring in a type expression.
allTVarsInTExp :: TypeExpr -> [Int]
allTVarsInTExp te = trTypeExpr (:) tcomb (.) forall te []
 where
  tcomb _ = foldr (.) id
  forall tvs texp = (map fst tvs ++) . texp

------------------------------------------------------------------------------
-- | Transforms a function declaration by adding type annotations to the
--   function's body.
addTypesInFunc ::
     (QName -> TypeExpr)  -- ^ map constructor name into its type
  -> (QName -> TypeExpr)  -- ^ map function name into its type
  -> FuncDecl             -- ^ function declaration to be transformed
  -> AFuncDecl TypeInfo   -- ^ transformed function
addTypesInFunc ctm ftm (Func qf ar vis ftype rule) =
  AFunc qf ar vis ftype
        (case rule of
           External s    -> AExternal s
           Rule args rhs ->
             ARule args
                   (addTypesInBody ctm ftm qf (zip args argtypes) rhs rhstype))
 where
  (argtypes,rhstype) = splitFuncType ar (stripForall ftype)

-- Add type annotations to an expression, i.e., each (sub)expression is
-- annotated with the list of typed variables visible in this expression
-- and the type of the expression.
addTypesInBody :: (QName -> TypeExpr) -> (QName -> TypeExpr) -> QName
               -> [(VarIndex,TypeExpr)] -> Expr -> TypeExpr -> AExpr TypeInfo
addTypesInBody ctypemap ftypemap _ typedvars expr resulttype =
  let (aexp0,rstate) = runState
                        (addTypedVars typedvars >> annExp expr resulttype)
                        (initTIState resulttype)
      -- apply the computed type substitution to the entire expression
      -- since in order to bind all (earlier introduced) type variables:
      aexp = mapAExpr (applyTS2TypeInfo (tiTSub rstate)) aexp0
  in -- checkAnnExpr ctypemap ftypemap qf aexp &> -- uncomment to check
     aexp
 where
  -- This auxiliary operation returns the type-annotated expression.
  -- The second parameter is the expected result type of the expression.
  annExp exp rt = do
    case exp of
      Var v         -> do vt <- getVarType v
                          if rt == vt
                            then returnAExpr rt $ \ti -> AVar ti v
                            else do
                              let tsub = unifyTExps rt vt
                              addTSub tsub
                              returnAExpr (applyTS tsub rt) $ \ti-> AVar ti v
      Lit l         -> if rt == typeOfLit l
                         then returnAExpr rt $ \ti -> ALit ti l
                              --returnAExpr (TCons (pre "Bool") []) $ \ti -> ALit ti l
                         else do
                           let tsub = unifyTExps rt (typeOfLit l)
                           addTSub tsub
                           returnAExpr (applyTS tsub rt) $ \ti -> ALit ti l
      Comb ct qn es -> annComb ct qn es rt
      Let bs e      -> annLet bs e rt
      Free vs e     -> do addTypedVars vs
                          aexp <- annExp e rt
                          returnAExpr rt $ \tie -> AFree tie vs aexp
      Or e1 e2      -> do aexp1 <- annExp e1 rt
                          aexp2 <- annExp e2 rt
                          returnAExpr rt $ \tie -> AOr tie aexp1 aexp2
      Case ct ce bs -> annCase ct ce bs rt
      Typed e t     -> do aexp <- annExp e rt
                          returnAExpr rt $ \tie -> ATyped tie aexp t

  annComb ct qn args rtype = do
    qntype <- getFreshTExp $ case ct of FuncCall       -> ftypemap qn
                                        ConsCall       -> ctypemap qn
                                        FuncPartCall _ -> ftypemap qn
                                        ConsPartCall _ -> ctypemap qn
    let (ats,rt) = splitFuncType (length args) qntype
    iats <- if rt == rtype
              then return ats
              else do let tsubst = unifyTExps rt rtype
                      addTSub tsubst
                      return $ map (applyTS tsubst) ats
    aes <- mapM (\ (e,te) -> annExp e te) (zip args iats)
    --aes <- mapM (\ (e,te) -> applyCurrentTSub te >>= annExp e) (zip args iats)
    returnAExpr rtype $ \ti -> AComb ti ct qn aes

  annLet bs e rtype = do
    addTypedVars (map (\ (v,t,_) -> (v,t)) bs)
    abs <- mapM (\ (v,t,be) -> annExp be t >>= \abe -> return (v,t,abe)) bs
    ae  <- annExp e rtype
    returnAExpr rtype $ \ti -> ALet ti abs ae

  annCase ct ce brs rtype = do
    -- if the discriminating expression is not a variable, we do not know
    -- its type so that we represent it by a fresh type variable which will
    -- be later instantiated by unifiying types.
    ctype <- case ce of Var v -> getVarType v
                        _     -> fmap TVar getFreshTVar
    ace <- annExp ce ctype
    ptype <- applyCurrentTSub ctype
    abs <- mapM (annBranch ptype) brs
    returnAExpr rtype $ \tie -> ACase tie ct ace abs
   where
    annBranch ptype (Branch pt be) = do
      --ctvs <- getTypedVars
      tvs  <- patVars ptype pt
      addTypedVars tvs
      aexp <- annExp be rtype
      --setTypedVars ctvs
      return $ ABranch pt aexp

    patVars _     (LPattern _)    = return []
    patVars ptype (Pattern qn vs) = do
      ctype <- getFreshTExp (ctypemap qn)
      let (ats,rt) = splitFuncType (length vs) ctype
          tsub     = unifyTExps rt ptype
      addTSub tsub
      return $ zip vs (map (applyTS tsub) ats)

------------------------------------------------------------------------------
-- Check the type annotations of an annotated expression.
-- Terminates with an error if something is not correct.
checkAnnExpr :: (QName -> TypeExpr) -> (QName -> TypeExpr) -> QName
             -> AExpr TypeInfo -> Bool
checkAnnExpr ctypemap ftypemap qf aexp = checkAExp aexp
 where
  checkError msg = error $ "Check error in function '" ++ snd qf ++ "':\n" ++
                           msg ++ "\n\n" ++ show aexp

  checkAExp exp =
    let checkMsg msg c = c ||
          (checkError $ "Annotated subexpression inconsistent:\n" ++ show exp ++
                        (if null msg then "" else  "\nProblem: " ++ show msg))
        checkET t1 t2 = checkMsg ("Types differ: " ++ show t1 ++ " " ++ show t2)
                                 (t1 == t2)
        TypeInfo tvs rt = annOfAExpr exp   in
    case exp of
      AVar _ v         -> maybe (checkError $ "Variable " ++ show v ++
                                 " not in list of typed variables!")
                                (\vt -> checkET vt rt)
                                (lookup v tvs)
      ALit _ l         -> checkET rt (typeOfLit l)
      AComb _ ct qn es -> let qntype = freshTE 1000 $
                                       case ct of FuncCall       -> ftypemap qn
                                                  ConsCall       -> ctypemap qn
                                                  FuncPartCall _ -> ftypemap qn
                                                  ConsPartCall _ -> ctypemap qn
                              ctype = foldr (FuncType . annTypeOfAExpr) rt es
                              tsub = unifyTExps qntype ctype -- check inst.
                           in checkMsg "comb type wrong" (length tsub >= 0) &&
                              --checkMsg "Unif is instantiation"
                              --   (all ((`elem` allTVarsInTExp qntype) . fst)
                              --        tsub) &&
                              all checkAExp es
      ALet _ bs e      -> checkET (annTypeOfAExpr e) rt &&
                          all (\ (v,t,be) ->
                                   checkMsg "not all let vars contained"
                                     ((v,t) `elem` annVarsOfAExpr e) &&
                                   checkAExp be &&
                                   checkET t (annTypeOfAExpr be)) bs &&
                          checkAExp e
      AFree _ vs e     -> checkET (annTypeOfAExpr e) rt &&
                          checkMsg "not all free vars contained"
                                   (all (`elem` annVarsOfAExpr e) vs) &&
                          checkAExp e
      AOr ta e1 e2     -> checkAExp e1 && checkAExp e2 &&
                          checkET ta (annOfAExpr e1) &&
                          checkET ta (annOfAExpr e2)
      ACase _ _ ce bs  -> checkAExp ce &&
                          all (\ (ABranch pt be) ->
                                 checkMsg "not all pattern vars contained"
                                   (all (`elem` map fst (annVarsOfAExpr be))
                                        (patVars pt)) &&
                                 checkAExp be &&
                                 checkET rt (annTypeOfAExpr be))
                              bs
      ATyped ta e t    -> checkET ta (annOfAExpr e) && checkET rt t &&
                          checkAExp e

  patVars (LPattern _)   = []
  patVars (Pattern _ vs) = vs

  annVarsOfAExpr = tiTypedVars . annOfAExpr
  annTypeOfAExpr = tiType . annOfAExpr

------------------------------------------------------------------------------
-- Type substitutions and their operations.

-- The representation of a type sbustitution.
type TSub = [(TVarIndex,TypeExpr)]

-- The empty (identity) type substitution.
idTSub :: TSub
idTSub = []

-- Applies a type substitution (first argument) to a type expression.
applyTS :: TSub -> TypeExpr -> TypeExpr
applyTS tvtexps te = subst te
 where
  subst texp = case texp of
    TVar v             -> maybe texp id (lookup v tvtexps)
    FuncType t1 t2     -> FuncType (subst t1) (subst t2)
    TCons tc tes       -> TCons tc (map subst tes)
    ForallType tvs fte -> ForallType tvs (subst fte)

-- Compose two type substitution such that `(compTSub s1 s2) v == s1 (s2 v)`
compTSub :: TSub -> TSub -> TSub
compTSub s1 s2 = s1 ++ map (\ (v,t) -> (v, applyTS s1 t)) s2

-- Computes a type substitution which unifies two type expressions.
-- It is assumed that the type expressions do not contain `ForallType`.
-- In a correct FlatCurry program, the unifier always exist so that an
-- error is raised if this is not the case.
--
-- Note that type expressions in FlatCurry are first order (due to historical
-- reasons) but the type expressions in Curry might include applications
-- of type variables to type expressions. This is represented in FlatCurry
-- by the type constructor `Prelude.Apply`. For instance, the type
-- application `(v1 v2)` is represented by the type expression
--
--     TCons ("Prelude","Apply") [TVar 1,TVar 2])
--
-- Moreover, there is a type constructor "(->)" to represent partial
-- applications of function types. For instance, the type expression
--
--     TCons ("Prelude","Apply") [TCons ("Prelude","(->)") [TVar 1],TVar 2]
--
-- is equivalent to
--
--     FuncType (TVar 1) (TVar 2)
--
-- These pecularities are considered in the subsequent unification procedure
-- (which is still not complete for arbitrary type expressions but sufficient
-- for type expressions where one side is sufficiently instantiated).
unifyTExps :: TypeExpr -> TypeExpr -> TSub
unifyTExps texp1 texp2 = unify [] texp1 texp2
 where
  unify tsub te1 te2 =
    unifyST tsub (simpTop (applyTS tsub te1)) (simpTop (applyTS tsub te2))

  unifyST tsub te1 te2 = case (te1, te2) of
      (TVar v, te) -> case te of
        TVar v' | v == v' -> tsub
                | v > v'  -> (v, TVar v') : tsub
                | v < v'  -> (v', TVar v) : tsub
        _                 -> (v,te) : tsub
      (_, TVar _) -> unifyST tsub te2 te1
      (FuncType at1 rt1, FuncType at2 rt2) ->
          unify (unify tsub at1 at2) rt1 rt2
      (TCons tc1 tes1, TCons tc2 tes2)
        | tc1 == tc2                       -> unifyTEs tsub (zip tes1 tes2)
        | tc1 == tcApply && tc2 == tcApply -> unifyTEs tsub (zip tes1 tes2)
        | tc1 == tcApply
         -> unifyTEs tsub (zip tes1 [TCons tc2 (init tes2), last tes2])
        | tc2 == tcApply
         -> unifyTEs tsub (zip tes2 [TCons tc1 (init tes1), last tes1])
      (TCons tca [TVar tf, te], FuncType fte1 fte2)
        | tca == tcApply -> unify ((tf, TCons tcFunc [fte1]):tsub) te fte2
      (FuncType fte1 fte2, TCons tca [TVar tf, te])
        | tca == tcApply -> unify ((tf, TCons tcFunc [fte1]):tsub) te fte2
      _ -> error $ "unifyTExps: types not unifiable:\n" ++
                   show texp1 ++ "\n" ++ show texp2 ++ "\n"

  unifyTEs tsub []              = tsub
  unifyTEs tsub ((te1,te2):tes) = unifyTEs (unify tsub te1 te2) tes

  -- Simplify top-level occurrences of specific type constructors:
  simpTop texp = case texp of
    TCons tc [TCons tc1 tes1, te2] | tc == tcApply
                                    -> simpTop (TCons tc1 (tes1 ++ [te2]))
    TCons tc [te1, te2]            | tc == tcFunc  -> FuncType te1 te2
    _                                              -> texp

  tcApply = pre "Apply"

  tcFunc = pre "(->)"

------------------------------------------------------------------------------
-- Some auxiliaries.

-- Transform name into Prelude-qualified name.
pre :: String -> QName
pre f = ("Prelude",f)

-- The type of a literal.
-- The type of a literal.
typeOfLit :: Literal -> TypeExpr
typeOfLit l = TCons (pre ltype) []
 where ltype = case l of Intc   _ -> "Int"
                         Floatc _ -> "Float"
                         Charc  _ -> "Char"

-- Split a function type into a list of argument types (where the length
-- is given as the first argument) and the remaining result type.
splitFuncType :: Int -> TypeExpr -> ([TypeExpr],TypeExpr)
splitFuncType ar te
  | ar < 0  = error "splitFuncType with negative arity"
  | ar == 0 = ([],te)
  | otherwise = case te of
                  FuncType t1 t2 -> let (ts,rt) = splitFuncType (ar-1) t2
                                    in (t1:ts, rt)
                  _ -> error "splitFuncType without function type"

-- Strip `ForallType` quantifications.
stripForall :: TypeExpr -> TypeExpr
stripForall texp = case texp of
  ForallType _ te  -> stripForall te
  FuncType t1 t2   -> FuncType (stripForall t1) (stripForall t2)
  TCons tc tes     -> TCons tc (map stripForall tes)
  TVar v           -> TVar v

-----------------------------------------------------------------------