content: {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE LambdaCase #-}
  
  -- |
  -- Module      : Jikka.Core.Convert.TypeInfer
  -- Description : does type inference. / 型推論を行います。
  -- Copyright   : (c) Kimiyuki Onaka, 2021
  -- License     : Apache License 2.0
  -- Maintainer  : kimiyuki95@gmail.com
  -- Stability   : experimental
  -- Portability : portable
  module Jikka.Core.Convert.TypeInfer
    ( run,
      runExpr,
      runRule,
  
      -- * internal types and functions
      Equation (..),
      formularizeProgram,
      sortEquations,
      mergeAssertions,
      Subst (..),
      subst,
      substDefault,
      solveEquations,
      substProgram,
    )
  where
  
  import Control.Arrow (second)
  import Control.Monad.State.Strict
  import Control.Monad.Writer.Strict (MonadWriter, censor, execWriterT, tell)
  import qualified Data.Map.Strict as M
  import Data.Monoid (Dual (..))
  import Jikka.Common.Alpha
  import Jikka.Common.Error
  import Jikka.Core.Format (formatExpr, formatToplevelExpr, formatType)
  import Jikka.Core.Language.BuiltinPatterns
  import Jikka.Core.Language.Expr
  import Jikka.Core.Language.FreeVars
  import Jikka.Core.Language.Lint
  import Jikka.Core.Language.NameCheck (namecheckExpr)
  import Jikka.Core.Language.TypeCheck (literalToType, typecheckExpr)
  import Jikka.Core.Language.Util
  
  data Hint
    = VarHint VarName
    | ExprHint Expr
    | ToplevelExprHint ToplevelExpr
    deriving (Eq, Ord, Show, Read)
  
  data Equation
    = TypeEquation Type Type [Hint]
    | TypeAssertion VarName Type
    deriving (Eq, Ord, Show, Read)
  
  type Eqns = Dual [Equation]
  
  consHint :: Hint -> Equation -> Equation
  consHint hint = \case
    TypeEquation t1 t2 hints -> TypeEquation t1 t2 (hint : hints)
    TypeAssertion x t -> TypeAssertion x t
  
  wrapHint :: MonadWriter Eqns m => Hint -> m a -> m a
  wrapHint hint = censor (fmap (map (consHint hint)))
  
  wrapErrorFromHint :: MonadError Error m => Hint -> m a -> m a
  wrapErrorFromHint = \case
    VarHint x -> wrapError' $ "around variable " ++ formatVarName x
    ExprHint e -> wrapError' $ "around expr " ++ summarize (formatExpr e)
    ToplevelExprHint e -> wrapError' $ "around toplevel expr " ++ summarize (formatToplevelExpr e)
    where
      summarize s = case lines s of
        (s : _ : _) -> s ++ " ..."
        _ -> s
  
  wrapErrorFromHints :: MonadError Error m => [Hint] -> m a -> m a
  wrapErrorFromHints hints = foldr (\hint f -> wrapErrorFromHint hint . f) id hints
  
  formularizeType :: MonadWriter Eqns m => Type -> Type -> m ()
  formularizeType t1 t2 = tell $ Dual [TypeEquation t1 t2 []]
  
  formularizeVarName :: MonadWriter Eqns m => VarName -> Type -> m ()
  formularizeVarName x t = tell $ Dual [TypeAssertion x t]
  
  formularizeExpr :: (MonadWriter Eqns m, MonadAlpha m, MonadError Error m) => Expr -> m Type
  formularizeExpr e = wrapHint (ExprHint e) $ case e of
    Var x -> do
      t <- genType
      formularizeVarName x t
      return t
    Lit lit -> case lit of
      LitBuiltin (Proj _) [] -> genType -- Proj may have a empty list.
      _ -> literalToType lit
    App f e -> do
      ret <- genType
      t <- formularizeExpr e
      formularizeExpr' f (FunTy t ret)
      return ret
    Lam x t body -> do
      formularizeVarName x t
      ret <- formularizeExpr body
      return $ FunTy t ret
    Let x t e1 e2 -> do
      formularizeVarName x t
      formularizeExpr' e1 t
      formularizeExpr e2
    Assert e1 e2 -> do
      formularizeExpr' e1 BoolTy
      formularizeExpr e2
  
  formularizeExpr' :: (MonadWriter Eqns m, MonadAlpha m, MonadError Error m) => Expr -> Type -> m ()
  formularizeExpr' e t = do
    t' <- formularizeExpr e
    wrapHint (ExprHint e) $ do
      formularizeType t t'
  
  formularizeToplevelExpr :: (MonadWriter Eqns m, MonadAlpha m, MonadError Error m) => ToplevelExpr -> m Type
  formularizeToplevelExpr e = wrapHint (ToplevelExprHint e) $ case e of
    ResultExpr e -> formularizeExpr e
    ToplevelLet x t e cont -> do
      formularizeVarName x t
      formularizeExpr' e t
      formularizeToplevelExpr cont
    ToplevelLetRec f args ret body cont -> do
      formularizeVarName f (curryFunTy (map snd args) ret)
      mapM_ (uncurry formularizeVarName) args
      formularizeExpr' body ret
      formularizeToplevelExpr cont
    ToplevelAssert e cont -> do
      formularizeExpr' e BoolTy
      formularizeToplevelExpr cont
  
  formularizeProgram :: (MonadAlpha m, MonadError Error m) => Program -> m [Equation]
  formularizeProgram prog = getDual <$> execWriterT (formularizeToplevelExpr prog)
  
  sortEquations :: [Equation] -> ([(Type, Type, [Hint])], [(VarName, Type)])
  sortEquations = go [] []
    where
      go eqns' assertions [] = (eqns', assertions)
      go eqns' assertions (eqn : eqns) = case eqn of
        TypeEquation t1 t2 hints -> go ((t1, t2, hints) : eqns') assertions eqns
        TypeAssertion x t -> go eqns' ((x, t) : assertions) eqns
  
  mergeAssertions :: [(VarName, Type)] -> [(Type, Type, [Hint])]
  mergeAssertions = go M.empty []
    where
      go _ eqns [] = eqns
      go gamma eqns ((x, t) : assertions) = case M.lookup x gamma of
        Nothing -> go (M.insert x t gamma) eqns assertions
        Just t' -> go gamma ((t, t', [VarHint x]) : eqns) assertions
  
  -- | `Subst` is type substituion. It's a mapping from type variables to their actual types.
  newtype Subst = Subst {unSubst :: M.Map TypeName Type}
    deriving (Eq, Ord, Show, Read)
  
  subst :: Subst -> Type -> Type
  subst sigma = \case
    VarTy x ->
      case M.lookup x (unSubst sigma) of
        Nothing -> VarTy x
        Just t -> subst sigma t
    IntTy -> IntTy
    BoolTy -> BoolTy
    ListTy t -> ListTy (subst sigma t)
    TupleTy ts -> TupleTy (map (subst sigma) ts)
    FunTy t ret -> FunTy (subst sigma t) (subst sigma ret)
    DataStructureTy ds -> DataStructureTy ds
  
  unifyTyVar :: (MonadState Subst m, MonadError Error m) => TypeName -> Type -> m ()
  unifyTyVar x t =
    if x `elem` freeTyVars t
      then throwInternalError $ "looped type equation " ++ formatTypeName x ++ " = " ++ formatType t
      else do
        modify' (Subst . M.insert x t . unSubst) -- This doesn't introduce the loop.
  
  unifyType :: (MonadState Subst m, MonadError Error m) => Type -> Type -> m ()
  unifyType t1 t2 = wrapError' ("failed to unify " ++ formatType t1 ++ " and " ++ formatType t2) $ do
    sigma <- get
    t1 <- return $ subst sigma t1 -- shadowing
    t2 <- return $ subst sigma t2 -- shadowing
    case (t1, t2) of
      _ | t1 == t2 -> return ()
      (VarTy x1, _) -> do
        unifyTyVar x1 t2
      (_, VarTy x2) -> do
        unifyTyVar x2 t1
      (ListTy t1, ListTy t2) -> do
        unifyType t1 t2
      (TupleTy ts1, TupleTy ts2) -> do
        if length ts1 == length ts2
          then mapM_ (uncurry unifyType) (zip ts1 ts2)
          else throwInternalError $ "different type ctors " ++ formatType t1 ++ " and " ++ formatType t2
      (FunTy t1 ret1, FunTy t2 ret2) -> do
        unifyType t1 t2
        unifyType ret1 ret2
      _ -> throwInternalError $ "different type ctors " ++ formatType t1 ++ " and " ++ formatType t2
  
  solveEquations :: MonadError Error m => [(Type, Type, [Hint])] -> m Subst
  solveEquations eqns = wrapError' "failed to solve type equations" $ do
    (`execStateT` Subst M.empty) $ do
      forM_ eqns $ \(t1, t2, hints) -> do
        wrapErrorFromHints hints $ do
          unifyType t1 t2
  
  -- | `substDefault` replaces all undetermined type variables with the given default type.
  substDefault :: Type -> Type -> Type
  substDefault t0 = \case
    VarTy _ -> TupleTy []
    IntTy -> IntTy
    BoolTy -> BoolTy
    ListTy t -> ListTy (substDefault t0 t)
    TupleTy ts -> TupleTy (map (substDefault t0) ts)
    FunTy t ret -> FunTy (substDefault t0 t) (substDefault t0 ret)
    DataStructureTy ds -> DataStructureTy ds
  
  subst' :: Maybe Type -> Subst -> Type -> Type
  subst' t0 sigma = maybe id substDefault t0 . subst sigma
  
  fixProj :: MonadError Error m => [(VarName, Type)] -> Expr -> m Expr
  fixProj env = \case
    Proj' [] i e -> do
      -- fix Proj with a empty list
      t <- typecheckExpr env e
      case t of
        TupleTy ts -> return $ Proj' ts i e
        _ -> throwInternalError $ "type of argument of proj must be a tuple: " ++ formatType t
    e -> return e
  
  substProgram :: MonadError Error m => Maybe Type -> Subst -> Program -> m Program
  substProgram t0 sigma = mapExprProgramM (mapSubExprM fixProj) . mapTypeProgram (subst' t0 sigma)
  
  substExpr :: MonadError Error m => Maybe Type -> Subst -> [(VarName, Type)] -> Expr -> m Expr
  substExpr t0 sigma env = mapSubExprM fixProj env . mapTypeExpr (subst' t0 sigma)
  
  -- | `run` does type inference.
  --
  -- * This assumes that program has no name conflicts.
  --
  -- Before:
  --
  -- > let f = fun y -> y
  -- > in let x = 1
  -- > in f(x + x)
  --
  -- After:
  --
  -- > let f: int -> int = fun y: int -> y
  -- > in let x: int = 1
  -- > in f(x + x)
  run :: (MonadAlpha m, MonadError Error m) => Program -> m Program
  run prog = wrapError' "Jikka.Core.Convert.TypeInfer.run" $ do
    precondition $ do
      ensureAlphaConverted prog
    eqns <- formularizeProgram prog
    let (eqns', assertions) = sortEquations eqns
    let eqns'' = mergeAssertions assertions
    sigma <- solveEquations (eqns' ++ eqns'')
    let t0 = Just UnitTy
    prog <- substProgram t0 sigma prog
    postcondition $ do
      ensureAlphaConverted prog
      ensureWellTyped prog
    return prog
  
  runExpr :: (MonadAlpha m, MonadError Error m) => [(VarName, Type)] -> Expr -> m Expr
  runExpr env e = wrapError' "Jikka.Core.Convert.TypeInfer.runExpr" $ do
    precondition $ do
      namecheckExpr env e
    eqns <- getDual <$> execWriterT (formularizeExpr e)
    let (eqns', assertions) = sortEquations eqns
    let eqns'' = mergeAssertions (env ++ assertions)
    sigma <- solveEquations (eqns' ++ eqns'')
    let t0 = Nothing -- don't use substDefault
    env <- return $ map (second (subst' t0 sigma)) env
    e <- substExpr t0 sigma env e
    postcondition $ do
      namecheckExpr env e
      typecheckExpr env e
    return e
  
  runRule :: (MonadAlpha m, MonadError Error m) => [(VarName, Type)] -> Expr -> Expr -> m ([(VarName, Type)], Expr, Expr)
  runRule args e1 e2 = wrapError' "Jikka.Core.Convert.TypeInfer.runRule" $ do
    precondition $ do
      -- Underscores are allowed for names, so we don't use namecheckExpr here.
      return ()
    eqns <- (getDual <$>) . execWriterT $ do
      t <- formularizeExpr e1
      formularizeExpr' e2 t
    let (eqns', assertions) = sortEquations eqns
    let eqns'' = mergeAssertions assertions
    sigma <- solveEquations (eqns' ++ eqns'')
    args <- return $ map (second (subst sigma)) args -- don't use substDefault
    e1 <- return $ mapTypeExpr (subst sigma) e1 -- don't use substDefault
    e2 <- return $ mapTypeExpr (subst sigma) e2 -- don't use substDefault
    postcondition $ do
      -- Underscores are allowed for names, so we don't use namecheckExpr here.
      -- Type variables can remain, so we don't use typecheckExpr here.
      return ()
    return (args, e1, e2)
  