fix UnitBox memory leak in UCE

src/main/java/soot/jimple/toolkits/scalar/UnreachableCodeEliminator.java

@@ -41,6 +41,7 @@
 import soot.Singletons;
 import soot.Trap;
 import soot.Unit;
+import soot.UnitBox;
 import soot.options.Options;
 import soot.toolkits.exceptions.PedanticThrowAnalysis;
 import soot.toolkits.exceptions.ThrowAnalysis;
@@ -64,6 +65,7 @@ public UnreachableCodeEliminator(ThrowAnalysis ta) {
     this.throwAnalysis = ta;
   }
 
+  @Override
   protected void internalTransform(Body body, String phaseName, Map<String, String> options) {
     final boolean verbose = Options.v().verbose();
     if (verbose) {
@@ -99,9 +101,16 @@ protected void internalTransform(Body body, String phaseName, Map<String, String
     // occur in practice, and certainly no in code generated from Java source.
     final Chain<Trap> traps = body.getTraps();
     for (Iterator<Trap> it = traps.iterator(); it.hasNext();) {
-      Trap trap = it.next();
-      if ((trap.getBeginUnit() == trap.getEndUnit()) || !reachable.contains(trap.getHandlerUnit())) {
+      final Trap trap = it.next();
+      UnitBox beginBox = trap.getBeginUnitBox();
+      UnitBox endBox = trap.getEndUnitBox();
+      UnitBox handlerBox = trap.getHandlerUnitBox();
+      if ((beginBox.getUnit() == endBox.getUnit()) || !reachable.contains(handlerBox.getUnit())) {
         it.remove();
+        // Cleanup UnitBox references that are no longer used
+        beginBox.getUnit().removeBoxPointingToThis(beginBox);
+        endBox.getUnit().removeBoxPointingToThis(endBox);
+        handlerBox.getUnit().removeBoxPointingToThis(handlerBox);
       }
     }
 

src/systemTest/java/soot/jimple/toolkit/scalar/UnreachableCodeEliminatorTest.java

@@ -0,0 +1,106 @@
+package soot.jimple.toolkit.scalar;
+
+/*-
+ * #%L
+ * Soot - a J*va Optimization Framework
+ * %%
+ * Copyright (C) 2020 Timothy Hoffman
+ * %%
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as
+ * published by the Free Software Foundation, either version 2.1 of the
+ * License, or (at your option) any later version.
+ * 
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Lesser Public License for more details.
+ * 
+ * You should have received a copy of the GNU General Lesser Public
+ * License along with this program.  If not, see
+ * <http://www.gnu.org/licenses/lgpl-2.1.html>.
+ * #L%
+ */
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+
+import org.junit.Assert;
+import org.junit.Test;
+import org.powermock.core.classloader.annotations.PowerMockIgnore;
+
+import soot.Body;
+import soot.SootMethod;
+import soot.Trap;
+import soot.Unit;
+import soot.UnitBox;
+import soot.jimple.toolkits.scalar.UnreachableCodeEliminator;
+import soot.options.Options;
+import soot.tagkit.CodeAttribute;
+import soot.tagkit.Tag;
+import soot.testing.framework.AbstractTestingFramework;
+
+@PowerMockIgnore({ "com.sun.org.apache.xerces.*", "javax.xml.*", "org.xml.*", "org.w3c.*" })
+public class UnreachableCodeEliminatorTest extends AbstractTestingFramework {
+
+  private static final String TEST_TARGET_CLASS = "soot.jimple.toolkits.scalar.UnreachableCodeEliminatorTestInput";
+
+  @Override
+  protected void setupSoot() {
+    // Skip JimpleBody UnreachableCodeEliminator so we can
+    // observe the effect of running it manually below.
+    Options.v().setPhaseOption("jb.uce", "enabled:false");
+
+    // Disable validation because it will give an error with UCE disabled.
+    Options.v().set_validate(false);
+  }
+
+  @Test
+  public void unitBoxConsistency() {
+    SootMethod target =
+        prepareTarget(methodSigFromComponents(TEST_TARGET_CLASS, "void", "unreachableTrap"), TEST_TARGET_CLASS);
+
+    Body body = target.retrieveActiveBody();
+
+    // There is 1 Trap before the optimization.
+    Assert.assertEquals(1, body.getTraps().size());
+
+    // Assert that the set of boxes obtained from all units equals the set
+    // obtained by checking Unit#getBoxesPointingToThis() on all units.
+    Assert.assertEquals(new HashSet<>(getAllUnitBoxes(body)), new HashSet<>(getAllBoxesPointingToUnits(body)));
+
+    UnreachableCodeEliminator.v().transform(body);
+
+    // There are no Traps after the optimization.
+    Assert.assertEquals(0, body.getTraps().size());
+
+    // Assert that the set of boxes obtained from all units equals the set
+    // obtained by checking Unit#getBoxesPointingToThis() on all units.
+    Assert.assertEquals(new HashSet<>(getAllUnitBoxes(body)), new HashSet<>(getAllBoxesPointingToUnits(body)));
+  }
+
+  private static List<UnitBox> getAllUnitBoxes(Body body) {
+    ArrayList<UnitBox> unitBoxList = new ArrayList<>();
+    for (Unit u : body.getUnits()) {
+      unitBoxList.addAll(u.getUnitBoxes());
+    }
+    for (Trap t : body.getTraps()) {
+      unitBoxList.addAll(t.getUnitBoxes());
+    }
+    for (Tag t : body.getTags()) {
+      if (t instanceof CodeAttribute) {
+        unitBoxList.addAll(((CodeAttribute) t).getUnitBoxes());
+      }
+    }
+    return unitBoxList;
+  }
+
+  private static List<UnitBox> getAllBoxesPointingToUnits(Body body) {
+    ArrayList<UnitBox> unitBoxList = new ArrayList<>();
+    for (Unit u : body.getUnits()) {
+      unitBoxList.addAll(u.getBoxesPointingToThis());
+    }
+    return unitBoxList;
+  }
+}

src/systemTest/targets/soot/jimple/toolkit/scalar/UnreachableCodeEliminatorTestInput.java

@@ -0,0 +1,41 @@
+package soot.jimple.toolkits.scalar;
+
+/*-
+ * #%L
+ * Soot - a J*va Optimization Framework
+ * %%
+ * Copyright (C) 2020 Timothy Hoffman
+ * %%
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as
+ * published by the Free Software Foundation, either version 2.1 of the
+ * License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Lesser Public License for more details.
+ *
+ * You should have received a copy of the GNU General Lesser Public
+ * License along with this program.  If not, see
+ * <http://www.gnu.org/licenses/lgpl-2.1.html>.
+ * #L%
+ */
+import java.io.PrintStream;
+
+/**
+ * @author Timothy Hoffman
+ */
+public class UnreachableCodeEliminatorTestInput {
+
+  public void unreachableTrap() {
+    final PrintStream out = System.out;
+    out.println("Before");
+    try {
+      int x = 2134;// dead code here allows the entire trap to be removed
+    } catch (RuntimeException ex) {
+      out.println("Handler");
+    }
+    out.println("After");
+  }
+}