Add back the not null checks in a better place.