From fcb21f6e632f526c79a5e3dab7b1f03ccbfd8657 Mon Sep 17 00:00:00 2001
From: aNNiMON <annimon119@gmail.com>
Date: Sat, 20 Apr 2024 00:00:17 +0300
Subject: [PATCH] Show AST before and after optimization

---
 .../src/main/java/com/annimon/ownlang/Main.java        | 10 ++++++----
 .../ownlang/parser/optimization/OptimizationStage.java |  2 ++
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/ownlang-desktop/src/main/java/com/annimon/ownlang/Main.java b/ownlang-desktop/src/main/java/com/annimon/ownlang/Main.java
index df1756e..e66d4b4 100644
--- a/ownlang-desktop/src/main/java/com/annimon/ownlang/Main.java
+++ b/ownlang-desktop/src/main/java/com/annimon/ownlang/Main.java
@@ -4,7 +4,7 @@ import com.annimon.ownlang.exceptions.OwnLangParserException;
 import com.annimon.ownlang.exceptions.StoppedException;
 import com.annimon.ownlang.parser.BeautifierStage;
 import com.annimon.ownlang.parser.Token;
-import com.annimon.ownlang.parser.ast.Statement;
+import com.annimon.ownlang.parser.ast.Node;
 import com.annimon.ownlang.parser.error.ParseErrorsFormatterStage;
 import com.annimon.ownlang.parser.linters.LinterStage;
 import com.annimon.ownlang.parser.optimization.OptimizationStage;
@@ -185,9 +185,11 @@ public final class Main {
                 }
             }
             if (options.showAst) {
-                Statement program = stagesData.get(ParserStage.TAG_PROGRAM);
-                System.out.println(program);
-                System.out.println(stagesData.getOrDefault(OptimizationStage.TAG_OPTIMIZATION_SUMMARY, ""));
+                System.out.println(stagesData.<Node>get(ParserStage.TAG_PROGRAM));
+                if (options.optimizationLevel > 0) {
+                    System.out.println(stagesData.getOrDefault(OptimizationStage.TAG_OPTIMIZATION_SUMMARY, ""));
+                    System.out.println(stagesData.<Node>get(OptimizationStage.TAG_OPTIMIZED_PROGRAM));
+                }
             }
             if (options.showMeasurements) {
                 System.out.println("=".repeat(25));
diff --git a/ownlang-parser/src/main/java/com/annimon/ownlang/parser/optimization/OptimizationStage.java b/ownlang-parser/src/main/java/com/annimon/ownlang/parser/optimization/OptimizationStage.java
index 0f7d1cd..c71ad85 100644
--- a/ownlang-parser/src/main/java/com/annimon/ownlang/parser/optimization/OptimizationStage.java
+++ b/ownlang-parser/src/main/java/com/annimon/ownlang/parser/optimization/OptimizationStage.java
@@ -7,6 +7,7 @@ import com.annimon.ownlang.stages.StagesData;
 public class OptimizationStage implements Stage<Node, Node> {
 
     public static final String TAG_OPTIMIZATION_SUMMARY = "optimizationSummary";
+    public static final String TAG_OPTIMIZED_PROGRAM = "optimizedProgram";
 
     private final int level;
     private final boolean summary;
@@ -48,6 +49,7 @@ public class OptimizationStage implements Stage<Node, Node> {
                 %s
                 """.formatted(iteration, optimization.summaryInfo())
             );
+            stagesData.put(TAG_OPTIMIZED_PROGRAM, result);
         }
 
         return result;