WASM: fix bugs in exception handling

core/src/main/java/org/teavm/backend/wasm/Example.java

@@ -37,6 +37,7 @@ public final class Example {
         testArrayCopy();
         testArrayIsObject();
         testIsAssignableFrom();
+        testExceptions();
         testBigInteger();
         testGC();
     }
@@ -186,6 +187,47 @@ public final class Example {
         System.out.println("GC complete");
     }
 
+    private static void testExceptions() {
+        try {
+            throwsException();
+        } catch (IllegalStateException e) {
+            System.out.println("Caught 1: " + e.getMessage());
+        }
+
+        int x = 0;
+        try {
+            x = 1;
+            doesNotThrow();
+            x = 2;
+            throwsException();
+            x = 3;
+        } catch (IllegalStateException e) {
+            System.out.println("Caught 2: " + e.getMessage());
+            System.out.println("x = " + x);
+        }
+
+        try {
+            throwsWithFinally();
+        } catch (IllegalStateException e) {
+            System.out.println("Caught 3: " + e.getMessage());
+        }
+    }
+
+    private static void doesNotThrow() {
+    }
+
+    private static void throwsWithFinally() {
+        try {
+            throwsException();
+        } finally {
+            System.out.println("Finally reached");
+        }
+    }
+
+    private static void throwsException() {
+        throw new IllegalStateException("exception message");
+    }
+
     private static Base instance(int index) {
         switch (index) {
             case 0:

core/src/main/java/org/teavm/backend/wasm/WasmRuntime.java

@@ -272,11 +272,16 @@ public final class WasmRuntime {
         return stackFrame.add(-size * 4);
     }
 
+    private static Address getExceptionHandlerPtr(Address stackFrame) {
+        int size = stackFrame.getInt();
+        return stackFrame.add(-size * 4 - 4);
+    }
+
     public static int getCallSiteId(Address stackFrame) {
-        return getStackRootPointer(stackFrame).getInt();
+        return getExceptionHandlerPtr(stackFrame).getInt();
     }
 
     public static void setExceptionHandlerId(Address stackFrame, int id) {
-        getStackRootPointer(stackFrame).putInt(id);
+        getExceptionHandlerPtr(stackFrame).putInt(id);
     }
 }

core/src/main/java/org/teavm/backend/wasm/WasmTarget.java

@@ -231,6 +231,9 @@ public class WasmTarget implements TeaVMTarget {
         dependencyChecker.linkMethod(new MethodReference(ExceptionHandling.class, "throwException",
                 Throwable.class, void.class), null).use();
 
+        dependencyChecker.linkMethod(new MethodReference(ExceptionHandling.class, "catchException",
+                Throwable.class), null).use();
+
         dependencyChecker.linkField(new FieldReference("java.lang.Object", "monitor"), null);
 
         ClassDependency runtimeClassDep = dependencyChecker.linkClass(RuntimeClass.class.getName(), null);
@@ -501,7 +504,7 @@ public class WasmTarget implements TeaVMTarget {
         tagExpression = new WasmLoadInt32(4, tagExpression, WasmInt32Subtype.INT32);
         body.add(new WasmSetLocal(subtypeVar, tagExpression));
 
-        Collections.sort(ranges, Comparator.comparingInt(range -> range.lower));
+        ranges.sort(Comparator.comparingInt(range -> range.lower));
 
         int lower = ranges.get(0).lower;
         int upper = ranges.get(ranges.size() - 1).upper;

core/src/main/java/org/teavm/backend/wasm/binary/BinaryWriter.java

@@ -107,8 +107,10 @@ public class BinaryWriter {
     public byte[] getData() {
         byte[] result = new byte[address];
         int offset = 0;
+        int index = 0;
         for (DataValue value : values) {
             offset = writeData(result, offset, value);
+            ++index;
         }
         return Arrays.copyOf(result, offset);
     }

core/src/main/java/org/teavm/backend/wasm/generate/WasmGenerationContext.java

@@ -42,7 +42,7 @@ public class WasmGenerationContext {
     private WasmStringPool stringPool;
     private Map<MethodReference, ImportedMethod> importedMethods = new HashMap<>();
     private List<WasmIntrinsic> intrinsics = new ArrayList<>();
-    private Map<MethodReference, WasmIntrinsic> intrinsicCache = new HashMap<>();
+    private Map<MethodReference, WasmIntrinsicHolder> intrinsicCache = new HashMap<>();
 
     public WasmGenerationContext(ClassReaderSource classSource, Diagnostics diagnostics,
             VirtualTableProvider vtableProvider, TagRegistry tagRegistry, WasmStringPool stringPool) {
@@ -58,9 +58,18 @@ public class WasmGenerationContext {
     }
 
     public WasmIntrinsic getIntrinsic(MethodReference method) {
-        return intrinsicCache.computeIfAbsent(method, key -> intrinsics.stream()
+        return intrinsicCache.computeIfAbsent(method, key -> new WasmIntrinsicHolder(intrinsics.stream()
                 .filter(intrinsic -> intrinsic.isApplicable(key))
-                .findFirst().orElse(null));
+                .findFirst().orElse(null)))
+                .value;
+    }
+
+    private static class WasmIntrinsicHolder {
+        WasmIntrinsic value;
+
+        public WasmIntrinsicHolder(WasmIntrinsic value) {
+            this.value = value;
+        }
     }
 
     public ImportedMethod getImportedMethod(MethodReference reference) {

core/src/main/java/org/teavm/backend/wasm/intrinsics/ExceptionHandlingIntrinsic.java

@@ -23,6 +23,9 @@ import org.teavm.backend.wasm.generate.CallSiteBinaryGenerator;
 import org.teavm.backend.wasm.generate.WasmClassGenerator;
 import org.teavm.backend.wasm.model.expression.WasmExpression;
 import org.teavm.backend.wasm.model.expression.WasmInt32Constant;
+import org.teavm.backend.wasm.model.expression.WasmIntBinary;
+import org.teavm.backend.wasm.model.expression.WasmIntBinaryOperation;
+import org.teavm.backend.wasm.model.expression.WasmIntType;
 import org.teavm.model.MethodReference;
 import org.teavm.model.lowlevel.CallSiteDescriptor;
 import org.teavm.runtime.ExceptionHandling;
@@ -59,6 +62,11 @@ public class ExceptionHandlingIntrinsic implements WasmIntrinsic {
         WasmInt32Constant constant = new WasmInt32Constant(0);
         constant.setLocation(invocation.getLocation());
         constants.add(constant);
-        return constant;
+
+        WasmExpression id = manager.generate(invocation.getArguments().get(0));
+        WasmExpression offset = new WasmIntBinary(WasmIntType.INT32, WasmIntBinaryOperation.SHL,
+                id, new WasmInt32Constant(3));
+
+        return new WasmIntBinary(WasmIntType.INT32, WasmIntBinaryOperation.ADD, constant, offset);
     }
 }

core/src/main/java/org/teavm/model/lowlevel/ExceptionHandlingShadowStackContributor.java

@@ -85,9 +85,25 @@ public class ExceptionHandlingShadowStackContributor {
             blockMapping[i] = i;
         }
 
+        List<Phi> allPhis = new ArrayList<>();
         int blockCount = program.basicBlockCount();
+        for (int i = 0; i < blockCount; ++i) {
+            allPhis.addAll(program.basicBlockAt(i).getPhis());
+        }
+
         for (int i = 0; i < blockCount; ++i) {
             BasicBlock block = program.basicBlockAt(i);
+
+            if (block.getExceptionVariable() != null) {
+                InvokeInstruction catchCall = new InvokeInstruction();
+                catchCall.setType(InvocationType.SPECIAL);
+                catchCall.setMethod(new MethodReference(ExceptionHandling.class, "catchException",
+                        Throwable.class));
+                catchCall.setReceiver(block.getExceptionVariable());
+                block.getInstructions().add(0, catchCall);
+                block.setExceptionVariable(null);
+            }
+
             int newIndex = contributeToBasicBlock(block);
             if (newIndex != i) {
                 blockMapping[i] = newIndex;
@@ -95,15 +111,12 @@ public class ExceptionHandlingShadowStackContributor {
             }
         }
 
-        for (int i = 0; i < blockCount; ++i) {
+        for (Phi phi : allPhis) {
-            BasicBlock block = program.basicBlockAt(i);
-            for (Phi phi : block.getPhis()) {
             for (Incoming incoming : phi.getIncomings()) {
                 int mappedSource = blockMapping[incoming.getSource().getIndex()];
                 incoming.setSource(program.basicBlockAt(mappedSource));
             }
         }
-        }
 
         return hasExceptionHandlers;
     }
@@ -130,6 +143,8 @@ public class ExceptionHandlingShadowStackContributor {
         }
 
         DefinitionExtractor defExtractor = new DefinitionExtractor();
+        List<BasicBlock> blocksToClearHandlers = new ArrayList<>();
+        blocksToClearHandlers.add(block);
 
         for (int i = 0; i < instructions.size(); ++i) {
             Instruction insn = instructions.get(i);
@@ -153,7 +168,7 @@ public class ExceptionHandlingShadowStackContributor {
                 } else {
                     next = program.createBasicBlock();
                     next.getTryCatchBlocks().addAll(ProgramUtils.copyTryCatches(block, program));
-                    block.getTryCatchBlocks().clear();
+                    blocksToClearHandlers.add(next);
 
                     List<Instruction> remainingInstructions = instructions.subList(i + 1, instructions.size());
                     List<Instruction> instructionsToMove = new ArrayList<>(remainingInstructions);
@@ -175,7 +190,7 @@ public class ExceptionHandlingShadowStackContributor {
                 }
                 block = next;
                 instructions = block.getInstructions();
-                i = 0;
+                i = -1;
             }
 
             insn.acceptVisitor(defExtractor);
@@ -185,7 +200,9 @@ public class ExceptionHandlingShadowStackContributor {
             }
         }
 
-        block.getTryCatchBlocks().clear();
+        for (BasicBlock blockToClear : blocksToClearHandlers) {
+            blockToClear.getTryCatchBlocks().clear();
+        }
 
         return block.getIndex();
     }
@@ -253,14 +270,14 @@ public class ExceptionHandlingShadowStackContributor {
         boolean defaultExists = false;
         int nextHandlerId = callSite.getId();
         for (TryCatchBlock tryCatch : block.getTryCatchBlocks()) {
-            if (tryCatch.getExceptionType() == null) {
-                defaultExists = true;
-                switchInsn.setDefaultTarget(tryCatch.getHandler());
-            } else {
             ExceptionHandlerDescriptor handler = new ExceptionHandlerDescriptor(++nextHandlerId,
                     tryCatch.getExceptionType());
             callSite.getHandlers().add(handler);
 
+            if (tryCatch.getExceptionType() == null) {
+                defaultExists = true;
+                switchInsn.setDefaultTarget(tryCatch.getHandler());
+            } else {
                 SwitchTableEntry catchEntry = new SwitchTableEntry();
                 catchEntry.setTarget(tryCatch.getHandler());
                 catchEntry.setCondition(handler.getId());

core/src/main/java/org/teavm/model/lowlevel/GCShadowStackContributor.java

@@ -356,6 +356,7 @@ public class GCShadowStackContributor {
             instructionsToAdd.add(slotConstant);
 
             InvokeInstruction registerInvocation = new InvokeInstruction();
+            registerInvocation.setLocation(callInstruction.getLocation());
             registerInvocation.setType(InvocationType.SPECIAL);
             registerInvocation.getArguments().add(slotVar);
             if (var >= 0) {

core/src/main/java/org/teavm/runtime/ExceptionHandling.java

@@ -28,7 +28,17 @@ public final class ExceptionHandling {
 
     public static native CallSite findCallSiteById(int id);
 
+    private static Throwable thrownException;
+
+    public static Throwable catchException() {
+        Throwable exception = thrownException;
+        thrownException = null;
+        return exception;
+    }
+
     public static void throwException(Throwable exception) {
+        thrownException = exception;
+
         RuntimeObject exceptionPtr = Address.ofObject(exception).toStructure();
         RuntimeClass exceptionClass = RuntimeClass.getClass(exceptionPtr);
         IsSupertypeFunction isExceptionSupertype = exceptionClass.isSupertypeOf;
@@ -40,7 +50,7 @@ public final class ExceptionHandling {
             ExceptionHandler handler = callSite.firstHandler;
 
             for (int i = 0; i < callSite.handlerCount; ++i) {
-                if (isExceptionSupertype.apply(handler.exceptionClass)) {
+                if (handler.exceptionClass == null || isExceptionSupertype.apply(handler.exceptionClass)) {
                     ShadowStack.setExceptionHandlerId(stackFrame, handler.id);
                     break stackLoop;
                 }