extend cases of unnecessary_null_checks/null_check_on_nullable_type_parameter

lib/src/rules/unnecessary_null_checks.dart

@@ -40,6 +40,10 @@ DartType? getExpectedType(PostfixExpression node) {
   var realNode =
       node.thisOrAncestorMatching((e) => e.parent is! ParenthesizedExpression);
   var parent = realNode?.parent;
+  var withAwait = parent is AwaitExpression;
+  if (withAwait) {
+    parent = parent!.parent;
+  }
 
   // in return value
   if (parent is ReturnStatement || parent is ExpressionFunctionBody) {
@@ -48,7 +52,27 @@ DartType? getExpectedType(PostfixExpression node) {
       return null;
     }
     var staticType = parentExpression.staticType;
-    return staticType is FunctionType ? staticType.returnType : null;
+    if (staticType is! FunctionType) {
+      return null;
+    }
+    if (withAwait || parentExpression.body.keyword?.lexeme == 'async') {
+      //  return type is obligatorily a Future<X>
+      return (staticType.returnType as ParameterizedType?)?.typeArguments.first;
+    } else {
+      return staticType.returnType;
+    }
+  }
+  // in yield value
+  if (parent is YieldStatement) {
+    var parentExpression = parent.thisOrAncestorOfType<FunctionExpression>();
+    if (parentExpression == null) {
+      return null;
+    }
+    var staticType = parentExpression.staticType;
+    if (staticType is! FunctionType) {
+      return null;
+    }
+    return (staticType.returnType as ParameterizedType).typeArguments.first;
   }
   // assignment
   if (parent is AssignmentExpression &&
@@ -71,6 +95,20 @@ DartType? getExpectedType(PostfixExpression node) {
     }
     return parentElement.parameters.first.type;
   }
+  // as member of list
+  if (parent is ListLiteral) {
+    return (parent.staticType as ParameterizedType?)?.typeArguments.first;
+  }
+  // as member of set
+  if (parent is SetOrMapLiteral && parent.isSet) {
+    return (parent.staticType as ParameterizedType?)?.typeArguments.first;
+  }
+  // as member of map
+  if (parent is MapLiteralEntry) {
+    var typeParameters =
+        (parent.parent! as SetOrMapLiteral).staticType as ParameterizedType?;
+    return typeParameters?.typeArguments[parent.key == node ? 0 : 1];
+  }
   // as parameter of function
   if (parent is NamedExpression) {
     realNode = parent;

test_data/rules/experiments/nnbd/rules/null_check_on_nullable_type_parameter.dart

@@ -21,6 +21,14 @@ T m30<T>(T? p) {
   t = p!; // LINT
   return t;
 }
+Future<T> m40<T extends Object?>(T? p) async => await p!; // LINT
+Future<List<T>> m41<T extends Object?>(T? p) async => await [p!]; // LINT
+List<T> m50<T>(T? p) => [p!]; // LINT
+Set<T> m60<T>(T? p) => {p!}; // LINT
+Map<String, T> m71<T>(T? p) => {'': p!}; // LINT
+Map<T, String> m72<T>(T? p) => {p!: ''}; // LINT
+Iterable<T> m80<T>(T? p) sync* {yield p!;} // LINT
+Stream<T> m90<T>(T? p) async* {yield p!;} // LINT
 class C<T> {
   late T t;
   m(T? p) {

test_data/rules/experiments/nnbd/rules/unnecessary_null_checks.dart

@@ -60,3 +60,16 @@ f5(int? p) {
   v1 ??= p!; // OK
 }
 
+Future<int?> f6(int? p) async => await p!; // LINT
+List<int?> f7(int? p) => [p!]; // LINT
+Set<int?> f8(int? p) => {p!}; // LINT
+Map<int?, String> f9(int? p) => {p!: ''}; // LINT
+Map<String, int?> f10(int? p) => {'': p!}; // LINT
+Iterable<int?> f11(int? p) sync* {yield p!;} // LINT
+Stream<int?> f12(int? p) async* {yield p!;} // LINT
+Future<void> f13(int? p) async {
+  var f = Future(() => p);
+  int? i;
+  i = await f!; // LINT
+}
+Future<int?> f14(int? p) async => p!; // LINT