fix[compiler]: preserve optional chaining when non-null access is behind a conditional guard inside a nested closure

compiler/packages/babel-plugin-react-compiler/src/HIR/CollectHoistablePropertyLoads.ts

@@ -448,8 +448,24 @@ function collectNonNullsInBlocks(
           const innerHoistables = assertNonNull(
             innerHoistableMap.get(innerFn.func.body.entry),
           );
+
+          // Collect the set of identifiers that are captured from outer scope
+          const capturedIdentifiers = new Set<IdentifierId>();
+          for (const place of innerFn.func.context) {
+            capturedIdentifiers.add(place.identifier.id);
+          }
+
+          // Only propagate non-null assumptions for identifiers that are NOT
+          // captured from the outer scope. This prevents incorrect nullability
+          // inference when the inner function is only called conditionally.
+          // For example:
+          //   if (currentDevice) { log(); }  // where log uses currentDevice.os
+          // We should NOT assume currentDevice is non-null in the outer scope
+          // just because it's used non-optionally inside log().
           for (const entry of innerHoistables.assumedNonNullObjects) {
-            assumedNonNullObjects.add(entry);
+            if (!capturedIdentifiers.has(entry.fullPath.identifier.id)) {
+              assumedNonNullObjects.add(entry);
+            }
           }
         }
       } else if (

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-nested-function-non-optional-access.expect.md

@@ -0,0 +1,85 @@
+
+## Input
+
+```javascript
+// Test that optional chaining is still removed when appropriate.
+// When a nested function is unconditionally called (e.g., as a JSX prop handler),
+// and the variable is accessed non-optionally, we should still optimize away the ?. 
+
+import {useState} from 'react';
+
+function Component({device}) {
+  const [count, setCount] = useState(0);
+
+  // This handler is unconditionally passed to onClick
+  const handleClick = () => {
+    console.log(device.type);  // this is safe to access directly
+    console.log(device.id);
+  };
+
+  return (
+    <div>
+      <button onClick={handleClick}>Click {count}</button>
+    </div>
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{device: {type: 'phone', id: 123}}],
+  isComponent: true,
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // Test that optional chaining is still removed when appropriate.
+// When a nested function is unconditionally called (e.g., as a JSX prop handler),
+// and the variable is accessed non-optionally, we should still optimize away the ?.
+
+import { useState } from "react";
+
+function Component(t0) {
+  const $ = _c(5);
+  const { device } = t0;
+  const [count] = useState(0);
+  let t1;
+  if ($[0] !== device) {
+    t1 = () => {
+      console.log(device.type);
+      console.log(device.id);
+    };
+    $[0] = device;
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  const handleClick = t1;
+  let t2;
+  if ($[2] !== count || $[3] !== handleClick) {
+    t2 = (
+      <div>
+        <button onClick={handleClick}>Click {count}</button>
+      </div>
+    );
+    $[2] = count;
+    $[3] = handleClick;
+    $[4] = t2;
+  } else {
+    t2 = $[4];
+  }
+  return t2;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ device: { type: "phone", id: 123 } }],
+  isComponent: true,
+};
+
+```
+
+### Eval output
+(kind: ok) <div><button>Click 0</button></div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-nested-function-non-optional-access.js

@@ -0,0 +1,27 @@
+// Test that optional chaining is still removed when appropriate.
+// When a nested function is unconditionally called (e.g., as a JSX prop handler),
+// and the variable is accessed non-optionally, we should still optimize away the ?. 
+
+import {useState} from 'react';
+
+function Component({device}) {
+  const [count, setCount] = useState(0);
+
+  // This handler is unconditionally passed to onClick
+  const handleClick = () => {
+    console.log(device.type);  // this is safe to access directly
+    console.log(device.id);
+  };
+
+  return (
+    <div>
+      <button onClick={handleClick}>Click {count}</button>
+    </div>
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{device: {type: 'phone', id: 123}}],
+  isComponent: true,
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-optional-chaining-conditional-closure.expect.md

@@ -0,0 +1,106 @@
+
+## Input
+
+```javascript
+// Test that optional chaining is preserved when used inside a closure
+// that is only conditionally called. The compiler should NOT strip the ?. 
+// operator even if the variable is accessed non-optionally inside a nested function,
+// because that function might only be called conditionally.
+//
+// Before fix (BROKEN): Compiler removes ?. and causes: TypeError: Cannot read properties of undefined
+// After fix (CORRECT): Compiler preserves ?. and code is safe
+
+import {useEffect} from 'react';
+
+function Component({devices}) {
+  useEffect(() => {
+    const currentDevice = devices?.[0];
+
+    // The function uses optional chaining, but the compiler was incorrectly
+    // inferring it's always non-null because it's called in a guarded context
+    const log = () => {
+      console.log(currentDevice?.type);
+      console.log(currentDevice?.os);
+      const match = currentDevice?.os?.match(/android|ios/i);
+      console.log(match);
+    };
+
+    // The function is only called conditionally (if currentDevice is truthy)
+    // This should NOT cause the compiler to assume currentDevice is non-null
+    // in the component render or effect setup
+    if (currentDevice) {
+      log();
+    }
+  }, [devices]);
+
+  return null;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{devices: []}],
+  isComponent: true,
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // Test that optional chaining is preserved when used inside a closure
+// that is only conditionally called. The compiler should NOT strip the ?.
+// operator even if the variable is accessed non-optionally inside a nested function,
+// because that function might only be called conditionally.
+//
+// Before fix (BROKEN): Compiler removes ?. and causes: TypeError: Cannot read properties of undefined
+// After fix (CORRECT): Compiler preserves ?. and code is safe
+
+import { useEffect } from "react";
+
+function Component(t0) {
+  const $ = _c(4);
+  const { devices } = t0;
+  let t1;
+  if ($[0] !== devices?.[0]) {
+    t1 = () => {
+      const currentDevice = devices?.[0];
+
+      const log = () => {
+        console.log(currentDevice?.type);
+        console.log(currentDevice?.os);
+        const match = currentDevice?.os?.match(/android|ios/i);
+        console.log(match);
+      };
+
+      if (currentDevice) {
+        log();
+      }
+    };
+    $[0] = devices?.[0];
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  let t2;
+  if ($[2] !== devices) {
+    t2 = [devices];
+    $[2] = devices;
+    $[3] = t2;
+  } else {
+    t2 = $[3];
+  }
+  useEffect(t1, t2);
+
+  return null;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ devices: [] }],
+  isComponent: true,
+};
+
+```
+
+### Eval output
+(kind: ok) null

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-optional-chaining-conditional-closure.js

@@ -0,0 +1,39 @@
+// Test that optional chaining is preserved when used inside a closure
+// that is only conditionally called. The compiler should NOT strip the ?. 
+// operator even if the variable is accessed non-optionally inside a nested function,
+// because that function might only be called conditionally.
+//
+// Before fix (BROKEN): Compiler removes ?. and causes: TypeError: Cannot read properties of undefined
+// After fix (CORRECT): Compiler preserves ?. and code is safe
+
+import {useEffect} from 'react';
+
+function Component({devices}) {
+  useEffect(() => {
+    const currentDevice = devices?.[0];
+
+    // The function uses optional chaining, but the compiler was incorrectly
+    // inferring it's always non-null because it's called in a guarded context
+    const log = () => {
+      console.log(currentDevice?.type);
+      console.log(currentDevice?.os);
+      const match = currentDevice?.os?.match(/android|ios/i);
+      console.log(match);
+    };
+
+    // The function is only called conditionally (if currentDevice is truthy)
+    // This should NOT cause the compiler to assume currentDevice is non-null
+    // in the component render or effect setup
+    if (currentDevice) {
+      log();
+    }
+  }, [devices]);
+
+  return null;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{devices: []}],
+  isComponent: true,
+};