Optimizing try/catch generator

teavm-core/src/main/java/org/teavm/javascript/Decompiler.java

@@ -47,6 +47,9 @@ public class Decompiler {
     private MethodNodeCache regularMethodCache;
     private Set<MethodReference> asyncMethods;
     private Set<MethodReference> splitMethods = new HashSet<>();
+    private List<TryCatchBookmark> tryCatchBookmarks = new ArrayList<>();
+    private Deque<Block> stack;
+    private Program program;
 
     public Decompiler(ClassHolderSource classSource, ClassLoader classLoader, Set<MethodReference> asyncMethods,
             Set<MethodReference> asyncFamilyMethods) {
@@ -70,10 +73,12 @@ public class Decompiler {
     }
 
     static class Block {
+        public Block parent;
         public final IdentifiedStatement statement;
         public final List<Statement> body;
         public final int end;
         public final int start;
+        public final List<TryCatchBookmark> tryCatches = new ArrayList<>();
 
         public Block(IdentifiedStatement statement, List<Statement> body, int start, int end) {
             this.statement = statement;
@@ -83,6 +88,14 @@ public class Decompiler {
         }
     }
 
+    static class TryCatchBookmark {
+        Block block;
+        int offset;
+        String exceptionType;
+        Integer exceptionVariable;
+        int exceptionHandler;
+    }
+
     public List<ClassNode> decompile(Collection<String> classNames) {
         List<String> sequence = new ArrayList<>();
         Set<String> visited = new HashSet<>();
@@ -286,7 +299,8 @@ public class Decompiler {
         loopGraph = new LoopGraph(this.graph);
         unflatCode();
         blockMap = new Block[program.basicBlockCount() * 2 + 1];
-        Deque<Block> stack = new ArrayDeque<>();
+        stack = new ArrayDeque<>();
+        this.program = program;
         BlockStatement rootStmt = new BlockStatement();
         rootStmt.setId("root");
         stack.push(new Block(rootStmt, rootStmt.getBody(), -1, -1));
@@ -299,6 +313,20 @@ public class Decompiler {
         currentNode = parentNode.getFirstChild();
         generator.async = async;
         for (int i = 0; i < this.graph.size(); ++i) {
+            int node = i < indexer.size() ? indexer.nodeAt(i) : -1;
+            int next = i + 1;
+            int head = loops[i];
+            if (head != -1 && loopSuccessors[head] == next) {
+                next = head;
+            }
+
+            if (node >= 0) {
+                generator.currentBlock = program.basicBlockAt(node);
+                int tmp = indexer.nodeAt(next);
+                generator.nextBlock = tmp >= 0 && next < indexer.size() ? program.basicBlockAt(tmp) : null;
+            }
+
+            List<TryCatchBookmark> inheritedBookmarks = new ArrayList<>();
             Block block = stack.peek();
             while (block.end == i) {
                 Block oldBlock = block;
@@ -310,26 +338,42 @@ public class Decompiler {
                         blockMap[mappedStart] = block;
                     }
                 }
+
+                for (int j = oldBlock.tryCatches.size() - 1; j >= 0; --j) {
+                    TryCatchBookmark bookmark = oldBlock.tryCatches.get(j);
+                    TryCatchStatement tryCatchStmt = new TryCatchStatement();
+                    tryCatchStmt.setExceptionType(bookmark.exceptionType);
+                    tryCatchStmt.setExceptionVariable(tryCatchStmt.getExceptionVariable());
+                    tryCatchStmt.getHandler().add(generator.generateJumpStatement(
+                            program.basicBlockAt(bookmark.exceptionHandler)));
+                    List<Statement> blockPart = oldBlock.body.subList(bookmark.offset, oldBlock.body.size());
+                    tryCatchStmt.getProtectedBody().addAll(blockPart);
+                    blockPart.clear();
+                    blockPart.add(tryCatchStmt);
+                    inheritedBookmarks.add(bookmark);
                 }
+                oldBlock.tryCatches.clear();
+            }
+
+            for (int j = inheritedBookmarks.size() - 1; j >= 0; --j) {
+                TryCatchBookmark bookmark = inheritedBookmarks.get(j);
+                bookmark.block = block;
+                bookmark.offset = block.body.size();
+                block.tryCatches.add(bookmark);
+            }
+
             while (parentNode.getEnd() == i) {
                 currentNode = parentNode.getNext();
                 parentNode = parentNode.getParent();
             }
             for (Block newBlock : createBlocks(i)) {
                 block.body.add(newBlock.statement);
+                newBlock.parent = block;
                 stack.push(newBlock);
                 block = newBlock;
             }
-            int node = i < indexer.size() ? indexer.nodeAt(i) : -1;
+
-            int next = i + 1;
-            int head = loops[i];
-            if (head != -1 && loopSuccessors[head] == next) {
-                next = head;
-            }
             if (node >= 0) {
-                generator.currentBlock = program.basicBlockAt(node);
-                int tmp = indexer.nodeAt(next);
-                generator.nextBlock = tmp >= 0 && next < indexer.size() ? program.basicBlockAt(tmp) : null;
                 generator.statements.clear();
                 InstructionLocation lastLocation = null;
                 NodeLocation nodeLocation = null;
@@ -351,18 +395,7 @@ public class Decompiler {
                     generator.statements.add(stmt);
                 }
 
-                for (TryCatchBlock tryCatch : generator.currentBlock.getTryCatchBlocks()) {
+                updateTryCatchBookmarks(generator, generator.currentBlock.getTryCatchBlocks());
-                    TryCatchStatement tryCatchStmt = new TryCatchStatement();
-                    tryCatchStmt.setExceptionType(tryCatch.getExceptionType());
-                    tryCatchStmt.setExceptionVariable(tryCatch.getExceptionVariable().getIndex());
-                    tryCatchStmt.getProtectedBody().addAll(generator.statements);
-                    generator.statements.clear();
-                    generator.statements.add(tryCatchStmt);
-                    Statement handlerStmt = generator.generateJumpStatement(tryCatch.getHandler());
-                    if (handlerStmt != null) {
-                        tryCatchStmt.getHandler().add(handlerStmt);
-                    }
-                }
                 block.body.addAll(generator.statements);
             }
         }
@@ -372,6 +405,72 @@ public class Decompiler {
         return result;
     }
 
+    private void updateTryCatchBookmarks(StatementGenerator generator, List<TryCatchBlock> tryCatchBlocks) {
+        tryCatchBlocks = new ArrayList<>(tryCatchBlocks);
+        Collections.reverse(tryCatchBlocks);
+
+        // Find which try catch blocks have remained since the previous basic block
+        int sz = Math.min(tryCatchBlocks.size(), tryCatchBookmarks.size());
+        int start;
+        for (start = 0; start < sz; ++start) {
+            TryCatchBlock tryCatch = tryCatchBlocks.get(start);
+            TryCatchBookmark bookmark = tryCatchBookmarks.get(start);
+            if (tryCatch.getHandler().getIndex() != bookmark.exceptionHandler) {
+                break;
+            }
+            if (!Objects.equals(tryCatch.getExceptionType(), bookmark.exceptionType)) {
+                break;
+            }
+            if (tryCatch.getExceptionVariable() != null && bookmark.exceptionVariable != null &&
+                    tryCatch.getExceptionVariable().getIndex() != bookmark.exceptionVariable.intValue()) {
+                break;
+            }
+        }
+
+        // Close old bookmarks
+        for (int i = tryCatchBookmarks.size() - 1; i >= start; --i) {
+            TryCatchBookmark bookmark = tryCatchBookmarks.get(i);
+            Block block = stack.peek();
+            while (block != bookmark.block) {
+                TryCatchStatement tryCatchStmt = new TryCatchStatement();
+                tryCatchStmt.setExceptionType(bookmark.exceptionType);
+                tryCatchStmt.setExceptionVariable(tryCatchStmt.getExceptionVariable());
+                tryCatchStmt.getHandler().add(generator.generateJumpStatement(
+                        program.basicBlockAt(bookmark.exceptionHandler)));
+                tryCatchStmt.getProtectedBody().addAll(block.body);
+                block.body.clear();
+                block.body.add(tryCatchStmt);
+                block = block.parent;
+            }
+
+            TryCatchStatement tryCatchStmt = new TryCatchStatement();
+            tryCatchStmt.setExceptionType(bookmark.exceptionType);
+            tryCatchStmt.setExceptionVariable(tryCatchStmt.getExceptionVariable());
+            tryCatchStmt.getHandler().add(generator.generateJumpStatement(
+                    program.basicBlockAt(bookmark.exceptionHandler)));
+            List<Statement> blockPart = block.body.subList(bookmark.offset, block.body.size());
+            tryCatchStmt.getProtectedBody().addAll(blockPart);
+            blockPart.clear();
+            blockPart.add(tryCatchStmt);
+
+            bookmark.block.tryCatches.remove(bookmark);
+        }
+
+        // Add new bookmarks
+        for (int i = start; i < tryCatchBlocks.size(); ++i) {
+            TryCatchBlock tryCatch = tryCatchBlocks.get(i);
+            TryCatchBookmark bookmark = new TryCatchBookmark();
+            bookmark.block = stack.peek();
+            bookmark.offset = bookmark.block.body.size();
+            bookmark.exceptionHandler = tryCatch.getHandler().getIndex();
+            bookmark.exceptionType = tryCatch.getExceptionType();
+            bookmark.exceptionVariable = tryCatch.getExceptionVariable() != null ?
+                    tryCatch.getExceptionVariable().getIndex() : null;
+            bookmark.block.tryCatches.add(bookmark);
+            tryCatchBookmarks.add(bookmark);
+        }
+    }
+
     private Set<NodeModifier> mapModifiers(Set<ElementModifier> modifiers) {
         Set<NodeModifier> result = EnumSet.noneOf(NodeModifier.class);
         if (modifiers.contains(ElementModifier.STATIC)) {