content: -- Copyright 2021 Google LLC
  --
  -- Use of this source code is governed by a BSD-style
  -- license that can be found in the LICENSE file or at
  -- https://developers.google.com/open-source/licenses/bsd
  
  {-# LANGUAGE UndecidableInstances #-}
  
  module Err (
    Err (..), Except (..), ToErr (..), PrintableErr (..),
    ParseErr (..), SyntaxErr (..), NameErr (..), TypeErr (..), MiscErr (..),
    Fallible (..), Catchable (..), catchErrExcept,
    HardFailM (..), runHardFail, throw,
    catchIOExcept, liftExcept, liftExceptAlt,
    ignoreExcept, getCurrentCallStack, printCurrentCallStack,
    ExceptT (..), rootSrcId, SrcId (..), assertEq, throwInternal,
    InferenceArgDesc, InfVarDesc (..), HasSrcId (..), getErrSrcId) where
  
  import Control.Exception hiding (throw)
  import Control.Applicative
  import Control.Monad
  import Control.Monad.Identity
  import Control.Monad.Writer.Strict
  import Control.Monad.State.Strict
  import Control.Monad.Reader
  import Data.Aeson (ToJSON)
  import Data.Coerce
  import Data.Hashable
  import Data.List (sort)
  import Data.Foldable (fold)
  import Data.Store (Store (..))
  import GHC.Stack
  import GHC.Generics
  
  import PPrint
  
  -- === source info ===
  
  -- XXX: 0 is reserved for the root The IDs are generated from left to right in
  -- parsing order, so IDs for lexemes are guaranteed to be sorted correctly.
  newtype SrcId = SrcId Int  deriving (Show, Eq, Ord, Generic)
  
  rootSrcId :: SrcId
  rootSrcId = SrcId 0
  
  class HasSrcId a where
    getSrcId :: a -> SrcId
  
  getErrSrcId :: Err -> Maybe SrcId
  getErrSrcId = \case
    SearchFailure _ -> Nothing
    InternalErr   _ -> Nothing
    ParseErr      _ -> Nothing
    SyntaxErr sid _ -> Just sid
    NameErr   sid _ -> Just sid
    TypeErr   sid _ -> Just sid
    RuntimeErr -> Nothing
    MiscErr _  -> Nothing
  
  -- === core errro type ===
  
  data Err =
     SearchFailure String      -- used as the identity for `Alternative` instances and for MonadFail.
   | InternalErr String
   | ParseErr  ParseErr
   | SyntaxErr SrcId SyntaxErr
   | NameErr   SrcId NameErr
   | TypeErr   SrcId TypeErr
   | RuntimeErr
   | MiscErr MiscErr
     deriving (Show, Eq)
  
  type MsgStr  = String
  type VarStr  = String
  type TypeStr = String
  
  data ParseErr =
    MiscParseErr MsgStr
    deriving (Show, Eq)
  
  data SyntaxErr =
     MiscSyntaxErr MsgStr
   | TopLevelArrowBinder
   | CantConstrainAnonBinders
   | UnexpectedBinder
   | OnlyUnaryWithoutParens
   | IllegalPattern
   | UnexpectedConstraint
   | ExpectedIdentifier String
   | UnexpectedEffectForm
   | UnexpectedMethodDef
   | BlockWithoutFinalExpr
   | UnexpectedGivenClause
   | ArgsShouldHaveParens
   | BadEqualSign
   | BadColon
   | ExpectedAnnBinder
   | BadField
   | BadPrefix VarStr
     deriving (Show, Eq)
  
  data NameErr =
     MiscNameErr MsgStr
   | UnboundVarErr VarStr    -- name of var
   | EscapedNameErr [VarStr] -- names
   | RepeatedPatVarErr VarStr
   | RepeatedVarErr VarStr
   | NotAnOrdinaryVar VarStr
   | NotADataCon VarStr
   | NotAClassName VarStr
   | NotAMethodName VarStr
   | AmbiguousVarErr VarStr [String]
   | VarDefErr VarStr
     deriving (Show, Eq)
  
  data TypeErr =
     MiscTypeErr MsgStr
   | CantSynthDict     TypeStr
   | CantSynthInfVars  TypeStr
   | NotASynthType     TypeStr
   | CantUnifySkolem
   | OccursCheckFailure VarStr TypeStr
   | UnificationFailure TypeStr TypeStr [VarStr]  -- expected, actual, inference vars
   | DisallowedEffects  String String  -- allowed, actual
   | InferEmptyTable
   | ArityErr        Int Int  -- expected, actual
   | PatternArityErr Int Int  -- expected, actual
   | SumTypeCantFail
   | PatTypeErr String String  -- expected type constructor (from pattern), actual type (from rhs)
   | EliminationErr String String  -- expected type constructor, actual type
   | IllFormedCasePattern
   | NotAMethod VarStr VarStr
   | DuplicateMethod VarStr
   | MissingMethod   VarStr
   | WrongArrowErr String String
   | AnnotationRequired
   | NotAUnaryConstraint TypeStr
   | InterfacesNoImplicitParams
   | RepeatedOptionalArgs [VarStr]
   | UnrecognizedOptionalArgs [VarStr] [VarStr]
   | NoFields TypeStr
   | TypeMismatch TypeStr TypeStr  -- TODO: should we merege this with UnificationFailure
   | InferHoleErr
   | InferDepPairErr
   | InferEmptyCaseEff
   | UnexpectedTerm String TypeStr
   | CantFindField VarStr TypeStr [VarStr]  -- field name, field type, known fields
   | TupleLengthMismatch Int Int
   | CantReduceType TypeStr
   | CantReduceDict
   | CantReduceDependentArg
   | AmbiguousInferenceVar VarStr TypeStr InfVarDesc
   | FFIResultTyErr    TypeStr
   | FFIArgTyNotScalar TypeStr
     deriving (Show, Eq)
  
  data MiscErr =
     MiscMiscErr MsgStr
   | ModuleImportErr VarStr
   | CantFindModuleSource VarStr
     deriving (Show, Eq)
  
  -- name of function, name of arg
  type InferenceArgDesc = (String, String)
  data InfVarDesc =
     ImplicitArgInfVar InferenceArgDesc
   | AnnotationInfVar String -- name of binder
   | TypeInstantiationInfVar String -- name of type
   | MiscInfVar
     deriving (Show, Generic, Eq, Ord)
  
  -- === ToErr class ===
  
  class ToErr a where
    toErr :: SrcId -> a -> Err
  
  instance ToErr SyntaxErr where toErr = SyntaxErr
  instance ToErr NameErr   where toErr = NameErr
  instance ToErr TypeErr   where toErr = TypeErr
  
  -- === Error messages ===
  
  class PrintableErr a where
    printErr :: a -> String
  
  instance PrintableErr Err where
    printErr = \case
      SearchFailure s -> "Internal search failure: " ++ s
      InternalErr s -> "Internal compiler error: " ++ s ++ "\n" ++
        "Please report this at github.com/google-research/dex-lang/issues\n"
      ParseErr    e -> "Parse error: "  ++ printErr e
      SyntaxErr _ e -> "Syntax error: " ++ printErr e
      NameErr   _ e -> "Name error: "   ++ printErr e
      TypeErr   _ e -> "Type error: "   ++ printErr e
      MiscErr     e -> "Error: "        ++ printErr e
      RuntimeErr    -> "Runtime error"
  
  instance PrintableErr ParseErr where
    printErr = \case
      MiscParseErr s -> s
  
  instance PrintableErr SyntaxErr where
    printErr = \case
      MiscSyntaxErr s -> s
      TopLevelArrowBinder ->
        "Arrow binder syntax <- not permitted at the top level, because the binding would have unbounded scope."
      CantConstrainAnonBinders -> "can't constrain anonymous binders"
      UnexpectedBinder -> "binder must be an identifier or `_`"
      OnlyUnaryWithoutParens ->"only unary constructors can form patterns without parens"
      IllegalPattern -> "illegal pattern"
      UnexpectedConstraint -> "unexpected constraint"
      ExpectedIdentifier ctx -> "expected " ++ ctx ++ " to be an identifier"
      UnexpectedEffectForm ->
        "unexpected effect form; expected one of `Read h`, `Accum h`, `State h`, `Except`, `IO`, "
         ++ "or the name of a user-defined effect."
      UnexpectedMethodDef -> "unexpected method definition. Expected `def` or `x = ...`."
      BlockWithoutFinalExpr -> "block must end in expression"
      UnexpectedGivenClause -> "unexpected `given` clause"
      ArgsShouldHaveParens -> "argument types should be in parentheses"
      BadEqualSign -> "equal sign must be used as a separator for labels or binders, not a standalone operator"
      BadColon ->
        "colon separates binders from their type annotations, is not a standalone operator.\n"
        ++ " If you are trying to write a dependent type, use parens: (i:Fin 4) => (..i)"
      ExpectedAnnBinder -> "expected an annotated binder"
      BadField -> "field must be a name or an integer"
      BadPrefix name -> "prefix (" ++ name ++ ") not legal as a bare expression"
  
  instance PrintableErr NameErr where
    printErr = \case
      MiscNameErr s -> s
      UnboundVarErr v     -> "variable not in scope: " ++ v
      EscapedNameErr vs   -> "leaked local variables: " ++ unwords vs
      RepeatedPatVarErr v -> "variable already defined within pattern: " ++ v
      RepeatedVarErr    v -> "variable already defined : " ++ v
      NotAnOrdinaryVar v -> "not an ordinary variable: " ++ v
      NotADataCon      v -> "not a data constructor: "   ++ v
      NotAClassName    v -> "not a class name: "         ++ v
      NotAMethodName   v -> "not a method name: "        ++ v
      -- we sort the lines to make the result a bit more deterministic for quine tests
      AmbiguousVarErr v defs ->
        "ambiguous occurrence: " ++ v ++ " is defined:\n"
        ++ unlines (sort defs)
      -- TODO: we see this message a lot. We should improve it by including more information.
      -- Ideally we'd provide a link to where the original error happened."
      VarDefErr v -> "error in (earlier) definition of variable: " ++ v
  
  instance PrintableErr TypeErr where
    printErr = \case
      MiscTypeErr s -> s
      FFIResultTyErr t    -> "FFI result type should be scalar or pair. Got: " ++ t
      FFIArgTyNotScalar t -> "FFI function arguments should be scalar. Got: " ++ t
      CantSynthDict t     -> "can't synthesize a class dictionary for: " ++ t
      CantSynthInfVars t  -> "can't synthesize a class dictionary for a type with inference vars: " ++ t
      NotASynthType t     -> "can't synthesize terms of type: " ++ t
      CantUnifySkolem     -> "can't unify with skolem vars"
      OccursCheckFailure v t -> "occurs check failure: " ++ v ++ " occurs in " ++ t
      DisallowedEffects r1 r2 -> "\nAllowed:   " ++ pprint r1 ++
                                 "\nRequested: " ++ pprint r2
      UnificationFailure t1 t2 vs -> "\nExpected: " ++ t1
                                  ++ "\nActual: " ++ t2 ++ case vs of
                                                            [] -> ""
                                                            _  -> "\n(Solving for: " ++ unwords vs ++ ")"
      InferEmptyTable       -> "can't infer type of empty table"
      ArityErr n1 n2 -> "wrong number of positional arguments provided. Expected " ++ show n1 ++ " but got " ++ show n2
      PatternArityErr n1 n2 -> "unexpected number of pattern binders. Expected "   ++ show n1 ++ " but got " ++ show n2
      SumTypeCantFail               -> "sum type constructor in can't-fail pattern"
      PatTypeErr patTy rhsTy        -> "pattern is for a " ++ patTy ++ "but we're matching against a " ++ rhsTy
      EliminationErr expected ty    -> "expected a " ++ expected ++ ". Got: " ++ ty
      IllFormedCasePattern          -> "case patterns must start with a data constructor or variant pattern"
      NotAMethod method className   -> "unexpected method: " ++ method ++ " is not a method of " ++ className
      DuplicateMethod method        -> "duplicate method: " ++ method
      MissingMethod   method        -> "missing method: "   ++ method
      WrongArrowErr expected actual -> "wrong arrow. Expected " ++ expected ++ " got " ++ actual
      AnnotationRequired            -> "type annotation or constraint required"
      NotAUnaryConstraint ty        -> "constraint should be a unary function. Got: " ++ ty
      InterfacesNoImplicitParams    -> "interfaces can't have implicit parameters"
      RepeatedOptionalArgs vs       -> "repeated names offered:" ++ unwords vs
      UnrecognizedOptionalArgs vs accepted -> "unrecognized named arguments: " ++ unwords vs
                                         ++ ". Should be one of: " ++ pprint accepted
      NoFields ty -> "can't get fields for type " ++ pprint ty
      TypeMismatch expected actual -> "\nExpected: " ++ expected ++
                                      "\nActual:   " ++ actual
      InferHoleErr -> "can't infer value of hole"
      InferDepPairErr -> "can't infer the type of a dependent pair; please annotate its type"
      InferEmptyCaseEff -> "can't infer empty case expressions"
      UnexpectedTerm term ty -> "unexpected " ++ term ++ ". Expected: " ++ ty
      CantFindField field fieldTy knownFields ->
        "can't resolve field " ++ field ++ " of type " ++ fieldTy ++
        "\nKnown fields are: " ++ unwords knownFields
      TupleLengthMismatch req actual -> do
          "tuple length mismatch. Expected: " ++ show req ++ " but got " ++ show actual
      CantReduceType  ty -> "Can't reduce type expression: " ++ ty
      CantReduceDict -> "Can't reduce dict"
      CantReduceDependentArg ->
        "dependent functions can only be applied to fully evaluated expressions. " ++
        "Bind the argument to a name before you apply the function."
      AmbiguousInferenceVar infVar ty desc -> case desc of
         AnnotationInfVar v ->
           "couldn't infer type of unannotated binder " <> v
         ImplicitArgInfVar (f, argName) ->
           "couldn't infer implicit argument `" <> argName <> "` of " <> f
         TypeInstantiationInfVar t ->
           "couldn't infer instantiation of type " <> t
         MiscInfVar ->
           "ambiguous type variable: " ++ infVar ++ ": " ++ ty
  
  instance PrintableErr MiscErr where
    printErr = \case
      MiscMiscErr s -> s
      ModuleImportErr v -> "couldn't import " ++ v
      CantFindModuleSource v ->
        "couldn't find a source file for module " ++ v ++
        "\nHint: Consider extending --lib-path"
  
  -- === monads and helpers ===
  
  class MonadFail m => Fallible m where
    throwErr :: Err -> m a
  
  class Fallible m => Catchable m where
    catchErr :: m a -> (Err -> m a) -> m a
  
  catchErrExcept :: Catchable m => m a -> m (Except a)
  catchErrExcept m = catchErr (Success <$> m) (\e -> return $ Failure e)
  
  catchSearchFailure :: Catchable m => m a -> m (Maybe a)
  catchSearchFailure m = (Just <$> m) `catchErr` \case
    SearchFailure _ -> return Nothing
    err -> throwErr err
  
  instance Fallible IO where
    throwErr errs = throwIO errs
    {-# INLINE throwErr #-}
  
  instance Catchable IO where
    catchErr cont handler =
      catchIOExcept cont >>= \case
        Success result -> return result
        Failure errs -> handler errs
  
  -- === ExceptT type ===
  
  newtype ExceptT m a = ExceptT { runExceptT :: m (Except a) }
  
  instance Monad m => Functor (ExceptT m) where
    fmap = liftM
    {-# INLINE fmap #-}
  
  instance Monad m => Applicative (ExceptT m) where
    pure = return
    {-# INLINE pure #-}
    liftA2 = liftM2
    {-# INLINE liftA2 #-}
  
  instance Monad m => Monad (ExceptT m) where
    return x = ExceptT $ return (Success x)
    {-# INLINE return #-}
    m >>= f = ExceptT $ runExceptT m >>= \case
      Failure errs -> return $ Failure errs
      Success x    -> runExceptT $ f x
    {-# INLINE (>>=) #-}
  
  instance Monad m => MonadFail (ExceptT m) where
    fail s = ExceptT $ return $ Failure $ SearchFailure s
    {-# INLINE fail #-}
  
  instance Monad m => Fallible (ExceptT m) where
    throwErr errs = ExceptT $ return $ Failure errs
    {-# INLINE throwErr #-}
  
  instance Monad m => Alternative (ExceptT m) where
    empty = throwErr $ SearchFailure ""
    {-# INLINE empty #-}
    m1 <|> m2 = do
      catchSearchFailure m1 >>= \case
        Nothing -> m2
        Just x -> return x
    {-# INLINE (<|>) #-}
  
  instance Monad m => Catchable (ExceptT m) where
    m `catchErr` f = ExceptT $ runExceptT m >>= \case
      Failure errs -> runExceptT $ f errs
      Success x    -> return $ Success x
    {-# INLINE catchErr #-}
  
  instance MonadState s m => MonadState s (ExceptT m) where
    get = lift get
    {-# INLINE get #-}
    put x = lift $ put x
    {-# INLINE put #-}
  
  instance MonadTrans ExceptT where
    lift m = ExceptT $ Success <$> m
    {-# INLINE lift #-}
  
  -- === Except type ===
  
  -- Except is isomorphic to `Either Err` but having a distinct type makes it
  -- easier to debug type errors.
  
  data Except a =
     Failure Err
   | Success a
     deriving (Show, Eq)
  
  instance Functor Except where
    fmap = liftM
    {-# INLINE fmap #-}
  
  instance Applicative Except where
    pure = return
    {-# INLINE pure #-}
    liftA2 = liftM2
    {-# INLINE liftA2 #-}
  
  instance Monad Except where
    return = Success
    {-# INLINE return #-}
    Failure errs >>= _ = Failure errs
    Success x >>= f = f x
    {-# INLINE (>>=) #-}
  
  instance Alternative Except where
    empty = throwErr $ SearchFailure ""
    {-# INLINE empty #-}
    m1 <|> m2 = do
      catchSearchFailure m1 >>= \case
        Nothing -> m2
        Just x -> return x
    {-# INLINE (<|>) #-}
  
  instance Catchable Except where
    Success ans `catchErr` _ = Success ans
    Failure errs `catchErr` f = f errs
    {-# INLINE catchErr #-}
  
  -- === HardFail ===
  
  -- Implements Fallible by crashing. Used in type querying when we want to avoid
  -- work by trusting decl annotations and skipping the checks.
  newtype HardFailM a =
    HardFailM { runHardFail' :: Identity a }
    -- We don't derive Functor, Applicative and Monad, because Identity doesn't
    -- use INLINE pragmas in its own instances, which unnecessarily inhibits optimizations.
  
  instance Functor HardFailM where
    fmap f (HardFailM (Identity x)) = HardFailM $ Identity $ f x
    {-# INLINE fmap #-}
  
  instance Applicative HardFailM where
    pure = HardFailM . Identity
    {-# INLINE pure #-}
    (<*>) = coerce
    {-# INLINE (<*>) #-}
    liftA2 = coerce
    {-# INLINE liftA2 #-}
  
  instance Monad HardFailM where
    (HardFailM (Identity x)) >>= k = k x
    {-# INLINE (>>=) #-}
    return = HardFailM . Identity
    {-# INLINE return #-}
  
  runHardFail :: HardFailM a -> a
  runHardFail m = runIdentity $ runHardFail' m
  {-# INLINE runHardFail #-}
  
  instance MonadFail HardFailM where
    fail s = error s
    {-# INLINE fail #-}
  
  instance Fallible HardFailM where
    throwErr errs = error $ pprint errs
    {-# INLINE throwErr #-}
  
  -- === convenience layer ===
  
  throw :: (ToErr e, Fallible m) => SrcId -> e -> m a
  throw sid e = throwErr $ toErr sid e
  {-# INLINE throw #-}
  
  getCurrentCallStack :: () -> Maybe [String]
  getCurrentCallStack () =
  #ifdef DEX_DEBUG
    case reverse (unsafePerformIO currentCallStack) of
      []    -> Nothing
      stack -> Just stack
  #else
    Nothing
  #endif
  {-# NOINLINE getCurrentCallStack #-}
  
  printCurrentCallStack :: Maybe [String] -> String
  printCurrentCallStack Nothing = "<no call stack available>"
  printCurrentCallStack (Just frames) = fold frames
  
  catchIOExcept :: MonadIO m => IO a -> m (Except a)
  catchIOExcept m = liftIO $ (liftM Success m) `catches`
    [ Handler \(e::Err)           -> return $ Failure e
    , Handler \(e::IOError)        -> return $ Failure $ MiscErr $ MiscMiscErr $ show e
    -- Propagate asynchronous exceptions like ThreadKilled; they are
    -- part of normal operation (of the live evaluation modes), not
    -- compiler bugs.
    , Handler \(e::AsyncException) -> liftIO $ throwIO e
    , Handler \(e::SomeException)  -> return $ Failure $ InternalErr $ show e
    ]
  
  liftExcept :: Fallible m => Except a -> m a
  liftExcept (Failure errs) = throwErr errs
  liftExcept (Success ans) = return ans
  {-# INLINE liftExcept #-}
  
  liftExceptAlt :: Alternative m => Except a -> m a
  liftExceptAlt = \case
    Success a -> pure a
    Failure _ -> empty
  {-# INLINE liftExceptAlt #-}
  
  ignoreExcept :: HasCallStack => Except a -> a
  ignoreExcept (Failure e) = error $ pprint e
  ignoreExcept (Success x) = x
  {-# INLINE ignoreExcept #-}
  
  assertEq :: (HasCallStack, Fallible m, Show a, Pretty a, Eq a) => a -> a -> String -> m ()
  assertEq x y s = if x == y then return ()
                             else throwInternal msg
    where msg = "assertion failure (" ++ s ++ "):\n"
                ++ pprint x ++ " != " ++ pprint y
  
  throwInternal :: (HasCallStack, Fallible m) => String -> m a
  throwInternal s = throwErr $ InternalErr $ s ++ "\n" ++ prettyCallStack callStack ++ "\n"
  
  instance (Monoid w, Fallible m) => Fallible (WriterT w m) where
    throwErr errs = lift $ throwErr errs
    {-# INLINE throwErr #-}
  
  instance Fallible [] where
    throwErr _ = []
    {-# INLINE throwErr #-}
  
  instance Fallible Maybe where
    throwErr _ = Nothing
    {-# INLINE throwErr #-}
  
  -- === instances ===
  
  instance Fallible Except where
    throwErr errs = Failure errs
    {-# INLINE throwErr #-}
  
  instance MonadFail Except where
    fail s = Failure $ SearchFailure s
    {-# INLINE fail #-}
  
  instance Exception Err
  
  instance Fallible m => Fallible (ReaderT r m) where
    throwErr errs = lift $ throwErr errs
    {-# INLINE throwErr #-}
  
  instance Catchable m => Catchable (ReaderT r m) where
    ReaderT f `catchErr` handler = ReaderT \r ->
      f r `catchErr` \e -> runReaderT (handler e) r
  
  instance Fallible m => Fallible (StateT s m) where
    throwErr errs = lift $ throwErr errs
    {-# INLINE throwErr #-}
  
  instance Catchable m => Catchable (StateT s m) where
    StateT f `catchErr` handler = StateT \s ->
      f s `catchErr` \e -> runStateT (handler e) s
  
  instance Pretty Err where
    pretty e = pretty $ printErr e
  
  instance ToJSON SrcId
  
  instance Store InfVarDesc
  instance Store SrcId
  
  instance Hashable InfVarDesc
  instance Hashable SrcId
  