content: {-# LANGUAGE TemplateHaskell #-}
  {-# LANGUAGE TupleSections   #-}
  
  module Infernu.Infer
      ( runTypeInference
      , test
      , getAnnotations
      , minifyVars
      , TypeError
      )
      where
  
  
  import           Control.Monad      (foldM, forM)
  import           Data.Map.Strict      (Map)
  import qualified Data.Map.Strict      as Map
  import           Data.Set           (Set)
  import qualified Data.Set           as Set
  import           Data.Char          (isUpper)
  
  import           Text.PrettyPrint.ANSI.Leijen (Pretty (..), align, text, (<+>), vsep, align, indent)
  
  import           Infernu.Prelude
  import qualified Infernu.Builtins   as Builtins
  import           Infernu.InferState
  import           Infernu.Lib        (safeLookup)
  import           Infernu.Log
  import           Infernu.Pretty()
  import           Infernu.Types
  import           Infernu.Source     (Source(..), TypeError(..))
  import           Infernu.Expr       (Exp(..), LitVal(..), EVarName(..), EPropName(..))
  import           Infernu.Unify      (unify, unifyAll, unifyPending, unifyPredsL, unifyl, tryMakeRow)
  
  
  
  getAnnotations :: Exp a -> [a]
  getAnnotations = foldr (:) []
  
  ----------------------------------------------------------------------
  
  closeRowList :: Bool -> Source -> TRowList Type -> Infer (TRowList Type)
  closeRowList unrollRec a   (TRowProp p t l) = TRowProp p t <$> closeRowList unrollRec a l
  closeRowList _         _   (TRowEnd _)      = return $ TRowEnd Nothing
  closeRowList unrollRec a r@(TRowRec tid ts) =
    if not unrollRec
    then return $ r
    else do qt <- unrollName a tid ts
            case qualType qt of
              Fix (TRow _ r') -> closeRowList False a r'
              _ -> error $ show $ text "Expected row type, got: " <+> pretty qt
  
  -- | Replaces a top-level open row type with the closed equivalent.
  -- >>> pretty $ closeRow (Fix $ TRow $ TRowProp "a" (schemeEmpty $ Fix $ TRow $ TRowProp "aa" (schemeEmpty $ Fix $ TBody TNumber) (TRowEnd (Just $ RowTVar 1))) (TRowEnd (Just $ RowTVar 2)))
  -- "{a: {aa: Number, ..b}}"
  -- >>> pretty $ closeRow (Fix $ TFunc [Fix $ TRow $ TRowProp "a" (schemeEmpty $ Fix $ TRow $ TRowProp "aa" (schemeEmpty $ Fix $ TBody TNumber) (TRowEnd Nothing)) (TRowEnd Nothing)] (Fix $ TBody TString))
  -- "{a: {aa: Number}}.(() -> String)"
  -- >>> pretty $ closeRow (Fix $ TFunc [Fix $ TRow $ TRowProp "a" (schemeEmpty $ Fix $ TRow $ TRowProp "a.a" (schemeEmpty $ Fix $ TBody TNumber) (TRowEnd (Just $ RowTVar 1))) (TRowEnd (Just $ RowTVar 2))] (Fix $ TBody TString))
  -- "{a: {a.a: Number, ..b}, ..c}.(() -> String)"
  closeRow :: Source -> Type -> Infer Type
  closeRow a (Fix (TRow l r)) = Fix . TRow (fmap (++"*") l) <$> closeRowList True a r
  closeRow _ t = return t
  
  ----------------------------------------------------------------------
  
  
  -- For efficiency reasons, types list is returned in reverse order.
  accumInfer :: TypeEnv -> [Exp Source] -> Infer [(QualType, Exp (Source, QualType))]
  accumInfer env =
    do traceLog $ text "accumInfer: env: " <+> pretty env
       foldM accumInfer' []
       where accumInfer' types expr =
               do (t, e) <- inferType env expr
                  return ((t,e):types)
  
  -- wrapInferError :: Exp Source -> Infer a -> Infer a
  -- wrapInferError exp act = wrapError format' s act
  --     where format' te = vsep [text "During type inference for:"
  --                             , indent 4 $ pretty exp
  --                             , text "Failed with:"
  --                             , indent 4 $ pretty te
  --                             ]
  --           s = head $ getAnnotations exp
  
  inferType  :: TypeEnv -> Exp Source -> Infer (QualType, Exp (Source, QualType))
  inferType env expr = do
    traceLog $ text ">> " <+> pretty expr <+> text " -- env: " <+> pretty env
  --  (t, e) <- wrapInferError expr $ inferType' env expr
    (t, e) <- inferType' env expr
    unifyPending
    s <- getMainSubst
    st <- getState
    traceLog $ text ">> " <+> pretty expr <+> text " -- inferred :: " <+> pretty (applySubst s t)
    traceLog $ indent 4 $ text "infer state: " <+> pretty st
    return (applySubst s t, fmap (applySubst s) e)
  
  inferType' :: TypeEnv -> Exp Source -> Infer (QualType, Exp (Source, QualType))
  inferType' _ (ELit a lit) = do
    let t = Fix $ TBody $ case lit of
                      LitNumber _ -> TNumber
                      LitBoolean _ -> TBoolean
                      LitString _ -> TString
                      LitRegex{} -> TRegex
                      LitUndefined -> TUndefined
                      LitNull -> TNull
                      LitEmptyThis -> TEmptyThis
    return (qualEmpty t, ELit (a, qualEmpty t) lit)
  inferType' env (EVar a n) =
    do t <- instantiateVar a n env
       return (t, EVar (a, t) n)
  inferType' env (EAbs a argNames e2) =
    do argTypes <- forM argNames (const $ Fix . TBody . TVar <$> freshFlex)
       env' <- foldM (\e (n, t) -> addVarScheme e n $ schemeEmpty t) env $ zip argNames argTypes
       (t1, e2') <- inferType env' e2
       pred' <- unifyPredsL a $ qualPred t1
       let t = TQual pred' $ Fix $ TFunc argTypes (qualType t1)
       return (t, EAbs (a, t) argNames e2')
  inferType' env (EApp a e1 eArgs) =
    do tvar <- Fix . TBody . TVar <$> freshFlex
       (t1, e1') <- inferType env e1
       traceLog $ text "EApp: Inferred type for func expr: " <+> pretty t1
       argsTE <- accumInfer env eArgs
       traceLog $ text "EApp: Inferred types for func args: " <+> pretty argsTE
       let rargsTE = reverse argsTE
           tArgs = map fst rargsTE
           eArgs' = map snd rargsTE
           preds = concatMap qualPred $ t1:tArgs
       unify a (Fix $ TFunc (map qualType tArgs) tvar) (qualType t1)
       traceLog $ text "Inferred preds: " <+> pretty preds
       tvar' <- do  pred' <- unifyPredsL a preds
                    tvarSubsted <- applyMainSubst tvar
                    return $ TQual pred' tvarSubsted
       traceLog $ text "Inferred func application: " <+> pretty tvar'
       return (tvar', EApp (a, tvar') e1' eArgs')
  inferType' env (ENew a e1 eArgs) =
    do (t1, e1') <- inferType env e1
       let label e' = case e' of
                          EVar _ n -> Just n
                          ELet _ _ _ e'' -> label e''
                          _ -> Nothing
       argsTE <- accumInfer env eArgs
       thisT <- Fix . TBody . TVar <$> freshFlex
       resT <- Fix . TBody . TVar <$> freshFlex
       let rargsTE = reverse argsTE
           tArgs = thisT : map (qualType . fst) rargsTE
           eArgs' = map snd rargsTE
           preds = concatMap qualPred $ t1 : map fst argsTE
       unify a (Fix $ TFunc tArgs resT) (qualType t1)
       -- constrain 'this' to be a row type:
       rowConstraintVar <- RowTVar <$> freshFlex
       unify a (Fix . TRow Nothing . TRowEnd $ Just rowConstraintVar) thisT
       -- close the row type
       resolvedThisT <- applyMainSubst thisT -- otherwise closeRow will not do what we want.
  --     unify a thisT (closeRow resolvedThisT)
       -- TODO: If the function returns a row type, it should be the resulting type; other it should be 'thisT'
       preds' <- unifyPredsL a preds
       let thisT' = TQual preds' thisTLabeled
           thisTLabeled = case resolvedThisT of
                              Fix (TRow Nothing rl) -> Fix $ TRow (label e1) rl
                              _ -> resolvedThisT
       return (thisT', ENew (a, thisT') e1' eArgs')
  inferType' env (ELet a n e1 e2) =
    do recType <- Fix . TBody . TVar <$> freshFlex
       recEnv <- addVarScheme env n $ schemeEmpty recType
       (t1, e1') <- inferType recEnv e1
       unify a (qualType t1) recType
       recType' <- applyMainSubst recType
       -- TODO hacky special treatment of names that start with uppercase letters. should perhaps be
       -- moved to translation layer.
       t1' <- case (isUpper $ head n, unFix recType') of
                  (True, TFunc (thisT:argsT) resT) ->
                    do closedThisT <- closeRow a thisT
                       traceLog $ vsep [ text "Closing 'this' row type: " <+> align (pretty thisT)
                                       , text "while inferring type of" <+> align (pretty n)
                                       , text "which has type: " <+> align (pretty recType')
                                       ]
                       traceLog $ text "\tClosed: " <+> align (pretty closedThisT)
                       unify a thisT closedThisT
                       return $ t1 { qualType = Fix $ TFunc (closedThisT:argsT) resT }
                  _ -> return $ t1
       (t', preds) <- generalize e1 env t1'
       env' <- addVarScheme env n t'
       (t2, e2') <- inferType env' e2
       preds' <- unifyPredsL a $ preds ++ concatMap qualPred [t1', t2]
       let resT = TQual preds' $ qualType t2
       return (resT, ELet (a, resT) n e1' e2')
  -- | Handling of mutable variable assignment.
  -- | Prevent mutable variables from being polymorphic.
  inferType' env (EPropAssign a objExpr prop expr1 expr2) =
    do (objT, objExpr') <- inferType env objExpr
       (rvalueT, expr1') <- inferType env expr1
       rowTailVar <- RowTVar <$> freshFlex
       -- Don't generalize prop assignments, because it causes inference of too-polymorphic types that
       -- the user may not have meant to use, and later cause errors when another assignment is not as
       -- polymorphic. In the future, type annotations could be used to make the prop assignment
       -- polymorphic.
       let rvalueSchemeFloated = TScheme [] TQual { qualPred = [], qualType = qualType rvalueT }
           rvalueRowType = Fix . TRow Nothing
                           . TRowProp (TPropGetName prop) rvalueSchemeFloated
                           . TRowProp (TPropSetName prop) rvalueSchemeFloated
                           $ TRowEnd (Just rowTailVar)
       unify a (qualType objT) rvalueRowType
       (expr2T, expr2') <- inferType env expr2 -- TODO what about the pred
       preds <- unifyPredsL a $ concatMap qualPred [objT, rvalueT, expr2T] -- TODO review
       let tRes = TQual preds $ qualType expr2T
       return (tRes, EPropAssign (a, tRes) objExpr' prop expr1' expr2')
  inferType' env (EArray a exprs) =
    do tv <- Fix . TBody . TVar <$> freshFlex
       te <- accumInfer env exprs
       let types = map (qualType . fst) te
       unifyl unify a $ zip (tv:types) types
       let t = qualEmpty $ Fix $ TCons TArray [tv]
       return (t, EArray (a,t) $ map snd te)
  inferType' env (ETuple a exprs) =
    do te <- accumInfer env exprs
       let t = TQual (concatMap (qualPred . fst) te) $ Fix . TCons TTuple . reverse $ map (qualType . fst) te
       return (t, ETuple (a,t) $ map snd te)
  inferType' env (EStringMap a exprs') =
    do let exprs = map snd exprs'
       elemType <- Fix . TBody . TVar <$> freshFlex
       te <- accumInfer env exprs
       let types = map (qualType . fst) te
       unifyAll a $ elemType:types
       allPreds <- unifyPredsL a . concatMap qualPred $ map fst te
       let t = TQual { qualPred = allPreds, qualType = Fix $ TCons TStringMap [elemType] }
       return (t, EStringMap (a,t) $ zip (map fst exprs') (map snd te))
  inferType' env (ERow a isOpen propExprs) =
    do te <- accumInfer env $ map snd propExprs
       endVar <- RowTVar <$> freshFlex
       let propNamesTypes = zip propExprs (reverse $ map fst te)
           rowEnd' = TRowEnd $ if isOpen then Just endVar else Nothing
           accumRowProp' (row, floatedPs') ((propName, propExpr), propType) =
             do (ts, floatedPs) <- generalize propExpr env propType
                -- TODO use unfloated predicates
                return (TRowProp (TPropGetName propName) ts $ TRowProp (TPropSetName propName) ts $ row, floatedPs' ++ floatedPs)
       (rowType', floatedPreds) <- foldM  accumRowProp' (rowEnd', []) propNamesTypes
       let rowType = TQual { qualPred = floatedPreds, qualType = Fix . TRow Nothing $ rowType' }
       return (rowType, ERow (a,rowType) isOpen $ zip (map fst propExprs) (map snd te))
  inferType' env (ECase a eTest eBranches) =
    do (eType, eTest') <- inferType env eTest
       infPatterns <- accumInfer env $ map (ELit a . fst) eBranches
       let patternTypes = map (qualType . fst) infPatterns
       unifyAll a $ qualType eType : patternTypes
       infBranches <- accumInfer env $ map snd eBranches
       let branchTypes = map (qualType . fst) infBranches
       unifyAll a branchTypes
       allPreds <- unifyPredsL a . concatMap qualPred $ eType : map fst infPatterns ++ map fst infBranches
       let tRes = TQual { qualPred = allPreds, qualType = qualType . fst $ head infBranches } -- TODO unsafe head
       return (tRes, ECase (a, tRes) eTest' $ zip (map fst eBranches) (map snd infBranches))
  inferType' env (EProp a eObj propName) =
    do (tObj, eObj') <- inferType env eObj
       -- This is a hack to make it easier to support higher-rank row properties.  We require to have
       -- determined the type of eObj "good enough" for us to find if it has the required property or
       -- not.  Otherwise we must assume the property is a monotype.
       let propTypeIfMono =
               do traceLog $ text "Failed to find prop: " <+> pretty propName <+> text " in type: " <+> pretty tObj
                  rowTail <- TRowEnd . Just . RowTVar <$> freshFlex
                  propTypeScheme <- schemeEmpty . Fix . TBody . TVar <$> freshFlex
                  unify a (Fix . TRow Nothing $ TRowProp (TPropGetName propName) propTypeScheme rowTail) (qualType tObj)
                  instantiate propTypeScheme
  
           propTypefromTRow tRowList =
               case Map.lookup (TPropGetName propName) . fst $ flattenRow tRowList of
                   Just knownPropTypeScheme -> instantiate knownPropTypeScheme
                   Nothing -> propTypeIfMono
  
       rowType <- case unFix (qualType tObj) of
                      TRow _ tRowList -> return $ Just tRowList
                      _ -> tryMakeRow $ unFix (qualType tObj)
       propType <-
           case rowType of
               Just tRowList -> propTypefromTRow tRowList
               _ -> propTypeIfMono
  
       return (propType, EProp (a,propType) eObj' propName)
  
  createEnv :: Map EVarName TypeScheme -> Infer (Map EVarName VarId)
  createEnv builtins = foldM addVarScheme' Map.empty $ Map.toList builtins
      where allTVars :: TypeScheme -> Set TVarName
            allTVars (TScheme qvars t) = freeTypeVars t `Set.union` Set.fromList qvars
  
            addVarScheme' :: Map EVarName VarId -> (EVarName, TypeScheme) -> Infer (Map EVarName VarId)
            addVarScheme' m (name, tscheme) =
              do allocNames <- forM (Set.toList $ allTVars tscheme)
                               $ \tvName -> (tvName,) <$> freshFlex
                 addVarScheme m name $ mapVarNames (safeLookup allocNames) tscheme
  
  
  typeInference :: Map EVarName TypeScheme -> Exp Source -> Infer (Exp (Source, QualType))
  typeInference builtins e =
    do env <- createEnv builtins
       (_t, e') <- inferType env e
       return e'
  
  ----------------------------------------------------------------------
  --
  -- | Mutable variable being assigned incompatible types:
  --
  -- >>> let p = emptySource
  -- >>> let fun args = EAbs p ("this":args)
  -- >>> let var = EVar p
  -- >>> let let' = ELet p
  -- >>> let tuple = ETuple p
  -- >>> let lit = ELit p
  -- >>> let app a b = EApp p a [lit LitUndefined, b]
  -- >>> let assign = EAssign p
  -- >>> let array = EArray p
  --
  -- x is known to have type forall a. a -> a, and to have been used in a context requiring bool -> bool (e.g. `x True`)
  --
  -- we now try to assign x := \y -> 2
  --
  -- This should fail because it "collapses" x to be Number -> Number which is not compatible with bool -> bool
  --
  -- >>> test $ let' "x" (fun ["z"] (var "z")) (let' "y" (tuple [app (var "x") (lit (LitNumber 2)), app (var "x") (lit (LitBoolean True))]) (assign "x" (fun ["y"] (lit (LitNumber 0))) (tuple [var "x", var "y"])))
  -- ":1:1*: Error: Could not unify: Number with Boolean"
  --
  -- The following should succeed because x is immutable and thus polymorphic:
  --
  -- >>> test $ let' "x" (fun ["z"] (var "z")) (let' "y" (tuple [app (var "x") (lit (LitNumber 2)), app (var "x") (lit (LitBoolean True))]) (tuple [var "x", var "y"]))
  -- "(c.(d -> d), (Number, Boolean))"
  --
  -- The following should fail because x is mutable and therefore a monotype:
  --
  -- >>> test $ let' "x" (fun ["z"] (var "z")) (let' "y" (tuple [app (var "x") (lit (LitNumber 2)), app (var "x") (lit (LitBoolean True))]) (assign "x" (fun ["z1"] (var "z1")) (tuple [var "x", var "y"])))
  -- ":1:1*: Error: Could not unify: Number with Boolean"
  --
  -- The following should also succeed because "x" is only ever used like this: (x True). The second assignment to x is: x := \z1 -> False, which is specific but matches the usage. Note that x's type is collapsed to: Boolean -> Boolean.
  --
  -- >>> test $ let' "x" (fun ["z"] (var "z")) (let' "y" (app (var "x") (lit (LitBoolean True))) (assign "x" (fun ["z1"] (lit (LitBoolean False))) (tuple [var "x", var "y"])))
  -- "((Boolean -> Boolean), Boolean)"
  --
  -- | Tests a setter for x being called with something more specific than x's original definition:
  -- >>> :{
  -- >>> test $ let'
  -- >>> "x" (fun ["a"] (var "a"))
  -- >>> (let' "setX"
  -- >>>    (fun ["v"]
  -- >>>             (let'
  -- >>>          "_" (assign "x" (var "v") (var "x")) (lit (LitBoolean False))))
  -- >>>    (let'
  -- >>>       "_" (app (var "setX") (fun ["a"] (lit (LitString "a"))))
  -- >>>       (app (var "x") (lit (LitBoolean True)))))
  -- >>> :}
  -- ":1:1*: Error: Could not unify: String with Boolean"
  --
  -- >>> test $ tuple [lit (LitBoolean True), lit (LitNumber 2)]
  -- "(Boolean, Number)"
  --
  -- >>> test $ let' "id" (fun ["x"] (var "x")) (assign "id" (fun ["y"] (var "y")) (var "id"))
  -- "a.(b -> b)"
  --
  -- >>> test $ let' "id" (fun ["x"] (var "x")) (assign "id" (lit (LitBoolean True)) (var "id"))
  -- ":1:1*: Error: Could not unify: a.(b -> b) with Boolean"
  --
  -- >>> test $ let' "x" (lit (LitBoolean True)) (assign "x" (lit (LitBoolean False)) (var "x"))
  -- "Boolean"
  --
  -- >>> test $ let' "x" (lit (LitBoolean True)) (assign "x" (lit (LitNumber 3)) (var "x"))
  -- ":1:1*: Error: Could not unify: Boolean with Number"
  --
  -- >>> test $ let' "x" (array [lit (LitBoolean True)]) (var "x")
  -- "[Boolean]"
  --
  -- >>> test $ let' "x" (array [lit $ LitBoolean True, lit $ LitBoolean False]) (var "x")
  -- "[Boolean]"
  --
  -- >>> test $ let' "x" (array []) (assign "x" (array []) (var "x"))
  -- "[a]"
  --
  -- >>> test $ let' "x" (array [lit $ LitBoolean True, lit $ LitNumber 2]) (var "x")
  -- ":1:1*: Error: Could not unify: Number with Boolean"
  --
  -- >>> test $ let' "id" (fun ["x"] (let' "y" (var "x") (var "y"))) (app (var "id") (var "id"))
  -- "c.(d -> d)"
  --
  -- >>> test $ let' "id" (fun ["x"] (let' "y" (var "x") (var "y"))) (app (app (var "id") (var "id")) (lit (LitNumber 2)))
  -- "Number"
  --
  -- >>> test $ let' "id" (fun ["x"] (app (var "x") (var "x"))) (var "id")
  -- ":1:1*: Error: Occurs check failed: a in (a -> b)"
  --
  -- >>> test $ fun ["m"] (let' "y" (var "m") (let' "x" (app (var "y") (lit (LitBoolean True))) (var "x")))
  -- "a.((Boolean -> b) -> b)"
  --
  -- >>> test $ app (lit (LitNumber 2)) (lit (LitNumber 2))
  -- ":1:1*: Error: Could not unify: Number with (Number -> a)"
  --
  -- EAssign tests
  -- >>> test $ let' "x" (fun ["y"] (lit (LitNumber 0))) (assign "x" (fun ["y"] (var "y")) (var "x"))
  -- "a.(Number -> Number)"
  --
  -- >>> test $ let' "x" (fun ["y"] (var "y")) (assign "x" (fun ["y"] (lit (LitNumber 0))) (var "x"))
  -- "a.(Number -> Number)"
  --
  -- >>> test $ let' "x" (fun ["y"] (var "y")) (tuple [app (var "x") (lit (LitNumber 2)), app (var "x") (lit (LitBoolean True))])
  -- "(Number, Boolean)"
  --
  -- >>> test $ let' "x" (fun ["y"] (var "y")) (app (var "x") (var "x"))
  -- "c.(d -> d)"
  --
  -- >>> test $ let' "x" (fun ["a"] (var "a")) (let' "getX" (fun ["v"] (var "x")) (let' "setX" (fun ["v"] (let' "_" (assign "x" (var "v") (var "x")) (lit (LitBoolean True)))) (let' "_" (app (var "setX") (fun ["a"] (lit (LitString "a")))) (var "getX"))))
  -- "e.(f -> d.(String -> String))"
  test :: Exp Source -> String
  test e = case runTypeInference e of
            Left err -> show $ pretty err
            Right expr -> show $ pretty $ snd . head . getAnnotations . minifyVars $ expr
  
  
  runTypeInference :: Exp Source -> Either TypeError (Exp (Source, QualType))
  runTypeInference e = runInfer $ typeInference Builtins.builtins e
  