diff --git a/CHANGELOG.md b/CHANGELOG.md
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,10 @@
 # Revision history for monadic-bang
 
+## 0.2.0.0 -- 2023-10-16
+
+* Added support for GHC 9.8
+* Prefix internal modules with `Internal`
+
 ## 0.1.1.0 -- 2023-07-10
 
 * Removed debug log message (thanks evincarofautumn!)
diff --git a/README.md b/README.md
--- a/README.md
+++ b/README.md
@@ -1,5 +1,7 @@
 # Monadic Bang
 
+[![Run Tests](https://github.com/JakobBruenker/monadic-bang/actions/workflows/haskell.yml/badge.svg?branch=main&event=push)](https://github.com/JakobBruenker/monadic-bang/actions/workflows/haskell.yml)
+
 This is a GHC Parser plugin for GHC 9.4 and above, intended to make monadic code within `do`-blocks more concise and nicer to work with. Works with HLS.
 
 This is heavily inspired by [Idris's !-notation](https://idris2.readthedocs.io/en/latest/tutorial/interfaces.html#notation), but with some [important differences](#comparison-with-idriss--notation).
diff --git a/monadic-bang.cabal b/monadic-bang.cabal
--- a/monadic-bang.cabal
+++ b/monadic-bang.cabal
@@ -14,7 +14,7 @@
 -- PVP summary:      +-+------- breaking API changes
 --                   | | +----- non-breaking API additions
 --                   | | | +--- code changes with no API change
-version:            0.1.1.0
+version:            0.2.1.0
 
 -- A short (one-line) description of the package.
 synopsis:           GHC plugin to desugar ! into do-notation
@@ -49,8 +49,9 @@
     CHANGELOG.md
     README.md
 
-tested-with:        GHC == 9.4.5
-                    GHC == 9.6.2
+tested-with:        GHC == 9.4.7
+                    GHC == 9.6.3
+                    GHC == 9.8.1
 
 source-repository head
     type:           git
@@ -60,12 +61,12 @@
     -- Modules exported by the library.
     exposed-modules:  MonadicBang
                       MonadicBang.Internal
-                      MonadicBang.Effect.Offer
-                      MonadicBang.Effect.Uniques
-                      MonadicBang.Effect.Writer.Discard
-                      MonadicBang.Options
-                      MonadicBang.Utils
-                      MonadicBang.Error
+                      MonadicBang.Internal.Effect.Offer
+                      MonadicBang.Internal.Effect.Uniques
+                      MonadicBang.Internal.Effect.Writer.Discard
+                      MonadicBang.Internal.Options
+                      MonadicBang.Internal.Utils
+                      MonadicBang.Internal.Error
 
     -- Modules included in this library but not exported.
     -- other-modules:
@@ -83,11 +84,11 @@
                       PatternSynonyms
 
     -- Other library packages from which modules are imported.
-    build-depends:    base >=4.17.0.0 && < 5,
-                      ghc >= 9.4 && < 9.7,
-                      containers ^>= 0.6.4.1,
-                      transformers >= 0.5.6.2 && < 0.7,
-                      fused-effects ^>= 1.1.1.2,
+    build-depends:    base >=4.17.0.0 && <4.20,
+                      ghc >=9.4 && <9.9,
+                      containers ^>=0.6.4.1 || ^>=0.7,
+                      transformers >=0.5.6.2 && <0.7,
+                      fused-effects ^>=1.1.1.2
 
     -- Directories containing source files.
     hs-source-dirs:   src
@@ -119,10 +120,9 @@
     -- Test dependencies.
     build-depends:    base,
                       ghc,
-                      ghc-boot >= 9.4 && < 9.7,
-                      ghc-paths ^>= 0.1.0.12,
+                      ghc-boot >=9.4 && <9.9,
+                      ghc-paths ^>=0.1.0.12,
                       transformers,
-                      exceptions ^>= 0.10,
                       monadic-bang
 
     ghc-options:      -Wall -Wcompat -plugin-package=monadic-bang
diff --git a/src/MonadicBang/Effect/Offer.hs b/src/MonadicBang/Effect/Offer.hs
deleted file mode 100644
--- a/src/MonadicBang/Effect/Offer.hs
+++ /dev/null
@@ -1,35 +0,0 @@
-{-# LANGUAGE BlockArguments #-}
-{-# LANGUAGE DerivingStrategies #-}
-{-# LANGUAGE UndecidableInstances #-}
-{-# LANGUAGE MonoLocalBinds #-}
-{-# LANGUAGE OverloadedRecordDot #-}
-
-module MonadicBang.Effect.Offer where
-
-import Control.Algebra
-import Control.Carrier.State.Strict
-import Data.Map.Strict (Map)
-import Data.Map.Strict qualified as M
-
--- | Offers a number of things that can be yoinked, but only once
-data Offer k v m a where
-  Yoink :: k -> Offer k v m (Maybe v)
-
-yoink :: Has (Offer k v) sig m => k -> m (Maybe v)
-yoink = send . Yoink
-
-newtype OfferC k v m a = OfferC {getOfferState :: StateC (Map k v) m a}
-  deriving newtype (Functor, Applicative, Monad)
-
--- Returns the result of the computation, along with the remaining offers
-runOffer :: Map k v -> OfferC k v m a -> m (Map k v, a)
-runOffer o (OfferC s) = runState o s
-
-instance (Algebra sig m, Ord k) => Algebra (Offer k v :+: sig) (OfferC k v m) where
-  alg hdl sig ctx = case sig of
-    L (Yoink k) -> OfferC do
-      (mv, remaining) <- M.updateLookupWithKey (\_ _ -> Nothing) k <$> get
-      put remaining
-      pure (mv <$ ctx)
-    R other -> OfferC (alg ((.getOfferState) . hdl) (R other) ctx)
-  {-# INLINE alg #-}
diff --git a/src/MonadicBang/Effect/Uniques.hs b/src/MonadicBang/Effect/Uniques.hs
deleted file mode 100644
--- a/src/MonadicBang/Effect/Uniques.hs
+++ /dev/null
@@ -1,41 +0,0 @@
-{-# LANGUAGE BlockArguments #-}
-{-# LANGUAGE DerivingStrategies #-}
-{-# LANGUAGE UndecidableInstances #-}
-{-# LANGUAGE MonoLocalBinds #-}
-{-# LANGUAGE OverloadedRecordDot #-}
-
-module MonadicBang.Effect.Uniques where
-
-import Control.Algebra
-import Control.Carrier.State.Strict
-import Control.Monad.IO.Class
-import Data.Functor
-import Data.Tuple
-
-import GHC.Types.Unique
-import GHC.Types.Unique.Supply
-
--- | Uniques provides arbitrarily many unique GHC Uniques
-data Uniques m a where
-  FreshUnique :: Uniques m Unique
-
-freshUnique :: Has Uniques sig m => m Unique
-freshUnique = send FreshUnique
-
-newtype UniquesC m a = UniquesC {getUniquesState :: StateC UniqSupply m a}
-  deriving newtype (Functor, Applicative, Monad)
-
--- | The "mask" (Char) supplied is purely cosmetic, making it easier to figure out where a Unique was born.
---
--- See Note [Uniques for wired-in prelude things and known masks] in GHC.Builtin.Uniques
-runUniquesIO :: MonadIO m => Char -> UniquesC m a -> m a
-runUniquesIO mask (UniquesC s) = flip evalState s =<< liftIO (mkSplitUniqSupply mask)
-
-runUniques :: Functor m => UniqSupply -> UniquesC m a -> m a
-runUniques uniqSupply (UniquesC s) = evalState uniqSupply s
-
-instance Algebra sig m => Algebra (Uniques :+: sig) (UniquesC m) where
-  alg hdl sig ctx = case sig of
-    L FreshUnique -> UniquesC . state $ fmap (ctx $>) . swap . takeUniqFromSupply
-    R other -> UniquesC (alg ((.getUniquesState) . hdl) (R other) ctx)
-  {-# INLINE alg #-} 
diff --git a/src/MonadicBang/Effect/Writer/Discard.hs b/src/MonadicBang/Effect/Writer/Discard.hs
deleted file mode 100644
--- a/src/MonadicBang/Effect/Writer/Discard.hs
+++ /dev/null
@@ -1,24 +0,0 @@
-{-# LANGUAGE DerivingStrategies #-}
-{-# LANGUAGE UndecidableInstances #-}
-{-# LANGUAGE BlockArguments #-}
-{-# LANGUAGE GADTs #-}
-
--- | A Writer carrier that discards any values it is told
-module MonadicBang.Effect.Writer.Discard where
-
-import Control.Algebra
-import Control.Effect.Writer
-
-newtype DiscardC w m a = DiscardC { evalDiscardC :: m a }
-  deriving newtype (Functor, Applicative, Monad)
-
-evalWriter :: (Monoid w, Algebra sig m) => DiscardC w m a -> m a
-evalWriter = evalDiscardC
-
-instance (Monoid w, Algebra sig m) => Algebra (Writer w :+: sig) (DiscardC w m) where
-  alg hdl sig ctx = DiscardC $ case sig of
-    L writer -> case writer of
-      Tell _ -> pure ctx
-      Listen m -> fmap (mempty,) <$> evalWriter (hdl (m <$ ctx))
-      Censor _ m -> evalWriter (hdl (m <$ ctx))
-    R other -> alg (evalDiscardC . hdl) other ctx
diff --git a/src/MonadicBang/Error.hs b/src/MonadicBang/Error.hs
deleted file mode 100644
--- a/src/MonadicBang/Error.hs
+++ /dev/null
@@ -1,39 +0,0 @@
-{-# LANGUAGE LambdaCase #-}
-{-# LANGUAGE CPP #-}
-
-module MonadicBang.Error where
-
-import Prelude hiding ((<>))
-
-import Control.Effect.Writer
-
-import GHC
-import GHC.Types.Error
-import GHC.Types.Name.Occurrence
-import GHC.Parser.Errors.Types
-import GHC.Utils.Outputable
-
-data Error = ErrOutOfScopeVariable OccName
-           | ErrBangOutsideOfDo
-
-type PsErrors = Writer (Messages PsError)
-
-customError :: Error -> PsError
-#if MIN_VERSION_ghc(9,6,0)
-customError = PsUnknownMessage . UnknownDiagnostic . \cases
-#else
-customError = PsUnknownMessage . \cases
-#endif
-  ErrBangOutsideOfDo -> DiagnosticMessage
-    { diagMessage = mkDecorated [text "Monadic ! outside of a 'do'-block is not allowed"]
-    , diagReason = ErrorWithoutFlag
-    , diagHints = [SuggestMissingDo]
-    }
-  (ErrOutOfScopeVariable name) -> DiagnosticMessage
-    { diagMessage = mkDecorated [text "The variable " <> quotes (ppr name) <> text " cannot be used inside of ! here, since its desugaring would escape its scope"]
-    , diagReason = ErrorWithoutFlag
-    , diagHints = [UnknownHint $ text "Maybe you meant to open a new 'do'-block after " <> ppr name <> text " has been bound?"]
-    }
-
-tellPsError :: Has PsErrors sig m => PsError -> SrcSpan -> m ()
-tellPsError err srcSpan = tell . singleMessage $ MsgEnvelope srcSpan neverQualify err SevError
diff --git a/src/MonadicBang/Internal.hs b/src/MonadicBang/Internal.hs
--- a/src/MonadicBang/Internal.hs
+++ b/src/MonadicBang/Internal.hs
@@ -44,14 +44,14 @@
 
 import GHC.Utils.Logger
 
-import MonadicBang.Effect.Offer
-import MonadicBang.Effect.Uniques
-import MonadicBang.Options
-import MonadicBang.Utils
-import MonadicBang.Error
+import MonadicBang.Internal.Effect.Offer
+import MonadicBang.Internal.Effect.Uniques
+import MonadicBang.Internal.Options
+import MonadicBang.Internal.Utils
+import MonadicBang.Internal.Error
+import MonadicBang.Internal.Effect.Writer.Discard
 
 import Data.Kind
-import MonadicBang.Effect.Writer.Discard
 
 -- We don't care about which file things are from, because the entire AST comes
 -- from the same module
diff --git a/src/MonadicBang/Internal/Effect/Offer.hs b/src/MonadicBang/Internal/Effect/Offer.hs
new file mode 100644
--- /dev/null
+++ b/src/MonadicBang/Internal/Effect/Offer.hs
@@ -0,0 +1,35 @@
+{-# LANGUAGE BlockArguments #-}
+{-# LANGUAGE DerivingStrategies #-}
+{-# LANGUAGE UndecidableInstances #-}
+{-# LANGUAGE MonoLocalBinds #-}
+{-# LANGUAGE OverloadedRecordDot #-}
+
+module MonadicBang.Internal.Effect.Offer where
+
+import Control.Algebra
+import Control.Carrier.State.Strict
+import Data.Map.Strict (Map)
+import Data.Map.Strict qualified as M
+
+-- | Offers a number of things that can be yoinked, but only once
+data Offer k v m a where
+  Yoink :: k -> Offer k v m (Maybe v)
+
+yoink :: Has (Offer k v) sig m => k -> m (Maybe v)
+yoink = send . Yoink
+
+newtype OfferC k v m a = OfferC {getOfferState :: StateC (Map k v) m a}
+  deriving newtype (Functor, Applicative, Monad)
+
+-- Returns the result of the computation, along with the remaining offers
+runOffer :: Map k v -> OfferC k v m a -> m (Map k v, a)
+runOffer o (OfferC s) = runState o s
+
+instance (Algebra sig m, Ord k) => Algebra (Offer k v :+: sig) (OfferC k v m) where
+  alg hdl sig ctx = case sig of
+    L (Yoink k) -> OfferC do
+      (mv, remaining) <- M.updateLookupWithKey (\_ _ -> Nothing) k <$> get
+      put remaining
+      pure (mv <$ ctx)
+    R other -> OfferC (alg ((.getOfferState) . hdl) (R other) ctx)
+  {-# INLINE alg #-}
diff --git a/src/MonadicBang/Internal/Effect/Uniques.hs b/src/MonadicBang/Internal/Effect/Uniques.hs
new file mode 100644
--- /dev/null
+++ b/src/MonadicBang/Internal/Effect/Uniques.hs
@@ -0,0 +1,41 @@
+{-# LANGUAGE BlockArguments #-}
+{-# LANGUAGE DerivingStrategies #-}
+{-# LANGUAGE UndecidableInstances #-}
+{-# LANGUAGE MonoLocalBinds #-}
+{-# LANGUAGE OverloadedRecordDot #-}
+
+module MonadicBang.Internal.Effect.Uniques where
+
+import Control.Algebra
+import Control.Carrier.State.Strict
+import Control.Monad.IO.Class
+import Data.Functor
+import Data.Tuple
+
+import GHC.Types.Unique
+import GHC.Types.Unique.Supply
+
+-- | Uniques provides arbitrarily many unique GHC Uniques
+data Uniques m a where
+  FreshUnique :: Uniques m Unique
+
+freshUnique :: Has Uniques sig m => m Unique
+freshUnique = send FreshUnique
+
+newtype UniquesC m a = UniquesC {getUniquesState :: StateC UniqSupply m a}
+  deriving newtype (Functor, Applicative, Monad)
+
+-- | The "mask" (Char) supplied is purely cosmetic, making it easier to figure out where a Unique was born.
+--
+-- See Note [Uniques for wired-in prelude things and known masks] in GHC.Builtin.Uniques
+runUniquesIO :: MonadIO m => Char -> UniquesC m a -> m a
+runUniquesIO mask (UniquesC s) = flip evalState s =<< liftIO (mkSplitUniqSupply mask)
+
+runUniques :: Functor m => UniqSupply -> UniquesC m a -> m a
+runUniques uniqSupply (UniquesC s) = evalState uniqSupply s
+
+instance Algebra sig m => Algebra (Uniques :+: sig) (UniquesC m) where
+  alg hdl sig ctx = case sig of
+    L FreshUnique -> UniquesC . state $ fmap (ctx $>) . swap . takeUniqFromSupply
+    R other -> UniquesC (alg ((.getUniquesState) . hdl) (R other) ctx)
+  {-# INLINE alg #-} 
diff --git a/src/MonadicBang/Internal/Effect/Writer/Discard.hs b/src/MonadicBang/Internal/Effect/Writer/Discard.hs
new file mode 100644
--- /dev/null
+++ b/src/MonadicBang/Internal/Effect/Writer/Discard.hs
@@ -0,0 +1,24 @@
+{-# LANGUAGE DerivingStrategies #-}
+{-# LANGUAGE UndecidableInstances #-}
+{-# LANGUAGE BlockArguments #-}
+{-# LANGUAGE GADTs #-}
+
+-- | A Writer carrier that discards any values it is told
+module MonadicBang.Internal.Effect.Writer.Discard where
+
+import Control.Algebra
+import Control.Effect.Writer
+
+newtype DiscardC w m a = DiscardC { evalDiscardC :: m a }
+  deriving newtype (Functor, Applicative, Monad)
+
+evalWriter :: (Monoid w, Algebra sig m) => DiscardC w m a -> m a
+evalWriter = evalDiscardC
+
+instance (Monoid w, Algebra sig m) => Algebra (Writer w :+: sig) (DiscardC w m) where
+  alg hdl sig ctx = DiscardC $ case sig of
+    L writer -> case writer of
+      Tell _ -> pure ctx
+      Listen m -> fmap (mempty,) <$> evalWriter (hdl (m <$ ctx))
+      Censor _ m -> evalWriter (hdl (m <$ ctx))
+    R other -> alg (evalDiscardC . hdl) other ctx
diff --git a/src/MonadicBang/Internal/Error.hs b/src/MonadicBang/Internal/Error.hs
new file mode 100644
--- /dev/null
+++ b/src/MonadicBang/Internal/Error.hs
@@ -0,0 +1,49 @@
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE CPP #-}
+
+module MonadicBang.Internal.Error where
+
+import Prelude hiding ((<>))
+
+import Control.Effect.Writer
+
+import GHC
+#if MIN_VERSION_ghc(9,8,0)
+import GHC.Utils.Error
+#endif
+import GHC.Types.Error
+import GHC.Types.Name.Occurrence
+import GHC.Parser.Errors.Types
+import GHC.Utils.Outputable
+
+data Error = ErrOutOfScopeVariable OccName
+           | ErrBangOutsideOfDo
+
+type PsErrors = Writer (Messages PsError)
+
+customError :: Error -> PsError
+#if MIN_VERSION_ghc(9,8,0)
+customError = PsUnknownMessage . mkUnknownDiagnostic . \cases
+#elif MIN_VERSION_ghc(9,6,0)
+customError = PsUnknownMessage . UnknownDiagnostic . \cases
+#else
+customError = PsUnknownMessage . \cases
+#endif
+  ErrBangOutsideOfDo -> DiagnosticMessage
+    { diagMessage = mkDecorated [text "Monadic ! outside of a 'do'-block is not allowed"]
+    , diagReason = ErrorWithoutFlag
+    , diagHints = [SuggestMissingDo]
+    }
+  (ErrOutOfScopeVariable name) -> DiagnosticMessage
+    { diagMessage = mkDecorated [text "The variable " <> quotes (ppr name) <> text " cannot be used inside of ! here, since its desugaring would escape its scope"]
+    , diagReason = ErrorWithoutFlag
+    , diagHints = [UnknownHint $ text "Maybe you meant to open a new 'do'-block after " <> ppr name <> text " has been bound?"]
+    }
+
+tellPsError :: Has PsErrors sig m => PsError -> SrcSpan -> m ()
+tellPsError err srcSpan = tell . singleMessage $
+#if MIN_VERSION_ghc(9,8,0)
+  mkErrorMsgEnvelope srcSpan neverQualify err
+#else
+  MsgEnvelope srcSpan neverQualify err SevError
+#endif
diff --git a/src/MonadicBang/Internal/Options.hs b/src/MonadicBang/Internal/Options.hs
new file mode 100644
--- /dev/null
+++ b/src/MonadicBang/Internal/Options.hs
@@ -0,0 +1,61 @@
+{-# LANGUAGE BlockArguments #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE OverloadedRecordDot #-}
+{-# LANGUAGE CPP #-}
+
+module MonadicBang.Internal.Options where
+
+import Control.Exception
+import Control.Algebra
+import Control.Carrier.State.Strict
+import Control.Effect.Throw
+import Control.Monad
+import Data.Bifunctor
+import Data.Bool
+import Data.List (intercalate, partition)
+
+import GHC
+import GHC.Plugins
+
+data Verbosity = DumpTransformed | Quiet
+
+data PreserveErrors = Preserve | Don'tPreserve
+
+data Options = MkOptions {verbosity :: Verbosity, preserveErrors :: PreserveErrors}
+
+#if MIN_VERSION_ghc(9,6,0)
+parseOptions :: Has (Throw ErrorCall) sig m => Located (HsModule GhcPs) -> [CommandLineOption] -> m Options
+#else
+parseOptions :: Has (Throw ErrorCall) sig m => Located HsModule -> [CommandLineOption] -> m Options
+#endif
+parseOptions mod' cmdLineOpts = do
+  (remaining, options) <- runState cmdLineOpts do
+    verbosity <- bool Quiet DumpTransformed <$> extractOpts verboseOpts
+    preserveErrors <- bool Don'tPreserve Preserve <$> extractOpts preserveErrorsOpts
+    pure $ MkOptions verbosity preserveErrors
+  unless (null remaining) . throwError . ErrorCall $
+    "Incorrect command line options for plugin MonadicBang, encountered in " ++ modName ++ modFile ++
+    "\n\tOptions that were supplied (via -fplugin-opt) are: " ++ intercalate ", " (map show cmdLineOpts) ++
+    "\n\tUnrecognized options: " ++ showOpts remaining ++
+    "\n\n\tUsage: [-ddump] [-preserve-errors]" ++
+    "\n" ++
+    "\n\t\t-ddump            Print the altered AST" ++
+    "\n\t\t-preserve-errors  Keep parse errors about ! outside of 'do' in their original form, rather then a more relevant explanation." ++
+    "\n\t\t                  This is mainly useful if another plugin expects those errors."
+  pure options
+
+  where
+    verboseOpts = ["-ddump"]
+    preserveErrorsOpts = ["-preserve-errors"]
+    extractOpts opt = do
+      (isOpt, opts') <- gets $ first (not . null) . partition (`elem` opt)
+      put opts'
+      pure isOpt
+
+    showOpts = intercalate ", " . map show
+
+    modFile = maybe "" ((" in file " ++) . unpackFS . srcSpanFile) $ toRealSrcSpan (getLoc mod')
+    modName = maybe "an unnamed module" (("module " ++) . moduleNameString . unLoc) $ (unLoc mod').hsmodName
+    toRealSrcSpan = \cases
+      (RealSrcSpan rss _) -> Just rss
+      (UnhelpfulSpan _) -> Nothing
diff --git a/src/MonadicBang/Internal/Utils.hs b/src/MonadicBang/Internal/Utils.hs
new file mode 100644
--- /dev/null
+++ b/src/MonadicBang/Internal/Utils.hs
@@ -0,0 +1,38 @@
+{-# LANGUAGE MonoLocalBinds #-}
+{-# LANGUAGE CPP #-}
+
+module MonadicBang.Internal.Utils where
+
+import Control.Monad.Trans.Maybe
+import Data.Monoid
+import GHC.Stack (withFrozenCallStack, HasCallStack)
+
+type DList a = Endo [a]
+
+-- | Handle a specific AST node
+type Handler m a = a -> m a
+
+-- | Try handling an AST node, but may fail (usually because the handler is not
+-- applicable)
+type Try m a = Handler (MaybeT m) a
+
+{-# INLINE fromDList #-}
+fromDList :: DList a -> [a]
+fromDList = appEndo ?? []
+
+{-# INLINE (??) #-}
+(??) :: Functor f => f (a -> b) -> a -> f b
+fs ?? x = ($ x) <$> fs
+
+#if MIN_VERSION_transformers(0,6,0)
+#else
+{-# INLINE hoistMaybe #-}
+hoistMaybe :: Applicative m => Maybe a -> MaybeT m a
+hoistMaybe = MaybeT . pure
+#endif
+
+panic :: HasCallStack => String -> a
+panic message = withFrozenCallStack $ error $
+  unlines ["MonadicBang panic:", message, "", submitReport]
+  where
+    submitReport = "This is likely a bug. Please submit a bug report at https://github.com/JakobBruenker/monadic-bang/issues"
diff --git a/src/MonadicBang/Options.hs b/src/MonadicBang/Options.hs
deleted file mode 100644
--- a/src/MonadicBang/Options.hs
+++ /dev/null
@@ -1,61 +0,0 @@
-{-# LANGUAGE BlockArguments #-}
-{-# LANGUAGE LambdaCase #-}
-{-# LANGUAGE OverloadedRecordDot #-}
-{-# LANGUAGE CPP #-}
-
-module MonadicBang.Options where
-
-import Control.Exception
-import Control.Algebra
-import Control.Carrier.State.Strict
-import Control.Effect.Throw
-import Control.Monad
-import Data.Bifunctor
-import Data.Bool
-import Data.List (intercalate, partition)
-
-import GHC
-import GHC.Plugins
-
-data Verbosity = DumpTransformed | Quiet
-
-data PreserveErrors = Preserve | Don'tPreserve
-
-data Options = MkOptions {verbosity :: Verbosity, preserveErrors :: PreserveErrors}
-
-#if MIN_VERSION_ghc(9,6,0)
-parseOptions :: Has (Throw ErrorCall) sig m => Located (HsModule GhcPs) -> [CommandLineOption] -> m Options
-#else
-parseOptions :: Has (Throw ErrorCall) sig m => Located HsModule -> [CommandLineOption] -> m Options
-#endif
-parseOptions mod' cmdLineOpts = do
-  (remaining, options) <- runState cmdLineOpts do
-    verbosity <- bool Quiet DumpTransformed <$> extractOpts verboseOpts
-    preserveErrors <- bool Don'tPreserve Preserve <$> extractOpts preserveErrorsOpts
-    pure $ MkOptions verbosity preserveErrors
-  unless (null remaining) . throwError . ErrorCall $
-    "Incorrect command line options for plugin MonadicBang, encountered in " ++ modName ++ modFile ++
-    "\n\tOptions that were supplied (via -fplugin-opt) are: " ++ intercalate ", " (map show cmdLineOpts) ++
-    "\n\tUnrecognized options: " ++ showOpts remaining ++
-    "\n\n\tUsage: [-ddump] [-preserve-errors]" ++
-    "\n" ++
-    "\n\t\t-ddump            Print the altered AST" ++
-    "\n\t\t-preserve-errors  Keep parse errors about ! outside of 'do' in their original form, rather then a more relevant explanation." ++
-    "\n\t\t                  This is mainly useful if another plugin expects those errors."
-  pure options
-
-  where
-    verboseOpts = ["-ddump"]
-    preserveErrorsOpts = ["-preserve-errors"]
-    extractOpts opt = do
-      (isOpt, opts') <- gets $ first (not . null) . partition (`elem` opt)
-      put opts'
-      pure isOpt
-
-    showOpts = intercalate ", " . map show
-
-    modFile = maybe "" ((" in file " ++) . unpackFS . srcSpanFile) $ toRealSrcSpan (getLoc mod')
-    modName = maybe "an unnamed module" (("module " ++) . moduleNameString . unLoc) $ (unLoc mod').hsmodName
-    toRealSrcSpan = \cases
-      (RealSrcSpan rss _) -> Just rss
-      (UnhelpfulSpan _) -> Nothing
diff --git a/src/MonadicBang/Utils.hs b/src/MonadicBang/Utils.hs
deleted file mode 100644
--- a/src/MonadicBang/Utils.hs
+++ /dev/null
@@ -1,39 +0,0 @@
-{-# LANGUAGE MonoLocalBinds #-}
-{-# LANGUAGE CPP #-}
-
-module MonadicBang.Utils where
-
-import Control.Monad.Trans.Maybe
-import Data.Monoid
-import GHC.Stack (withFrozenCallStack, HasCallStack)
-
-type DList a = Endo [a]
-
--- | Handle a specific AST node
-type Handler m a = a -> m a
-
--- | Try handling an AST node, but may fail (usually because the handler is not
--- applicable)
-type Try m a = Handler (MaybeT m) a
-
-{-# INLINE fromDList #-}
-fromDList :: DList a -> [a]
-fromDList = appEndo ?? []
-
-{-# INLINE (??) #-}
-(??) :: Functor f => f (a -> b) -> a -> f b
-fs ?? x = ($ x) <$> fs
-
-#if MIN_VERSION_ghc(9,6,0)
-#else
--- This is included in transformers 0.6, but that can't be used together with ghc 9.4
-{-# INLINE hoistMaybe #-}
-hoistMaybe :: Applicative m => Maybe a -> MaybeT m a
-hoistMaybe = MaybeT . pure
-#endif
-
-panic :: HasCallStack => String -> a
-panic message = withFrozenCallStack $ error $
-  unlines ["MonadicBang panic:", message, "", submitReport]
-  where
-    submitReport = "This is likely a bug. Please submit a bug report at https://github.com/JakobBruenker/monadic-bang/issues"
diff --git a/test/MonadicBang/Test.hs b/test/MonadicBang/Test.hs
--- a/test/MonadicBang/Test.hs
+++ b/test/MonadicBang/Test.hs
@@ -1,13 +1,51 @@
+{-# LANGUAGE CPP #-}
 {-# LANGUAGE BlockArguments #-}
 {-# LANGUAGE LexicalNegation #-}
 {-# LANGUAGE MonadComprehensions #-}
+{-# LANGUAGE ViewPatterns #-}
 
 module Main (main) where
 
+import Prelude hiding ((<>))
+#if MIN_VERSION_ghc(9,6,0)
+#else
+import Control.Applicative (liftA2)
+#endif
+import Data.Foldable
+import Data.Traversable
+import System.IO
+
+import MonadicBang.Test.Utils
 import MonadicBang.Test.ShouldPass
 import MonadicBang.Test.ShouldFail
 
+import GHC.Utils.Outputable
+import GHC.Utils.Ppr (Mode(PageMode))
+import System.Exit (exitFailure)
+
 main :: IO ()
 main = do
-  shouldPass
-  shouldFail
+  (numFailures, numFailedSuites) <- liftA2 (,) sum (length . filter (> 0)) <$> for suites \suite -> do
+    failures <- runSuite suite
+    for_ failures \failure -> putSDoc stderr $ vcat [space, prettyFail failure]
+    pure $ length failures
+  if numFailures == 0
+    then putStrLn "All tests passed!"
+    else do
+      putSDoc stdout $
+        space $+$ plural' numFailures "test" <+> text "in" <+> int numFailedSuites <> char '/' <> plural' (length suites) "suite" <+> text "failed."
+      exitFailure
+  where
+    plural' :: Int -> String -> SDoc
+    plural' n (text -> s) = int n <+> case n of
+      1 -> s
+      _ -> s <> char 's'
+
+putSDoc :: Handle -> SDoc -> IO ()
+putSDoc = printSDocLn defaultSDocContext (PageMode True)
+
+suites :: [TestType]
+suites =
+  [ shouldPass
+  , shouldFail
+  ]
diff --git a/test/MonadicBang/Test/ShouldFail.hs b/test/MonadicBang/Test/ShouldFail.hs
--- a/test/MonadicBang/Test/ShouldFail.hs
+++ b/test/MonadicBang/Test/ShouldFail.hs
@@ -3,7 +3,7 @@
 module MonadicBang.Test.ShouldFail (shouldFail) where
 
 import MonadicBang.Test.Utils
-import MonadicBang.Error
+import MonadicBang.Internal.Error
 
 import GHC.Types.Name.Occurrence
 
diff --git a/test/MonadicBang/Test/ShouldPass.hs b/test/MonadicBang/Test/ShouldPass.hs
--- a/test/MonadicBang/Test/ShouldPass.hs
+++ b/test/MonadicBang/Test/ShouldPass.hs
@@ -16,6 +16,7 @@
 
 import MonadicBang.Test.Utils
 import MonadicBang.Test.Utils.QualifiedDo qualified as QualifiedDo
+import Control.Monad.IO.Class
 
 shouldPass :: Test
 shouldPass = do
@@ -39,7 +40,7 @@
   confusing
   qualifiedDo
 
-getA, getB, getC :: IO String
+getA, getB, getC :: MonadIO m => m String
 getA = pure "a"
 getB = pure "b"
 getC = pure "c"
diff --git a/test/MonadicBang/Test/Utils.hs b/test/MonadicBang/Test/Utils.hs
--- a/test/MonadicBang/Test/Utils.hs
+++ b/test/MonadicBang/Test/Utils.hs
@@ -4,6 +4,9 @@
 #if MIN_VERSION_ghc(9,6,0)
 {-# LANGUAGE ScopedTypeVariables #-}
 #endif
+{-# LANGUAGE NoFieldSelectors #-}
+{-# LANGUAGE OverloadedRecordDot #-}
+{-# LANGUAGE ViewPatterns #-}
 
 module MonadicBang.Test.Utils where
 
@@ -11,6 +14,8 @@
 import Data.Foldable
 import Data.Function
 
+import Control.Monad.Trans.Writer.CPS
+
 import GHC.Stack
 
 import GHC
@@ -20,49 +25,76 @@
 import GHC.Utils.Outputable hiding ((<>))
 
 import MonadicBang.Test.Utils.RunGhcParser
+import MonadicBang.Internal.Utils
+import Data.Monoid
 
--- TODO: This should use a Writer to collect all errors
-type Test = HasCallStack => IO ()
+data FailType
+  = forall a . Show a => IncorrectResult { expectedValue :: a, actualValue :: a }
+  | forall a . Outputable a => Didn'tFail { expectedFails :: [PsMessage], actualValue :: a }
+  | FailedIncorrectly { expectedFails :: [PsMessage], actualFails :: [GhcMessage] }
 
-assertEq :: (HasCallStack, Show a, Eq a) => a -> a -> IO ()
--- We don't care about seeing where the `error` call itself happens in the
--- call stack, so we freeze it
-assertEq expected actual = when (expected /= actual) $ withFrozenCallStack do
-  error $ "Expected " <> show expected <> ", but got " <> show actual
+data Fail = MkFail { error :: FailType, callStack :: CallStack }
 
+type TestType = WriterT (DList Fail) IO ()
+
+type Test = HasCallStack => TestType
+
+runSuite :: TestType -> IO [Fail]
+runSuite test = fromDList <$> execWriterT test
+
+prettyFail :: Fail -> SDoc
+prettyFail failure = vcat
+  [ case failure.error of
+      IncorrectResult{ expectedValue, actualValue } -> vcat
+        [ text "Expected: " <+> text (show expectedValue)
+        , text "but got:  " <+> text (show actualValue)
+        ]
+      Didn'tFail{ expectedFails } -> vcat
+        [ text "Expected failure with"
+        , nest 2 $ diagnosticsSDoc expectedFails
+        , text "but execution succeeded"
+        ]
+      FailedIncorrectly{ expectedFails, actualFails } -> vcat
+        [ text "Expected failure with"
+        , nest 2 $ diagnosticsSDoc expectedFails
+        , text "but execution failed with these errors instead:"
+        , nest 2 $ diagnosticsSDoc actualFails
+        ]
+  , text "at" <+> text (prettyCallStack failure.callStack)
+  ]
+  where
+    diagnosticsSDoc diags = vcat (map (vcat . unDecorated . diagMsg) diags)
+
+recordFail :: HasCallStack => FailType -> TestType
+recordFail err = tell . Endo . (:) $ MkFail err callStack
+
+assertEq :: (HasCallStack, Show a, Eq a) => a -> a -> TestType
+assertEq expected actual = when (expected /= actual) $
+  withFrozenCallStack $ recordFail $ IncorrectResult expected actual
+
 sdocEq :: SDoc -> SDoc -> Bool
 sdocEq = (==) `on` showSDocUnsafe
 
-assertFailWith :: (HasCallStack, Outputable a) => [PsMessage] -> Either SourceError a -> IO ()
+assertFailWith :: (HasCallStack, Outputable a) => [PsMessage] -> Either SourceError a -> TestType
 assertFailWith expected = \case
-  Right result -> withFrozenCallStack $ error . showSDocUnsafe $
-    text "\n    Expected failure with" $$
-    diagnosticsSDoc expected $$
-    text "    but execution succeeded with this result:" $$
-    ppr result
-  Left err -> unless sameErrors do
-    error . showSDocUnsafe $
-      text "\n    Expected failure with" $$
-      diagnosticsSDoc expected $$
-      text "    but execution failed with these errors instead:" $$
-      diagnosticsSDoc errMsgs
+  Right result -> withFrozenCallStack $ recordFail $ Didn'tFail expected result
+  Left err -> unless sameFails do
+    withFrozenCallStack $ recordFail $ FailedIncorrectly expected errMsgs
     where
       errMsgs = toList (srcErrorMessages err)
       toPsMessage = \case
         GhcPsMessage m -> Just m
         _ -> Nothing
       listEq eq xs ys = and $ zipWith eq xs ys
-      sameErrors = maybe False (((listEq . listEq) sdocEq `on` map (unDecorated . diagMsg)) expected) $ traverse toPsMessage errMsgs
-  where
-    diagnosticsSDoc diags = vcat (map (vcat . unDecorated . diagMsg) diags)
+      sameFails = maybe False (((listEq . listEq) sdocEq `on` map (unDecorated . diagMsg)) expected) $ traverse toPsMessage errMsgs
 
-    diagMsg :: forall a . Diagnostic a => a -> DecoratedSDoc
+diagMsg :: forall a . Diagnostic a => a -> DecoratedSDoc
 #if MIN_VERSION_ghc(9,6,0)
-    diagMsg = diagnosticMessage (defaultDiagnosticOpts @a)
+diagMsg = diagnosticMessage (defaultDiagnosticOpts @a)
 #else
-    diagMsg = diagnosticMessage
+diagMsg = diagnosticMessage
 #endif
 
-assertParseFailWith :: HasCallStack => [PsMessage] -> String -> IO ()
-assertParseFailWith expected source = withFrozenCallStack do
+assertParseFailWith :: HasCallStack => [PsMessage] -> String -> TestType
+assertParseFailWith expected source = withFrozenCallStack $
   assertFailWith expected . fmap pm_parsed_source =<< parseGhc source