1 files changed, 101 insertions, 1 deletions

diff --git a/lib/AST/NSAPI.cpp b/lib/AST/NSAPI.cpp
index 3dc750a81787..2749100e14a2 100644
--- a/lib/AST/NSAPI.cpp
+++ b/lib/AST/NSAPI.cpp
@@ -27,7 +27,10 @@ IdentifierInfo *NSAPI::getNSClassId(NSClassIdKindKind K) const {
     "NSMutableArray",
     "NSDictionary",
     "NSMutableDictionary",
-    "NSNumber"
+    "NSNumber",
+    "NSMutableSet",
+    "NSCountedSet",
+    "NSMutableOrderedSet"
   };
 
   if (!ClassIds[K])
@@ -124,6 +127,25 @@ Selector NSAPI::getNSArraySelector(NSArrayMethodKind MK) const {
       Sel = Ctx.Selectors.getSelector(2, KeyIdents);
       break;
     }
+    case NSMutableArr_addObject:
+      Sel = Ctx.Selectors.getUnarySelector(&Ctx.Idents.get("addObject"));
+      break;
+    case NSMutableArr_insertObjectAtIndex: {
+      IdentifierInfo *KeyIdents[] = {
+        &Ctx.Idents.get("insertObject"),
+        &Ctx.Idents.get("atIndex")
+      };
+      Sel = Ctx.Selectors.getSelector(2, KeyIdents);
+      break;
+    }
+    case NSMutableArr_setObjectAtIndexedSubscript: {
+      IdentifierInfo *KeyIdents[] = {
+        &Ctx.Idents.get("setObject"),
+        &Ctx.Idents.get("atIndexedSubscript")
+      };
+      Sel = Ctx.Selectors.getSelector(2, KeyIdents);
+      break;
+    }
     }
     return (NSArraySelectors[MK] = Sel);
   }
@@ -209,6 +231,22 @@ Selector NSAPI::getNSDictionarySelector(
       Sel = Ctx.Selectors.getSelector(2, KeyIdents);
       break;
     }
+    case NSMutableDict_setObjectForKeyedSubscript: {
+      IdentifierInfo *KeyIdents[] = {
+        &Ctx.Idents.get("setObject"),
+        &Ctx.Idents.get("forKeyedSubscript")
+      };
+      Sel = Ctx.Selectors.getSelector(2, KeyIdents);
+      break;
+    }
+    case NSMutableDict_setValueForKey: {
+      IdentifierInfo *KeyIdents[] = {
+        &Ctx.Idents.get("setValue"),
+        &Ctx.Idents.get("forKey")
+      };
+      Sel = Ctx.Selectors.getSelector(2, KeyIdents);
+      break;
+    }
     }
     return (NSDictionarySelectors[MK] = Sel);
   }
@@ -227,6 +265,63 @@ NSAPI::getNSDictionaryMethodKind(Selector Sel) {
   return None;
 }
 
+Selector NSAPI::getNSSetSelector(NSSetMethodKind MK) const {
+  if (NSSetSelectors[MK].isNull()) {
+    Selector Sel;
+    switch (MK) {
+    case NSMutableSet_addObject:
+      Sel = Ctx.Selectors.getUnarySelector(&Ctx.Idents.get("addObject"));
+      break;
+    case NSOrderedSet_insertObjectAtIndex: {
+      IdentifierInfo *KeyIdents[] = {
+        &Ctx.Idents.get("insertObject"),
+        &Ctx.Idents.get("atIndex")
+      };
+      Sel = Ctx.Selectors.getSelector(2, KeyIdents);
+      break;
+    }
+    case NSOrderedSet_setObjectAtIndex: {
+      IdentifierInfo *KeyIdents[] = {
+        &Ctx.Idents.get("setObject"),
+        &Ctx.Idents.get("atIndex")
+      };
+      Sel = Ctx.Selectors.getSelector(2, KeyIdents);
+      break;
+    }
+    case NSOrderedSet_setObjectAtIndexedSubscript: {
+      IdentifierInfo *KeyIdents[] = {
+        &Ctx.Idents.get("setObject"),
+        &Ctx.Idents.get("atIndexedSubscript")
+      };
+      Sel = Ctx.Selectors.getSelector(2, KeyIdents);
+      break;
+    }
+    case NSOrderedSet_replaceObjectAtIndexWithObject: {
+      IdentifierInfo *KeyIdents[] = {
+        &Ctx.Idents.get("replaceObjectAtIndex"),
+        &Ctx.Idents.get("withObject")
+      };
+      Sel = Ctx.Selectors.getSelector(2, KeyIdents);
+      break;
+    }
+    }
+    return (NSSetSelectors[MK] = Sel);
+  }
+
+  return NSSetSelectors[MK];
+}
+
+Optional<NSAPI::NSSetMethodKind>
+NSAPI::getNSSetMethodKind(Selector Sel) {
+  for (unsigned i = 0; i != NumNSSetMethods; ++i) {
+    NSSetMethodKind MK = NSSetMethodKind(i);
+    if (Sel == getNSSetSelector(MK))
+      return MK;
+  }
+
+  return None;
+}
+
 Selector NSAPI::getNSNumberLiteralSelector(NSNumberLiteralMethodKind MK,
                                            bool Instance) const {
   static const char *ClassSelectorName[NumNSNumberLiteralMethods] = {
@@ -410,6 +505,11 @@ StringRef NSAPI::GetNSIntegralKind(QualType T) const {
   return StringRef();
 }
 
+bool NSAPI::isMacroDefined(StringRef Id) const {
+  // FIXME: Check whether the relevant module macros are visible.
+  return Ctx.Idents.get(Id).hasMacroDefinition();
+}
+
 bool NSAPI::isObjCTypedef(QualType T,
                           StringRef name, IdentifierInfo *&II) const {
   if (!Ctx.getLangOpts().ObjC1)