content: {-# LANGUAGE DeriveGeneric #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE PatternGuards #-}
  {-# LANGUAGE MonadComprehensions #-}
  {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  
  module Language.Haskell.Exference.Core.Types
    ( TVarId
    , QualifiedName(..)
    , HsType (..)
    , HsTypeOffset (..)
    , Subst (..)
    , Substs
    , HsTypeClass (..)
    , HsInstance (..)
    , HsConstraint (..)
    , StaticClassEnv (..)
    , QueryClassEnv ( qClassEnv_env
                    , qClassEnv_constraints
                    , qClassEnv_inflatedConstraints
                    , qClassEnv_varConstraints )
    , constraintApplySubsts
    , inflateHsConstraints
    , applySubst
    , applySubsts
    -- , typeParser
    , containsVar
    , showVar
    , showTypedVar
    , mkQueryClassEnv
    , addQueryClassEnv
    , freeVars
    , showHsConstraint
    , TypeVarIndex
    , showHsType
  #if !MIN_VERSION_base(4,8,0)
    , Alt (..)
  #endif
    )
  where
  
  
  
  import Data.Char ( ord, chr, isLower, isUpper, toLower )
  import Data.List ( intercalate, intersperse )
  import Data.Foldable ( fold, foldMap )
  import Control.Applicative ( (<$>), (<*>), (*>), (<*) )
  import Data.Maybe ( maybeToList, fromMaybe )
  import Data.Monoid ( Monoid(..), Any(..) )
  import Control.Monad ( liftM, liftM2, MonadPlus )
  import Control.Applicative ( Applicative, Alternative, empty, (<|>) )
  import Control.Arrow ( first )
  import Data.Orphans ()
  
  import qualified Data.Set as S
  import qualified Data.IntSet as IntSet
  import qualified Data.Map.Strict as M
  import qualified Data.IntMap.Strict as IntMap
  import qualified Data.List as L
  
  import Text.ParserCombinators.Parsec hiding (State, (<|>))
  import Text.ParserCombinators.Parsec.Char
  
  import Language.Haskell.Exts.Syntax ( Name (..) )
  
  import Control.DeepSeq.Generics
  import Control.DeepSeq
  import GHC.Generics
  import Data.Data ( Data )
  import Data.Typeable ( Typeable )
  import Control.Monad.Trans.MultiState
  import Safe
  
  import Debug.Hood.Observe
  import Debug.Trace
  
  
  
  #if !MIN_VERSION_base(4,8,0)
  newtype Alt f a = Alt {getAlt :: f a}
    deriving (Generic, Generic1, Read, Show, Eq, Ord, Num, Enum,
              Monad, MonadPlus, Applicative, Alternative, Functor)
  
  instance forall f a . Alternative f => Monoid (Alt f a) where
          mempty = Alt empty
          (Alt x) `mappend` (Alt y) = Alt (x <|> y)
  #endif
  
  type TVarId = Int
  data Subst  = Subst {-# UNPACK #-} !TVarId !HsType
  type Substs = IntMap.IntMap HsType
  
  data QualifiedName
    = QualifiedName [String] String
    | ListCon
    | TupleCon Int
    | Cons
    deriving (Eq, Ord, Generic, Data, Typeable)
  
  data HsType = TypeVar      {-# UNPACK #-} !TVarId
              | TypeConstant {-# UNPACK #-} !TVarId
                -- like TypeCons, for exference-internal purposes.
              | TypeCons     QualifiedName
              | TypeArrow    !HsType !HsType
              | TypeApp      !HsType !HsType
              | TypeForall   [TVarId] [HsConstraint] !HsType
    deriving (Ord, Eq, Generic, Data, Typeable)
  
  data HsTypeOffset = HsTypeOffset !HsType {-# UNPACK #-} !Int
  
  type TypeVarIndex = M.Map Name Int
  
  data HsTypeClass = HsTypeClass
    { tclass_name :: QualifiedName
    , tclass_params :: [TVarId]
    , tclass_constraints :: [HsConstraint]
    }
    deriving (Eq, Show, Ord, Generic, Data, Typeable)
  
  data HsInstance = HsInstance
    { instance_constraints :: [HsConstraint]
    , instance_tclass :: HsTypeClass
    , instance_params :: [HsType]
    }
    deriving (Eq, Show, Ord, Generic, Data, Typeable)
  
  data HsConstraint = HsConstraint
    { constraint_tclass :: HsTypeClass
    , constraint_params :: [HsType]
    }
    deriving (Eq, Ord, Generic, Data, Typeable)
  
  data StaticClassEnv = StaticClassEnv
    { sClassEnv_tclasses :: [HsTypeClass]
    , sClassEnv_instances :: M.Map QualifiedName [HsInstance]
    }
    deriving (Show, Generic, Data, Typeable)
  
  data QueryClassEnv = QueryClassEnv
    { qClassEnv_env :: StaticClassEnv
    , qClassEnv_constraints :: S.Set HsConstraint
    , qClassEnv_inflatedConstraints :: S.Set HsConstraint
    , qClassEnv_varConstraints :: IntMap.IntMap (S.Set HsConstraint)
    }
    deriving (Generic)
  
  instance NFData QualifiedName  where rnf = genericRnf
  instance NFData HsType         where rnf = genericRnf
  instance NFData HsTypeClass    where rnf = genericRnf
  instance NFData HsInstance     where rnf = genericRnf
  instance NFData HsConstraint   where rnf = genericRnf
  instance NFData StaticClassEnv where rnf = genericRnf
  instance NFData QueryClassEnv  where rnf = genericRnf
  
  instance Show QualifiedName where
    show (QualifiedName ns n) = if    length n >= 2
                                   && head n == '('
                                   && last n == ')'
                                then "(" ++ intercalate "." (ns ++ [tail n])
                                else        intercalate "." (ns ++ [n])
    show ListCon              = "[]"
    show (TupleCon 0)         = "()"
    show (TupleCon i)         = "(" ++ replicate (i-1) ',' ++ ")"
    show Cons                 = "(:)"
  
  instance Show HsType where
    showsPrec _ (TypeVar i) = showString $ showVar i
    showsPrec _ (TypeConstant i) = showString $ "C" ++ showVar i
    showsPrec d (TypeCons s) = showsPrec d s
    showsPrec d (TypeArrow t1 t2) =
      showParen (d> -2) $ showsPrec (-1) t1 . showString " -> " . showsPrec (-1) t2
    showsPrec d (TypeApp t1 t2) =
      showParen (d> -1) $ showsPrec 0 t1 . showString " " . showsPrec 0 t2
    showsPrec d (TypeForall [] [] t) = showsPrec d t
    showsPrec d (TypeForall is cs t) =
      showParen (d>0)
      $ showString ("forall " ++ intercalate ", " (showVar <$> is) ++ " . ")
      . showParen True (\x -> foldr (++) x $ intersperse ", " $ map show cs)
      . showString " => "
      . showsPrec (-2) t
  
  showHsType :: TypeVarIndex -> HsType -> String
  showHsType convMap t = h 0 t ""
   where
    h :: Int -> HsType -> ShowS
    h _ (TypeVar i)      = showString
                         $ maybe "badNameInternalError"
                                 (\(Ident n, _) -> n)
                         $ L.find ((i ==) .  snd)
                         $ M.toList convMap
    h _ (TypeConstant i) = showString
                         $ maybe "badNameInternalError"
                                 (\(Ident n, _) -> n)
                         $ L.find ((i ==) .  snd)
                         $ M.toList convMap
    h _ (TypeCons s) = shows s
    h d (TypeArrow t1 t2) =
      showParen (d> -2) $ t1Shows . showString " -> " . t2Shows
      where
        t1Shows = h (-1) t1
        t2Shows = h (-1) t2
    h d (TypeApp t1 t2) =
      showParen (d> -1) $ t1Shows . showString " " . t2Shows
      where
        t1Shows = h 0 t1
        t2Shows = h 0 t2
    h d (TypeForall [] [] ty) = h d ty
    h d (TypeForall is cs ty) =
      showParen (d>0)
        $ showString ("forall " ++ intercalate ", " (showVar <$> is) ++ " . ")
        . showParen True (\x -> foldr (++) x $ intersperse ", " $ map show cs)
        . showString " => "
        . tShows
      where
        tShows = h (-2) ty
  
  instance Observable HsType where
    observer x = observeOpaque (show x) x
  
  -- instance Read HsType where
  --   readsPrec _ = maybeToList . parseType
  
  instance Show HsConstraint where
    show (HsConstraint c ps) = unwords $ show (tclass_name c) : map show ps
  
  showHsConstraint :: TypeVarIndex
                   -> HsConstraint
                   -> String
  showHsConstraint convMap (HsConstraint c ps) =
    unwords $ show name : tyStrs  
   where
    name = tclass_name c
    tyStrs = showHsType convMap <$> ps
    
  
  instance Show QueryClassEnv where
    show (QueryClassEnv _ cs _ _) = "(QueryClassEnv _ " ++ show cs ++ " _)"
  instance Observable HsConstraint where
    observer x = observeOpaque (show x) x
  
  instance Observable QueryClassEnv where
    observer x = observeOpaque (show x) x
  
  instance Observable HsInstance where
    observer x = observeOpaque (show x) x
  
  filterHsConstraintsByVarId :: TVarId
                             -> S.Set HsConstraint
                             -> S.Set HsConstraint
  filterHsConstraintsByVarId i = S.filter
                               $ any (containsVar i) . constraint_params
  
  containsVar :: TVarId -> HsType -> Bool
  containsVar i = S.member i . freeVars
  
  mkQueryClassEnv :: StaticClassEnv -> [HsConstraint] -> QueryClassEnv
  mkQueryClassEnv sClassEnv constrs = addQueryClassEnv constrs $ QueryClassEnv {
    qClassEnv_env = sClassEnv,
    qClassEnv_constraints = S.empty,
    qClassEnv_inflatedConstraints = S.empty,
    qClassEnv_varConstraints = IntMap.empty
  }
  
  addQueryClassEnv :: [HsConstraint] -> QueryClassEnv -> QueryClassEnv
  addQueryClassEnv constrs env = env {
    qClassEnv_constraints = csSet,
    qClassEnv_inflatedConstraints = inflateHsConstraints csSet,
    qClassEnv_varConstraints = helper constrs
  }
    where
      csSet = S.fromList constrs `S.union` qClassEnv_constraints env
      helper :: [HsConstraint] -> IntMap.IntMap (S.Set HsConstraint)
      helper cs =
        let ids :: IntSet.IntSet
            ids = IntSet.fromList $ S.toList $ fold $ freeVars <$> (constraint_params =<< cs)
        in IntMap.fromSet (flip filterHsConstraintsByVarId
                          $ inflateHsConstraints csSet) ids
  
  inflateHsConstraints :: S.Set HsConstraint -> S.Set HsConstraint
  inflateHsConstraints = inflate (S.fromList . f)
    where
      f :: HsConstraint -> [HsConstraint]
      f (HsConstraint (HsTypeClass _ ids constrs) ps) =
        map (snd . constraintApplySubsts (IntMap.fromList $ zip ids ps)) constrs
  
  -- uses f to find new elements. adds these new elements, and recursively
  -- tried to find even more elements. will not terminate if there are cycles
  -- in the application of f
  inflate :: (Ord a, Show a) => (a -> S.Set a) -> S.Set a -> S.Set a
  inflate f = fold . takeWhile (not . S.null) . iterate (foldMap f)
  
  constraintApplySubst :: Subst -> HsConstraint -> HsConstraint
  constraintApplySubst s (HsConstraint c ps) =
    HsConstraint c $ map (applySubst s) ps
  
  -- returns if any change was necessary,
  -- plus the (potentially changed) constraint
  -- constraintApplySubst' :: Subst -> HsConstraint -> (Bool, HsConstraint)
  -- constraintApplySubst' s (HsConstraint c ps) =
  --   let applied = map (applySubst' s) ps
  --   in (any fst applied, HsConstraint c $ snd <$> applied)
  
  -- returns if any change was necessary,
  -- plus the (potentially changed) constraint
  {-# INLINE constraintApplySubsts #-}
  constraintApplySubsts :: Substs -> HsConstraint -> (Any, HsConstraint)
  constraintApplySubsts ss c
    | IntMap.null ss = return c
    | HsConstraint cl ps <- c =
      HsConstraint cl <$> mapM (applySubsts ss) ps
  
  showVar :: TVarId -> String
  showVar 0 = "v0"
  showVar i | i<27      = [chr (ord 'a' + i - 1)]
            | otherwise = "t"++show (i-27)
  
  showTypedVar :: forall m
                . ( MonadMultiState (M.Map TVarId HsType) m )
               => TVarId
               -> m String
  showTypedVar i = do
    m <- mGet
    fromJustNote "missing collectVarTypes before showTypedVar"
      $ h <$> M.lookup i m
   where
    -- h t | traceShow (i, t) False = undefined
    h TypeVar{}          = return $ showVar i
    h TypeConstant{}     = return $ showVar i
    h (TypeCons qName) = do
      return $ case qName of
        QualifiedName _ (c:_) -> toLower c : show i
        QualifiedName{}       -> showVar i
        ListCon               -> showVar i ++ "s"
        TupleCon{}            -> showVar i
        Cons                  -> showVar i
    h TypeArrow{}        = return $ "f" ++ show i
    h (TypeApp t _)      = h t
    h (TypeForall _ _ t) = h t
  
  -- parseType :: _ => String -> m (Maybe (HsType, String))
  -- parseType s = either (const Nothing) Just
  --             $ runParser (    (,)
  --                          <$> typeParser
  --                          <*> many anyChar)
  --                         ()
  --                         ""
  --                         s
  -- 
  -- typeParser :: forall m . (_) => Parser (m HsType)
  -- typeParser = parseAll
  --   where
  --     parseAll :: Parser (m HsType)
  --     parseAll = parseUn >>= parseBin
  --     parseUn :: Parser (m HsType) -- TODO: forall
  --     parseUn = spaces *> (
  --             try (TypeCons . QualifiedName [] <$> ((:) <$> satisfy isUpper <*> many alphaNum))
  --         <|> try ((TypeVar . (\x -> x - ord 'a') . ord) <$> satisfy isLower)
  --         <|>     (char '(' *> parseAll <* char ')')
  --       )
  --     parseBin :: HsType -> Parser HsType
  --     parseBin left =
  --         try (    try (TypeArrow left <$> (spaces *> string "->" *> parseAll))
  --              <|>     ((TypeApp   left <$> (space *> parseUn)) >>= parseBin)
  --              )
  --         <|>
  --         (spaces *> return left)
  
  applySubst :: Subst -> HsType -> HsType
  applySubst (Subst i t) v@(TypeVar j) = if i==j then t else v
  applySubst _ c@(TypeConstant _) = c
  applySubst _ c@(TypeCons _)     = c
  applySubst s (TypeArrow t1 t2)  = TypeArrow (applySubst s t1) (applySubst s t2)
  applySubst s (TypeApp t1 t2)    = TypeApp (applySubst s t1) (applySubst s t2)
  applySubst s@(Subst i _) f@(TypeForall js cs t) = if i `elem` js
    then f
    else TypeForall js (constraintApplySubst s <$> cs) (applySubst s t)
  
  applySubsts :: Substs -> HsType -> (Any, HsType)
  applySubsts s v@(TypeVar i)      = fromMaybe (return v)
                                    $ (,) (Any True) <$> IntMap.lookup i s
  applySubsts _ c@(TypeConstant _) = return c
  applySubsts _ c@(TypeCons _)     = return c
  applySubsts s (TypeArrow t1 t2)  = liftM2 TypeArrow (applySubsts s t1) (applySubsts s t2)
  applySubsts s (TypeApp t1 t2)    = liftM2 TypeApp   (applySubsts s t1) (applySubsts s t2)
  applySubsts s (TypeForall js cs t) = liftM2 (TypeForall js) 
    (sequence $ constraintApplySubsts s <$> cs)
    (applySubsts (foldr IntMap.delete s js) t)
  
  freeVars :: HsType -> S.Set TVarId
  freeVars (TypeVar i)         = S.singleton i
  freeVars (TypeConstant _)    = S.empty
  freeVars (TypeCons _)        = S.empty
  freeVars (TypeArrow t1 t2)   = S.union (freeVars t1) (freeVars t2)
  freeVars (TypeApp t1 t2)     = S.union (freeVars t1) (freeVars t2)
  freeVars (TypeForall is _ t) = foldr S.delete (freeVars t) is
  
  -- instance Monoid w => Monad ((,) w) where
  --   return = (,) mempty
  --   (w,x) >>= f = first (mappend w) (f x)
  