From edee70cf31d3ec24896c8884dda512188059b203 Mon Sep 17 00:00:00 2001
From: Irmen de Jong <irmen@razorvine.net>
Date: Fri, 15 May 2020 00:16:53 +0200
Subject: [PATCH] use new api for ast mods in unused code remover

---
 compiler/res/version.txt                      |  2 +-
 compiler/src/prog8/compiler/Main.kt           |  1 +
 .../src/prog8/optimizer/UnusedCodeRemover.kt  | 37 +++++--------------
 3 files changed, 12 insertions(+), 28 deletions(-)

diff --git a/compiler/res/version.txt b/compiler/res/version.txt
index cd5ac039d..ceb0059e7 100644
--- a/compiler/res/version.txt
+++ b/compiler/res/version.txt
@@ -1 +1 @@
-2.0
+2.1-SNAPSHOT
diff --git a/compiler/src/prog8/compiler/Main.kt b/compiler/src/prog8/compiler/Main.kt
index b5beca928..1e371b85d 100644
--- a/compiler/src/prog8/compiler/Main.kt
+++ b/compiler/src/prog8/compiler/Main.kt
@@ -174,6 +174,7 @@ private fun optimizeAst(programAst: Program, errors: ErrorReporter) {
 
     val remover = UnusedCodeRemover()
     remover.visit(programAst)
+    remover.applyModifications()
 }
 
 private fun postprocessAst(programAst: Program, errors: ErrorReporter, compilerOptions: CompilationOptions) {
diff --git a/compiler/src/prog8/optimizer/UnusedCodeRemover.kt b/compiler/src/prog8/optimizer/UnusedCodeRemover.kt
index 3c62a764d..9e9397421 100644
--- a/compiler/src/prog8/optimizer/UnusedCodeRemover.kt
+++ b/compiler/src/prog8/optimizer/UnusedCodeRemover.kt
@@ -1,55 +1,38 @@
 package prog8.optimizer
 
-import prog8.ast.Module
+import prog8.ast.Node
 import prog8.ast.Program
-import prog8.ast.processing.IAstModifyingVisitor
+import prog8.ast.processing.AstWalker
+import prog8.ast.processing.IAstModification
 import prog8.ast.statements.Block
-import prog8.ast.statements.Subroutine
 
 
-internal class UnusedCodeRemover: IAstModifyingVisitor {
+internal class UnusedCodeRemover: AstWalker() {
 
-
-    override fun visit(program: Program) {
+    override fun before(program: Program, parent: Node): Iterable<IAstModification> {
         val callgraph = CallGraph(program)
+        val removals = mutableListOf<IAstModification>()
 
         // remove all subroutines that aren't called, or are empty
-        val removeSubroutines = mutableSetOf<Subroutine>()
         val entrypoint = program.entrypoint()
         program.modules.forEach {
             callgraph.forAllSubroutines(it) { sub ->
                 if (sub !== entrypoint && !sub.keepAlways && (sub.calledBy.isEmpty() || (sub.containsNoCodeNorVars() && !sub.isAsmSubroutine)))
-                    removeSubroutines.add(sub)
+                    removals.add(IAstModification.Remove(sub, sub.definingScope() as Node))
             }
         }
 
-        if (removeSubroutines.isNotEmpty()) {
-            removeSubroutines.forEach {
-                it.definingScope().remove(it)
-            }
-        }
-
-        val removeBlocks = mutableSetOf<Block>()
         program.modules.flatMap { it.statements }.filterIsInstance<Block>().forEach { block ->
             if (block.containsNoCodeNorVars() && "force_output" !in block.options())
-                removeBlocks.add(block)
-        }
-
-        if (removeBlocks.isNotEmpty()) {
-            removeBlocks.forEach { it.definingScope().remove(it) }
+                removals.add(IAstModification.Remove(block, block.definingScope() as Node))
         }
 
         // remove modules that are not imported, or are empty (unless it's a library modules)
-        val removeModules = mutableSetOf<Module>()
         program.modules.forEach {
             if (!it.isLibraryModule && (it.importedBy.isEmpty() || it.containsNoCodeNorVars()))
-                removeModules.add(it)
+                removals.add(IAstModification.Remove(it, it.parent))        // TODO does removing modules work like this?
         }
 
-        if (removeModules.isNotEmpty()) {
-            program.modules.removeAll(removeModules)
-        }
-
-        super.visit(program)
+        return removals
     }
 }