Extend try-catch blocks to cover inserted probes.

2 files changed, 76 insertions, 0 deletions

diff --git a/org.jacoco.core.test/src/org/jacoco/core/internal/flow/MethodProbesAdapterTest.java b/org.jacoco.core.test/src/org/jacoco/core/internal/flow/MethodProbesAdapterTest.java
index ea329632..5590323f 100644
--- a/org.jacoco.core.test/src/org/jacoco/core/internal/flow/MethodProbesAdapterTest.java
+++ b/org.jacoco.core.test/src/org/jacoco/core/internal/flow/MethodProbesAdapterTest.java
@@ -312,6 +312,58 @@ public class MethodProbesAdapterTest implements IProbeIdGenerator {
 		expectedVisitor.visitTableSwitchInsn(0, 1, label, labels);
 	}
 
+	@Test
+	public void testVisitTryCatchBlockNoProbe() {
+		Label start = new Label();
+		Label end = new Label();
+		Label handler = new Label();
+
+		adapter.visitTryCatchBlock(start, end, handler, "java/lang/Exception");
+		adapter.visitLabel(start);
+
+		expectedVisitor.visitTryCatchBlock(start, end, handler, "java/lang/Exception");
+		expectedVisitor.visitLabel(start);
+	}
+
+	@Test
+	public void testVisitTryCatchBlockWithProbe() {
+		Label target = new Label();
+		LabelInfo.setSuccessor(target);
+		LabelInfo.setTarget(target);
+		Label end = new Label();
+		Label handler = new Label();
+		Label start = new Label();
+
+		adapter.visitTryCatchBlock(target, end, handler, "java/lang/Exception");
+		adapter.visitLabel(target);
+
+		expectedVisitor.visitTryCatchBlock(start, end, handler, "java/lang/Exception");
+		expectedVisitor.visitLabel(start);
+		expectedVisitor.visitProbe(1000);
+		expectedVisitor.visitLabel(target);
+	}
+
+	@Test
+	public void testVisitMultipleTryCatchBlocksWithProbe() {
+		Label target = new Label();
+		LabelInfo.setSuccessor(target);
+		LabelInfo.setTarget(target);
+		Label end = new Label();
+		Label handler1 = new Label();
+		Label handler2 = new Label();
+		Label start = new Label();
+
+		adapter.visitTryCatchBlock(target, end, handler1, "java/lang/Exception");
+		adapter.visitTryCatchBlock(target, end, handler2, "java/io/IOException");
+		adapter.visitLabel(target);
+
+		expectedVisitor.visitTryCatchBlock(start, end, handler1, "java/lang/Exception");
+		expectedVisitor.visitTryCatchBlock(start, end, handler2, "java/io/IOException");
+		expectedVisitor.visitLabel(start);
+		expectedVisitor.visitProbe(1000);
+		expectedVisitor.visitLabel(target);
+	}
+
 	// === IProbeIdGenerator ===
 
 	public int nextId() {
diff --git a/org.jacoco.core/src/org/jacoco/core/internal/flow/MethodProbesAdapter.java b/org.jacoco.core/src/org/jacoco/core/internal/flow/MethodProbesAdapter.java
index 6d44bc5a..2fe7e315 100644
--- a/org.jacoco.core/src/org/jacoco/core/internal/flow/MethodProbesAdapter.java
+++ b/org.jacoco.core/src/org/jacoco/core/internal/flow/MethodProbesAdapter.java
@@ -17,6 +17,9 @@ import org.objectweb.asm.MethodVisitor;
 import org.objectweb.asm.Opcodes;
 import org.objectweb.asm.commons.AnalyzerAdapter;
 
+import java.util.HashMap;
+import java.util.Map;
+
 /**
  * Adapter that creates additional visitor events for probes to be inserted into
  * a method.
@@ -29,6 +32,8 @@ public final class MethodProbesAdapter extends MethodVisitor {
 
 	private AnalyzerAdapter analyzer;
 
+	private Map<Label, Label> tryCatchProbeLabels = new HashMap<Label, Label>();
+
 	/**
 	 * Create a new adapter instance.
 	 * 
@@ -56,8 +61,27 @@ public final class MethodProbesAdapter extends MethodVisitor {
 	}
 
 	@Override
+	public void visitTryCatchBlock(Label start, final Label end,
+			final Label handler, final String type) {
+		// If a probe will be inserted before the start label, we'll need to use a
+		// different label for the try-catch block.
+		if (tryCatchProbeLabels.containsKey(start)) {
+			start = tryCatchProbeLabels.get(start);
+		} else if (LabelInfo.isMultiTarget(start) && LabelInfo.isSuccessor(start)) {
+			Label probeLabel = new Label();
+			LabelInfo.setSuccessor(probeLabel);
+			tryCatchProbeLabels.put(start, probeLabel);
+			start = probeLabel;
+		}
+		probesVisitor.visitTryCatchBlock(start, end, handler, type);
+	}
+
+	@Override
 	public void visitLabel(final Label label) {
 		if (LabelInfo.isMultiTarget(label) && LabelInfo.isSuccessor(label)) {
+			if (tryCatchProbeLabels.containsKey(label)) {
+				probesVisitor.visitLabel(tryCatchProbeLabels.get(label));
+			}
 			probesVisitor.visitProbe(idGenerator.nextId());
 		}
 		probesVisitor.visitLabel(label);