Fix issue #2358 Make Map.keySet() handle multiple argments to @KeyFor

[Gaoyan1999]

Summary

Enhance @KeyFor propagation for Map.keySet() by modifying addComputedTypeAnnotations in KeyForAnnotatedTypeFactory.

Problem

The element type of Set returned by Map.keySet() only carried @KeyFor for the current variable (e.g., "sharedCounts1") and did not inherit @KeyFor values from the Map’s key type argument (e.g., "sharedBooks").

    Map<@KeyFor({"sharedBooks"}) String, Integer> sharedBooks = new HashMap<>();
    Map<@KeyFor({"sharedBooks"}) String, Integer> sharedCounts1 = new HashMap<>();
    Set<@KeyFor({"sharedBooks", "sharedCounts1"}) String> otherChars1 = sharedCounts1.keySet();

Solution

Override addComputedTypeAnnotations in KeyForAnnotatedTypeFactory. When the tree is a Map.keySet() call, read @KeyFor values from the receiver Map’s key type argument and from the current Set element type, merge the names, and update the keySet() return type’s element with the combined @KeyFor.

Reference

#2358

[coderabbitai]

📝 Walkthrough

Walkthrough

The changes modify KeyFor annotation handling in the Checker Framework's nullness checker. The main file is updated to merge KeyFor annotations from a Map's key type into the return type of keySet() invocations. A new method override handles type argument extraction and annotation merging through a private helper method. Additionally, new imports are introduced to support this logic. A test file is updated by removing a skip annotation and adding error diagnostic markers to specific test methods, enabling test execution with expected error annotations.

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 40.00% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

📜 Recent review details

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between d14f924 and c86cce7.

⛔ Files ignored due to path filters (1)

• checker/tests/nullness/KeyForMultiple.class is excluded by !**/*.class

📒 Files selected for processing (2)

• checker/src/main/java/org/checkerframework/checker/nullness/KeyForAnnotatedTypeFactory.java (4 hunks)
• checker/tests/nullness/KeyForMultiple.java (2 hunks)

🧰 Additional context used

🧠 Learnings (2)

📓 Common learnings

Learnt from: mernst
Repo: typetools/checker-framework PR: 7354
File: annotation-file-utilities/tests/LocalMultipleManyMethodsShifted.java:14-14
Timestamp: 2025-11-02T02:18:00.514Z
Learning: In the Checker Framework repository, when doing refactoring PRs (such as preferring isEmpty() over size() comparisons), test files should not be changed. Tests should remain stable to preserve their exact patterns for annotation processing and bytecode verification purposes.

📚 Learning: 2025-10-22T20:41:01.112Z

Learnt from: kelloggm
Repo: typetools/checker-framework PR: 7166
File: checker/src/main/java/org/checkerframework/checker/collectionownership/CollectionOwnershipAnnotatedTypeFactory.java:128-136
Timestamp: 2025-10-22T20:41:01.112Z
Learning: In checker/src/main/java/org/checkerframework/checker/collectionownership/CollectionOwnershipAnnotatedTypeFactory.java, prefer instance fields over static state for analysis caches (e.g., maps keyed by Tree/Block). There is at most one ATF per running checker, so static mutable fields are discouraged.

Applied to files:

• checker/src/main/java/org/checkerframework/checker/nullness/KeyForAnnotatedTypeFactory.java

🧬 Code graph analysis (1)

checker/src/main/java/org/checkerframework/checker/nullness/KeyForAnnotatedTypeFactory.java (3)

framework/src/main/java/org/checkerframework/framework/type/AnnotatedTypeMirror.java (2)

• AnnotatedTypeMirror (64-2780)
• AnnotatedDeclaredType (948-1237)

javacutil/src/main/java/org/checkerframework/javacutil/TreeUtils.java (1)

• TreeUtils (116-2967)

javacutil/src/main/java/org/checkerframework/javacutil/AnnotationUtils.java (1)

• AnnotationUtils (51-1512)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (7)

• GitHub Check: typetools.checker-framework (typecheck_part2_jdk25)
• GitHub Check: typetools.checker-framework (misc_jdk25)
• GitHub Check: typetools.checker-framework (typecheck_part1_jdk25)
• GitHub Check: typetools.checker-framework (nonjunit_jdk25)
• GitHub Check: typetools.checker-framework (junit_jdk25)
• GitHub Check: typetools.checker-framework (inference_part1_jdk25)
• GitHub Check: typetools.checker-framework (inference_part2_jdk25)

🔇 Additional comments (2)

checker/src/main/java/org/checkerframework/checker/nullness/KeyForAnnotatedTypeFactory.java (1)

252-326: Great integration of keySet qualifier merging.

Hooking into addComputedTypeAnnotations keeps the propagation localized, and the helper cleanly unions receiver and return @KeyFor sets without disturbing other qualifiers. Nicely done.

checker/tests/nullness/KeyForMultiple.java (1)

15-33: Thanks for tightening the regression test.

The added diagnostics confirm the new propagation path and guard against future regressions.

---

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 and usage tips.}

[Gaoyan1999]

Hi, @mernst @smillst, I would appreciate it if you could take a look when you have a chance, thanks in advance!