Model Map.remove as nullable

[codingkiddo]

Fixes #1591.

This PR models java.util.Map.remove(Object) as returning nullable, similar to the existing model for Map.get(Object).

Map.remove(Object) can return null when there is no mapping for the key, so dereferencing its result should produce a NullAway warning.

Changes:

• Add java.util.Map.remove(Object) to the always-nullable library return models
• Add a regression test for dereferencing map.remove(key)

Tested with:

• ./gradlew :nullaway:test --tests "com.uber.nullaway.FrameworkTests.defaultLibraryModelsMapRemove"
• ./gradlew spotlessCheck

Summary by CodeRabbit

• Bug Fixes
  • Improved null-safety modeling for Map.remove(...) so its return value is treated as potentially null, helping prevent unsafe dereferences after removing entries.
• Tests
  • Added a regression test that compiles a small example using map.remove(...).toString() and verifies the expected null-safety diagnostic.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

Walkthrough

LibraryModelsHandler now treats java.util.Map.remove(Object) as always returning @Nullable. FrameworkTests adds a test that calls Map.remove("key").toString() and expects a NullAway diagnostic for dereferencing a nullable expression.

Suggested reviewers

• lazaroclapp

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title is concise and directly describes the main change: modeling Map.remove as nullable.
Linked Issues check	✅ Passed	The change models java.util.Map.remove(Object) as nullable and adds a regression test for the dereference warning, matching #1591.
Out of Scope Changes check	✅ Passed	The PR stays focused on the Map.remove nullability model and its test, with no unrelated changes.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[subhramit]

@codingkiddo your commits will be squashed anyway at the time of merge, so you can avoid force-pushing. Incremental changes are easier to track via non-squashed commits.

[msridhar]

Thanks for this change! While I agree the change is correct, it is going to be a bit disruptive in terms of generating new errors for users (see the failed integration tests). I think at the least, we need containsKey support before landing this. We have support for reasoning that if containsKey returns true for some key in a Map, then the get call for that same key will return @NonNull. I think we want the same support for containsKey and remove, to handle cases like:

https://github.com/junit-team/junit-framework/blob/c32e9b138ceffc28d2e733ee2e2b1d73d2c1f914/junit-jupiter-engine/src/main/java/org/junit/jupiter/engine/descriptor/ClassTemplateTestDescriptor.java#L267-L269

Some of the relevant code in NullAway for this support is:

NullAway/nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPath.java

Lines 188 to 191 in f40a9fa

	if (isMapGet(ASTHelpers.getSymbol(node.getTree()), state)) {
	return fromMapGetCall(node, state, apContext);
	}
	return fromVanillaMethodCall(node, apContext);

I think basically, we want to handle Map.get and Map.remove calls the same way. @codingkiddo can you look into that?

(Aside: with JSpecify mode, we'll want to separately reason about maps with a @Nullable value type, and assume they can still return null even if containsKey returns true. But that is for later.)

[codingkiddo]

Thanks for this change! While I agree the change is correct, it is going to be a bit disruptive in terms of generating new errors for users (see the failed integration tests). I think at the least, we need containsKey support before landing this. We have support for reasoning that if containsKey returns true for some key in a Map, then the get call for that same key will return @NonNull. I think we want the same support for containsKey and remove, to handle cases like:

https://github.com/junit-team/junit-framework/blob/c32e9b138ceffc28d2e733ee2e2b1d73d2c1f914/junit-jupiter-engine/src/main/java/org/junit/jupiter/engine/descriptor/ClassTemplateTestDescriptor.java#L267-L269

Some of the relevant code in NullAway for this support is:

NullAway/nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPath.java

Lines 188 to 191 in f40a9fa

	if (isMapGet(ASTHelpers.getSymbol(node.getTree()), state)) {
	return fromMapGetCall(node, state, apContext);
	}
	return fromVanillaMethodCall(node, apContext);

I think basically, we want to handle Map.get and Map.remove calls the same way. @codingkiddo can you look into that?

(Aside: with JSpecify mode, we'll want to separately reason about maps with a @Nullable value type, and assume they can still return null even if containsKey returns true. But that is for later.)

Thanks for the detailed pointers. I updated the PR so Map.remove(key) is represented like Map.get(key) in access-path handling, allowing the existing containsKey(key) reasoning to apply to remove(key) as well.

I also added regression tests for the containsKey-guarded safe case and the different-key unsafe case.

[msridhar]

@codingkiddo thanks for the updates! I've made some further changes (with help of Codex) to handle cases like:

if (map.containsKey(key)) {
  map.remove(key).toString(); // intended to be safe
  map.remove(key).toString(); // should warn: key was removed
  map.get(key).toString();    // should warn: key was removed
}

This required slightly more involved changes than what you did. I've gone through the failures in the integration tests and the warnings look valid and should be fairly easy to fix. So, I will likely land this, and then we'll try to cut a release soonish for users.

[codecov]

Codecov Report

❌ Patch coverage is 92.30769% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 88.10%. Comparing base (addfe41) to head (ef0426c).

Files with missing lines	Patch %	Lines
...llaway/dataflow/AccessPathNullnessPropagation.java	75.00%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff            @@
##             master    #1623   +/-   ##
=========================================
  Coverage     88.10%   88.10%           
- Complexity     3046     3050    +4     
=========================================
  Files           105      105           
  Lines         10181    10192   +11     
  Branches       2062     2066    +4     
=========================================
+ Hits           8970     8980   +10     
  Misses          572      572           
- Partials        639      640    +1     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[codingkiddo]

@codingkiddo thanks for the updates! I've made some further changes (with help of Codex) to handle cases like:

if (map.containsKey(key)) {
  map.remove(key).toString(); // intended to be safe
  map.remove(key).toString(); // should warn: key was removed
  map.get(key).toString();    // should warn: key was removed
}

This required slightly more involved changes than what you did. I've gone through the failures in the integration tests and the warnings look valid and should be fairly easy to fix. So, I will likely land this, and then we'll try to cut a release soonish for users.

Thanks for the follow-up and for handling the more involved dataflow cases.

That makes sense — after remove(key), the previous containsKey(key) fact should no longer make later remove(key) or get(key) calls safe. I appreciate the explanation and the additional refinement.

Glad the remaining integration-test warnings look valid. Thanks for moving this forward!