1 files changed, 45 insertions, 0 deletions
diff --git a/lib/StaticAnalyzer/Checkers/BasicObjCFoundationChecks.cpp b/lib/StaticAnalyzer/Checkers/BasicObjCFoundationChecks.cpp
index 955e79ae46..3dda6d04ca 100644
--- a/lib/StaticAnalyzer/Checkers/BasicObjCFoundationChecks.cpp
+++ b/lib/StaticAnalyzer/Checkers/BasicObjCFoundationChecks.cpp
@@ -716,6 +716,47 @@ void ObjCLoopChecker::checkPostStmt(const ObjCForCollectionStmt *FCS,
   C.addTransition(State);
 }
 
+namespace {
+/// \class ObjCNonNilReturnValueChecker
+/// \brief The checker restricts the return values of APIs known to never
+/// return nil.
+class ObjCNonNilReturnValueChecker
+  : public Checker<check::PostObjCMessage> {
+    mutable bool Initialized;
+    mutable Selector ObjectAtIndex;
+    mutable Selector ObjectAtIndexedSubscript;
+public:
+  void checkPostObjCMessage(const ObjCMethodCall &M, CheckerContext &C) const;
+};
+}
+
+void ObjCNonNilReturnValueChecker::checkPostObjCMessage(const ObjCMethodCall &M,
+                                                        CheckerContext &C)
+                                                         const {
+  ProgramStateRef State = C.getState();
+
+  if (!Initialized) {
+    ASTContext &Ctx = C.getASTContext();
+    ObjectAtIndex = GetUnarySelector("objectAtIndex", Ctx);
+    ObjectAtIndexedSubscript = GetUnarySelector("objectAtIndexedSubscript", Ctx);
+  }
+
+  // Check the receiver type.
+  if (const ObjCInterfaceDecl *Interface = M.getReceiverInterface()) {
+    FoundationClass Cl = findKnownClass(Interface);
+    if (Cl == FC_NSArray || Cl == FC_NSOrderedSet) {
+      Selector Sel = M.getSelector();
+      if (Sel == ObjectAtIndex || Sel == ObjectAtIndexedSubscript) {
+        // Go ahead and assume the value is non-nil.
+        SVal Val = State->getSVal(M.getOriginExpr(), C.getLocationContext());
+        if (!isa<DefinedOrUnknownSVal>(Val))
+          return;
+        State = State->assume(cast<DefinedOrUnknownSVal>(Val), true);
+        C.addTransition(State);
+      }
+    }
+  }
+}
 
 //===----------------------------------------------------------------------===//
 // Check registration.
@@ -744,3 +785,7 @@ void ento::registerVariadicMethodTypeChecker(CheckerManager &mgr) {
 void ento::registerObjCLoopChecker(CheckerManager &mgr) {
   mgr.registerChecker<ObjCLoopChecker>();
 }
+
+void ento::registerObjCNonNilReturnValueChecker(CheckerManager &mgr) {
+  mgr.registerChecker<ObjCNonNilReturnValueChecker>();
+}