content: module Typing.Expr (i_expr) where
  
  import Typing.Constraint
  import Typing.Env
  import Typing.Substitution
  import Typing.Subtyping
  import Typing.TypeError
  import Typing.Types
  import Typing.Util
  
  import Absyn.Base
  import Absyn.Meta
  import qualified Absyn.Untyped as U
  import qualified Absyn.Typed as T
  import Util.Error
  
  import Control.Monad (when, zipWithM)
  import Data.List (union)
  
  i_expr :: U.Expr -> Tc (T.Expr, Type)
  i_expr (meta :< Literal lit) =
    return (meta :< Literal lit, i_lit lit)
  
  i_expr (meta :< Ident [i]) = do
    ty <- lookupValue i
    return (meta :< Ident [i], ty)
  
  -- TODO: Clear this up - should be handled by the renamer now
  i_expr (_ :< Ident (_:_)) = undefined
  i_expr (_ :< Ident []) = undefined
  
  i_expr (meta :< ParenthesizedExpr expr) = do
    (expr', ty) <- i_expr expr
    return (meta :< ParenthesizedExpr expr', ty)
  
  i_expr (meta :< BinOp _ _ lhs op rhs) = do
    tyOp@(Fun gen _ _) <- lookupValue op
    (lhs', lhsTy) <- i_expr lhs
    (rhs', rhsTy) <- i_expr rhs
    (retType', typeArgs)  <- inferTyArgs [lhsTy, rhsTy] tyOp
    constraintArgs <- inferConstraintArgs gen [] typeArgs
    return (meta :< BinOp constraintArgs [] lhs' (op, tyOp) rhs', retType')
  
  i_expr (meta :< Match expr cases) = do
    (expr', ty) <- i_expr expr
    (cases', casesTy) <- unzip <$> mapM (i_case ty) cases
    let retTy = case casesTy of
                  [] -> void
                  x:xs -> foldl (\/) x xs
    return (meta :< Match expr' cases', retTy)
  
  i_expr (meta :< Call fn constraintArgs types []) =
    i_expr (meta :< Call fn constraintArgs types [emptySpanFromPos (spanEnd meta) :< VoidExpr])
      where
        emptySpanFromPos pos = SourceSpan { spanStart = pos, spanEnd = pos }
  
  i_expr (meta :< Call fn _ types args) = do
    (fn', tyFn) <- i_expr fn
    let tyFn' = normalizeFnType tyFn
    (tyFn''@(Fun gen _ retType), skippedVars) <- adjustFnType (null types) args tyFn'
    (retType', args', typeArgs)  <-
          case (tyFn'', types) of
            (Fun (_:_) params _, []) -> do
              (_, argsTy) <- unzip <$> mapM i_expr args
              (retType, typeArgs) <- inferTyArgs argsTy tyFn''
              let s = zipSubst (map fst gen) typeArgs
              let params' = map (applySubst s) params
              args' <- zipWithM instSubtype args params'
              return (retType, args', typeArgs)
            (Fun gen params _, _) -> do
              types' <- mapM resolveType types
              let s = zipSubst (map fst gen) types'
              let params' = map (applySubst s) params
              args' <- zipWithM instSubtype args params'
              return (applySubst s retType, args', types')
            _ -> undefined
    constraintArgs <- inferConstraintArgs gen skippedVars typeArgs
    return (meta :< Call fn' constraintArgs types args', retType')
  
  i_expr (meta :< Record fields) = do
    (exprs, types) <- mapM (i_expr . snd) fields >>= return . unzip
    let labels = map fst fields
    let fieldsTy = zip labels types
    let recordTy = Rec fieldsTy
    let record = Record (zip labels exprs)
    return (meta :< record, recordTy)
  
  i_expr (meta :< FieldAccess expr field) = do
    (expr', ty) <- i_expr expr
    let
        aux :: Type -> [(String, Type)] -> Tc (T.Expr, Type)
        aux ty r = case lookup field r of
                  Nothing -> throwError $ UnknownField ty field
                  Just t -> return (meta :< FieldAccess expr' field, t)
    case ty of
      Rec r -> aux ty r
      Cls _ -> do
        vars <- lookupInstanceVars ty
        aux ty vars
      _ -> throwError . GenericError $ "Expected a record, but found value of type " ++ show ty
  
  i_expr (meta :< If ifCond ifBody elseBody) = do
    (ifCond', ty) <- i_expr ifCond
    ty <:! bool
    (ifBody', ifTy) <- i_body ifBody
    (elseBody', elseTy) <- i_body elseBody
    return (meta :< If ifCond' ifBody' elseBody', ifTy \/ elseTy)
  
  i_expr (meta :< List _ items) = do
    (items', itemsTy) <- unzip <$> mapM i_expr items
    (ty, itemTy) <- case itemsTy of
                      [] -> do
                        nilTy <- lookupValue "Nil"
                        return (nilTy, Bot)
                      x:xs ->
                        let ty = foldl (\/) x xs
                         in return (list ty, ty)
    return (meta :< List (Just itemTy) items', ty)
  
  i_expr (meta :< FnExpr fn) = do
    (fn', ty) <- i_fn fn
    return $ (meta :< FnExpr fn', ty)
  
  i_expr (meta :< Negate _ expr) = do
    (expr', ty) <- i_expr expr
    intf <- lookupInterface "Std.Number"
    constrArgs <- boundsCheck ty intf
    return (meta :< Negate constrArgs expr', ty)
  
  -- Expressions generated during type checking
  i_expr (meta :< VoidExpr) = return (meta :< VoidExpr, void)
  
  i_expr (_ :< TypeCall {}) = undefined
  
  instSubtype :: U.Expr -> Type -> Tc T.Expr
  instSubtype arg@(meta :< _) ty = do
    (arg', argTy) <- i_expr arg
    arg'' <- case (argTy, ty) of
               (Fun gen@(_:_) _ _, Fun [] _ _) -> do
                 typeArgs <- inferTyAbs argTy ty
                 constraintArgs <- inferConstraintArgs gen [] typeArgs
                 return $ meta :< TypeCall arg' constraintArgs
               _ -> do
                 argTy <:! ty
                 return arg'
    return arg''
  
  inferConstraintArgs :: [BoundVar] -> [BoundVar] -> [Type] -> Tc [ConstraintArg]
  inferConstraintArgs gen skippedVars typeArgs = do
    concat <$> zipWithM findConstrArgs gen typeArgs'
      where
        typeArgs' = zipWith findHoles gen typeArgs
  
        findHoles var ty =
          if var `elem` skippedVars
             then mkHole var
             else ty
  
        findConstrArgs (_, []) tyArg = do
          return [CAType tyArg]
  
        findConstrArgs (_, bounds) tyArg = do
          concat <$> mapM (boundsCheck tyArg) bounds
  
  -- Checks if `t1` implements the interface in `t2`
  -- returns a list with a single ConstraintArg indicating
  --  how to satisfy the bounds. See `ConstraintArg` for
  --  an explanation on the kinds of ConstraintArgs
  boundsCheck :: Type -> Intf -> Tc [ConstraintArg]
  boundsCheck t1 t2@(Intf name _ _) = do
    args <- boundsCheck' t1 t2
    if null args
       then throwError $ MissingImplementation name t1
       else return args
  
  boundsCheck' :: Type -> Intf -> Tc [ConstraintArg]
  boundsCheck' v@(Var _ bounds) intf = do
    return $ if intf `elem` bounds
                then [CABound v intf]
                else []
  
  boundsCheck' (TyApp ty args) intf@(Intf name _ _) = do
    implementations <- lookupImplementations name
    case lookup ty implementations of
      Nothing -> return []
      Just vars -> do
        let aux arg (_, bounds) =
              concat <$> mapM (boundsCheck arg) bounds
        args <- concat <$> zipWithM aux args vars
        return [CAPoly ty intf args]
  
  boundsCheck' (Forall params ty) intf =
    boundsCheck' (params \\ ty) intf
  
  boundsCheck' (TyAbs params ty) intf =
    boundsCheck' (params \\ ty) intf
  
  boundsCheck' Bot intf =
    return [CABound Bot intf]
  
  boundsCheck' ty intf@(Intf name _ _) = do
    implementations <- lookupImplementations name
    case lookup ty implementations of
      Just [] -> return [CABound ty intf]
      _ -> return []
  
  normalizeFnType :: Type -> Type
  normalizeFnType (Fun gen params (Fun [] params' retTy)) =
    normalizeFnType (Fun gen (params ++ params') retTy)
  normalizeFnType ty = ty
  
  adjustFnType :: Bool -> [a] -> Type -> Tc (Type, [BoundVar])
  adjustFnType allowHoles args fn@(Fun gen params retType) = do
    let lArgs = length args
    case compare lArgs (length params) of
      EQ -> return (fn, [])
      LT ->
        let headArgs = take lArgs params
            tailArgs = drop lArgs params
            skippedGen = filter aux gen
            aux (v, _) = allowHoles && v `elem` (fv $ Fun [] tailArgs retType) && v `notElem` (foldl union [] $ map fv headArgs)
         in return (Fun gen headArgs $ Fun skippedGen tailArgs retType, skippedGen)
      GT -> throwError ArityMismatch
  adjustFnType _ _ ty = throwError . GenericError $ "Expected a function, found " ++ show ty
  
  i_lit :: Literal -> Type
  i_lit (Integer _) = int
  i_lit (Float _) = float
  i_lit (Char _) = char
  i_lit (String _) = string
  
  i_case :: Type -> U.Case -> Tc (T.Case, Type)
  i_case ty (meta :< Case pattern caseBody) = do
    m <- startMarker
    pattern' <- c_pattern ty pattern
    endMarker m
  
    (caseBody', ty) <- i_body caseBody
  
    clearMarker m
  
    return (meta :< Case pattern' caseBody', ty)
  
  c_pattern :: Type -> U.Pattern -> Tc T.Pattern
  c_pattern _ (meta :< PatDefault) =
    return $ meta :< PatDefault
  
  c_pattern ty (meta :< PatLiteral l) = do
    let litTy = i_lit l
    litTy <:! ty
    return $ meta :< PatLiteral l
  
  c_pattern ty (meta :< PatVar v) = do
    insertValue v ty
    return $ meta :< PatVar v
  
  c_pattern ty@(Rec tyFields) (meta :< PatRecord fields) = do
    fields' <- mapM aux fields
    return $ meta :< PatRecord fields'
      where
        aux (key, pat) = do
          case lookup key tyFields of
            Just ty -> do
              pat' <- c_pattern ty pat
              return (key, pat')
            Nothing ->
              throwError . GenericError $ "Matching against field `" ++ key ++ "`, which is not included in the type of the value being matched, `" ++ show ty ++ "`"
  
  c_pattern ty (_ :< PatRecord _) = do
    throwError . GenericError $ "Using a record pattern, but value being matched has type `" ++ show ty ++ "`"
  
  c_pattern ty (meta :< PatList pats rest) = do
    itemTy <- getItemTy ty
    pats' <- mapM (c_pattern itemTy) pats
    rest' <- case rest of
               NoRest -> return NoRest
               DiscardRest -> return DiscardRest
               NamedRest n -> do
                 insertValue n ty
                 return (NamedRest n)
    return $ meta :< PatList pats' rest'
    where
      getItemTy (Forall _ (TyApp (Con "List") _)) =
        return Top
  
      getItemTy (TyApp (Con "List") [ty]) =
        return ty
  
      getItemTy _ =
        throwError . GenericError $ "Using a list pattern, but value being matched has type `" ++ show ty ++ "`"
  
  c_pattern ty (meta :< PatCtor name vars) = do
    ctorTy <- lookupValue name
    let (fnTy, params, retTy) = case ctorTy of
                              fn@(Fun [] params retTy) -> (fn, params, retTy)
                              fn@(Fun gen params retTy) -> (fn, params, Forall (map fst gen) retTy)
                              t -> (Fun [] [] t, [], t)
    when (length vars /= length params) (throwError ArityMismatch)
    retTy <:! ty
    let substs = case (retTy, ty) of
                   (Forall gen _, TyApp _ args) -> zipSubst gen args
                   _ -> emptySubst
    let params' = map (applySubst substs) params
    vars' <- zipWithM c_pattern params' vars
    return $ meta :< PatCtor (name, fnTy) vars'
  