Prove non-nullness of `get` for iterated keys

I have an interface which has map-like functionality, but does not implement Java's Map interface.
The map interface also implements Iterable<Object>; it iterates over the keys of the map
I'd like to use this in the body of an enhanced loop (see below), but without an assertion, and use get to retrieve values for the iterated keys, and without an [ERROR] from the Checker Framework.
Is that at all possible and could you provide pointers where to start or examples to learn from? I tried haphazardly to sprinkle some @KeyFors here and there, but with a lack of fully understanding what I'm doing it could take a while before I hit the right spots ;-)
I understand we might use an "Entry Iterator" and avoid to have to solve this problem in the first place, but I'm really just interested in learning how to teach the Checker Framework about the semantic relation between a key iterator and an @Nullable get method.

Here's a minimal working example:

import org.checkerframework.checker.nullness.qual.Nullable;

interface IMap extends Iterable<Object> {
    @Nullable Object get(Object o);
    IMap put(Object key, Object value); // immutable put
    IMap empty();

    default IMap remove(Object key) {
       IMap tmp = empty();

       for (Object k : this) {
           if (!k.equals(key)) {
               tmp.put(k, get(k)); // get(k) is always non-null because of the key iterator
           }
       }

       return tmp;
    }
}

class Map implements IMap {
   java.util.Map<Object, Object> contents = new java.util.HashMap<>();

   public Map() { }

   private Map(java.util.Map<Object, Object> contents) {
      this.contents = contents;   
   }

   @Override
   public @Nullable Object get(Object key) {
     return contents.get(key);
   }

   @Override
   public IMap empty() {
       return new Map();
   }

   @Override
   public IMap put(Object key, Object value) {
      java.util.Map<Object, Object> newContents = new java.util.HashMap<>();
      newContents.putAll(contents);
      newContents.put(key, value);
      return new Map(newContents);
   }

   @Override
   public java.util.Iterator<Object> iterator() {
      return contents.keySet().iterator();
   }
}

Solution

The Nullness Checker is warning you that the specifications (the type annotations) are inconsistent with the code itself.

Nullness problem

The key problem with your code is here:

tmp.put(k, get(k))

and the error message is:

error: [argument.type.incompatible] incompatible types in argument.
               tmp.put(k, get(k)); // get(k) is always non-null because of the key iterator
                             ^
  found   : @Initialized @Nullable Object
  required: @Initialized @NonNull Object

Here are the two specifications that are incompatible:

put requires a non-null second argument (recall that @NonNull is the default):

   public IMap put(Object key, Object value) { ... }

get might return null at any time, and clients have no way to know when the return value might be non-null:

   @Nullable Object get(Object o);

If you want to state that the return value of a method is nullable in general, but is non-null in certain situations, then you need to use a conditional postcondition such as @EnsuresNonNullIf.

That said, the Nullness Checker has special handling for Map.get. Your code doesn't use that, because you don't have a method that overrides java.util.Map.get (though it does have a class named Map that has nothing to do with java.util.Map).

If you want special-case handling for IMap.get, then either:

your class should extend java.util.Map, or
you should extend the Nullness Checker to recognize your class.

Map key problem

could you provide pointers where to start or examples to learn from?

I suggest starting with the Checker Framework Manual. It has lots of explanations and examples. You should read at least the Map Key Checker chapter. It links to further documentation, such as Javadoc for @KeyFor.

I tried haphazardly to sprinkle some @KeyFors here and there, but with a lack of fully understanding what I'm doing it could take a while before I hit the right spots ;-)

Please don't do that! That way lies suffering. The manual tells you not to do that; instead, think first and write specifications that describe your code.

Here are three @KeyFor annotations that you culd write:

interface IMap extends Iterable<@KeyFor("this") Object> {
...
    default IMap remove(@KeyFor("this") Object key) {
...
    @SuppressWarnings("keyfor") // a key for `contents` is a key for this object
    public java.util.Iterator<@KeyFor("this") Object> iterator() {

These annotations state, respectively:

The iterator returns keys for this object.
Clients must pass a key for this object.
The iterator returns keys for this object. I suppressed a warning because this object acts as a wrapper for a contained object, and I don't recall that the Checker Framework has a way to say, "This object is a wrapper around a field and each of its methods has the same properties as the methods of that field."

The result type-checks without problems (except the nullness one noted in the first section of this answer):

import org.checkerframework.checker.nullness.qual.KeyFor;
import org.checkerframework.checker.nullness.qual.NonNull;
import org.checkerframework.checker.nullness.qual.Nullable;

interface IMap extends Iterable<@KeyFor("this") Object> {
    @Nullable Object get(Object o);
    IMap put(Object key, Object value); // immutable put
    IMap empty();

    default IMap remove(@KeyFor("this") Object key) {
        IMap tmp = empty();

        for (Object k : this) {
            if (!k.equals(key)) {
                tmp.put(k, get(k)); // get(k) is always non-null because of the key iterator
            }
        }

        return tmp;
    }
}

class Map implements IMap {
    java.util.Map<Object, Object> contents = new java.util.HashMap<>();

    public Map() {}

    private Map(java.util.Map<Object, Object> contents) {
        this.contents = contents;
    }

    @Override
    public @Nullable Object get(Object key) {
        return contents.get(key);
    }

    @Override
    public IMap empty() {
        return new Map();
    }

    @Override
    public IMap put(Object key, Object value) {
        java.util.Map<Object, Object> newContents = new java.util.HashMap<>();
        newContents.putAll(contents);
        newContents.put(key, value);
        return new Map(newContents);
    }

    @Override
    @SuppressWarnings("keyfor") // a key for `contents` is a key for this object
    public java.util.Iterator<@KeyFor("this") Object> iterator() {
        return contents.keySet().iterator();
    }
}