1 files changed, 13 insertions, 8 deletions
diff --git a/src/Main.hs b/src/Main.hs
index 9330df5..f92d6c6 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -5,11 +5,12 @@ import Control.Monad.Writer
 import Data.Aeson (encode)
 import qualified Data.ByteString.Lazy.Char8 as BL
 import qualified Data.Text.IO
-import GHC.IO.Exception (ExitCode (ExitFailure))
+import GHC.IO.Exception (ExitCode (ExitFailure, ExitSuccess))
 import Language.Fiddle.Ast
 import Language.Fiddle.Compiler (coloredFormat, compile_, printDiagnostic)
 import Language.Fiddle.Compiler.Stage0
 import Language.Fiddle.Compiler.Stage1
+import Language.Fiddle.Compiler.Stage2
 import Language.Fiddle.GenericTree (ToGenericSyntaxTree (toGenericSyntaxTree))
 import qualified Language.Fiddle.Parser
 import qualified Language.Fiddle.Tokenizer
@@ -23,14 +24,18 @@ main = do
   case argv of
     [filePath] -> do
       text <- Data.Text.IO.readFile filePath
-      let (diags, ma) = compile_ $ toStage2 =<< toStage1 =<< toStage0 filePath text
+      let (diags, ma) = compile_ $ toStage3 =<< toStage2 =<< toStage1 =<< toStage0 filePath text
+      ec <- 
+        case ma of
+          Just ast -> do
+            putStrLn $ BL.unpack $ encode $ toGenericSyntaxTree ast
+            return ExitSuccess
+          Nothing -> do
+            putStrLn "\x1b[1;31mCompilation Failed\x1b[0m"
+            return (ExitFailure 1)
+
       forM_ diags printDiagnostic
-      case ma of
-        Just ast -> do
-          putStrLn $ BL.unpack $ encode $ toGenericSyntaxTree ast
-        Nothing -> do
-          putStrLn "\x1b[1;31mCompilation Failed\x1b[0m"
-          exitWith (ExitFailure 1)
+      exitWith ec
     _ -> do
       putStrLn "Wrong Args"
       exitWith (ExitFailure 2)