1 files changed, 36 insertions, 7 deletions

diff --git a/src/Language/Fiddle/Compiler.hs b/src/Language/Fiddle/Compiler.hs
index c029765..049d533 100644
--- a/src/Language/Fiddle/Compiler.hs
+++ b/src/Language/Fiddle/Compiler.hs
@@ -6,6 +6,7 @@ module Language.Fiddle.Compiler where
 
 import Control.Arrow
 import Control.Monad (when)
+import Control.Monad.Identity (Identity)
 import Control.Monad.RWS
 import Control.Monad.Trans.Maybe
 import Control.Monad.Writer
@@ -21,6 +22,15 @@ data Level = Error | Warning | Info
 
 data Diagnostic = Diagnostic Level String SourceSpan
 
+emitDiagnosticError :: String -> Commented SourceSpan -> Compile a ()
+emitDiagnosticError str a = tell [Diagnostic Error str (unCommented a)]
+
+emitDiagnosticWarning :: String -> Commented SourceSpan -> Compile a ()
+emitDiagnosticWarning str a = tell [Diagnostic Warning str (unCommented a)]
+
+emitDiagnosticInfo :: String -> Commented SourceSpan -> Compile a ()
+emitDiagnosticInfo str a = tell [Diagnostic Info str (unCommented a)]
+
 newtype Compile s a = Compile (MaybeT (RWS () [Diagnostic] s) a)
   deriving newtype (Functor, Applicative, Monad)
 
@@ -103,13 +113,13 @@ pureCompilationPhase ::
   (CompilationStage stageFrom) =>
   ( FiddleUnit
       stageFrom
-      (StageFunctor stageFrom)
+      Identity
       (StageAnnotation stageFrom) ->
     Compile
       ()
       ( FiddleUnit
           stageTo
-          (StageFunctor stageTo)
+          Identity
           (StageAnnotation stageTo)
       )
   ) ->
@@ -138,7 +148,7 @@ data CompilationPhase stageFrom stageTo where
       -- only time a side effect may be performed.
       ioAction ::
         privateFlags ->
-        TreeType FiddleUnit Parsed ->
+        FiddleUnit Parsed Identity (StageAnnotation Parsed) ->
         IO ([Diagnostic], Maybe privateState),
       -- | 'nextStage' is the function that transforms a 'FiddleUnit' from
       -- the current stage ('stageFrom') to the next stage ('stageTo'). It
@@ -149,13 +159,13 @@ data CompilationPhase stageFrom stageTo where
         privateState ->
         FiddleUnit
           stageFrom
-          (StageFunctor stageFrom)
+          Identity
           (StageAnnotation stageFrom) ->
         Compile
           ()
           ( FiddleUnit
               stageTo
-              (StageFunctor stageTo)
+              Identity
               (StageAnnotation stageTo)
           )
     } ->
@@ -198,11 +208,11 @@ thenPhase
 execCompilationPipelineWithCmdline ::
   CompilationPhase Parsed s' ->
   Parser
-    ( FiddleUnit Parsed (StageFunctor Parsed) (StageAnnotation Parsed) ->
+    ( FiddleUnit Parsed Identity (StageAnnotation Parsed) ->
       IO
         ( [Diagnostic],
           Maybe
-            ( FiddleUnit s' (StageFunctor s') (StageAnnotation s')
+            ( FiddleUnit s' Identity (StageAnnotation s')
             )
         )
     )
@@ -216,3 +226,22 @@ execCompilationPipelineWithCmdline
             Nothing -> return (diags, Nothing)
       )
       flagsParser
+
+squeezeDiagnostics :: (Alter t) => t (Either [Diagnostic]) a -> Compile () (t Identity a)
+squeezeDiagnostics ast = do
+  _ <-
+    alter
+      ( \case
+          (Left l) -> tell l >> return (Left l)
+          r -> return r
+      )
+      return
+      ast
+
+  case squeeze ast of
+    (Left _) -> compilationFailure
+    (Right a) -> return a
+
+resolveOrFail :: Either [Diagnostic] a -> Compile s a
+resolveOrFail (Right a) = return a
+resolveOrFail (Left l) = tell l >> compilationFailure