src/jdk.internal.vm.compiler/share/classes/org.graalvm.compiler.core.test/src/org/graalvm/compiler/core/test/ea/EATestBase.java
changeset 52578 7dd81e82d083
parent 50858 2d3e99a72541
child 52910 583fd71c47d6
--- a/src/jdk.internal.vm.compiler/share/classes/org.graalvm.compiler.core.test/src/org/graalvm/compiler/core/test/ea/EATestBase.java	Thu Nov 15 21:05:47 2018 +0100
+++ b/src/jdk.internal.vm.compiler/share/classes/org.graalvm.compiler.core.test/src/org/graalvm/compiler/core/test/ea/EATestBase.java	Thu Nov 15 09:04:07 2018 -0800
@@ -24,6 +24,8 @@
 
 package org.graalvm.compiler.core.test.ea;
 
+import static org.graalvm.compiler.graph.iterators.NodePredicates.isA;
+
 import java.util.List;
 
 import org.graalvm.compiler.core.test.GraalCompilerTest;
@@ -33,6 +35,7 @@
 import org.graalvm.compiler.nodes.StructuredGraph.AllowAssumptions;
 import org.graalvm.compiler.nodes.java.NewArrayNode;
 import org.graalvm.compiler.nodes.java.NewInstanceNode;
+import org.graalvm.compiler.nodes.virtual.AllocatedObjectNode;
 import org.graalvm.compiler.nodes.virtual.CommitAllocationNode;
 import org.graalvm.compiler.phases.common.CanonicalizerPhase;
 import org.graalvm.compiler.phases.common.DeadCodeEliminationPhase;
@@ -139,6 +142,10 @@
      *            iteration
      */
     protected void testEscapeAnalysis(String snippet, JavaConstant expectedConstantResult, boolean iterativeEscapeAnalysis) {
+        testEscapeAnalysis(snippet, expectedConstantResult, iterativeEscapeAnalysis, 0);
+    }
+
+    protected void testEscapeAnalysis(String snippet, JavaConstant expectedConstantResult, boolean iterativeEscapeAnalysis, int expectedAllocationCount) {
         prepareGraph(snippet, iterativeEscapeAnalysis);
         if (expectedConstantResult != null) {
             for (ReturnNode returnNode : returnNodes) {
@@ -146,9 +153,11 @@
                 Assert.assertEquals(expectedConstantResult, returnNode.result().asConstant());
             }
         }
-        int newInstanceCount = graph.getNodes().filter(NewInstanceNode.class).count() + graph.getNodes().filter(NewArrayNode.class).count() +
-                        graph.getNodes().filter(CommitAllocationNode.class).count();
-        Assert.assertEquals(0, newInstanceCount);
+        int newInstanceCount = graph.getNodes().filter(isA(NewInstanceNode.class).or(NewArrayNode.class).or(AllocatedObjectNode.class)).count();
+        Assert.assertEquals("Expected allocation count does not match", expectedAllocationCount, newInstanceCount);
+        if (expectedAllocationCount == 0) {
+            Assert.assertTrue("Unexpected CommitAllocationNode", graph.getNodes().filter(CommitAllocationNode.class).isEmpty());
+        }
     }
 
     @SuppressWarnings("try")