Show AST before and after optimization

Author: aNNiMON

ownlang-desktop/src/main/java/com/annimon/ownlang/Main.java

@@ -4,7 +4,7 @@
 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 @@ private static void run(RunOptions options) {
                 }
             }
             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));

ownlang-parser/src/main/java/com/annimon/ownlang/parser/optimization/OptimizationStage.java

@@ -7,6 +7,7 @@
 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 Node perform(StagesData stagesData, Node input) {
                 %s
                 """.formatted(iteration, optimization.summaryInfo())
             );
+            stagesData.put(TAG_OPTIMIZED_PROGRAM, result);
         }
 
         return result;