Omit null checks on return values when they're known to be not null

src/hotspot/share/ci/bcEscapeAnalyzer.cpp

@@ -287,8 +287,21 @@ void BCEscapeAnalyzer::invoke(StateInfo &state, Bytecodes::Code code, ciMethod*
   }
   if (skip_callee) {
     TRACE_BCEA(3, tty->print_cr("[EA] skipping method %s::%s", holder->name()->as_utf8(), target->name()->as_utf8()));
-    for (i = 0; i < arg_size; i++) {
-      set_method_escape(state.raw_pop());
+    // If the callee is a constructor (`<init>`) and if the object (passed as
+    // the last argument to `<init>`) is freshly allocated, then the object is
+    // being initialized and is not escaping.
+    if (code == Bytecodes::_invokespecial && target->is_object_initializer()) {
+      for (i = 0; i < arg_size; i++) {
+        ArgumentMap arg = state.raw_pop();
+        bool fresh_object = i == arg_size - 1 && arg.contains_allocated();
+        if (!fresh_object) {
+          set_method_escape(arg);
+        }
+      }
+    } else {
+      for (i = 0; i < arg_size; i++) {
+        set_method_escape(state.raw_pop());
+      }
     }
     _unknown_modified = true;  // assume the worst since we don't analyze the called method
     return;

src/hotspot/share/opto/doCall.cpp

@@ -22,6 +22,7 @@
  *
  */
 
+#include "ci/bcEscapeAnalyzer.hpp"
 #include "ci/ciCallSite.hpp"
 #include "ci/ciMethodHandle.hpp"
 #include "ci/ciSymbols.hpp"
@@ -524,6 +525,52 @@ static bool check_call_consistency(JVMState* jvms, CallGenerator* cg) {
 }
 #endif // ASSERT
 
+static bool return_is_not_null(ciMethod* callee, CallNode* call, Compile* C, PhaseGVN* gvn) {
+  if (callee == nullptr) {
+    return false;
+  }
+
+  BCEscapeAnalyzer* bcea = callee->get_bcea();
+  if (bcea == nullptr) {
+    return false;
+  }
+
+  bcea->copy_dependencies(C->dependencies());
+  if (bcea->is_return_allocated()) {
+    // The callee is known to return a not null value.
+    return true;
+  }
+
+  if (bcea->is_return_local()) {
+    // The callee returns one of its input arguments, so we find all
+    // args that may be returned and if they all are known to be not
+    // null, then we mark the return value as not null as well.
+
+    if (call == nullptr) {
+      return false;
+    }
+
+    // In addition to checking `is_return_local()`, we also need to check if we
+    // found at least one argument that may be returned because of conservative
+    // fallbacks in BCEA.  If we don't, then removing the null check can cause
+    // JVM crashes.
+    bool found = false;
+    uint arg_count = call->tf()->domain()->cnt() - TypeFunc::Parms;
+    for (uint idx = 0; idx < arg_count; idx++) {
+      if (bcea->is_arg_returned(idx)) {
+        found = true;
+        Node* arg = call->in(idx + TypeFunc::Parms);
+        if (gvn->type(arg)->maybe_null()) {
+          return false;
+        }
+      }
+    }
+    return found;
+  }
+
+  return false;
+}
+
 //------------------------------do_call----------------------------------------
 // Handle your basic call.  Inline if we can & want to, else just setup call.
 void Parse::do_call() {
@@ -798,6 +845,12 @@ void Parse::do_call() {
     }
     BasicType ct = ctype->basic_type();
     if (is_reference_type(ct)) {
+      // If the callee is not inlined, check whether bytecode escape analysis
+      // can prove that the return value is not null.  If callee is inlined,
+      // then C2 can infer whether the callee returns a non-null value.
+      if (!cg->is_inline() && return_is_not_null(cg->method(), cg->call_node(), C, &_gvn)) {
+        cast_not_null(peek());
+      }
       record_profiled_return_for_speculation();
     }
   }

test/hotspot/jtreg/compiler/escapeAnalysis/TestReturnNotNull.java

@@ -0,0 +1,155 @@
+/*
+ * Copyright (c) 2026, Oracle and/or its affiliates. All rights reserved.
+ * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+ *
+ * This code is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU General Public License version 2 only, as
+ * published by the Free Software Foundation.
+ *
+ * This code 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 Public License
+ * version 2 for more details (a copy is included in the LICENSE file that
+ * accompanied this code).
+ *
+ * You should have received a copy of the GNU General Public License version
+ * 2 along with this work; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
+ *
+ * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
+ * or visit www.oracle.com if you need additional information or have any
+ * questions.
+ */
+
+package compiler.escapeAnalysis;
+
+import compiler.lib.ir_framework.*;
+
+/**
+ * @test
+ * @summary Test that C2 eliminates null checks on return values when possible.
+ * @library /test/lib /
+ * @run driver compiler.escapeAnalysis.TestReturnNotNull
+ */
+public class TestReturnNotNull {
+    static Object leaked;
+    static int sideEffect = 5;
+
+    @DontInline
+    static Object allocateObject() {
+        return new Object();
+    }
+
+    @DontInline
+    static int[] allocateArray(int len) {
+        return new int[len];
+    }
+
+    static class Builder {
+        int value;
+
+        @DontInline
+        Builder setValue(int v) {
+            this.value = v;
+            return this;
+        }
+
+        @DontInline
+        int getValue() {
+            return value;
+        }
+    }
+
+    @DontInline
+    static Object identity(Object o) {
+        return o;
+    }
+
+    @DontInline
+    static Object allocateOnBothBranches(boolean flag) {
+        if (flag) {
+            return new Object();
+        } else {
+            return new Object();
+        }
+    }
+
+    @DontInline
+    static Object maybeNull(boolean flag) {
+        return flag ? new Object() : null;
+    }
+
+    @DontInline
+    static Object allocateButLeaks() {
+        Object o = new Object();
+        leaked = o;
+        return o;
+    }
+
+    @Test
+    @IR(failOn = IRNode.NULL_CHECK, phase = CompilePhase.FINAL_CODE)
+    static int testAllocateObject() {
+        Object o = allocateObject();
+        return o.hashCode();
+    }
+
+    @Test
+    @IR(failOn = IRNode.NULL_CHECK, phase = CompilePhase.FINAL_CODE)
+    static int testAllocateArray() {
+        int[] a = allocateArray(10);
+        return a.length;
+    }
+
+    @Test
+    @IR(failOn = IRNode.NULL_CHECK, phase = CompilePhase.FINAL_CODE)
+    static int testAllocateOnBothBranches() {
+        Object o = allocateOnBothBranches(sideEffect > 0);
+        return o.hashCode();
+    }
+
+    @Test
+    @IR(failOn = IRNode.NULL_CHECK, phase = CompilePhase.FINAL_CODE)
+    static int testReturnsThis() {
+        Builder b = new Builder();
+        return b.setValue(42).getValue();
+    }
+
+    @Test
+    @IR(failOn = IRNode.NULL_CHECK, phase = CompilePhase.FINAL_CODE)
+    static int testReturnsThisChained() {
+        Builder b = new Builder();
+        return b.setValue(1).setValue(2).getValue();
+    }
+
+    @Test
+    @IR(failOn = IRNode.NULL_CHECK, phase = CompilePhase.FINAL_CODE)
+    static int testIdentityOnNonNull() {
+        Object o = identity(new Object());
+        return o.hashCode();
+    }
+
+    @Test
+    @IR(counts = {IRNode.NULL_CHECK, "1"}, phase = CompilePhase.FINAL_CODE)
+    static int testMaybeNull() {
+        Object o = maybeNull(sideEffect > 0);
+        return o.hashCode();
+    }
+
+    @Test
+    @IR(counts = {IRNode.NULL_CHECK, "1"}, phase = CompilePhase.FINAL_CODE)
+    static int testAllocateButLeaks() {
+        Object o = allocateButLeaks();
+        return o.hashCode();
+    }
+
+    @Test
+    @IR(counts = {IRNode.NULL_CHECK, "1"}, phase = CompilePhase.FINAL_CODE)
+    static int testIdentityOnMaybeNull() {
+        Object o = identity(maybeNull(sideEffect > 0));
+        return o.hashCode();
+    }
+
+    public static void main(String[] args) {
+        TestFramework.runWithFlags("-XX:-TieredCompilation");
+    }
+}