CursorInfo: Teach SemaLocResolver to recognize the start of an expression, if the location cannot be further refined to other kind, e.g. a reference. rdar://32749670

Author: nkcsgexi

include/swift/IDE/Utils.h

@@ -148,6 +148,7 @@ enum class SemaTokenKind {
   Invalid,
   ValueRef,
   ModuleRef,
+  ExprStart,
   StmtStart,
 };
 
@@ -164,6 +165,7 @@ struct SemaToken {
   DeclContext *DC = nullptr;
   Type ContainerType;
   Stmt *TrailingStmt = nullptr;
+  Expr *TrailingExpr = nullptr;
 
   SemaToken() = default;
   SemaToken(ValueDecl *ValueD, TypeDecl *CtorTyRef, ExtensionDecl *ExtTyRef,
@@ -175,6 +177,8 @@ struct SemaToken {
                                                Mod(Mod), Loc(Loc) { }
   SemaToken(Stmt *TrailingStmt) : Kind(SemaTokenKind::StmtStart),
                                   TrailingStmt(TrailingStmt) {}
+  SemaToken(Expr* TrailingExpr) : Kind(SemaTokenKind::ExprStart),
+                                  TrailingExpr(TrailingExpr) {}
   bool isValid() const { return !isInvalid(); }
   bool isInvalid() const { return Kind == SemaTokenKind::Invalid; }
 };
@@ -191,6 +195,7 @@ class SemaLocResolver : public SourceEntityWalker {
   SourceManager &getSourceMgr() const;
 private:
   bool walkToExprPre(Expr *E) override;
+  bool walkToExprPost(Expr *E) override;
   bool walkToDeclPre(Decl *D, CharSourceRange Range) override;
   bool walkToDeclPost(Decl *D) override;
   bool walkToStmtPre(Stmt *S) override;
@@ -211,6 +216,7 @@ class SemaLocResolver : public SourceEntityWalker {
                   SourceLoc Loc, bool IsRef, Type Ty = Type());
   bool tryResolve(ModuleEntity Mod, SourceLoc Loc);
   bool tryResolve(Stmt *St);
+  bool tryResolve(Expr *Exp);
   bool visitSubscriptReference(ValueDecl *D, CharSourceRange Range,
                                bool IsOpenBracket) override;
 };

lib/IDE/SwiftSourceDocInfo.cpp

@@ -108,6 +108,14 @@ bool SemaLocResolver::tryResolve(Stmt *St) {
   return false;
 }
 
+bool SemaLocResolver::tryResolve(Expr *Exp) {
+  if (Exp->getStartLoc() == LocToResolve) {
+    SemaTok = { Exp };
+    return true;
+  }
+  return false;
+}
+
 bool SemaLocResolver::visitSubscriptReference(ValueDecl *D, CharSourceRange Range,
                                               bool IsOpenBracket) {
   // We should treat both open and close brackets equally
@@ -189,6 +197,12 @@ bool SemaLocResolver::walkToExprPre(Expr *E) {
   return true;
 }
 
+bool SemaLocResolver::walkToExprPost(Expr *E) {
+  if (isDone())
+    return false;
+  return !tryResolve(E);
+}
+
 bool SemaLocResolver::visitCallArgName(Identifier Name, CharSourceRange Range,
                                        ValueDecl *D) {
   if (isDone())

tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp

@@ -1181,6 +1181,7 @@ static void resolveCursor(SwiftLangSupport &Lang,
         }
         return;
       }
+      case SemaTokenKind::ExprStart:
       case SemaTokenKind::StmtStart: {
         Receiver({});
         return;