diff --git a/lib/StaticAnalyzer/Checkers/IvarInvalidationChecker.cpp b/lib/StaticAnalyzer/Checkers/IvarInvalidationChecker.cpp
index 80cb58d..527470a 100644
--- a/lib/StaticAnalyzer/Checkers/IvarInvalidationChecker.cpp
+++ b/lib/StaticAnalyzer/Checkers/IvarInvalidationChecker.cpp
@@ -227,15 +227,23 @@
   }
 
   // If interface, check all parent protocols and super.
-  // TODO: Visit all categories in case the invalidation method is declared in
-  // a category.
-  if (const ObjCInterfaceDecl *InterfaceD = dyn_cast<ObjCInterfaceDecl>(D)) {
+  if (const ObjCInterfaceDecl *InterfD = dyn_cast<ObjCInterfaceDecl>(D)) {
+
+    // Visit all protocols.
     for (ObjCInterfaceDecl::protocol_iterator
-        I = InterfaceD->protocol_begin(),
-        E = InterfaceD->protocol_end(); I != E; ++I) {
+        I = InterfD->protocol_begin(),
+        E = InterfD->protocol_end(); I != E; ++I) {
       containsInvalidationMethod(*I, OutInfo);
     }
-    containsInvalidationMethod(InterfaceD->getSuperClass(), OutInfo);
+
+    // Visit all categories in case the invalidation method is declared in
+    // a category.
+    for (const ObjCCategoryDecl *I = InterfD->getFirstClassExtension(); I;
+        I = I->getNextClassExtension()) {
+      containsInvalidationMethod(I, OutInfo);
+    }
+
+    containsInvalidationMethod(InterfD->getSuperClass(), OutInfo);
     return;
   }
 
@@ -249,7 +257,7 @@
     return;
   }
 
-  llvm_unreachable("One of the casts above should have succeeded.");
+  return;
 }
 
 bool IvarInvalidationChecker::trackIvar(const ObjCIvarDecl *Iv,
diff --git a/test/Analysis/objc_invalidation.m b/test/Analysis/objc_invalidation.m
index f21bfdf..15bcabe 100644
--- a/test/Analysis/objc_invalidation.m
+++ b/test/Analysis/objc_invalidation.m
@@ -41,6 +41,13 @@
 @interface Invalidation1Class <Invalidation1>
 @end
 
+@interface ClassWithInvalidationMethodInCategory <NSObject>
+@end
+
+@interface ClassWithInvalidationMethodInCategory ()
+- (void) invalidate __attribute__((annotate("objc_instance_variable_invalidator")));
+@end
+
 @interface SomeInvalidationImplementingObject: NSObject <Invalidation3, Invalidation2> {
   SomeInvalidationImplementingObject *ObjA; // invalidation in the parent
 }
@@ -107,6 +114,7 @@
 @implementation SomeSubclassInvalidatableObject{
   @private
   SomeInvalidationImplementingObject *Ivar5;
+  ClassWithInvalidationMethodInCategory *Ivar13;
 }
 
 @synthesize Prop7 = _propIvar;
@@ -156,6 +164,7 @@
  // expected-warning@-5 {{Instance variable _Ivar3 needs to be invalidated}}
  // expected-warning@-6 {{Instance variable _Ivar4 needs to be invalidated}}
  // expected-warning@-7 {{Instance variable Ivar5 needs to be invalidated or set to nil}}
+// expected-warning@-8 {{Instance variable Ivar13 needs to be invalidated or set to nil}}
 @end
 
 // Example, where the same property is inherited through