2 files changed, 60 insertions, 20 deletions

diff --git a/clang/lib/StaticAnalyzer/Checkers/NullabilityChecker.cpp b/clang/lib/StaticAnalyzer/Checkers/NullabilityChecker.cpp
index bb86ea401df..01c7287c97f 100644
--- a/clang/lib/StaticAnalyzer/Checkers/NullabilityChecker.cpp
+++ b/clang/lib/StaticAnalyzer/Checkers/NullabilityChecker.cpp
@@ -366,24 +366,20 @@ static bool checkPreconditionViolation(ProgramStateRef State, ExplodedNode *N,
   if (!D)
     return false;
 
-  if (const auto *BlockD = dyn_cast<BlockDecl>(D)) {
-    if (checkParamsForPreconditionViolation(BlockD->parameters(), State,
-                                            LocCtxt)) {
-      if (!N->isSink())
-        C.addTransition(State->set<PreconditionViolated>(true), N);
-      return true;
-    }
+  ArrayRef<ParmVarDecl*> Params;
+  if (const auto *BD = dyn_cast<BlockDecl>(D))
+    Params = BD->parameters();
+  else if (const auto *FD = dyn_cast<FunctionDecl>(D))
+    Params = FD->parameters();
+  else if (const auto *MD = dyn_cast<ObjCMethodDecl>(D))
+    Params = MD->parameters();
+  else
     return false;
-  }
 
-  if (const auto *FuncDecl = dyn_cast<FunctionDecl>(D)) {
-    if (checkParamsForPreconditionViolation(FuncDecl->parameters(), State,
-                                            LocCtxt)) {
-      if (!N->isSink())
-        C.addTransition(State->set<PreconditionViolated>(true), N);
-      return true;
-    }
-    return false;
+  if (checkParamsForPreconditionViolation(Params, State, LocCtxt)) {
+    if (!N->isSink())
+      C.addTransition(State->set<PreconditionViolated>(true), N);
+    return true;
   }
   return false;
 }
@@ -484,16 +480,20 @@ void NullabilityChecker::checkPreStmt(const ReturnStmt *S,
   if (!RetSVal)
     return;
 
+  QualType RequiredRetType;
   AnalysisDeclContext *DeclCtxt =
       C.getLocationContext()->getAnalysisDeclContext();
-  const FunctionType *FuncType = DeclCtxt->getDecl()->getFunctionType();
-  if (!FuncType)
+  const Decl *D = DeclCtxt->getDecl();
+  if (auto *MD = dyn_cast<ObjCMethodDecl>(D))
+    RequiredRetType = MD->getReturnType();
+  else if (auto *FD = dyn_cast<FunctionDecl>(D))
+    RequiredRetType = FD->getReturnType();
+  else
     return;
 
   NullConstraint Nullness = getNullConstraint(*RetSVal, State);
 
-  Nullability RequiredNullability =
-      getNullabilityAnnotation(FuncType->getReturnType());
+  Nullability RequiredNullability = getNullabilityAnnotation(RequiredRetType);
 
   // If the returned value is null but the type of the expression
   // generating it is nonnull then we will suppress the diagnostic.
diff --git a/clang/test/Analysis/nullability_nullonly.mm b/clang/test/Analysis/nullability_nullonly.mm
index 56b3f9e1449..6479d67bda6 100644
--- a/clang/test/Analysis/nullability_nullonly.mm
+++ b/clang/test/Analysis/nullability_nullonly.mm
@@ -1,5 +1,21 @@
 // RUN: %clang_cc1 -analyze -analyzer-checker=core,nullability.NullPassedToNonnull,nullability.NullReturnedFromNonnull -verify %s
 
+#define nil 0
+#define BOOL int
+
+@protocol NSObject
++ (id)alloc;
+- (id)init;
+@end
+
+@protocol NSCopying
+@end
+
+__attribute__((objc_root_class))
+@interface
+NSObject<NSObject>
+@end
+
 int getRandom();
 
 typedef struct Dummy { int val; } Dummy;
@@ -85,3 +101,27 @@ Dummy *_Nonnull testDefensiveInlineChecks(Dummy * p) {
     takesNonnull(p);
   return p;
 }
+
+
+@interface SomeClass : NSObject
+@end
+
+@implementation SomeClass (MethodReturn)
+- (SomeClass * _Nonnull)testReturnsNilInNonnull {
+  SomeClass *local = nil;
+  return local; // expected-warning {{Null is returned from a function that is expected to return a non-null value}}
+}
+
+- (SomeClass * _Nonnull)testReturnsCastSuppressedNilInNonnull {
+  SomeClass *local = nil;
+  return (SomeClass * _Nonnull)local; // no-warning
+}
+
+- (SomeClass * _Nonnull)testReturnsNilInNonnullWhenPreconditionViolated:(SomeClass * _Nonnull) p {
+  SomeClass *local = nil;
+  if (!p) // Pre-condition violated here.
+    return local; // no-warning
+  else
+    return p; // no-warning
+}
+@end