[CodeCompletion] Hide literals that don't match the type context in SourceKit

If there is a type context, hide literal suggesetions that don't match
it, unless they are keywords and we have filtered to their names.

Incidentally fix an output buffering issue when combining filtering with
the -raw flag in complete-test.

Part of rdar://problem/23865118

Author: benlangmuir

include/swift/IDE/CodeCompletion.h

@@ -693,6 +693,7 @@ class CodeCompletionContext {
 public:
   CodeCompletionCache &Cache;
   CompletionKind CodeCompletionKind = CompletionKind::None;
+  bool HasExpectedTypeRelation = false;
 
   CodeCompletionContext(CodeCompletionCache &Cache)
       : Cache(Cache) {}

lib/IDE/CodeCompletion.cpp

@@ -1546,9 +1546,14 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
   }
 
   void setExpectedTypes(ArrayRef<Type> Types) {
-    ExpectedTypes = Types;
+    ExpectedTypes.reserve(Types.size());
+    for (auto T : Types)
+      if (T)
+        ExpectedTypes.push_back(T);
   }
 
+  bool hasExpectedTypes() const { return !ExpectedTypes.empty(); }
+
   bool needDot() const {
     return NeedLeadingDot;
   }
@@ -4777,6 +4782,8 @@ void CodeCompletionCallbacksImpl::doneParsing() {
     Lookup.RequestedCachedResults.reset();
   }
 
+  CompletionContext.HasExpectedTypeRelation = Lookup.hasExpectedTypes();
+
   deliverCompletionResults();
 }
 

test/IDE/complete_value_literals.swift

@@ -1,3 +1,4 @@
+// RUN: %target-swift-ide-test -code-completion -source-filename=%s -code-completion-token=NO_CONTEXT_0 | FileCheck %s -check-prefix=NO_CONTEXT_0
 // RUN: %target-swift-ide-test -code-completion -source-filename=%s -code-completion-token=NIL_0 | FileCheck %s -check-prefix=NIL_0
 // RUN: %target-swift-ide-test -code-completion -source-filename=%s -code-completion-token=NIL_1 | FileCheck %s -check-prefix=NIL_1
 // RUN: %target-swift-ide-test -code-completion -source-filename=%s -code-completion-token=NIL_2 | FileCheck %s -check-prefix=NIL_2
@@ -25,6 +26,22 @@
 // RUN: %target-swift-ide-test -code-completion -source-filename=%s -code-completion-token=IMAGE_0 | FileCheck %s -check-prefix=IMAGE_0
 // RUN: %target-swift-ide-test -code-completion -source-filename=%s -code-completion-token=IMAGE_1 | FileCheck %s -check-prefix=IMAGE_1
 
+func testAll0() {
+  // Not type context.
+  let x = #^NO_CONTEXT_0^#
+// NO_CONTEXT_0-DAG: Begin completions
+// NO_CONTEXT_0-DAG: Literal[Integer]/None:              0[#Int#];
+// NO_CONTEXT_0-DAG: Literal[Boolean]/None:              true[#Bool#];
+// NO_CONTEXT_0-DAG: Literal[Boolean]/None:              false[#Bool#];
+// NO_CONTEXT_0-DAG: Literal[Nil]/None:                  nil;
+// NO_CONTEXT_0-DAG: Literal[String]/None:               "{#(abc)#}"[#String#];
+// NO_CONTEXT_0-DAG: Literal[Array]/None:                [{#(values)#}][#Array#];
+// NO_CONTEXT_0-DAG: Literal[Dictionary]/None:           [{#(key)#}: {#(value)#}][#Dictionary#];
+// NO_CONTEXT_0-DAG: Literal[_Color]/None:               [#Color({#colorLiteralRed: Float#}, {#green: Float#}, {#blue: Float#}, {#alpha: Float#})#];
+// NO_CONTEXT_0-DAG: Literal[_Image]/None:               [#Image({#imageLiteral: String#})#];
+// NO_CONTEXT_0: End completions
+}
+
 struct MyNil1: NilLiteralConvertible {
   init(nilLiteral: ()) {}
 }

test/SourceKit/CodeComplete/complete_literals.swift

@@ -33,7 +33,46 @@ func foo(x: Int) {
 // RUN: %complete-test -tok=EXPR1 %s -raw | FileCheck %s -check-prefix=LITERALS
 // RUN: %complete-test -tok=EXPR2 %s -raw | FileCheck %s -check-prefix=LITERALS
 // RUN: %complete-test -tok=EXPR3 %s -raw | FileCheck %s -check-prefix=LITERALS
-let x = #^EXPR1^#
-x + #^EXPR2^#
+let x1 = #^EXPR1^#
+x1 + #^EXPR2^#
 if #^EXPR3^# { }
-foo(#^EXPR3^#)
+
+// RUN: %complete-test -tok=EXPR4 %s -raw | FileCheck %s -check-prefix=LITERAL_INT
+foo(#^EXPR4^#)
+// LITERAL_INT-NOT: source.lang.swift.literal
+// LITERAL_INT: key.kind: source.lang.swift.literal.integer
+// LITERAL_INT-NOT: source.lang.swift.literal
+
+// RUN: %complete-test -tok=EXPR5 %s -raw | FileCheck %s -check-prefix=LITERAL_TUPLE
+let x2: (String, Int) = #^EXPR5^#
+// LITERAL_TUPLE-NOT: source.lang.swift.literal
+// LITERAL_TUPLE: key.kind: source.lang.swift.literal.tuple
+// LITERAL_TUPLE-NOT: source.lang.swift.literal
+
+// RUN: %complete-test -tok=EXPR6 %s -raw | FileCheck %s -check-prefix=LITERAL_NO_TYPE
+// When there is a type context that doesn't match, we should see no literals
+// except the keywords and they should be prioritized like keywords not
+// literals.
+struct Opaque {}
+let x3: Opaque = #^EXPR6,,colo,ni,tru,fals^#
+// LITERAL_NO_TYPE-LABEL: Results for filterText:  [
+// LITERAL_NO_TYPE-NOT: source.lang.swift.literal
+// LITERAL_NO_TYPE: ]
+// LITERAL_NO_TYPE-LABEL: Results for filterText: colo [
+// LITERAL_NO_TYPE-NOT: source.lang.swift.literal
+// LITERAL_NO_TYPE: ]
+// LITERAL_NO_TYPE-LABEL: Results for filterText: ni [
+// LITERAL_NO_TYPE-NOT: source.lang.swift.literal
+// LITERAL_NO_TYPE: key.kind: source.lang.swift.literal.nil
+// LITERAL_NO_TYPE-NOT: source.lang.swift.literal
+// LITERAL_NO_TYPE: ]
+// LITERAL_NO_TYPE-LABEL: Results for filterText: tru [
+// LITERAL_NO_TYPE-NOT: source.lang.swift.literal
+// LITERAL_NO_TYPE: key.kind: source.lang.swift.literal.boolean
+// LITERAL_NO_TYPE-NOT: source.lang.swift.literal
+// LITERAL_NO_TYPE: ]
+// LITERAL_NO_TYPE-LABEL: Results for filterText: fals [
+// LITERAL_NO_TYPE-NOT: source.lang.swift.literal
+// LITERAL_NO_TYPE: key.kind: source.lang.swift.literal.boolean
+// LITERAL_NO_TYPE-NOT: source.lang.swift.literal
+// LITERAL_NO_TYPE: ]

tools/SourceKit/lib/SwiftLang/CodeCompletionOrganizer.cpp

@@ -188,12 +188,13 @@ bool SourceKit::CodeCompletion::addCustomCompletions(
 class CodeCompletionOrganizer::Impl {
   std::unique_ptr<Group> rootGroup;
   CompletionKind completionKind;
+  bool completionHasExpectedTypes;
 
   void groupStemsRecursive(Group *group, bool recurseIntoNewGroups,
                            StringRef(getStem)(StringRef));
 
 public:
-  Impl(CompletionKind kind);
+  Impl(CompletionKind kind, bool hasExpectedTypes);
 
   void addCompletionsWithFilter(ArrayRef<Completion *> completions,
                                 StringRef filterText, Options options,
@@ -249,8 +250,9 @@ class CodeCompletionOrganizer::Impl {
 //===----------------------------------------------------------------------===//
 
 CodeCompletionOrganizer::CodeCompletionOrganizer(const Options &options,
-                                                 CompletionKind kind)
-    : impl(*new Impl(kind)), options(options) {}
+                                                 CompletionKind kind,
+                                                 bool hasExpectedTypes)
+    : impl(*new Impl(kind, hasExpectedTypes)), options(options) {}
 CodeCompletionOrganizer::~CodeCompletionOrganizer() { delete &impl; }
 
 void CodeCompletionOrganizer::preSortCompletions(
@@ -383,8 +385,8 @@ static std::unique_ptr<Result> make_result(Completion *result) {
 // CodeCompletionOrganizer::Impl implementation
 //===----------------------------------------------------------------------===//
 
-CodeCompletionOrganizer::Impl::Impl(CompletionKind kind)
-    : completionKind(kind) {
+CodeCompletionOrganizer::Impl::Impl(CompletionKind kind, bool hasExpectedTypes)
+    : completionKind(kind), completionHasExpectedTypes(hasExpectedTypes) {
   assert(!rootGroup && "initialized twice");
   rootGroup = make_group("");
 }
@@ -516,7 +518,8 @@ void CodeCompletionOrganizer::Impl::addCompletionsWithFilter(
         }
         if (completion->getExpectedTypeRelation() >= Completion::Convertible ||
             (completion->getKind() == Completion::Literal &&
-             completionKind != CompletionKind::StmtOrExpr))
+             completionKind != CompletionKind::StmtOrExpr &&
+             !completionHasExpectedTypes))
           break;
 
         if (completion->getKind() == Completion::Keyword &&
@@ -543,6 +546,16 @@ void CodeCompletionOrganizer::Impl::addCompletionsWithFilter(
     if (rules.hideCompletion(completion))
       continue;
 
+    // Hide literals other than the ones that are also keywords if they don't
+    // match the expected types.
+    if (completion->getKind() == Completion::Literal &&
+        completionHasExpectedTypes &&
+        completion->getExpectedTypeRelation() < Completion::Convertible &&
+        completion->getLiteralKind() !=
+            CodeCompletionLiteralKind::BooleanLiteral &&
+        completion->getLiteralKind() != CodeCompletionLiteralKind::NilLiteral)
+      continue;
+
     bool match = false;
     if (options.fuzzyMatching && filterText.size() >= options.minFuzzyLength) {
       match = pattern.matchesCandidate(completion->getName());
@@ -641,7 +654,7 @@ enum class ResultBucket {
 };
 } // end anonymous namespace
 
-static ResultBucket getResultBucket(Item &item) {
+static ResultBucket getResultBucket(Item &item, bool hasExpectedTypes) {
   if (isa<Group>(item))
     return ResultBucket::Normal; // FIXME: take best contained result.
   auto *completion = cast<Result>(item).value;
@@ -655,7 +668,15 @@ static ResultBucket getResultBucket(Item &item) {
 
   switch (completion->getKind()) {
   case Completion::Literal:
-    return matchesType ? ResultBucket::LiteralTypeMatch : ResultBucket::Literal;
+    if (matchesType) {
+      return ResultBucket::LiteralTypeMatch;
+    } else if (!hasExpectedTypes) {
+      return ResultBucket::Literal;
+    } else {
+      // When we have type context, we still show literals that are keywords,
+      // but we treat them as keywords instead of literals for prioritization.
+      return ResultBucket::Normal;
+    }
   case Completion::Keyword:
     return isHighPriorityKeyword(completion->getKeywordKind())
                ? ResultBucket::HighPriorityKeyword
@@ -721,13 +742,14 @@ static int compareLiterals(Item &a_, Item &b_) {
   return 0;
 }
 
-static void sortRecursive(const Options &options, Group *group) {
+static void sortRecursive(const Options &options, Group *group,
+                          bool hasExpectedTypes) {
   // Sort all of the subgroups first, and fill in the bucket for each result.
   auto &contents = group->contents;
   double best = -1.0;
   for (auto &item : contents) {
     if (Group *g = dyn_cast<Group>(item.get())) {
-      sortRecursive(options, g);
+      sortRecursive(options, g, hasExpectedTypes);
     } else {
       Result *r = cast<Result>(item.get());
       item->finalScore = combinedScore(options, item->matchScore, r->value);
@@ -749,40 +771,40 @@ static void sortRecursive(const Options &options, Group *group) {
     return;
   }
 
-  llvm::array_pod_sort(contents.begin(), contents.end(), [](const std::unique_ptr<Item> *a_, const std::unique_ptr<Item> *b_) {
-    Item &a = **a_;
-    Item &b = **b_;
+  std::sort(contents.begin(), contents.end(), [=](const std::unique_ptr<Item> &a_, const std::unique_ptr<Item> &b_) {
+    Item &a = *a_;
+    Item &b = *b_;
 
-    auto bucketA = getResultBucket(a);
-    auto bucketB = getResultBucket(b);
+    auto bucketA = getResultBucket(a, hasExpectedTypes);
+    auto bucketB = getResultBucket(b, hasExpectedTypes);
     if (bucketA < bucketB)
-      return 1;
+      return false;
     else if (bucketB < bucketA)
-      return -1;
+      return true;
 
     // Special internal orderings.
     switch (bucketA) {
     case ResultBucket::HighPriorityKeyword:
-      return compareHighPriorityKeywords(a, b);
+      return compareHighPriorityKeywords(a, b) < 0;
     case ResultBucket::Literal:
     case ResultBucket::LiteralTypeMatch:
-      return compareLiterals(a, b);
+      return compareLiterals(a, b) < 0;
     default:
       break;
     }
 
     // "Normal" order.
     if (a.finalScore < b.finalScore)
-      return 1;
+      return false;
     else if (b.finalScore < a.finalScore)
-      return -1;
+      return true;
 
-    return compareResultName(a, b);
+    return compareResultName(a, b) < 0;
   });
 }
 
 void CodeCompletionOrganizer::Impl::sort(Options options) {
-  sortRecursive(options, rootGroup.get());
+  sortRecursive(options, rootGroup.get(), completionHasExpectedTypes);
 }
 
 void CodeCompletionOrganizer::Impl::groupStemsRecursive(

tools/SourceKit/lib/SwiftLang/CodeCompletionOrganizer.h

@@ -75,7 +75,8 @@ class CodeCompletionOrganizer {
   Impl &impl;
   const Options &options;
 public:
-  CodeCompletionOrganizer(const Options &options, CompletionKind kind);
+  CodeCompletionOrganizer(const Options &options, CompletionKind kind,
+                          bool hasExpectedTypes);
   ~CodeCompletionOrganizer();
 
   static void

tools/SourceKit/lib/SwiftLang/SwiftCompletion.cpp

@@ -697,6 +697,10 @@ CompletionKind CodeCompletion::SessionCache::getCompletionKind() {
   llvm::sys::ScopedLock L(mtx);
   return completionKind;
 }
+bool CodeCompletion::SessionCache::getCompletionHasExpectedTypes() {
+  llvm::sys::ScopedLock L(mtx);
+  return completionHasExpectedTypes;
+}
 const CodeCompletion::FilterRules &
 CodeCompletion::SessionCache::getFilterRules() {
   llvm::sys::ScopedLock L(mtx);
@@ -928,7 +932,8 @@ static void transformAndForwardResults(
   };
 
   CodeCompletion::CodeCompletionOrganizer organizer(
-      options, session->getCompletionKind());
+      options, session->getCompletionKind(),
+      session->getCompletionHasExpectedTypes());
 
   bool hasEarlyInnerResults =
       session->getCompletionKind() == CompletionKind::PostfixExpr;
@@ -1060,10 +1065,12 @@ void SwiftLangSupport::codeCompleteOpen(
   }
 
   CompletionKind completionKind = CompletionKind::None;
+  bool hasExpectedTypes = false;
 
   SwiftCodeCompletionConsumer swiftConsumer(
       [&](ArrayRef<CodeCompletionResult *> results, SwiftCompletionInfo &info) {
         completionKind = info.completionContext->CodeCompletionKind;
+        hasExpectedTypes = info.completionContext->HasExpectedTypeRelation;
         completions =
             extendCompletions(results, sink, info, nameToPopularity, CCOpts);
       });
@@ -1097,7 +1104,7 @@ void SwiftLangSupport::codeCompleteOpen(
   std::vector<std::string> argsCopy(extendedArgs.begin(), extendedArgs.end());
   SessionCacheRef session{new SessionCache(
       std::move(sink), std::move(bufferCopy), std::move(argsCopy),
-      completionKind, std::move(filterRules))};
+      completionKind, hasExpectedTypes, std::move(filterRules))};
   session->setSortedCompletions(std::move(completions));
 
   if (!CCSessions.set(name, offset, session)) {

tools/SourceKit/lib/SwiftLang/SwiftLangSupport.h

@@ -137,22 +137,26 @@ class SessionCache : public ThreadSafeRefCountedBase<SessionCache> {
   CompletionSink sink;
   std::vector<Completion *> sortedCompletions;
   CompletionKind completionKind;
+  bool completionHasExpectedTypes;
   FilterRules filterRules;
   llvm::sys::Mutex mtx;
 
 public:
   SessionCache(CompletionSink &&sink,
                std::unique_ptr<llvm::MemoryBuffer> &&buffer,
                std::vector<std::string> &&args, CompletionKind completionKind,
-               FilterRules filterRules)
+               bool hasExpectedTypes, FilterRules filterRules)
       : buffer(std::move(buffer)), args(std::move(args)), sink(std::move(sink)),
-        completionKind(completionKind), filterRules(std::move(filterRules)) {}
+        completionKind(completionKind),
+        completionHasExpectedTypes(hasExpectedTypes),
+        filterRules(std::move(filterRules)) {}
   void setSortedCompletions(std::vector<Completion *> &&completions);
   ArrayRef<Completion *> getSortedCompletions();
   llvm::MemoryBuffer *getBuffer();
   ArrayRef<std::string> getCompilerArgs();
   const FilterRules &getFilterRules();
   CompletionKind getCompletionKind();
+  bool getCompletionHasExpectedTypes();
 };
 typedef RefPtr<SessionCache> SessionCacheRef;
 

tools/SourceKit/tools/complete-test/complete-test.cpp

@@ -537,6 +537,7 @@ static void printResponse(sourcekitd_response_t resp, bool raw, bool structure,
   }
   ResponsePrinter p(llvm::outs(), 4, indentation, structure);
   p.printResponse(resp);
+  llvm::outs().flush();
 }
 
 static std::unique_ptr<llvm::MemoryBuffer>
@@ -770,9 +771,11 @@ static int handleTestInvocation(TestOptions &options) {
             return true;
           }
           llvm::outs() << "Results for filterText: " << prefix << " [\n";
+          llvm::outs().flush();
           printResponse(response, options.rawOutput, options.structureOutput,
                         /*indentation*/ 4);
           llvm::outs() << "]\n";
+          llvm::outs().flush();
           return false;
         });
     if (isError)