{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}

-- |
-- Module      : Jikka.RestrictedPython.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.RestrictedPython.Convert.TypeInfer
  ( run,

    -- * internal types and functions
    Equation (..),
    formularizeProgram,
    sortEquations,
    mergeAssertions,
    Subst (..),
    subst,
    solveEquations,
    mapTypeProgram,
  )
where

import Control.Arrow (second)
import Control.Monad.Reader
import Control.Monad.State.Strict
import Control.Monad.Writer.Strict
import qualified Data.Map.Strict as M
import Jikka.Common.Alpha
import Jikka.Common.Error
import Jikka.RestrictedPython.Format (formatType)
import Jikka.RestrictedPython.Language.Builtin
import Jikka.RestrictedPython.Language.Expr
import Jikka.RestrictedPython.Language.Util

data Equation
  = TypeEquation Type Type (Maybe Loc)
  | TypeAssertion VarName' Type
  deriving (Eq, Ord, Show, Read)

type Eqns = Dual [Equation]

formularizeType :: MonadWriter Eqns m => Type -> Type -> Maybe Loc -> m ()
formularizeType t1 t2 location = tell $ Dual [TypeEquation t1 t2 location]

formularizeVarName :: MonadWriter Eqns m => VarName' -> Type -> m ()
formularizeVarName x t = tell $ Dual [TypeAssertion x t]

formularizeTarget :: (MonadWriter Eqns m, MonadAlpha m) => Target' -> m Type
formularizeTarget x0 = case value' x0 of
  SubscriptTrg f index -> do
    t <- genType
    tf <- formularizeTarget f
    formularizeType tf (ListTy t) (loc' x0)
    tindex <- formularizeExpr index
    formularizeType tindex IntTy (loc' x0)
    return t
  NameTrg x -> do
    t <- genType
    formularizeVarName x t
    return t
  TupleTrg xs -> do
    TupleTy <$> mapM formularizeTarget xs

formularizeTarget' :: (MonadWriter Eqns m, MonadAlpha m) => Target' -> Type -> m ()
formularizeTarget' x0 t = do
  t' <- formularizeTarget x0
  formularizeType t t' (loc' x0)

formularizeExpr :: (MonadWriter Eqns m, MonadAlpha m) => Expr' -> m Type
formularizeExpr e0 = case value' e0 of
  BoolOp e1 _ e2 -> do
    formularizeExpr' e1 BoolTy
    formularizeExpr' e2 BoolTy
    return BoolTy
  BinOp e1 _ e2 -> do
    formularizeExpr' e1 IntTy
    formularizeExpr' e2 IntTy
    return IntTy
  UnaryOp op e -> do
    let t' = if op == Not then BoolTy else IntTy
    formularizeExpr' e t'
    return t'
  Lambda args body -> do
    mapM_ (uncurry formularizeVarName) args
    ret <- genType
    formularizeExpr' body ret
    return $ CallableTy (map snd args) ret
  IfExp e1 e2 e3 -> do
    formularizeExpr' e1 BoolTy
    t <- formularizeExpr e2
    formularizeExpr' e3 t
    return t
  ListComp e comp -> do
    let Comprehension x iter pred = comp
    te <- formularizeExpr e
    tx <- formularizeTarget x
    formularizeExpr' iter (ListTy tx)
    case pred of
      Nothing -> return ()
      Just pred -> formularizeExpr' pred BoolTy
    return $ ListTy te
  Compare e1 (CmpOp' op t) e2 -> do
    formularizeExpr' e1 t
    formularizeExpr' e2 (if op == In || op == NotIn then ListTy t else t)
    return BoolTy
  Call f args -> do
    ts <- mapM formularizeExpr args
    ret <- genType
    formularizeExpr' f (CallableTy ts ret)
    return ret
  Constant const ->
    return $ case const of
      ConstNone -> NoneTy
      ConstInt _ -> IntTy
      ConstBool _ -> BoolTy
      ConstBuiltin b -> typeBuiltin b
  Attribute e x -> do
    let (t1, t2) = typeAttribute (value' x)
    formularizeExpr' e t1
    return t2
  Subscript e1 e2 -> do
    t <- genType
    formularizeExpr' e1 (ListTy t)
    formularizeExpr' e2 IntTy
    return t
  Starred e -> do
    t <- genType
    formularizeExpr' e (ListTy t)
    return t -- because @*xs@ and @y@ has the same type in @[*xs, y]@
  Name x -> do
    t <- genType
    formularizeVarName x t
    return t
  List t es -> do
    forM_ es $ \e -> do
      formularizeExpr' e t
    return $ ListTy t
  Tuple es -> TupleTy <$> mapM formularizeExpr es
  SubscriptSlice e from to step -> do
    t' <- genType
    formularizeExpr' e (ListTy t')
    let formularize = \case
          Nothing -> return ()
          Just e -> formularizeExpr' e IntTy
    formularize from
    formularize to
    formularize step
    return (ListTy t')

formularizeExpr' :: (MonadWriter Eqns m, MonadAlpha m) => Expr' -> Type -> m ()
formularizeExpr' e0 t = do
  t' <- formularizeExpr e0
  formularizeType t t' (loc' e0)

formularizeStatement :: (MonadWriter Eqns m, MonadAlpha m) => Type -> Statement -> m ()
formularizeStatement ret = \case
  Return e -> do
    t <- formularizeExpr e
    formularizeType t ret (loc' e)
  AugAssign x _ e -> do
    formularizeTarget' x IntTy
    formularizeExpr' e IntTy
  AnnAssign x t e -> do
    formularizeTarget' x t
    formularizeExpr' e t
  For x e body -> do
    t <- formularizeTarget x
    formularizeExpr' e (ListTy t)
    mapM_ (formularizeStatement ret) body
  If e body1 body2 -> do
    formularizeExpr' e BoolTy
    mapM_ (formularizeStatement ret) body1
    mapM_ (formularizeStatement ret) body2
  Assert e -> do
    formularizeExpr' e BoolTy
  Expr' e -> do
    formularizeExpr' e SideEffectTy

formularizeToplevelStatement :: (MonadWriter Eqns m, MonadAlpha m) => ToplevelStatement -> m ()
formularizeToplevelStatement = \case
  ToplevelAnnAssign x t e -> do
    formularizeVarName x t
    formularizeExpr' e t
  ToplevelFunctionDef f args ret body -> do
    mapM_ (uncurry formularizeVarName) args
    formularizeVarName f (CallableTy (map snd args) ret)
    mapM_ (formularizeStatement ret) body
  ToplevelAssert e -> do
    formularizeExpr' e BoolTy

formularizeProgram :: MonadAlpha m => Program -> m [Equation]
formularizeProgram prog = getDual <$> execWriterT (mapM_ formularizeToplevelStatement prog)

sortEquations :: [Equation] -> ([(Type, Type, Maybe Loc)], [(VarName', Type)])
sortEquations = go [] []
  where
    go eqns' assertions [] = (eqns', assertions)
    go eqns' assertions (eqn : eqns) = case eqn of
      TypeEquation t1 t2 loc -> go ((t1, t2, loc) : eqns') assertions eqns
      TypeAssertion x t -> go eqns' ((x, t) : assertions) eqns

mergeAssertions :: [(VarName', Type)] -> [(Type, Type, Maybe Loc)]
mergeAssertions = go M.empty []
  where
    go _ eqns [] = eqns
    go gamma eqns ((x, t) : assertions) = case M.lookup (value' x) gamma of
      Nothing -> go (M.insert (value' x) t gamma) eqns assertions
      Just t' -> go gamma ((t, t', loc' 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}

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)
  CallableTy ts ret -> CallableTy (map (subst sigma) ts) (subst sigma ret)
  StringTy -> StringTy
  SideEffectTy -> SideEffectTy

unifyTyVar :: (MonadState Subst m, MonadError Error m) => TypeName -> Type -> m ()
unifyTyVar x t =
  if x `elem` freeTyVars t
    then throwTypeError $ "type equation loops: " ++ formatType (VarTy 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 = 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 throwTypeError $ "type " ++ formatType t1 ++ " is not type " ++ formatType t2
    (CallableTy args1 ret1, CallableTy args2 ret2) -> do
      if length args1 == length args2
        then mapM_ (uncurry unifyType) (zip args1 args2)
        else throwTypeError $ "type " ++ formatType t1 ++ " is not type " ++ formatType t2
      unifyType ret1 ret2
    _ -> throwTypeError $ "type " ++ formatType t1 ++ " is not type " ++ formatType t2

solveEquations :: MonadError Error m => [(Type, Type, Maybe Loc)] -> m Subst
solveEquations eqns = wrapError' "failed to solve type equations" $ do
  flip execStateT (Subst M.empty) $ do
    errs <- forM eqns $ \(t1, t2, loc) -> do
      (Right <$> unifyType t1 t2) `catchError` \err -> do
        sigma <- get
        t1 <- return $ subst sigma t1 -- shadowing
        t2 <- return $ subst sigma t2 -- shadowing
        return $ Left (maybe id WithLocation loc (WithWrapped ("failed to unify type " ++ formatType t1 ++ " and type " ++ formatType t2) err))
    reportErrors errs

mapTypeConstant :: (Type -> Type) -> Constant -> Constant
mapTypeConstant f = \case
  ConstNone -> ConstNone
  ConstInt n -> ConstInt n
  ConstBool p -> ConstBool p
  ConstBuiltin b -> ConstBuiltin (mapTypeBuiltin f b)

mapTypeTarget :: (Type -> Type) -> Target' -> Target'
mapTypeTarget f = fmap $ \case
  SubscriptTrg x index -> SubscriptTrg (mapTypeTarget f x) (mapTypeExpr f index)
  NameTrg x -> NameTrg x
  TupleTrg xs -> TupleTrg (map (mapTypeTarget f) xs)

mapTypeExpr :: (Type -> Type) -> Expr' -> Expr'
mapTypeExpr f = mapSubExpr go
  where
    go = fmap $ \case
      Lambda args body -> Lambda (map (second f) args) (go body)
      ListComp e (Comprehension x iter pred) -> ListComp (go e) (Comprehension (mapTypeTarget f x) (go iter) (fmap go pred))
      Compare e1 (CmpOp' op t) e2 -> Compare (go e1) (CmpOp' op (f t)) (go e2)
      Constant const -> Constant (mapTypeConstant f const)
      Attribute e a -> Attribute (go e) (mapTypeAttribute f <$> a)
      List t es -> List (f t) (map go es)
      e -> e

mapTypeStatement :: (Type -> Type) -> Statement -> Statement
mapTypeStatement f = \case
  Return e -> Return (mapTypeExpr f e)
  AugAssign x op e -> AugAssign (mapTypeTarget f x) op (mapTypeExpr f e)
  AnnAssign x t e -> AnnAssign (mapTypeTarget f x) (f t) (mapTypeExpr f e)
  For x iter body -> For (mapTypeTarget f x) (mapTypeExpr f iter) (map (mapTypeStatement f) body)
  If pred body1 body2 -> If (mapTypeExpr f pred) (map (mapTypeStatement f) body1) (map (mapTypeStatement f) body2)
  Assert e -> Assert (mapTypeExpr f e)
  Expr' e -> Expr' (mapTypeExpr f e)

mapTypeToplevelStatement :: (Type -> Type) -> ToplevelStatement -> ToplevelStatement
mapTypeToplevelStatement f = \case
  ToplevelAnnAssign x t e -> ToplevelAnnAssign x (f t) (mapTypeExpr f e)
  ToplevelFunctionDef g args ret body -> ToplevelFunctionDef g (map (second f) args) (f ret) (map (mapTypeStatement f) body)
  ToplevelAssert e -> ToplevelAssert (mapTypeExpr f e)

mapTypeProgram :: (Type -> Type) -> Program -> Program
mapTypeProgram f prog = map (mapTypeToplevelStatement f) prog

-- | `substUnit` replaces all undetermined type variables with the unit type.
substUnit :: Type -> Type
substUnit = \case
  VarTy _ -> NoneTy
  IntTy -> IntTy
  BoolTy -> BoolTy
  ListTy t -> ListTy (substUnit t)
  TupleTy ts -> TupleTy (map substUnit ts)
  CallableTy ts ret -> CallableTy (map substUnit ts) (substUnit ret)
  StringTy -> StringTy
  SideEffectTy -> SideEffectTy

-- | `subst'` does `subst` and replaces all undetermined type variables with the unit type.
subst' :: Subst -> Type -> Type
subst' sigma = substUnit . subst sigma

-- | `run` infers types of given programs.
--
-- As the interface, you can understand this function does the following:
--
-- 1. Finds a type environment \(\Gamma\) s.t. for all statement \(\mathrm{stmt}\) in the given program, \(\Gamma \vdash \mathrm{stmt}\) holds, and
-- 2. Annotates each variable in the program using the \(\Gamma\).
--
-- In its implementation, this is just something like a Hindley-Milner type inference.
--
-- == Requirements
--
-- * There must be no name conflicts in given programs. They must be alpha-converted. (`Jikka.RestrictedPython.Convert.Alpha`)
-- * All names must be resolved. (`Jikka.RestrictedPython.Convert.ResolveBuiltin`)
run :: (MonadAlpha m, MonadError Error m) => Program -> m Program
run prog = wrapError' "Jikka.RestrictedPython.Convert.TypeInfer" $ do
  eqns <- formularizeProgram prog
  let (eqns', assertions) = sortEquations eqns
  let eqns'' = mergeAssertions assertions
  sigma <- solveEquations (eqns' ++ eqns'')
  return $ mapTypeProgram (subst' sigma) prog