1 files changed, 33 insertions, 47 deletions
diff --git a/src/Language/Fiddle/Compiler.hs b/src/Language/Fiddle/Compiler.hs
index 8d8d65c..b523a78 100644
--- a/src/Language/Fiddle/Compiler.hs
+++ b/src/Language/Fiddle/Compiler.hs
@@ -1,6 +1,9 @@
 module Language.Fiddle.Compiler where
 
+import Control.Monad (when)
+import Control.Monad.RWS (RWS, RWST, evalRWS, mapRWS, runRWS)
 import Control.Monad.State
+import Control.Monad.Trans.Maybe
 import Control.Monad.Writer
 import Data.Default
 import Language.Fiddle.Ast
@@ -12,65 +15,39 @@ data Level = Error | Warning | Info
 data Diagnostic = Diagnostic Level String SourceSpan
 
 -- Compilation monad. Has diagnostics. Optionally produces a value.
-newtype Compile s a = Compile (s -> (s, [Diagnostic], Maybe a))
+-- newtype Compile s a = Compile (s -> (s, [Diagnostic], Maybe a))
 
-instance Functor (Compile s) where
-  fmap fn (Compile cfn) = Compile $ \s ->
-    let (s', d', ma) = cfn s in (s, d', fmap fn ma)
+newtype Compile s a = Compile (MaybeT (RWS () [Diagnostic] s) a)
+  deriving (Functor, Applicative, Monad)
 
-instance Applicative (Compile s) where
-  (<*>) mfn ma = do
-    fn <- mfn
-    fn <$> ma
-
-  pure = return
-
-instance Monad (Compile s) where
-  return a = Compile (,[],Just a)
-
-  -- m a -> (a -> m b) -> m b
-  (>>=) (Compile cfn) fn = Compile $ \s ->
-    let (s', diags, ma) = cfn s
-     in case ma of
-          Nothing -> (s', diags, Nothing)
-          Just a ->
-            let (Compile cfn') = fn a
-                (s'', diags', mb) = cfn' s'
-             in (s'', diags ++ diags', mb)
+compilationFailure :: Compile s a
+compilationFailure = Compile $ MaybeT (return Nothing)
 
 instance MonadWriter [Diagnostic] (Compile s) where
-  tell diag = Compile (,diag,Just ())
-
-  listen (Compile fn) = Compile $ \s ->
-    let (s', diags, ma) = fn s in (s', diags, (,diags) <$> ma)
-
-  -- Not really "correctly" implemented, but I suspect this function will not be
-  -- used very much.
-  pass (Compile fn) = Compile $ \s ->
-    let (s', diags, mafn) = fn s
-     in case mafn of
-          Just (a, fn) -> (s', fn diags, Just a)
-          Nothing -> (s', diags, Nothing)
+  tell = Compile . tell
+  listen (Compile fn) = Compile $ listen fn
+  pass (Compile fn) = Compile $ pass fn
 
 instance MonadState s (Compile s) where
-  get = Compile $ \s -> (s, [], Just s)
-
-  put s = Compile $ const (s, [], Just ())
-
-hoistMaybe :: Maybe a -> Compile s a
-hoistMaybe ma = Compile (,[],ma)
+  get = Compile get
+  put s = Compile $ put s
+  state fn = Compile $ state fn
 
 -- Runs a sub-compilation routine with the given state, but discards the
 -- resulting state in favor of the original state.
-subCompile :: s' -> Compile s' a -> Compile s a
-subCompile s' (Compile fn) = Compile $ \s ->
-  let (_, diags, ma) = fn s' in (s, diags, ma)
+subCompile :: s' -> Compile s' a -> Compile s (s', a)
+subCompile s' (Compile mtrws) = Compile $ do
+  let (a, s, w) = runRWS (runMaybeT mtrws) () s'
+  tell w
+  MaybeT $ return $ fmap (s,) a
+
+-- mapMaybeT (mapRWS (\(a, _, w) -> (a, s', w))) mtrws
 
 -- Saves the state, runs the routine, then restores the state.
 pushState :: Compile s a -> Compile s a
 pushState cp = do
   s <- get
-  subCompile s cp
+  snd <$> subCompile s cp
 
 -- Runs a compilation routine. It produces diagnostics and maybe a result.
 -- Generally if the diagnostics contain an error, the result will be Nothing,
@@ -80,12 +57,15 @@ pushState cp = do
 -- from returning something even if the diagnostics contain errors, but it
 -- generally wouldn't make much sense for this to be the case.
 compile :: Compile s a -> s -> ([Diagnostic], Maybe a)
-compile (Compile fn) initState =
-  let (_, d, ma) = fn initState in (d, ma)
+compile (Compile fn) initState = do
+  let (a, _, w) = runRWS (runMaybeT fn) () initState in (w, a)
 
 compile_ :: (Default s) => Compile s a -> ([Diagnostic], Maybe a)
 compile_ c = compile c def
 
+hoistMaybe :: Maybe a -> Compile s a
+hoistMaybe = Compile . MaybeT . return
+
 newtype DiagnosticFormat = DiagnosticFormat (Diagnostic -> String)
 
 coloredFormat :: DiagnosticFormat
@@ -117,3 +97,9 @@ diagnosticToString (DiagnosticFormat f) = f
 printDiagnostic :: Diagnostic -> IO ()
 printDiagnostic d =
   putStrLn (diagnosticToString coloredFormat d)
+
+fromMayberOrFail :: SourceSpan -> String -> Maybe a -> Compile s a
+fromMayberOrFail sourceSpan err Nothing = do
+  tell [Diagnostic Error err sourceSpan]
+  compilationFailure
+fromMayberOrFail _ _ (Just a) = return a