content: {-# LANGUAGE PatternGuards #-}
  module ISyntaxCheck(iGetKind,
                      tCheckIPackage,
                      tCheckIModule) where
  
  import qualified Data.Map as M
  import qualified Data.Set as S
  import qualified EquivalenceClass as EC
  
  import ErrorUtil(internalError)
  import Flags(Flags)
  import PPrint hiding(empty)
  import Position
  import Id
  import PreIds
  import ISyntax
  import ISyntaxUtil
  import SymTab(SymTab, mustFindClass, findSClass)
  
  import Pred
  import CType
  import TCMisc
  import TIMonad
  
  import IOUtil(progArgs)
  import Util(tracep, fromJustOrErr)
  
  doTraceEqWitnesses :: Bool
  doTraceEqWitnesses = "-trace-eq-witnesses" `elem` progArgs
  
  trace_witness :: String -> a -> a
  trace_witness = tracep doTraceEqWitnesses
  
  -----
  
  eqType :: Flags -> SymTab -> Env -> IType -> IType -> Bool
  eqType flags symt r (ITForAll i k t) (ITForAll i' k' t') =
      k == k' &&
      eqType flags symt r t (tSubst i' (ITVar i) t')
  eqType flags symt r t t' = eqType0 flags symt r t t'
  
  eqType0 :: Flags -> SymTab -> Env -> IType -> IType -> Bool
  eqType0 flags symt r@(E _ _ eqs _) t t' =
      --trace ("eqType0 " ++ ppReadable(t,t')) $
      t == t' ||
      EC.isEqual eqs t t' ||
      eqType1 flags symt r t t' ||
      -- try satisfying NumEq
      -- XXX this is wasted work when the kind is not numeric
      eqTypeFinal flags symt r t t'
  
  eqType1 :: Flags -> SymTab -> Env -> IType -> IType -> Bool
  
  -- hack to make commutative type constructors work
  eqType1 flags symt r (ITAp (ITAp (ITAp (ITCon i _ _) t1) t2) t3)
                       (ITAp (ITAp (ITAp (ITCon i' _ _) t1') t2') t3')
      | (i == i') && (hasIdProp i IdPCommutativeTCon) =
      eqType0 flags symt r t3 t3' &&
      (((eqType0 flags symt r t1 t1') && (eqType0 flags symt r t2 t2')) ||
       ((eqType0 flags symt r t1 t2') && (eqType0 flags symt r t2 t1')))
  
  -- look for arrows and add any dictionaries
  eqType1 flags symt r (ITAp (ITAp (ITCon tc _ _) tA) tB)
                       (ITAp (ITAp (ITCon tc' _ _) tA') tB')
      | (tc == idArrow noPosition) && (tc' == idArrow noPosition) =
      let r2 = addDict symt tA r
      in  eqType0 flags symt r tA tA' && eqType0 flags symt r2 tB tB'
  
  eqType1 flags symt r (ITAp f a) (ITAp f' a') =
      eqType0 flags symt r f f' && eqType0 flags symt r a a'
  eqType1 _ _ _ (ITVar i) (ITVar i') = i == i'
  eqType1 _ _ _ (ITCon i _ _) (ITCon i' _ _) = i == i'
  eqType1 _ _ _ (ITNum n) (ITNum n') = n == n'
  eqType1 _ _ _ _ _ = False
  
  -- Decide if two (numeric) types are equal by creating a NumEq proviso
  -- in CSyntax and applying "satisfy".
  eqTypeFinal :: Flags -> SymTab -> Env -> IType -> IType -> Bool
  eqTypeFinal flags symt e t1 t2
      -- Attempt to save time by weeding out cases without TAdd, SizeOf, etc
      -- (since there's no use trying to equate "n" and "m", for instance)
      -- XXX can we also weed out when the kind is not numeric?
      | isITAp t1 || isITAp t2 =
      let numEqCls = mustFindClass symt (CTypeclass idNumEq)
          (e', t1') = convType e t1
          (E _ _ _ (PredEnv _ m s), t2') = convType e' t2
          --satisfyEq :: TI ([VPred], [CDefl])
          satisfyEq = do
            eqs <- mapM mkEPred (S.toList s)
            addBoundTVs (M.elems m)
            addExplPreds eqs
            vp <- mkVPredFromPred [] (IsIn numEqCls [t1', t2'])
            satisfy eqs [vp]
      in  case (fst $ runTI flags False symt satisfyEq) of
            Right ([],_) -> True
            res -> --trace("eqTypeFinal: not satisfied: " ++ ppReadable res) $
                   False
   where isITAp (ITAp _ _) = True
         isITAp _          = False
  eqTypeFinal _ _ _ _ _ = False
  
  -------
  
  assert :: (PPrint e, Show t) => Bool -> String -> e -> t -> a -> a
  assert True  s e t x = x
  assert False s e t x = internalError ("assert failed: " ++ s ++ "\n" ++ ppReadable e ++ "\n" ++ (show t))
  --assert False s e x = internalError ("assert failed: " ++ s ++ "\n" ++ ppDebug e)
  
  tCheck :: Flags -> SymTab -> Env -> IExpr a -> IType
  tCheck flags symt r ec@(ILam i t e) =
      -- assert (kCheckErr r t == IKStar) "ILam" (ec, kCheckErr r t) $
          itFun t (tCheck flags symt (addT symt i t r) e)
  tCheck flags symt r ec@(IAps f0 ts [a]) =
          let f = iAps f0 ts [] in
          case tCheck flags symt r f of
          ITAp (ITAp arr at') rt | arr == itArrow ->
              let at = tCheck flags symt r a
              in  --trace(ppReadable((f,tCheck r f),(a,at))) $
                  assert (eqType flags symt r at at') "IAp"
                      (r, ec, a, (at, at') {-, (f,ft),(a,at)-}) (at, at') rt
          tt -> internalError ("tCheck IAp: " ++ ppReadable(ec, f, tt))
  tCheck flags symt r (IAps f ts (e:es)) =
      tCheck flags symt r (IAps (IAps f ts [e]) [] es)
  tCheck flags symt r (IVar i) = findT i r
  tCheck flags symt r (ILAM i k e) =
      ITForAll i k (tCheck flags symt (addK i k r) e)
  tCheck flags symt r ec@(IAps e [t] []) =
          case tCheck flags symt r e of
          ITForAll i k rt ->
              let kt = kCheckErr r t
                  rt'= tSubst i t rt
              in  --trace ("tCheck " ++ ppReadable ((e,et),(t,kt))) $
                  assert (k == kt) "IAP" (ec, (i,k,rt), kt) (k, kt) rt'
          tt -> internalError ("tCheck IAP: " ++ ppReadable (ec, tt))
  tCheck flags symt r (IAps f (t:ts) []) =
      tCheck flags symt r (IAps (IAps f [t] []) ts [])
  tCheck flags symt r (ICon c ic) = iConType ic
  tCheck flags symt r (IAps f [] []) =
      {- trace ("tCheck " ++ show f) $ -}
      tCheck flags symt r f
  tCheck flags symt r (IRefT t _ _) = t
  --tCheck _ _ _ e = internalError ("no match in tCheck: " ++ ppReadable e)
  
  kCheck :: Env -> IType -> Maybe IKind
  kCheck r (ITForAll i k t) = do
    kr <- kCheck (addK i k r) t
    return (IKFun k kr)
  kCheck r tc@(ITAp f a) = do
    kf <- kCheck r f
    case kf of
      IKFun ka kr ->
          do ka' <- kCheck r a
             assert (ka' == ka) "ITAp" (tc, (f, kf), (ka', ka)) (ka', ka) $ return kr
      k -> internalError ("kCheck " ++ ppReadable (tc, k))
  kCheck r (ITVar i) = findK i r
  kCheck r (ITCon _ k _) = return k
  kCheck r (ITNum _) = return IKNum
  kCheck r (ITStr _) = return IKStr
  
  iGetKind :: IType -> Maybe IKind
  iGetKind = kCheck emptyEnv
  
  kCheckErr :: Env -> IType -> IKind
  kCheckErr r t = fj $ kCheck r t
    where fj = fromJustOrErr ("findK: " ++ ppReadable (r, t))
  
  tCheckIPackage :: Flags -> SymTab -> IPackage a -> Bool
  tCheckIPackage flags symt (IPackage pi _ _ ds) =
      let r  = emptyEnv
          defOK (IDef i t e _) =
              let t' = (tCheck flags symt r e)
              in  assert (eqType flags symt r t' t) "defOK1"
                      (i,e,(t,t')) (t, t') True
      in  all defOK ds
  
  tCheckIModule :: Flags -> SymTab -> IModule a -> Bool
  tCheckIModule flags symt (IModule { imod_type_args  = iks,
                                      imod_local_defs = ds,
                                      imod_rules      = rs,
                                      imod_interface  = ifc }) =
          let r = foldr (\ (i, k) r -> addK i k r) emptyEnv iks
              defOK (IDef i t e _) =
                  let t' = (tCheck flags symt r e)
                  in  assert (eqType flags symt r t' t) "defOK2"
                          (i,e,(t,t')) (t, t') True
              ifcOK (IEFace i _ maybe_e maybe_r _ _) =
                         (case maybe_e of
                            Just (e,t) -> defOK (IDef i t e [])
                            _ -> True)
                         &&
                         (case maybe_r of
                            Just rs -> rulesOK rs
                            _ -> True)
  
              rulesOK (IRules sps rs) = all ruleOK rs
              ruleOK (IRule { irule_pred = p , irule_body = a }) =
                  let tp = tCheck flags symt r p
                      ta = tCheck flags symt r a
                  in
                      assert (eqType flags symt r tp itBit1) "ruleOK p"
                          (p, tp) (p, tp) True &&
                      assert (eqType flags symt r ta itAction) "ruleOK a"
                          (a, ta) (p, tp) True
          in  all defOK ds && rulesOK rs && all ifcOK ifc
  
  -------
  
  data Env = E (M.Map Id IType) (M.Map Id IKind) (EC.EquivClasses IType) PredEnv
  
  emptyEnv :: Env
  emptyEnv = E M.empty M.empty EC.empty emptyPredEnv
  
  addDict :: SymTab -> IType -> Env -> Env
  addDict symt t e@(E tm km eqs ps) = E tm km eqs' ps'
    where new_eqs = case t of
                      (ITAp (ITAp (ITCon i _ _) t1) t2)
                           | i == idLog   -> [(ITAp iTLog t1, t2)]
                           | i == idBits  -> [(ITAp iTSizeOf t1, t2)]
                           | i == idNumEq -> [(t1, t2)]
                      (ITAp (ITAp (ITAp (ITCon i _ _) t1) t2) t3)
                           -- XXX should we also equate T#(t2,t1)
                           | i == idAdd -> [(ITAp (ITAp iTAdd t1) t2, t3)]
                           | i == idMax -> [(ITAp (ITAp iTMax t1) t2, t3)]
                           | i == idMin -> [(ITAp (ITAp iTMin t1) t2, t3)]
                           | i == idMul -> [(ITAp (ITAp iTMul t1) t2, t3)]
                           | i == idDiv -> [(ITAp (ITAp iTDiv t1) t2, t3)]
                      otherwise -> []
          eqs' = trace_witness ("num eq witnesses: " ++ ppReadable new_eqs) $
                 let eqFn (x,y) ec = EC.equate x y ec
                 in  foldr eqFn eqs new_eqs
          ps' = addPred symt e t
  
  addT :: SymTab -> Id -> IType -> Env -> Env
  addT symt i t (E tm km eqs ps) = addDict symt t $ E (M.insert i t tm) km eqs ps
  
  addK :: Id -> IKind -> Env -> Env
  addK i k (E tm km eqs ps) = E tm (M.insert i k km) eqs ps
  
  findT :: Id -> Env -> IType
  findT i (E tm _ _ _) =
      case M.lookup i tm of
          Just t -> t
          Nothing -> internalError ("ISyntaxCheck.findT " ++ ppString i ++ "\n" ++ ppReadable (M.toList tm))
  
  findK :: Id -> Env -> Maybe IKind
  findK i (E _ km _ _) = M.lookup i km
  
  instance PPrint Env where
      pPrint d _ (E tm km eqs ps) =
          text "Env" <+>
          (pPrint d 0 (M.toList tm) $$
           pPrint d 0 (M.toList km) $$
           pPrint d 0 (EC.classes eqs) $$
           text "PredEnv")
  
  ------
  
  data PredEnv = PredEnv Int (M.Map Id TyVar) (S.Set Pred.Pred)
  
  emptyPredEnv :: PredEnv
  emptyPredEnv = PredEnv 0 M.empty S.empty
  
  convType :: Env -> IType -> (Env, Type)
  convType _ (ITForAll i k t) = internalError ("convType: ITForAll " ++ ppReadable (i, k, t))
  convType e (ITAp t1 t2) =
      let (e', t1') = convType e t1
          (e'', t2') = convType e' t2
      in  (e'', TAp t1' t2')
  convType e@(E tm km eqs (PredEnv n m s)) (ITVar i) =
      case (M.lookup i m) of
        Just tyvar -> (e, TVar tyvar)
        Nothing ->
            let k = fromJustOrErr ("findK: " ++ ppReadable (e, i)) $ findK i e
                tyvar = TyVar i n (iToCK k)
                e' = E tm km eqs (PredEnv (n+1) (M.insert i tyvar m) s)
            in  (e', TVar tyvar)
  convType e (ITCon i k s) = (e, TCon (TyCon i (Just (iToCK k)) s))
  convType e (ITNum n) = (e, TCon (TyNum n noPosition))
  convType e (ITStr s) = (e, TCon (TyStr s noPosition))
  
  convTypes :: Env -> [IType] -> (Env, [Type])
  convTypes e [] = (e, [])
  convTypes e (t:ts) = let (e', t') = convType e t
                           (e'', ts') = convTypes e' ts
                       in  (e'', t':ts')
  
  addPred :: SymTab -> Env -> IType -> PredEnv
  addPred symt e@(E _ _ _ ps) t =
      let
          res =
              case (splitITAp t) of
                (ITCon i _ _, as) | (Just cls) <- findSClass symt (CTypeclass i)
                  -> let (e', as') = convTypes e as
                     in  Just (e', IsIn cls as')
                _ -> Nothing
      in
          case res of
            Just (E _ _ _ (PredEnv n m s), new_pred) ->
                PredEnv n m (S.insert new_pred s)
            Nothing -> ps
  
  ------
  