Add basic typo correction for unqualified lookup.

There's a lot of room for better QoI / performance here.

Author: rjmccall

include/swift/AST/DiagnosticsSema.def

@@ -549,6 +549,8 @@ ERROR(use_undeclared_type,none,
       "use of undeclared type %0", (Identifier))
 ERROR(use_undeclared_type_did_you_mean,none,
       "use of undeclared type %0; did you mean to use '%1'?", (Identifier, StringRef))
+NOTE(note_typo_candidate,none,
+     "did you mean '%0'?", (StringRef))
 NOTE(note_remapped_type,none,
      "did you mean to use '%0'?", (StringRef))
 ERROR(identifier_init_failure,none,

include/swift/AST/NameLookup.h

@@ -158,6 +158,22 @@ class VisibleDeclConsumer {
   virtual void foundDecl(ValueDecl *VD, DeclVisibilityKind Reason) = 0;
 };
 
+/// An implementation of VisibleDeclConsumer that's built from a lambda.
+template <class Fn>
+class LambdaDeclConsumer : public VisibleDeclConsumer {
+  Fn Callback;
+public:
+  LambdaDeclConsumer(Fn &&callback) : Callback(std::move(callback)) {}
+
+  void foundDecl(ValueDecl *VD, DeclVisibilityKind reason) {
+    Callback(VD, reason);
+  }
+};
+template <class Fn>
+LambdaDeclConsumer<Fn> makeDeclConsumer(Fn &&callback) {
+  return LambdaDeclConsumer<Fn>(std::move(callback));
+}
+
 /// A consumer that inserts found decls into an externally-owned SmallVector.
 class VectorDeclConsumer : public VisibleDeclConsumer {
   virtual void anchor() override;

include/swift/Basic/TopCollection.h

@@ -0,0 +1,189 @@
+//===--- TopCollection.h - A size-limiting top-N collection -----*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2016 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See http://swift.org/LICENSE.txt for license information
+// See http://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines the TopCollection class, a data structure which,
+// given a size limit, keeps the best-scoring (i.e. lowest) N values
+// added to it.
+//
+// The current implementation of this is only suited for small values
+// of maxSize.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_BASIC_TOPCOLLECTION_H
+#define SWIFT_BASIC_TOPCOLLECTION_H
+
+#include "swift/Basic/LLVM.h"
+#include "llvm/ADT/SmallVector.h"
+
+namespace swift {
+
+template <class ScoreType, class T, unsigned InlineCapacity = 16>
+class TopCollection {
+public:
+  struct Entry {
+    ScoreType Score;
+    T Value;
+  };
+
+private:
+  SmallVector<Entry, InlineCapacity> Data;
+
+  unsigned MaxSize;
+  unsigned EndOfAccepted = 0;
+public:
+  explicit TopCollection(unsigned maxSize) : MaxSize(maxSize) {
+    assert(maxSize > 0 && "creating collection with a maximum size of 0?");
+    Data.reserve(maxSize);
+  }
+
+  // The invariants work fine with these.
+  TopCollection(const TopCollection &other) = default;
+  TopCollection &operator=(const TopCollection &other) = default;
+
+  // We need to clear EndOfAccepted to preserve the invariants here.
+  TopCollection(TopCollection &&other)
+    : Data(std::move(other.Data)),
+      MaxSize(other.MaxSize),
+      EndOfAccepted(other.EndOfAccepted) {
+    other.EndOfAccepted = 0;
+  }
+  TopCollection &operator=(TopCollection &&other) {
+    Data = std::move(other.Data);
+    MaxSize = other.MaxSize;
+    EndOfAccepted = other.EndOfAccepted;
+    other.EndOfAccepted = 0;
+    return *this;
+  }
+
+  ~TopCollection() {}
+
+  bool empty() const {
+    return EndOfAccepted == 0;
+  }
+
+  size_t size() const {
+    return EndOfAccepted;
+  }
+
+  using iterator = const Entry *;
+  iterator begin() const { return Data.begin(); }
+  iterator end() const { return Data.begin() + EndOfAccepted; }
+
+  /// Return a score beyond which scores are uninteresting.  Inserting
+  /// a value with this score will never change the collection.
+  ScoreType getMinUninterestingScore(ScoreType defaultBound) const {
+    assert(EndOfAccepted <= MaxSize);
+    assert(EndOfAccepted <= Data.size());
+
+    // If we've accepted as many values as we can, then all scores up (and
+    // including) that value are interesting.
+    if (EndOfAccepted == MaxSize)
+      return Data[EndOfAccepted - 1].Score + 1;
+
+    // Otherwise, if there are values in the collection that we've rejected,
+    // any score up to that is still interesting.
+    if (EndOfAccepted != Data.size())
+      return Data[EndOfAccepted].Score;
+
+    // Otherwise, use the default bound.
+    return defaultBound;
+  }
+
+  /// Try to add a scored value to the collection.
+  ///
+  /// \return true if the insertion was successful
+  bool insert(ScoreType score, T &&value) {
+    assert(EndOfAccepted <= MaxSize);
+    assert(EndOfAccepted <= Data.size());
+
+    // Find the index of the last entry whose score is larger than 'score'.
+    auto i = EndOfAccepted;
+    while (i != 0 && score < Data[i - 1].Score)
+      --i;
+
+    assert(0 <= i && i <= EndOfAccepted);
+    assert(i == 0 || score >= Data[i - 1].Score);
+
+    // If we skipped anything, it's definitely safe to insert.
+    if (i != EndOfAccepted) {
+      // fall through
+
+    // If there's a tie with the previous thing, we might have to declare
+    // an existing tier off-bounds.
+    } else if (i != 0 && score == Data[i - 1].Score) {
+      // Only if there isn't space to expand the tier.
+      if (i == MaxSize) {
+        for (--i; i != 0; --i) {
+          if (Data[i - 1].Score != Data[i].Score)
+            break;
+        }
+        EndOfAccepted = i;
+        return false;
+      }
+
+    // Otherwise, there's a non-trivial prefix that the new element has a
+    // strictly higher score than.  We can still insert if there's space.
+    } else {
+      // Don't insert if there's no room.
+      if (i == MaxSize)
+        return false;
+
+      // Don't insert if we're at least as high as things we've previously
+      // rejected.
+      if (i != Data.size() && score >= Data[i].Score)
+        return false;
+    }
+
+    // We don't care about any of the actual values after EndOfAccepted
+    // *except* that we need to remember the minimum value following
+    // EndOfAccepted if that's less than MaxSize so that we continue to
+    // drop values with that score.
+    //
+    // Note that all of the values between EndOfAccepted and MaxSize
+    // should have the same score, beause otherwise there's a tier we
+    // shouldn't have marked dead.
+
+    // Just overwrite the next element instead of inserting if possible.
+    if (i == EndOfAccepted && i != Data.size()) {
+      Data[i].Score = score;
+      Data[i].Value = std::move(value);
+    } else {
+      if (Data.size() == MaxSize)
+        Data.pop_back();
+      Data.insert(Data.begin() + i, { score, std::move(value) });
+    }
+
+    EndOfAccepted++;
+    assert(EndOfAccepted <= Data.size());
+    assert(EndOfAccepted <= MaxSize);
+    return true;
+  }
+
+  /// Drop any values which a score more than the given value from the
+  /// minimum score.
+  template <class Range>
+  void filterMaxScoreRange(Range difference) {
+    if (EndOfAccepted < 2) return;
+    for (unsigned i = 1; i != EndOfAccepted; ++i) {
+      if (Data[i].Score > Data[0].Score + difference) {
+        EndOfAccepted = i;
+        return;
+      }
+    }
+
+  }
+};
+
+} // end namespace swift
+
+#endif // SWIFT_BASIC_CLUSTEREDBITVECTOR_H

lib/Sema/TypeCheckConstraints.cpp

@@ -450,11 +450,28 @@ resolveDeclRefExpr(UnresolvedDeclRefExpr *UDRE, DeclContext *DC) {
     // If we failed lookup of an operator, check to see it to see if it is
     // because two operators are juxtaposed e.g. (x*-4) that needs whitespace.
     // If so, emit specific diagnostics for it.
-    if (!diagnoseOperatorJuxtaposition(UDRE, DC, *this)) {
-      diagnose(Loc, diag::use_unresolved_identifier, Name,
-               UDRE->getName().isOperator())
-        .highlight(UDRE->getSourceRange());
+    if (diagnoseOperatorJuxtaposition(UDRE, DC, *this)) {
+      return new (Context) ErrorExpr(UDRE->getSourceRange());
+    }
+
+    // TODO: Name will be a compound name if it was written explicitly as
+    // one, but we should also try to propagate labels into this.
+    DeclNameLoc nameLoc = UDRE->getNameLoc();
+
+    performTypoCorrection(DC, UDRE->getRefKind(), Name, Loc, LookupOptions,
+                          Lookup);
+
+    diagnose(Loc, diag::use_unresolved_identifier, Name, Name.isOperator())
+      .highlight(UDRE->getSourceRange());
+
+    // Note all the correction candidates.
+    for (auto &result : Lookup) {
+      noteTypoCorrection(Name, nameLoc, result);
     }
+
+    // TODO: consider recovering from here.  We may want some way to suppress
+    // downstream diagnostics, though.
+
     return new (Context) ErrorExpr(UDRE->getSourceRange());
   }
 

lib/Sema/TypeCheckNameLookup.cpp

@@ -17,6 +17,7 @@
 //===----------------------------------------------------------------------===//
 #include "TypeChecker.h"
 #include "swift/AST/NameLookup.h"
+#include "swift/Basic/TopCollection.h"
 #include <algorithm>
 
 using namespace swift;
@@ -415,3 +416,126 @@ LookupResult TypeChecker::lookupConstructors(DeclContext *dc, Type type,
                                              NameLookupOptions options) {
   return lookupMember(dc, type, Context.Id_init, options);
 }
+
+enum : unsigned {
+  /// Never consider a candidate that's this distance away or worse.
+  UnreasonableCallEditDistance = 8,
+
+  /// Don't consider candidates that score worse than the given distance
+  /// from the best candidate.
+  MaxCallEditDistanceFromBestCandidate = 1
+};
+
+static unsigned getCallEditDistance(DeclName argName, DeclName paramName,
+                                    unsigned maxEditDistance) {
+  // TODO: consider arguments.
+  // TODO: maybe ignore certain kinds of missing / present labels for the
+  //   first argument label?
+  // TODO: word-based rather than character-based?
+  StringRef argBase = argName.getBaseName().str();
+  StringRef paramBase = paramName.getBaseName().str();
+
+  unsigned distance = argBase.edit_distance(paramBase, maxEditDistance);
+
+  // Bound the distance to UnreasonableCallEditDistance.
+  if (distance >= maxEditDistance ||
+      distance > (paramBase.size() + 2) / 3) {
+    return UnreasonableCallEditDistance;
+  }
+
+  return distance;
+}
+
+static bool isPlausibleTypo(DeclRefKind refKind, DeclName typedName,
+                            ValueDecl *candidate) {
+  // Ignore anonymous declarations.
+  if (!candidate->hasName())
+    return false;
+
+  // An operator / identifier mismatch is never a plausible typo.
+  auto fn = dyn_cast<FuncDecl>(candidate);
+  if (typedName.isOperator() != (fn && fn->isOperator()))
+    return false;
+  if (!typedName.isOperator())
+    return true;
+
+  // TODO: honor ref kind?  This is trickier than it sounds because we
+  // may not have processed attributes and types on the candidate yet.
+  return true;
+}
+
+void TypeChecker::performTypoCorrection(DeclContext *DC, DeclRefKind refKind,
+                                        DeclName targetDeclName,
+                                        SourceLoc nameLoc,
+                                        NameLookupOptions lookupOptions,
+                                        LookupResult &result,
+                                        unsigned maxResults) {
+  // Fill in a collection of the most reasonable entries.
+  TopCollection<unsigned, ValueDecl*> entries(maxResults);
+  auto consumer = makeDeclConsumer([&](ValueDecl *decl,
+                                       DeclVisibilityKind reason) {
+    // Never match an operator with an identifier or vice-versa; this is
+    // not a plausible typo.
+    if (!isPlausibleTypo(refKind, targetDeclName, decl))
+      return;
+
+    // Don't suggest a variable within its own initializer.
+    if (auto var = dyn_cast<VarDecl>(decl)) {
+      if (auto binding = var->getParentPatternBinding()) {
+        if (!binding->isImplicit() &&
+            Context.SourceMgr.rangeContainsTokenLoc(binding->getSourceRange(),
+                                                    nameLoc))
+          return;
+      }
+    }
+
+    // Don't waste time computing edit distances that are more than
+    // the worst in our collection.
+    unsigned maxDistance =
+      entries.getMinUninterestingScore(UnreasonableCallEditDistance);
+
+    unsigned distance =
+      getCallEditDistance(targetDeclName, decl->getFullName(), maxDistance);
+
+    // Ignore values that are further than a reasonable distance.
+    if (distance >= UnreasonableCallEditDistance)
+      return;
+
+    entries.insert(distance, std::move(decl));
+  });
+
+  lookupVisibleDecls(consumer, DC, nullptr, /*top level*/ true, nameLoc);
+
+  // Impose a maximum distance from the best score.
+  entries.filterMaxScoreRange(MaxCallEditDistanceFromBestCandidate);
+
+  for (auto &entry : entries)
+    result.add({ entry.Value, nullptr });
+}
+
+static InFlightDiagnostic
+diagnoseTypoCorrection(TypeChecker &tc, DeclNameLoc loc, ValueDecl *decl) {
+  if (auto var = dyn_cast<VarDecl>(decl)) {
+    // Suggest 'self' at the use point instead of pointing at the start
+    // of the function.
+    if (var->isSelfParameter())
+      return tc.diagnose(loc.getBaseNameLoc(), diag::note_typo_candidate,
+                         decl->getName().str());
+  }
+
+  return tc.diagnose(decl, diag::note_typo_candidate, decl->getName().str());
+}
+
+void TypeChecker::noteTypoCorrection(DeclName writtenName, DeclNameLoc loc,
+                                     const LookupResult::Result &suggestion) {
+  auto decl = suggestion.Decl;
+  auto &&diagnostic = diagnoseTypoCorrection(*this, loc, decl);
+
+  DeclName declName = decl->getFullName();
+
+  if (writtenName.getBaseName() != declName.getBaseName())
+    diagnostic.fixItReplace(loc.getBaseNameLoc(), declName.getBaseName().str());
+
+  // TODO: add fix-its for typo'ed argument labels.  This is trickier
+  // because of the reordering rules.
+}