-- | Functions relating to skolemization used during typechecking
module Language.PureScript.TypeChecker.Skolems
  ( newSkolemConstant
  , introduceSkolemScope
  , newSkolemScope
  , skolemize
  , skolemizeTypesInValue
  , skolemEscapeCheck
  ) where

import Prelude

import Control.Monad.Error.Class (MonadError(..))
import Control.Monad.State.Class (MonadState(..), gets, modify)
import Data.Foldable (traverse_)
import Data.Functor.Identity (Identity(), runIdentity)
import Data.Set (Set, fromList, notMember)
import Data.Text (Text)
import Language.PureScript.AST (Binder(..), ErrorMessageHint(..), Expr(..), SourceAnn, SourceSpan, everythingWithContextOnValues, everywhereWithContextOnValuesM, nonEmptySpan)
import Language.PureScript.Crash (internalError)
import Language.PureScript.Errors (ErrorMessage(..), MultipleErrors, SimpleErrorMessage(..), positionedError, singleError)
import Language.PureScript.Traversals (defS)
import Language.PureScript.TypeChecker.Monad (CheckState(..))
import Language.PureScript.Types (SkolemScope(..), SourceType, Type(..), everythingOnTypes, everywhereOnTypesM, replaceTypeVars)

-- | Generate a new skolem constant
newSkolemConstant :: MonadState CheckState m => m Int
newSkolemConstant = do
  s <- gets checkNextSkolem
  modify $ \st -> st { checkNextSkolem = s + 1 }
  return s

-- | Introduce skolem scope at every occurrence of a ForAll
introduceSkolemScope :: MonadState CheckState m => Type a -> m (Type a)
introduceSkolemScope = everywhereOnTypesM go
  where
  go (ForAll ann vis ident mbK ty Nothing) = ForAll ann vis ident mbK ty <$> (Just <$> newSkolemScope)
  go other = return other

-- | Generate a new skolem scope
newSkolemScope :: MonadState CheckState m => m SkolemScope
newSkolemScope = do
  s <- gets checkNextSkolemScope
  modify $ \st -> st { checkNextSkolemScope = s + 1 }
  return $ SkolemScope s

-- | Skolemize a type variable by replacing its instances with fresh skolem constants
skolemize :: a -> Text -> Maybe (Type a) -> Int -> SkolemScope -> Type a -> Type a
skolemize ann ident mbK sko scope = replaceTypeVars ident (Skolem ann ident mbK sko scope)

-- | This function skolemizes type variables appearing in any type signatures or
-- 'DeferredDictionary' placeholders. These type variables are the only places
-- where scoped type variables can appear in expressions.
skolemizeTypesInValue :: SourceAnn -> Text -> Maybe SourceType -> Int -> SkolemScope -> Expr -> Expr
skolemizeTypesInValue ann ident mbK sko scope =
    runIdentity . onExpr'
  where
    onExpr' :: Expr -> Identity Expr
    (_, onExpr', _, _, _, _) = everywhereWithContextOnValuesM [] defS onExpr onBinder defS defS defS

    onExpr :: [Text] -> Expr -> Identity ([Text], Expr)
    onExpr sco (DeferredDictionary c ts)
      | ident `notElem` sco = return (sco, DeferredDictionary c (map (skolemize ann ident mbK sko scope) ts))
    onExpr sco (TypedValue check val ty)
      | ident `notElem` sco = return (sco ++ peelTypeVars ty, TypedValue check val (skolemize ann ident mbK sko scope ty))
    onExpr sco (VisibleTypeApp val ty)
      | ident `notElem` sco = return (sco ++ peelTypeVars ty, VisibleTypeApp val (skolemize ann ident mbK sko scope ty))
    onExpr sco other = return (sco, other)

    onBinder :: [Text] -> Binder -> Identity ([Text], Binder)
    onBinder sco (TypedBinder ty b)
      | ident `notElem` sco = return (sco ++ peelTypeVars ty, TypedBinder (skolemize ann ident mbK sko scope ty) b)
    onBinder sco other = return (sco, other)

    peelTypeVars :: SourceType -> [Text]
    peelTypeVars (ForAll _ _ i _ ty _) = i : peelTypeVars ty
    peelTypeVars _ = []

-- | Ensure skolem variables do not escape their scope
--
-- Every skolem variable is created when a 'ForAll' type is skolemized.
-- This determines the scope of that skolem variable, which is copied from
-- the 'SkolemScope' field of the 'ForAll' constructor.
--
-- This function traverses the tree top-down, and collects any 'SkolemScope's
-- introduced by 'ForAll's. If a 'Skolem' is encountered whose 'SkolemScope' is
-- not in the current list, then we have found an escaped skolem variable.
skolemEscapeCheck :: MonadError MultipleErrors m => Expr -> m ()
skolemEscapeCheck (TypedValue False _ _) = return ()
skolemEscapeCheck expr@TypedValue{} =
    traverse_ (throwError . singleError) (toSkolemErrors expr)
  where
    toSkolemErrors :: Expr -> [ErrorMessage]
    (_, toSkolemErrors, _, _, _) = everythingWithContextOnValues (mempty, Nothing) [] (<>) def go def def def

    def s _ = (s, [])

    go :: (Set SkolemScope, Maybe SourceSpan)
       -> Expr
       -> ((Set SkolemScope, Maybe SourceSpan), [ErrorMessage])
    go (scopes, _) (PositionedValue ss _ _) = ((scopes, Just ss), [])
    go (scopes, ssUsed) val@(TypedValue _ _ ty) =
        ( (allScopes, ssUsed)
        , [ ErrorMessage (maybe id ((:) . positionedError) ssUsed [ ErrorInExpression val ]) $
              EscapedSkolem name (nonEmptySpan ssBound) ty
          | (ssBound, name, scope) <- collectSkolems ty
          , notMember scope allScopes
          ]
        )
      where
        -- Any new skolem scopes introduced by universal quantifiers
        newScopes :: [SkolemScope]
        newScopes = collectScopes ty

        -- All scopes, including new scopes
        allScopes :: Set SkolemScope
        allScopes = fromList newScopes <> scopes

        -- Collect any scopes appearing in quantifiers at the top level
        collectScopes :: SourceType -> [SkolemScope]
        collectScopes (ForAll _ _ _ _ t (Just sco)) = sco : collectScopes t
        collectScopes ForAll{} = internalError "skolemEscapeCheck: No skolem scope"
        collectScopes _ = []

        -- Collect any skolem variables appearing in a type
        collectSkolems :: SourceType -> [(SourceAnn, Text, SkolemScope)]
        collectSkolems = everythingOnTypes (++) collect where
          collect (Skolem ss name _ _ scope) = [(ss, name, scope)]
          collect _ = []
    go scos _ = (scos, [])
skolemEscapeCheck _ = internalError "skolemEscapeCheck: untyped value"