{-# LANGUAGE FlexibleContexts #-}

module Language.Dickinson.ScopeCheck ( checkScope
                                     , checkScopeExpr
                                     , checkScopeDecl
                                     ) where

import           Control.Applicative              (Alternative, (<|>))
import           Control.Monad.Except             (MonadError)
import           Control.Monad.State              (State, evalState, get, modify)
import           Data.Foldable                    (asum, traverse_)
import qualified Data.IntSet                      as IS
import           Language.Dickinson.Check.Pattern
import           Language.Dickinson.Error
import           Language.Dickinson.Name
import           Language.Dickinson.Type
import           Language.Dickinson.Unique

type CheckM = State IS.IntSet

runCheckM :: CheckM a -> a
runCheckM = flip evalState IS.empty

insertName :: Name a -> CheckM ()
insertName (Name _ (Unique i) _) = modify (IS.insert i)

deleteName :: Name a -> CheckM ()
deleteName (Name _ (Unique i) _) = modify (IS.delete i)

-- | Checks that there are not any identifiers that aren't in scope; needs to run
-- after the renamer
checkScope :: [Declaration a] -> Maybe (DickinsonError a)
checkScope = runCheckM . checkDickinson

checkScopeExpr :: MonadError (DickinsonError a) m => Expression a -> m ()
checkScopeExpr = maybeThrow . runCheckM . checkExpr

checkScopeDecl :: MonadError (DickinsonError a) m => Declaration a -> m ()
checkScopeDecl = maybeThrow . runCheckM . checkDecl

checkDickinson :: [Declaration a] -> CheckM (Maybe (DickinsonError a))
checkDickinson d = traverse_ insDecl d *> mapSumM checkDecl d

insDecl :: Declaration a -> CheckM ()
insDecl (Define _ n _)   = insertName n
insDecl (TyDecl _ n tys) = insertName n *> traverse_ insertName tys

checkDecl :: Declaration a -> CheckM (Maybe (DickinsonError a))
checkDecl (Define _ _ e) = checkExpr e
checkDecl TyDecl{}       = pure Nothing

checkType :: DickinsonTy a -> CheckM (Maybe (DickinsonError a))
checkType TyText{}         = pure Nothing
checkType (TyFun _ ty ty') = (<|>) <$> checkType ty <*> checkType ty'
checkType (TyTuple _ tys)  = mapSumM checkType tys
checkType (TyNamed l n@(Name _ (Unique k) _)) = do
    b <- get
    if k `IS.member` b
        then pure Nothing
        else pure $ Just (UnfoundType l n)

checkExpr :: Expression a -> CheckM (Maybe (DickinsonError a))
checkExpr Literal{}      = pure Nothing
checkExpr StrChunk{}     = pure Nothing
checkExpr (Apply _ e e') = (<|>) <$> checkExpr e <*> checkExpr e'
checkExpr (Interp _ es)  = mapSumM checkExpr es
checkExpr (MultiInterp _ es)  = mapSumM checkExpr es
checkExpr (Choice _ brs) = mapSumM checkExpr (snd <$> brs)
checkExpr (Concat _ es)  = mapSumM checkExpr es
checkExpr (Tuple _ es)   = mapSumM checkExpr es
checkExpr (Flatten _ e)  = checkExpr e
checkExpr (Annot _ e ty) = (<|>) <$> checkExpr e <*> checkType ty
checkExpr (Lambda _ n _ e) = do
    insertName n
    checkExpr e <* deleteName n
checkExpr (Constructor _ n@(Name _ (Unique i) l)) = do
    b <- get
    if i `IS.member` b
        then pure Nothing
        else pure $ Just (UnfoundConstructor l n)
checkExpr (Var _ n@(Name _ (Unique i) l)) = do
    b <- get
    if i `IS.member` b
        then pure Nothing
        else pure $ Just (UnfoundName l n)
checkExpr (Let _ bs e) = do
    let ns = fst <$> bs
    traverse_ insertName ns
    (<|>) <$> checkExpr e <*> mapSumM checkExpr (snd <$> bs)
        <* traverse_ deleteName ns
checkExpr (Match _ e p e') =
    ((<|>) <$> checkExpr e) <*> do
        let ns = traversePattern p
        traverse_ insertName ns
        checkExpr e' <* traverse_ deleteName ns

mapSumM :: (Traversable t, Alternative f, Applicative m) => (a -> m (f b)) -> t a -> m (f b)
mapSumM = (fmap asum .) . traverse