content: {- |
  Module      :  ./HasCASL/TypeCheck.hs
  Description :  type checking terms and program equations
  Copyright   :  (c) Christian Maeder and Uni Bremen 2003
  License     :  GPLv2 or higher, see LICENSE.txt
  
  Maintainer  :  Christian.Maeder@dfki.de
  Stability   :  experimental
  Portability :  portable
  
  type inference based on
  
  <http://www.cs.fiu.edu/~smithg/papers/>
  
  Principal Type Schemes for Functional Programs with Overloading and
  Subtyping, Geoffrey S. Smith, Science of Computer Programming 23(2-3),
  pp. 197-226, December 1994
  -}
  
  module HasCASL.TypeCheck
      ( typeCheck
      , resolveTerm
      ) where
  
  import HasCASL.Unify
  import HasCASL.VarDecl
  import HasCASL.As
  import HasCASL.Builtin
  import HasCASL.FoldType
  import HasCASL.Le
  import HasCASL.PrintLe
  import HasCASL.AsUtils
  import HasCASL.MixAna
  import HasCASL.TypeAna
  import HasCASL.MapTerm
  import HasCASL.FoldTerm
  import HasCASL.Constrain
  import HasCASL.ProgEq
  import HasCASL.MinType
  
  import qualified Data.Map as Map
  import qualified Data.Set as Set
  import Common.Id
  import Common.GlobalAnnotations
  import Common.Result
  import Common.DocUtils
  import Common.Lib.State
  
  import Data.Maybe (catMaybes)
  import Control.Monad (when, unless)
  
  substTerm :: Subst -> Term -> Term
  substTerm s = mapTerm (id, subst s)
  
  -- | mixfix and type resolution
  resolveTerm :: Type -> Term -> State Env (Maybe Term)
  resolveTerm mt trm = do
      mtrm <- resolve trm
      case mtrm of
        Nothing -> return Nothing
        Just t -> typeCheck mt t
  
  -- | get a constraint from a type argument instantiated with a type
  mkConstraint :: (Type, VarKind) -> Constrain
  mkConstraint (ty, vk) = case vk of
      MissingKind -> error "mkConstraint"
      VarKind k -> Kinding ty k
      Downset super -> Subtyping ty super
  
  instantiate :: [Type] -> TypeScheme
              -> State Env (Maybe (Type, [(Type, VarKind)]))
  instantiate tys sc@(TypeScheme tArgs t _) =
      if null tys || length tys /= length tArgs then
           if null tys then fmap Just $ toEnvState $ freshInst sc
           else do
               addDiags [mkDiag Hint ("for type scheme '" ++
                   showDoc t "' wrong length of instantiation list") tys]
               return Nothing
          else let s = Map.fromList $ zip [-1, -2 ..] tys
               in return $ Just
                  (substGen s t, zip tys $ map (substTypeArg s) tArgs)
  
  instOpInfo :: [Type] -> OpInfo
             -> State Env (Maybe (Type, [Type], Constraints, OpInfo))
  instOpInfo tys oi = do
       m <- instantiate tys $ opType oi
       return $ case m of
         Just (ty, cl) ->
             Just (ty, map fst cl, Set.fromList $ map mkConstraint cl, oi)
         Nothing -> Nothing
  
  {- This function has the problem that the type of an earlier component may
     restrict the types of following components too much. -}
  checkList :: Bool -> [Maybe Type] -> [Term]
            -> State Env [(Subst, Constraints, [Type], [Term])]
  checkList isP mtys trms = case (mtys, trms) of
      (ty : rty, trm : rt) -> do
        fts0 <- inferWithMaybeType isP ty trm
        fts <- reduce True fts0
        combs <- mapM ( \ (sf, cs, tyf, tf) -> do
                        vs <- gets localVars
                        putLocalVars $ substVarTypes sf vs
                        rts <- checkList isP (map (fmap (subst sf)) rty) rt
                        putLocalVars vs
                        return $ map ( \ (sr, cr, tys, tts) ->
                               (compSubst sf sr,
                                substC sr cs `joinC` cr,
                                subst sr tyf : tys,
                                       tf : tts)) rts) fts
        return $ concat combs
      _ -> return [(eps, noC, [], [])]
  
  -- | reduce a substitution
  reduceR :: Bool -> Env -> (Subst, Constraints, Type, Term)
          -> State Int (Result (Subst, Constraints, Type, Term))
  reduceR doMono te (s, cr, ty, tr) = do
      Result ds0 mc <- shapeRelAndSimplify True te cr
        $ if doMono then Just ty else Nothing
      let ds = map (improveDiag tr) ds0
      return $ case mc of
          Nothing -> Result ds Nothing
          Just (s1, qs) -> let
            s2 = compSubst s s1
            in Result ds $ Just
               (s2, qs, subst s1 ty, substTerm s2 tr)
  
  -- | reduce a substitution
  reduce :: Bool -> [(Subst, Constraints, Type, Term)]
         -> State Env [(Subst, Constraints, Type, Term)]
  reduce doMono alts = do
         te <- get
         combs <- mapM (\ alt -> do
           Result ds mc <- toEnvState $ reduceR doMono te alt
           addDiags ds
           case mc of
             Nothing -> return []
             Just q -> return [q]) alts
         return $ concat combs
  
  checkForUninstantiatedVars :: GlobalAnnos -> Term -> Range -> State Env ()
  checkForUninstantiatedVars ga t p = let
    tys = getAllVarTypes t
    in unless (null tys) $ addDiags
      [(mkDiag Error ("in term '" ++ showGlobalDoc ga t
                      "'\n are uninstantiated type variables")
        $ Set.toList $ Set.unions
        $ map (Set.fromList . freeTVarIds) tys)
       {diagPos = p}]
  
  simplifyTypedTerms :: Env -> Term -> Term
  simplifyTypedTerms e = foldTerm mapRec
    { foldTypedTerm = \ _ nt q ty ps ->
        let ntyped = TypedTerm nt q ty ps
            ityped = TypedTerm nt Inferred ty ps
        in case getTypeOf nt of
          Nothing -> ntyped
          Just ty2 -> let isSubT = lesserType e ty2 ty in
            case q of
              InType | isSubT -> unitTerm trueId ps
              _ -> case nt of
                TypedTerm nt2 q2 _ _
                 | q2 == AsType && q /= InType && lesserType e ty ty2
                    -> TypedTerm nt2 q2 ty ps
                 | q == AsType && elem q2 [OfType, Inferred] && isSubT
                     -> ityped
                 | otherwise -> ntyped
                _ -> if q == AsType && isSubT then ityped else ntyped }
  
  -- | type checking a term
  typeCheck :: Type -> Term -> State Env (Maybe Term)
  typeCheck exTy trm =
      do alts0 <- inferWithMaybeType False (Just exTy) trm
         alts <- reduce True alts0
         te <- get
         let p = getRange trm
             ga = globAnnos te
         case typeNub te q2p alts of
           [] -> do
             addDiags [mkNiceDiag ga Error "no typing for" trm]
             return Nothing
           [(_, rcs, ty, t)] -> do
             unless (Set.null rcs)
                   $ addDiags [(mkDiag Error ("in term '"
                               ++ showGlobalDoc ga t "' of type '"
                               ++ showDoc ty "'\n unresolved constraints")
                                   rcs) {diagPos = p}]
             checkForUninstantiatedVars ga t p
             return $ Just $ simplifyTypedTerms te t
           falts -> do
                 addDiags [Diag Error
                           ("ambiguous typings\n " ++
                            showSepList ("\n " ++)
                            ( \ (n, t) -> shows n . (". " ++) . showDoc t)
                            (zip [1 .. (5 :: Int)] $ map ( \ (_, _, _, t) ->
                                            t) falts) "")
                              p]
                 return Nothing
  
  freshTypeVar :: Term -> State Env Type
  freshTypeVar t =
      do (var, c) <- toEnvState $ freshVar $ Id [] [] $ getRange t
         return $ TypeName var rStar c
  
  substVarTypes :: Subst -> Map.Map Id VarDefn -> Map.Map Id VarDefn
  substVarTypes s = Map.map ( \ (VarDefn t) -> VarDefn $ subst s t)
  
  warnEmpty :: Maybe Type -> Term -> [a] -> State Env ()
  warnEmpty mt trm res = do
      ga <- gets globAnnos
      when (null res) $ addDiags [mkNiceDiag ga Hint ("untypeable term" ++
        case mt of
          Nothing -> ""
          Just ty -> " (with type: " ++ showGlobalDoc ga ty ")") trm]
  
  -- | infer type of application, consider lifting for lazy types
  inferAppl :: Bool -> Range -> Term -> Term
            -> State Env [(Subst, Constraints, Type, Term)]
  inferAppl isP ps t1 t2 = do
      ops <- infer isP t1
      warnEmpty Nothing t1 ops
      vs <- gets localVars
      e <- get
      combs <- mapM ( \ (sf, cf, funty, tf) -> do
          (ok, sfty, frty, sub) <- case getTypeAppl funty of
                            (topTy, [paty, prty]) |
                              lesserType e topTy $ toFunType PFunArr ->
                                  return (True, Just paty, prty, False)
                            (topTy, [prty]) |
                              lesserType e topTy lazyTypeConstr ->
                                  return (True, Just unitType, prty, False)
                            (TypeName _ _ c, []) | c > 0 -> do
                              rty <- freshTypeVar t1
                              return (True, Nothing, rty, True)
                            _ -> return (False, Nothing, funty, False)
          if ok then do
                  putLocalVars $ substVarTypes sf vs
                  args <- inferWithMaybeType isP sfty t2
                  warnEmpty sfty t2 args
                  putLocalVars vs
                  return $ map ( \ (sa, ca, arty, ta) ->
                          let nTy = subst sa frty in
                            (compSubst sf sa, (if sub then
                                 insertC (Subtyping (subst sa funty)
                                  $ mkFunArrType arty PFunArr nTy) else id)
                                $ joinC ca $ substC sa cf, nTy,
                              TypedTerm (ApplTerm tf ta ps)
                              Inferred nTy ps)) args
              else return []) ops
      reduce False $ concat combs
  
  getAllVarTypes :: Term -> [Type]
  getAllVarTypes = filter (not . null . freeTVars) . getAllTypes
  
  mkTypedTerm :: Term -> Type -> Term
  mkTypedTerm trm ty = case trm of
      TupleTerm ts ps | not (null ts) -> let
        n = length ts
        (topTy, tArgs) = getTypeAppl ty
        in if n > 1 && topTy == toProdType n ps
               && length tArgs == n then
        TupleTerm (zipWith mkTypedTerm ts tArgs) ps
        else TypedTerm trm Inferred ty ps
      LetTerm br eqs inTrm ps -> LetTerm br eqs (mkTypedTerm inTrm ty) ps
      QuantifiedTerm quant decls t ps ->
          QuantifiedTerm quant decls (mkTypedTerm t ty) ps
      _ -> TypedTerm trm Inferred ty $ getRange trm
  
  -- | efficiently infer type of a monomorphic tuple term
  inferWithMaybeType :: Bool -> Maybe Type -> Term
                     -> State Env [(Subst, Constraints, Type, Term)]
  inferWithMaybeType isP mt trm = case (trm, mt) of
      (TupleTerm ts@(_ : _ : _) ps, Just ty) -> case getTypeAppl ty of
          (TypeName i _ _, argTys@(_ : _ : _)) | isProductId i ->
            if length ts == length argTys then
                if all (null . freeTVars) argTys then do
                   -- remaining type variables would not become instantiated
                  ls <- checkList isP (map Just argTys) ts
                  return $ map ( \ (su, cs, tys, trms) ->
                    ( su, cs, mkProductTypeWithRange tys ps
                    , mkTupleTerm trms ps)) ls
                else inferWithMaybeTypeAux isP mt trm
             else return [] -- fail for tuples of different lengths
          _ -> inferWithMaybeTypeAux isP mt trm
      _ -> inferWithMaybeTypeAux isP mt trm
  
  -- | infer type of term (or a pattern if the Bool is True)
  inferWithMaybeTypeAux :: Bool -> Maybe Type -> Term
                     -> State Env [(Subst, Constraints, Type, Term)]
  inferWithMaybeTypeAux isP mt trm = do
    rs <- infer isP trm
    te <- get
    case mt of
      Nothing -> return rs
      Just inTy -> do
          combs <- mapM (\ q@(s, c, ty, t) -> let nTy = subst s inTy in
              if ty == nTy then return [q] else do
              Result ds mc <- toEnvState $ reduceR False te
                (s, insertC (Subtyping ty nTy) c, nTy, mkTypedTerm t nTy)
              case mc of
                Nothing -> do
                  addDiags ds
                  return []
                Just alt -> return [alt]) rs
          return $ concat combs
  
  -- | infer type of term (or a pattern if the Bool is True)
  infer :: Bool -> Term -> State Env [(Subst, Constraints, Type, Term)]
  infer isP trm = do
      e <- get
      let tm = typeMap e
          bs = binders e
          vs = localVars e
          ga = globAnnos e
      case trm of
          qv@(QualVar (VarDecl _ ty _ _)) -> return [(eps, noC, ty, qv)]
          QualOp br i sc tys k ps -> do
              ms <- instOpInfo tys OpInfo { opType = sc
                                          , opAttrs = Set.empty
                                          , opDefn = NoOpDefn br }
              return $ case ms of
                  Nothing -> []
                  Just (ty, inst, cs, _) ->
                      let qv = TypedTerm (QualOp br i sc inst k ps)
                               Inferred ty ps
                      in [(eps, cs, ty, qv)]
          ResolvedMixTerm i tys ts ps -> case (Map.lookup i bs, ts) of
            (Just j, pat : rt@(_ : _)) -> case reverse rt of
              lt : ft -> infer isP $ ResolvedMixTerm j tys
                  (reverse $ LambdaTerm [pat] Partial lt ps : ft) ps
              [] -> error "ResolvedMixTerm: binder"
            _ ->
              if null ts then case Map.lookup i vs of
                 Just (VarDefn t) ->
                   infer isP $ QualVar $ VarDecl i t Other ps
                 Nothing -> do
                      insts <- mapM (instOpInfo tys) $ getMinAssumps e i
                      let ls = map ( \ (ty, is, cs, oi) ->
                                (eps, ty, is, cs, oi)) $ catMaybes insts
                      -- possibly fresh variable
                      vl <- if isP && null tys && null ls
                            && (isSimpleId i || i == simpleIdToId uTok) then do
                          vty <- freshTypeVar trm
                          return
                            [(eps, noC, vty, QualVar $ VarDecl i vty Other ps)]
                        else do
                          when (null ls) $
                            addDiags [mkDiag Hint "no type found for" i]
                          return []
                      return $ vl ++ map
                        ( \ (s, ty, is, cs, oi) ->
                          let od = opDefn oi
                              br = case od of
                                   NoOpDefn v -> v
                                   Definition v _ -> v
                                   _ -> Op
                              ik = if null tys then Infer else UserGiven
                         in (s, cs, ty, case opType oi of
                             sc@(TypeScheme [] _ _) ->
                                    QualOp br (PolyId i [] ps) sc [] ik ps
                             sc -> TypedTerm (QualOp br (PolyId i [] ps)
                                              sc is ik ps)
                                         Inferred ty ps)) ls
              else inferAppl isP ps (ResolvedMixTerm i tys [] ps)
                   $ mkTupleTerm ts ps
          ApplTerm t1 t2 ps -> inferAppl isP ps t1 t2
          TupleTerm ts ps -> if null ts then return [(eps, noC, unitType, trm)]
            else do
              ls <- checkList isP (map (const Nothing) ts) ts
              return $ map ( \ (su, cs, tys, trms) ->
                (su, cs, mkProductTypeWithRange tys ps, mkTupleTerm trms ps)) ls
          TypedTerm t qual ty ps ->
              case qual of
                  InType -> do
                      vTy <- freshTypeVar t
                      rs <- infer False t
                      return $ map ( \ (s, cs, typ, tr) ->
                             let sTy = subst s ty in
                                 ( s, insertC (Subtyping sTy vTy)
                                       $ insertC (Subtyping typ vTy) cs
                                 , unitType, TypedTerm tr qual sTy ps)) rs
                  AsType -> do
                      vTy <- freshTypeVar t
                      rs <- infer False t
                      return $ map ( \ (s, cs, typ, tr) ->
                          let sTy = subst s ty in
                                  ( s, insertC (Subtyping sTy vTy)
                                         $ insertC (Subtyping typ vTy) cs
                                  , sTy, TypedTerm tr qual sTy ps)) rs
                  _ -> do
                      let decl = case t of
                            ResolvedMixTerm _ tys ts _
                                | isP && null tys && null ts && qual == OfType
                                -> True
                            _ -> False
                      rs <- inferWithMaybeType isP
                             (if decl then Nothing else Just ty) t
                      return $ map ( \ (s, cs, _, tr) ->
                            let sTy = subst s ty in
                                  (s, cs, sTy, case tr of
                                   QualVar (VarDecl vp _ po _)
                                       | decl -- shadow
                                       -> QualVar (VarDecl vp sTy po ps)
                                   _ -> if (qual == Inferred || case tr of
                                          QualVar _ -> True
                                          QualOp {} -> True
                                          TypedTerm _ OfType _ _ -> True
                                          _ -> False)
                                          && maybe False (eqStrippedType sTy)
                                                 (getTypeOf tr)
                                        then tr
                                        else TypedTerm tr qual sTy ps)) rs
          QuantifiedTerm quant decls t ps -> do
              mapM_ addGenVarDecl decls
              rs <- inferWithMaybeType False (Just $ mkLazyType unitType) t
              putLocalVars vs
              putTypeMap tm
              return $ map ( \ (s, cs, typ, tr) ->
                (s, cs, typ, QuantifiedTerm quant decls tr ps)) rs
          LambdaTerm pats part resTrm ps -> do
              ls <- checkList True (map (const Nothing) pats) pats
              rs <- mapM ( \ ( s, cs, patys, nps) -> do
                         mapM_ (addLocalVar True) $ concatMap extractVars nps
                         es <- infer False resTrm
                         putLocalVars vs
                         return $ map ( \ (s2, cr, rty, rtm) ->
                                        let s3 = compSubst s s2
                                            typ = getFunType (subst s rty)
                                               part $ map (subst s2) patys
                                        in
                                        (s3, joinC (substC s2 cs) $ substC s cr
                                        , typ,
                                         TypedTerm
                                         (LambdaTerm nps part rtm ps)
                                         Inferred typ ps)) es) ls
              return $ concat rs
          CaseTerm ofTrm eqs ps -> do
              ts <- infer False ofTrm
              rty <- freshTypeVar trm
              when (null ts) $ addDiags [mkNiceDiag ga Hint
                                        "unresolved of-term in case" ofTrm]
              rs <- mapM ( \ (s1, cs, oty, otrm) -> do
                   es <- inferCaseEqs oty (subst s1 rty) eqs
                   return $ map ( \ (s2, cr, _, ty, nes) ->
                                  (compSubst s1 s2,
                                   substC s2 cs `joinC` cr, ty,
                                   TypedTerm (CaseTerm otrm nes ps)
                                   Inferred ty ps)) es) ts
              return $ concat rs
          LetTerm br eqs inTrm ps -> do
              es <- inferLetEqs eqs
              rs <- mapM ( \ (s1, cs, _, nes) -> do
                 mapM_ (addLocalVar True) $ concatMap
                           ( \ (ProgEq p _ _) -> extractVars p) nes
                 ts <- infer False inTrm
                 return $ map ( \ (s2, cr, ty, nt) ->
                                (compSubst s1 s2,
                                 substC s2 cs `joinC` cr,
                                 ty, LetTerm br nes nt ps)) ts) es
              putLocalVars vs
              return $ concat rs
          AsPattern (VarDecl v _ ok qs) pat ps -> do
             pats <- infer True pat
             return $ map ( \ (s1, cs, t1, p1) -> (s1, cs, t1,
                            AsPattern (VarDecl v t1 ok qs) p1 ps)) pats
          _ -> do ty <- freshTypeVar trm
                  addDiags [mkNiceDiag ga Error "unexpected term" trm]
                  return [(eps, noC, ty, trm)]
  
  inferLetEqs :: [ProgEq] -> State Env [(Subst, Constraints, [Type], [ProgEq])]
  inferLetEqs es = do
      let pats = map (\ (ProgEq p _ _) -> p) es
          trms = map (\ (ProgEq _ t _) -> t) es
          qs = map (\ (ProgEq _ _ q) -> q) es
      do vs <- gets localVars
         newPats <- checkList True (map (const Nothing) pats) pats
         combs <- mapM ( \ (sf, pcs, tys, pps) -> do
               mapM_ (addLocalVar True) $ concatMap extractVars pps
               newTrms <- checkList False (map Just tys) trms
               return $ map ( \ (sr, tcs, tys2, tts ) ->
                            (compSubst sf sr,
                             joinC tcs $ substC sr pcs, tys2,
                             zipWith3 (ProgEq . substTerm sr)
                                 pps tts qs)) newTrms) newPats
         putLocalVars vs
         return $ concat combs
  
  inferCaseEq :: Type -> Type -> ProgEq
              -> State Env [(Subst, Constraints, Type, Type, ProgEq)]
  inferCaseEq pty tty (ProgEq pat trm ps) = do
     pats1 <- inferWithMaybeType True (Just pty) pat
     e <- get
     let pats = filter ( \ (_, _, _, p) -> isPat e p) pats1
         ga = globAnnos e
     when (null pats)
       $ addDiags [mkNiceDiag ga Hint "unresolved case pattern" pat]
     vs <- gets localVars
     es <- mapM ( \ (s, cs, ty, p) -> do
                  mapM_ (addLocalVar True) $ extractVars p
                  ts <- inferWithMaybeType False (Just $ subst s tty) trm
                  putLocalVars vs
                  return $ map ( \ (st, cr, tyt, t) ->
                         (compSubst s st,
                          substC st cs `joinC` cr,
                          subst st ty, tyt,
                          ProgEq p t ps)) ts) pats
     return $ concat es
  
  inferCaseEqs :: Type -> Type -> [ProgEq]
              -> State Env [(Subst, Constraints, Type, Type, [ProgEq])]
  inferCaseEqs pty tTy [] = return [(eps, noC, pty, tTy, [])]
  inferCaseEqs pty tty (eq : eqs) = do
    fts <- inferCaseEq pty tty eq
    rs <- mapM (\ (s, cs, pty1, tty1, ne) -> do
                rts <- inferCaseEqs pty1 tty1 eqs
                return $ map ( \ (s2, cr, pty2, tty2, nes) ->
                               (compSubst s s2,
                                substC s2 cs `joinC` cr,
                                pty2, tty2, ne : nes)) rts) fts
    return $ concat rs
  