From b182f7e6934c340e06df98d2d5fc18a0e9514002 Mon Sep 17 00:00:00 2001
From: Irmen de Jong <irmen@razorvine.net>
Date: Fri, 13 Mar 2020 02:31:53 +0100
Subject: [PATCH] optimizer removes unreachable code following call to exit()

---
 compiler/src/prog8/ast/processing/AstChecker.kt  |  2 +-
 .../src/prog8/optimizer/StatementOptimizer.kt    | 16 +++++++++++++++-
 docs/source/todo.rst                             |  1 -
 examples/test.p8                                 | 14 +++++++++++---
 4 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/compiler/src/prog8/ast/processing/AstChecker.kt b/compiler/src/prog8/ast/processing/AstChecker.kt
index 8629df5ce..6f5ba59d3 100644
--- a/compiler/src/prog8/ast/processing/AstChecker.kt
+++ b/compiler/src/prog8/ast/processing/AstChecker.kt
@@ -1049,7 +1049,7 @@ internal class AstChecker(private val program: Program,
     private fun visitStatements(statements: List<Statement>) {
         for((index, stmt) in statements.withIndex()) {
             if(stmt is FunctionCallStatement && stmt.target.nameInSource.last()=="exit") {
-                if(index < statements.size-1) {
+                if(index < statements.lastIndex) {
                     printWarning("unreachable code", statements[index+1].position, "exit call above never returns")
                 }
             }
diff --git a/compiler/src/prog8/optimizer/StatementOptimizer.kt b/compiler/src/prog8/optimizer/StatementOptimizer.kt
index d48d91461..f1ee57e9f 100644
--- a/compiler/src/prog8/optimizer/StatementOptimizer.kt
+++ b/compiler/src/prog8/optimizer/StatementOptimizer.kt
@@ -116,6 +116,7 @@ internal class StatementOptimizer(private val program: Program) : IAstModifyingV
             return NopStatement.insteadOf(subroutine)
         }
 
+        visitStatements(subroutine.statements)
         return subroutine
     }
 
@@ -564,6 +565,7 @@ internal class StatementOptimizer(private val program: Program) : IAstModifyingV
         if(linesToRemove.isNotEmpty()) {
             linesToRemove.reversed().forEach{scope.statements.removeAt(it)}
         }
+        visitStatements(scope.statements)
         return super.visit(scope)
     }
 
@@ -571,11 +573,23 @@ internal class StatementOptimizer(private val program: Program) : IAstModifyingV
         // remove duplicate labels
         val stmts = label.definingScope().statements
         val startIdx = stmts.indexOf(label)
-        if(startIdx<(stmts.size-1) && stmts[startIdx+1] == label)
+        if(startIdx< stmts.lastIndex && stmts[startIdx+1] == label)
             return NopStatement.insteadOf(label)
 
         return super.visit(label)
     }
+
+
+    private fun visitStatements(statements: MutableList<Statement>) {
+        // remove all statements following the call to exit()
+        val exitCallIndex = statements.indexOfFirst { it is FunctionCallStatement && it.target.nameInSource.last()=="exit" }
+        if(exitCallIndex>=0) {
+            while(exitCallIndex < statements.lastIndex) {
+                statements.removeAt(statements.lastIndex)
+            }
+        }
+    }
+
 }
 
 
diff --git a/docs/source/todo.rst b/docs/source/todo.rst
index 8029585d8..dcf09e870 100644
--- a/docs/source/todo.rst
+++ b/docs/source/todo.rst
@@ -3,7 +3,6 @@ TODO
 ====
 
 - option to load library files from a directory instead of the embedded ones
-- code optimizer: remove all statements following an exit() function call
 - vector inc/dec/add/sub/mul/div...?
         arrayvar++ / arrayvar-- / arrayvar += 2 / arrayvar -= 2 / arrayvar *= 3 / arrayvar /= 3
 
diff --git a/examples/test.p8 b/examples/test.p8
index 55daaa391..d18dbf90a 100644
--- a/examples/test.p8
+++ b/examples/test.p8
@@ -21,9 +21,17 @@ main {
         c64scr.print("sp2:")
         print_stackpointer()
         exit(65)
-        sub3()          ; TODO  remove the code in optimizer
-        sub3()          ; TODO  remove the code in optimizer
-        sub3()          ; TODO  remove the code in optimizer
+        sub3()
+        sub3()
+        sub3()
+        sub3()
+        sub3()
+        sub3()
+        sub3()
+        sub3()
+        sub3()
+        sub3()
+        sub3()
     }
 
     sub sub3() {