CompilationPhases are categories.

2 files changed, 167 insertions, 62 deletions

diff --git a/Main.hs b/Main.hs
index 34ef997..97be0a3 100644
--- a/Main.hs
+++ b/Main.hs
@@ -20,10 +20,12 @@ import Options.Applicative
 import qualified System.Environment as System
 import System.Exit (exitWith)
 import System.IO (stderr, hPutStrLn)
+import Control.Category ((>>>))
 
 newtype IntermediateAst = IntermediateAst (GenericSyntaxTree Identity (Maybe Value))
   deriving (Typeable)
 
+-- | The total compilation pipeline.
 compilationPipeline ::
   ( FilePath ->
     IO
@@ -50,6 +52,7 @@ compilationPipeline parse compile =
     >>> (dumpPhase "qualified" >>> consistencyCheckPhase)
     >>> dumpPhase "checked"
 
+-- | Compilation phase that just emits an ast artifact.
 dumpPhase ::
   forall stage.
   ( Typeable stage,
@@ -90,7 +93,7 @@ parseGlobalFlags :: Parser GlobalFlags
 parseGlobalFlags =
   GlobalFlags
     <$> argument str (metavar "INPUT" <> help "Input file")
-    <*> ((\b -> if b then jsonFormat else coloredFormat) <$> switch
+    <*> ((\b -> if b then jsonDiagnosticFormat else coloredFormat) <$> switch
       ( long "diagnostics-as-json"
           <> help "Dump diagnostics in JSON format."
       ))
diff --git a/src/Language/Fiddle/Compiler.hs b/src/Language/Fiddle/Compiler.hs
index 812d4f6..571c7b0 100644
--- a/src/Language/Fiddle/Compiler.hs
+++ b/src/Language/Fiddle/Compiler.hs
@@ -1,80 +1,145 @@
 {-# LANGUAGE DerivingStrategies #-}
+{-# LANGUAGE FunctionalDependencies #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE RankNTypes #-}
 
 module Language.Fiddle.Compiler where
 
 import Control.Arrow
+import Control.Category
 import Control.Monad (when)
 import Control.Monad.Identity (Identity)
 import Control.Monad.RWS
 import Control.Monad.Trans.Maybe
 import Control.Monad.Writer
+import Data.Aeson (FromJSON, ToJSON (toJSON), encode)
+import qualified Data.ByteString.Lazy.Char8 as BL
 import Data.Default
+import Data.Kind (Type)
 import Data.Typeable
+import GHC.Generics (Generic)
 import Language.Fiddle.Ast
 import Language.Fiddle.Types
 import Options.Applicative
 import System.IO (hPutStrLn, stderr)
-import qualified Data.ByteString.Lazy.Char8 as BL
 import Text.Parsec (sourceColumn, sourceLine, sourceName)
-import Data.Aeson (FromJSON, ToJSON (toJSON), encode)
-import GHC.Generics (Generic)
-
+import Prelude hiding (id, (.))
+
+-- | Represents the severity level of a diagnostic message. It indicates
+-- how critical the diagnostic is.
+--
+-- - 'Error': Indicates a serious problem that causes the compilation to fail,
+--   meaning the resulting AST will not be passed to the backend.
+-- - 'Warning': Represents a less severe issue, suggesting that something may
+--   need attention but does not halt the compilation.
+-- - 'Info': Used for informational messages that do not indicate a problem.
 data Level = Error | Warning | Info
   deriving (Eq, Ord, Show, Read, Enum, ToJSON, FromJSON, Generic)
 
-data Diagnostic = Diagnostic {
-  diagnosticLevel :: Level,
-  diagnosticMesasge :: String,
-  diagnosticLocation :: SourceSpan
+-- | A general representation of a compilation artifact. This can be any
+-- typeable value that is generated or collected during the compilation
+-- process. The 'Artifact' type is used to carry various outputs or
+-- intermediate results produced by the compiler.
+data Artifact where
+  Artifact :: forall t. (Typeable t) => t -> Artifact
+
+-- | Represents a diagnostic message that may be emitted during compilation.
+-- A 'Diagnostic' contains the severity level, a descriptive message, and
+-- the location in the source code where the issue was found.
+data Diagnostic = Diagnostic
+  { -- | The severity of the diagnostic.
+    diagnosticLevel :: Level,
+    -- | The message describing the issue.
+    diagnosticMessage :: String,
+    -- | The location in the source code.
+    diagnosticLocation :: SourceSpan
   }
   deriving (Generic, ToJSON, FromJSON, Typeable)
 
+-- | Emits a diagnostic with 'Error' severity, indicating a critical issue
+-- that will cause the compilation to fail. The diagnostic includes a message
+-- and the location where the error occurred.
 emitDiagnosticError :: String -> Commented SourceSpan -> Compile a ()
 emitDiagnosticError str a = tell [Diagnostic Error str (unCommented a)]
 
+-- | Emits a diagnostic with 'Warning' severity, indicating a potential issue
+-- that does not prevent compilation but may require attention. Includes a
+-- message and location information.
 emitDiagnosticWarning :: String -> Commented SourceSpan -> Compile a ()
 emitDiagnosticWarning str a = tell [Diagnostic Warning str (unCommented a)]
 
+-- | Emits a diagnostic with 'Info' severity, used to provide informational
+-- messages during compilation. Includes a message and the relevant source
+-- location.
 emitDiagnosticInfo :: String -> Commented SourceSpan -> Compile a ()
 emitDiagnosticInfo str a = tell [Diagnostic Info str (unCommented a)]
 
-data Artifact where
-  Artifact :: forall t. (Typeable t) => t -> Artifact
+-- | Emits a compilation artifact, which is a general output or data produced
+-- during the compilation process. Artifacts are collected for later use, such
+-- as generating output files or debugging information.
+emitArtifact :: Artifact -> Compile s ()
+emitArtifact art = Compile $ tell ([], [art])
 
+-- | The primary monad for managing the compilation process. 'Compile' handles
+-- emitting diagnostics and collecting artifacts while allowing the state to be
+-- modified. The 'MaybeT' layer is used to represent potential compilation
+-- failure.
 newtype Compile s a = Compile (MaybeT (RWS () ([Diagnostic], [Artifact]) s) a)
   deriving newtype (Functor, Applicative, Monad)
 
+-- | Signals a compilation failure, terminating the current compilation process.
+-- No further diagnostics will be emitted. This should be used sparingly and
+-- only when it is not possible to continue compilation.
 compilationFailure :: Compile s a
 compilationFailure = Compile $ MaybeT (return Nothing)
 
+-- | Provides a 'MonadWriter' instance for the 'Compile' monad, allowing
+-- diagnostics to be collected and manipulated during the compilation process.
+-- The diagnostics are treated as the output to be written.
 instance MonadWriter [Diagnostic] (Compile s) where
+  -- \| Add diagnostics to the compilation output.
   tell t = Compile $ tell (t, [])
+
+  -- \| Listen for diagnostics emitted during a computation, returning the result
+  -- along with the collected diagnostics.
   listen (Compile fn) = Compile $ do
     (a, (ds, _)) <- listen fn
     return (a, ds)
+
+  -- \| Modify the collected diagnostics using the given function.
   pass (Compile mfun) = Compile $ pass (second first <$> mfun)
 
+-- | Provides a 'MonadState' instance for the 'Compile' monad, allowing
+-- manipulation of the compiler's state during the compilation process.
 instance MonadState s (Compile s) where
+  -- \| Get the current state.
   get = Compile get
+
+  -- \| Set a new state.
   put s = Compile $ put s
+
+  -- \| Apply a state transformation function.
   state fn = Compile $ state fn
 
+-- | Run a compilation routine under a temporary state. After the routine
+-- completes, the original state is restored, and the result along with the
+-- modified temporary state is returned.
 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
 
-emitArtifact :: Artifact -> Compile s ()
-emitArtifact art = Compile $ tell ([], [art])
-
+-- | Save the current state and run a provided compilation routine with a
+-- duplicated state. Once the routine finishes, the original state is restored.
 pushState :: Compile s a -> Compile s a
 pushState cp = do
   s <- get
   snd <$> subCompile s cp
 
+-- | Runs a compilation routine with the given initial state and returns the
+-- diagnostics, artifacts, and result. If there are any 'Error' diagnostics,
+-- the result will be 'Nothing'.
 compile :: Compile s a -> s -> ([Diagnostic], [Artifact], Maybe a)
 compile (Compile fn) initState = do
   let (a, _, (w, as)) = runRWS (runMaybeT fn) () initState
@@ -82,84 +147,103 @@ compile (Compile fn) initState = do
   where
     hasError = any (\(Diagnostic e _ _) -> e == Error)
 
+-- | Runs a compilation routine with a default initial state, using the 'Default'
+-- type class to provide the initial state.
 compile_ :: (Default s) => Compile s a -> ([Diagnostic], [Artifact], Maybe a)
 compile_ c = compile c def
 
+-- | Converts a 'Maybe' value into the 'Compile' monad. If the value is 'Nothing',
+-- the compilation fails.
 hoistMaybe :: Maybe a -> Compile s a
 hoistMaybe = Compile . MaybeT . return
 
-newtype DiagnosticFormat = DiagnosticFormat ([Diagnostic] -> String)
+-- | Represents a format for displaying diagnostics, providing a function
+-- to convert a list of diagnostics into a string.
+newtype DiagnosticFormat = DiagnosticFormat
+  { diagnosticsToString :: [Diagnostic] -> String
+  }
 
-jsonFormat :: DiagnosticFormat
-jsonFormat = DiagnosticFormat $ BL.unpack . encode . toJSON
+-- | Diagnostic format that outputs diagnostics as a JSON string.
+jsonDiagnosticFormat :: DiagnosticFormat
+jsonDiagnosticFormat = DiagnosticFormat $ BL.unpack . encode . toJSON
 
+-- | Diagnostic format that outputs diagnostics in a colored format using ANSI
+-- escape codes. The colors indicate the diagnostic level (Error, Warning, Info).
 coloredFormat :: DiagnosticFormat
-coloredFormat = DiagnosticFormat $ \diags -> unlines $ map (\(Diagnostic level message (SourceSpan pos1 pos2)) ->
-  execWriter $ do
-    case level of
-      Error -> tell "\x1b[01;31mError "
-      Warning -> tell "\x1b[01;33mWarn "
-      Info -> tell "\x1b[01;37mInfo "
-
-    tell "\x1b[0m"
-    tell (sourceName pos1)
-    tell "("
-    tellPos pos1
-    when (pos2 /= pos1) $ do
-      tell "-"
-      tellPos pos2
-    tell "): "
-    tell (unwords $ words message)
-  ) diags
+coloredFormat = DiagnosticFormat $ \diags ->
+  unlines $
+    map
+      ( \(Diagnostic level message (SourceSpan pos1 pos2)) ->
+          execWriter $ do
+            case level of
+              Error -> tell "\x1b[01;31mError "
+              Warning -> tell "\x1b[01;33mWarn "
+              Info -> tell "\x1b[01;37mInfo "
+
+            tell "\x1b[0m"
+            tell (sourceName pos1)
+            tell "("
+            tellPos pos1
+            when (pos2 /= pos1) $ do
+              tell "-"
+              tellPos pos2
+            tell "): "
+            tell (unwords $ words message)
+      )
+      diags
   where
+    -- \| Helper function to format a 'SourcePos' as a line:column string.
     tellPos pos = do
       tell (show $ sourceLine pos)
       tell ":"
       tell (show $ sourceColumn pos)
 
-diagnosticsToString :: DiagnosticFormat -> [Diagnostic] -> String
-diagnosticsToString (DiagnosticFormat f) = f
-
+-- | Prints a list of diagnostics to 'stderr' using the specified
+-- diagnostic format.
 printDiagnostics :: DiagnosticFormat -> [Diagnostic] -> IO ()
 printDiagnostics fmt d = hPutStrLn stderr (diagnosticsToString fmt d)
 
+-- | Prints a single diagnostic to 'stderr' using the 'coloredFormat'
+-- to visually indicate the diagnostic level (Error, Warning, Info).
 printDiagnostic :: Diagnostic -> IO ()
 printDiagnostic d =
   hPutStrLn stderr (diagnosticsToString coloredFormat [d])
 
+-- | Converts a 'Maybe' value into the 'Compile' monad, emitting an
+-- 'Error' diagnostic and causing compilation to fail if the value is 'Nothing'.
+-- If the value is 'Just a', it is returned.
 fromMayberOrFail :: SourceSpan -> String -> Maybe a -> Compile s a
 fromMayberOrFail sourceSpan err Nothing = do
   tell [Diagnostic Error err sourceSpan]
   compilationFailure
 fromMayberOrFail _ _ (Just a) = return a
 
+-- | Creates a pure compilation phase. The provided function is used to
+-- transform the 'FiddleUnit' from one stage to another within the compilation
+-- pipeline, without requiring any additional state or side effects.
 pureCompilationPhase ::
-  (CompilationStage stageFrom) =>
   ( FiddleUnit
       stageFrom
       Identity
-      (StageAnnotation stageFrom) ->
+      (Commented SourceSpan) ->
     Compile
       ()
       ( FiddleUnit
           stageTo
           Identity
-          (StageAnnotation stageTo)
+          (Commented SourceSpan)
       )
   ) ->
   CompilationPhase stageFrom stageTo
 pureCompilationPhase fn =
   CompilationPhase (pure ()) (\_ _ -> return ([], Just ())) (\_ _ -> fn)
 
+-- | Emits a list of compilation artifacts, adding them to the compilation
+-- output. Artifacts are user-defined data that can be generated during
+-- the compilation process.
 emitArtifacts :: [Artifact] -> Compile s ()
 emitArtifacts arts = Compile $ tell ([], arts)
 
--- data IOActionExtraData = IOActionExtraData
---   { parseFile :: FilePath -> IO (TreeType FiddleUnit Parsed),
---     stage3Compile :: TreeType FiddleUnit Parsed ->
---       TreeType FiddleUnit Checked
---   }
-
 -- | 'CompilationPhase' represents a phase in the compilation process.
 -- It consists of an IO action that performs necessary side effects or state
 -- preparations before the next stage, and a function that transforms the
@@ -187,17 +271,21 @@ data CompilationPhase stageFrom stageTo where
         FiddleUnit
           stageFrom
           Identity
-          (StageAnnotation stageFrom) ->
+          (Commented SourceSpan) ->
         Compile
           ()
           ( FiddleUnit
               stageTo
               Identity
-              (StageAnnotation stageTo)
+              (Commented SourceSpan)
           )
     } ->
     CompilationPhase stageFrom stageTo
 
+instance Category CompilationPhase where
+  id = pureCompilationPhase return
+  (.) = flip thenPhase
+
 -- | 'thenPhase' composes two 'CompilationPhase' stages into a single pipeline
 -- phase. It combines their IO actions and applies each stage in sequence.
 thenPhase ::
@@ -224,23 +312,20 @@ thenPhase
           compile2 f2 s2 secondStage
       )
 
--- | Infix operator for 'thenPhase' to chain compilation phases.
-(>>>) :: CompilationPhase stage1 stage2 -> CompilationPhase stage2 stage3 -> CompilationPhase stage1 stage3
-(>>>) = thenPhase
-
 -- | 'execCompilationPipeline' executes a full compilation pipeline starting
 -- from the 'Parsed' phase. It performs the IO action of the first phase and
 -- then invokes the compilation function for the remaining stages. It returns
 -- a tuple containing diagnostics and an optional final 'FiddleUnit'.
 execCompilationPipelineWithCmdline ::
+  (StageAnnotation Parsed ~ Commented SourceSpan) =>
   CompilationPhase Parsed s' ->
   Parser
-    ( FiddleUnit Parsed Identity (StageAnnotation Parsed) ->
+    ( FiddleUnit Parsed Identity (Commented SourceSpan) ->
       IO
         ( [Diagnostic],
           [Artifact],
           Maybe
-            ( FiddleUnit s' Identity (StageAnnotation s')
+            ( FiddleUnit s' Identity (Commented SourceSpan)
             )
         )
     )
@@ -250,26 +335,43 @@ execCompilationPipelineWithCmdline
       ( \opts ast -> do
           (diags, ms) <- ioAct opts ast
           case ms of
-            Just s -> return $ (\(a, b, c) -> (diags ++ a, b, c)) $ compile_ $ rest opts s ast
+            Just s ->
+              return $
+                (\(a, b, c) -> (diags ++ a, b, c)) $
+                  compile_ $
+                    rest opts s ast
             Nothing -> return (diags, [], Nothing)
       )
       flagsParser
 
+-- | Extracts all embedded diagnostics from the abstract syntax tree (AST)
+-- and emits them as compilation warnings or errors. This function inspects
+-- each subelement in the AST for diagnostic messages. If any subelement
+-- contains errors (represented as a 'Left' value), the compilation fails.
+-- Otherwise, the AST is returned with all deferred diagnostics resolved.
 squeezeDiagnostics :: (Alter t) => t (Either [Diagnostic]) a -> Compile () (t Identity a)
 squeezeDiagnostics ast = do
+  -- Iterate over the AST, emitting any diagnostics found.
   _ <-
     alter
       ( \case
-          (Left l) -> tell l >> return (Left l)
-          r -> return r
+          (Left l) -> tell l >> return (Left l) -- Emit diagnostics if present.
+          r -> return r -- Otherwise, leave the subelement unchanged.
       )
       return
       ast
 
+  -- Attempt to transform the AST into a non-deferred form.
+  -- If any subelement contained errors, the transformation will fail.
   case squeeze ast of
-    (Left _) -> compilationFailure
-    (Right a) -> return a
-
+    (Left _) -> compilationFailure -- Fail compilation if there are unresolved diagnostics.
+    (Right a) -> return a -- Return the resolved AST if no errors are found.
+
+-- | Checks the result of a computation and either returns the result or
+-- emits any associated diagnostics and fails the compilation. If the result
+-- is a 'Right' value, the contained value is returned. If it is a 'Left'
+-- value containing diagnostics, they are emitted and the compilation is
+-- terminated.
 resolveOrFail :: Either [Diagnostic] a -> Compile s a
-resolveOrFail (Right a) = return a
-resolveOrFail (Left l) = tell l >> compilationFailure
+resolveOrFail (Right a) = return a -- Return the value if no diagnostics.
+resolveOrFail (Left l) = tell l >> compilationFailure -- Emit diagnostics and fail if errors are present.