From d6bf0c390e0587e39b5c35d3ec9196f6905ef878 Mon Sep 17 00:00:00 2001
From: Sei Lisa <sei-lisa@email.fake>
Date: Mon, 1 Jun 2020 20:49:48 +0200
Subject: [PATCH] Fix two issues in the dead code removal module

1. When the last statement of a function is a RETURN statement which is syntactically required, it could still be deleted.
2. The child of a RETURN statement could be removed if the statement was not executed.

This commit fixes both issues.

Bug report and test case provided by @Tonaie. Fixes #14.
---
 lslopt/lsldeadcode.py                         | 31 +++++++++++++------
 .../regression.suite/state-in-udf-6.skp       |  1 +
 2 files changed, 22 insertions(+), 10 deletions(-)
 create mode 100644 unit_tests/regression.suite/state-in-udf-6.skp

diff --git a/lslopt/lsldeadcode.py b/lslopt/lsldeadcode.py
index 35deb55..49ee4a6 100644
--- a/lslopt/lsldeadcode.py
+++ b/lslopt/lsldeadcode.py
@@ -386,7 +386,7 @@ class deadcode(object):
 
         return False
 
-    def CleanNode(self, curnode):
+    def CleanNode(self, curnode, isFnDef = False):
         """Recursively checks if the children are used, deleting those that are
         not.
         """
@@ -402,14 +402,25 @@ class deadcode(object):
             node = curnode.ch[index]
 
             if not hasattr(node, 'X'):
-                if curnode.ch[index].nt == 'JUMP':
-                    # Decrease label reference count
-                    scope = curnode.ch[index].scope
-                    name = curnode.ch[index].name
-                    assert self.symtab[scope][name]['ref'] > 0
-                    self.symtab[scope][name]['ref'] -= 1
-                del curnode.ch[index]
-                continue
+                deleted = curnode.ch[index]
+                # Don't delete the child of a RETURN statement!
+                if curnode.nt == 'RETURN':
+                    pass
+                # If it's a RETURN as the last statement of a block in a
+                # function definition, don't delete it (a hacky workaround for
+                # issue #14).
+                elif (deleted.nt == 'RETURN' and index == len(curnode.ch) - 1
+                      and isFnDef):
+                    pass
+                else:
+                    del curnode.ch[index]
+                    if deleted.nt == 'JUMP':
+                        # Decrease label reference count
+                        scope = deleted.scope
+                        name = deleted.name
+                        assert self.symtab[scope][name]['ref'] > 0
+                        self.symtab[scope][name]['ref'] -= 1
+                    continue
 
             nt = node.nt
 
@@ -484,7 +495,7 @@ class deadcode(object):
                     child[1] = nr(nt='CONST', X=True, SEF=True, t='integer',
                         value=0)
 
-            self.CleanNode(node)
+            self.CleanNode(node, isFnDef = (curnode.nt == 'FNDEF'))
             index += 1
 
 
diff --git a/unit_tests/regression.suite/state-in-udf-6.skp b/unit_tests/regression.suite/state-in-udf-6.skp
new file mode 100644
index 0000000..c79470d
--- /dev/null
+++ b/unit_tests/regression.suite/state-in-udf-6.skp
@@ -0,0 +1 @@
+Skipped because our fix to #14 brings back the RETURN statement.