[SourceKit] Preparation for solver-based cursor info

include/swift/AST/Stmt.h

@@ -171,6 +171,13 @@ class BraceStmt final : public Stmt,
   SourceLoc getStartLoc() const;
   SourceLoc getEndLoc() const;
 
+  SourceLoc getContentStartLoc() const;
+  SourceLoc getContentEndLoc() const;
+  /// The range of the brace statement without the braces.
+  SourceRange getContentRange() const {
+    return {getContentStartLoc(), getContentEndLoc()};
+  }
+
   bool empty() const { return getNumElements() == 0; }
   unsigned getNumElements() const { return Bits.BraceStmt.NumElements; }
 

include/swift/IDE/Utils.h

@@ -128,71 +128,200 @@ enum class CursorInfoKind {
   StmtStart,
 };
 
+/// Base class of more specialized \c ResolvedCursorInfos that also represents
+/// and \c Invalid cursor info.
+/// Subclasses of \c ResolvedCursorInfo cannot add new stored properies because
+/// \c ResolvedCursorInfo is being passed around as its base class and thus any
+/// properties in subclasses would get lost.
 struct ResolvedCursorInfo {
+protected:
   CursorInfoKind Kind = CursorInfoKind::Invalid;
   SourceFile *SF = nullptr;
   SourceLoc Loc;
-  ValueDecl *ValueD = nullptr;
-  TypeDecl *CtorTyRef = nullptr;
-  ExtensionDecl *ExtTyRef = nullptr;
-  /// Declarations that were shadowed by \c ValueD using a shorthand syntax that
-  /// names both the newly declared variable and the referenced variable by the
-  /// same identifier in the source text. This includes shorthand closure
-  /// captures (`[foo]`) and shorthand if captures
-  /// (`if let foo {`).
-  /// Decls that are shadowed using shorthand syntax should be reported as
-  /// additional cursor info results.
-  SmallVector<ValueDecl *, 2> ShorthandShadowedDecls;
-  ModuleEntity Mod;
-  bool IsRef = true;
-  bool IsKeywordArgument = false;
-  Type Ty;
-  Type ContainerType;
-  Stmt *TrailingStmt = nullptr;
-  Expr *TrailingExpr = nullptr;
-  /// It this is a ref, whether it is "dynamic". See \c ide::isDynamicRef.
-  bool IsDynamic = false;
-  /// If this is a dynamic ref, the types of the base (multiple in the case of
-  /// protocol composition).
-  SmallVector<NominalTypeDecl *, 1> ReceiverTypes;
 
+  // Technically, these structs could form a union (because only one of them is
+  // active at a time). But I had issues with C++ complaining about copy
+  // constructors and gave up. At the moment it's only wasting 3 words for non
+  // ValueRef data.
+  struct {
+    ValueDecl *ValueD = nullptr;
+    TypeDecl *CtorTyRef = nullptr;
+    ExtensionDecl *ExtTyRef = nullptr;
+    bool IsRef = true;
+    Type Ty;
+    Type ContainerType;
+    bool IsKeywordArgument = false;
+    /// It this is a ref, whether it is "dynamic". See \c ide::isDynamicRef.
+    bool IsDynamic = false;
+    /// If this is a dynamic ref, the types of the base (multiple in the case of
+    /// protocol composition).
+    SmallVector<NominalTypeDecl *> ReceiverTypes;
+    /// Declarations that were shadowed by \c ValueD using a shorthand syntax
+    /// that names both the newly declared variable and the referenced variable
+    /// by the same identifier in the source text. This includes shorthand
+    /// closure captures (`[foo]`) and shorthand if captures
+    /// (`if let foo {`).
+    /// Decls that are shadowed using shorthand syntax should be reported as
+    /// additional cursor info results.
+    SmallVector<ValueDecl *> ShorthandShadowedDecls;
+  } ValueRefInfo;
+  struct {
+    ModuleEntity Mod;
+  } ModuleRefInfo;
+  struct {
+    Expr *TrailingExpr = nullptr;
+  } ExprStartInfo;
+  struct {
+    Stmt *TrailingStmt = nullptr;
+  } StmtStartInfo;
+
+public:
   ResolvedCursorInfo() = default;
   ResolvedCursorInfo(SourceFile *SF) : SF(SF) {}
 
-  ValueDecl *typeOrValue() { return CtorTyRef ? CtorTyRef : ValueD; }
+  CursorInfoKind getKind() const { return Kind; }
+
+  SourceFile *getSourceFile() const { return SF; }
+
+  SourceLoc getLoc() const { return Loc; }
+  void setLoc(SourceLoc Loc) { this->Loc = Loc; }
 
   friend bool operator==(const ResolvedCursorInfo &lhs,
                          const ResolvedCursorInfo &rhs) {
     return lhs.SF == rhs.SF &&
       lhs.Loc.getOpaquePointerValue() == rhs.Loc.getOpaquePointerValue();
   }
 
-  void setValueRef(ValueDecl *ValueD, TypeDecl *CtorTyRef,
-                   ExtensionDecl *ExtTyRef, bool IsRef,
-                   Type Ty, Type ContainerType) {
+  bool isValid() const { return !isInvalid(); }
+  bool isInvalid() const { return Kind == CursorInfoKind::Invalid; }
+};
+
+struct ResolvedValueRefCursorInfo : public ResolvedCursorInfo {
+  // IMPORTANT: Don't add stored properties here. See comment on
+  // ResolvedCursorInfo.
+
+  ResolvedValueRefCursorInfo() = default;
+  explicit ResolvedValueRefCursorInfo(const ResolvedCursorInfo &Base,
+                                      ValueDecl *ValueD, TypeDecl *CtorTyRef,
+                                      ExtensionDecl *ExtTyRef, bool IsRef,
+                                      Type Ty, Type ContainerType)
+      : ResolvedCursorInfo(Base) {
+    assert(Base.getKind() == CursorInfoKind::Invalid &&
+           "Can only specialize from invalid");
     Kind = CursorInfoKind::ValueRef;
-    this->ValueD = ValueD;
-    this->CtorTyRef = CtorTyRef;
-    this->ExtTyRef = ExtTyRef;
-    this->IsRef = IsRef;
-    this->Ty = Ty;
-    this->ContainerType = ContainerType;
+    ValueRefInfo.ValueD = ValueD;
+    ValueRefInfo.CtorTyRef = CtorTyRef;
+    ValueRefInfo.ExtTyRef = ExtTyRef;
+    ValueRefInfo.IsRef = IsRef;
+    ValueRefInfo.Ty = Ty;
+    ValueRefInfo.ContainerType = ContainerType;
+  }
+
+  ValueDecl *getValueD() const { return ValueRefInfo.ValueD; }
+  void setValueD(ValueDecl *ValueD) { ValueRefInfo.ValueD = ValueD; }
+
+  ExtensionDecl *getExtTyRef() const { return ValueRefInfo.ExtTyRef; }
+
+  TypeDecl *getCtorTyRef() const { return ValueRefInfo.CtorTyRef; }
+
+  bool isRef() const { return ValueRefInfo.IsRef; }
+  void setIsRef(bool IsRef) { ValueRefInfo.IsRef = IsRef; }
+
+  Type getType() const { return ValueRefInfo.Ty; }
+
+  Type getContainerType() const { return ValueRefInfo.ContainerType; }
+  void setContainerType(Type Ty) { ValueRefInfo.ContainerType = Ty; }
+
+  bool isKeywordArgument() const { return ValueRefInfo.IsKeywordArgument; }
+  void setIsKeywordArgument(bool IsKeywordArgument) {
+    ValueRefInfo.IsKeywordArgument = IsKeywordArgument;
+  }
+
+  bool isDynamic() const { return ValueRefInfo.IsDynamic; }
+  void setIsDynamic(bool IsDynamic) { ValueRefInfo.IsDynamic = IsDynamic; }
+
+  ArrayRef<NominalTypeDecl *> getReceiverTypes() const {
+    return ValueRefInfo.ReceiverTypes;
+  }
+  void setReceiverTypes(const SmallVector<NominalTypeDecl *> &ReceiverTypes) {
+    ValueRefInfo.ReceiverTypes = ReceiverTypes;
+  }
+
+  ArrayRef<ValueDecl *> getShorthandShadowedDecls() const {
+    return ValueRefInfo.ShorthandShadowedDecls;
+  };
+  void setShorthandShadowedDecls(
+      const SmallVector<ValueDecl *> &ShorthandShadowedDecls) {
+    ValueRefInfo.ShorthandShadowedDecls = ShorthandShadowedDecls;
+  };
+
+  ValueDecl *typeOrValue() {
+    return ValueRefInfo.CtorTyRef ? ValueRefInfo.CtorTyRef
+                                  : ValueRefInfo.ValueD;
   }
-  void setModuleRef(ModuleEntity Mod) {
+
+  static bool classof(const ResolvedCursorInfo *Info) {
+    return Info->getKind() == CursorInfoKind::ValueRef;
+  }
+};
+
+struct ResolvedModuleRefCursorInfo : public ResolvedCursorInfo {
+  // IMPORTANT: Don't add stored properties here. See comment on
+  // ResolvedCursorInfo.
+
+  ResolvedModuleRefCursorInfo(const ResolvedCursorInfo &Base, ModuleEntity Mod)
+      : ResolvedCursorInfo(Base) {
+    assert(Base.getKind() == CursorInfoKind::Invalid &&
+           "Can only specialize from invalid");
     Kind = CursorInfoKind::ModuleRef;
-    this->Mod = Mod;
+    ModuleRefInfo.Mod = Mod;
   }
-  void setTrailingStmt(Stmt *TrailingStmt) {
-    Kind = CursorInfoKind::StmtStart;
-    this->TrailingStmt = TrailingStmt;
+
+  ModuleEntity getMod() const { return ModuleRefInfo.Mod; }
+
+  static bool classof(const ResolvedCursorInfo *Info) {
+    return Info->getKind() == CursorInfoKind::ModuleRef;
   }
-  void setTrailingExpr(Expr* TrailingExpr) {
+};
+
+struct ResolvedExprStartCursorInfo : public ResolvedCursorInfo {
+  // IMPORTANT: Don't add stored properties here. See comment on
+  // ResolvedCursorInfo.
+
+  ResolvedExprStartCursorInfo(const ResolvedCursorInfo &Base,
+                              Expr *TrailingExpr)
+      : ResolvedCursorInfo(Base) {
+    assert(Base.getKind() == CursorInfoKind::Invalid &&
+           "Can only specialize from invalid");
     Kind = CursorInfoKind::ExprStart;
-    this->TrailingExpr = TrailingExpr;
+    ExprStartInfo.TrailingExpr = TrailingExpr;
   }
 
-  bool isValid() const { return !isInvalid(); }
-  bool isInvalid() const { return Kind == CursorInfoKind::Invalid; }
+  Expr *getTrailingExpr() const { return ExprStartInfo.TrailingExpr; }
+
+  static bool classof(const ResolvedCursorInfo *Info) {
+    return Info->getKind() == CursorInfoKind::ExprStart;
+  }
+};
+
+struct ResolvedStmtStartCursorInfo : public ResolvedCursorInfo {
+  // IMPORTANT: Don't add stored properties here. See comment on
+  // ResolvedCursorInfo.
+
+  ResolvedStmtStartCursorInfo(const ResolvedCursorInfo &Base,
+                              Stmt *TrailingStmt)
+      : ResolvedCursorInfo(Base) {
+    assert(Base.getKind() == CursorInfoKind::Invalid &&
+           "Can only specialize from invalid");
+    Kind = CursorInfoKind::StmtStart;
+    StmtStartInfo.TrailingStmt = TrailingStmt;
+  }
+
+  Stmt *getTrailingStmt() const { return StmtStartInfo.TrailingStmt; }
+
+  static bool classof(const ResolvedCursorInfo *Info) {
+    return Info->getKind() == CursorInfoKind::StmtStart;
+  }
 };
 
 void simple_display(llvm::raw_ostream &out, const ResolvedCursorInfo &info);

include/swift/Parse/CodeCompletionCallbacks.h

@@ -254,7 +254,7 @@ class CodeCompletionCallbacks {
 
   /// Signals that the AST for the all the delayed-parsed code was
   /// constructed.  No \c complete*() callbacks will be done after this.
-  virtual void doneParsing() = 0;
+  virtual void doneParsing(SourceFile *SrcFile) = 0;
 };
 
 /// A factory to create instances of \c CodeCompletionCallbacks.

include/swift/Sema/IDETypeChecking.h

@@ -164,6 +164,11 @@ namespace swift {
       constraints::SolutionApplicationTarget &target, bool needsPrecheck,
       llvm::function_ref<void(const constraints::Solution &)> callback);
 
+  /// Thunk around \c TypeChecker::resolveDeclRefExpr to make it available to
+  /// \c swift::ide
+  Expr *resolveDeclRefExpr(UnresolvedDeclRefExpr *UDRE, DeclContext *Context,
+                         bool replaceInvalidRefsWithErrors);
+
   LookupResult
   lookupSemanticMember(DeclContext *DC, Type ty, DeclName name);
 
@@ -339,12 +344,17 @@ namespace swift {
   /// these shorthand shadows.
   /// The first element in the pair is the implicitly declared variable and the
   /// second variable is the shadowed one.
+  /// If a \c DeclContext is passed, it is used to resolve any
+  /// \c UnresolvedDeclRef that a shorthand shadow may refer to.
   SmallVector<std::pair<ValueDecl *, ValueDecl *>, 1>
-  getShorthandShadows(CaptureListExpr *CaptureList);
+  getShorthandShadows(CaptureListExpr *CaptureList, DeclContext *DC = nullptr);
 
   /// Same as above but for shorthand `if let foo {` syntax.
+  /// If a \c DeclContext is passed, it is used to resolve any
+  /// \c UnresolvedDeclRef that a shorthand shadow may refer to.
   SmallVector<std::pair<ValueDecl *, ValueDecl *>, 1>
-  getShorthandShadows(LabeledConditionalStmt *CondStmt);
+  getShorthandShadows(LabeledConditionalStmt *CondStmt,
+                      DeclContext *DC = nullptr);
 }
 
 #endif

lib/AST/Stmt.cpp

@@ -160,6 +160,17 @@ SourceLoc BraceStmt::getStartLoc() const {
   if (LBLoc) {
     return LBLoc;
   }
+  return getContentStartLoc();
+}
+
+SourceLoc BraceStmt::getEndLoc() const {
+  if (RBLoc) {
+    return RBLoc;
+  }
+  return getContentEndLoc();
+}
+
+SourceLoc BraceStmt::getContentStartLoc() const {
   for (auto elt : getElements()) {
     if (auto loc = elt.getStartLoc()) {
       return loc;
@@ -168,10 +179,7 @@ SourceLoc BraceStmt::getStartLoc() const {
   return SourceLoc();
 }
 
-SourceLoc BraceStmt::getEndLoc() const {
-  if (RBLoc) {
-    return RBLoc;
-  }
+SourceLoc BraceStmt::getContentEndLoc() const {
   for (auto elt : llvm::reverse(getElements())) {
     if (auto loc = elt.getEndLoc()) {
       return loc;

lib/IDE/CodeCompletion.cpp

@@ -287,7 +287,7 @@ class CodeCompletionCallbacksImpl : public CodeCompletionCallbacks {
   void completeTypeAttrBeginning() override;
   void completeOptionalBinding() override;
 
-  void doneParsing() override;
+  void doneParsing(SourceFile *SrcFile) override;
 
 private:
   void addKeywords(CodeCompletionResultSink &Sink, bool MaybeFuncBody);
@@ -1552,7 +1552,7 @@ static void undoSingleExpressionReturn(DeclContext *DC) {
   }
 }
 
-void CodeCompletionCallbacksImpl::doneParsing() {
+void CodeCompletionCallbacksImpl::doneParsing(SourceFile *SrcFile) {
   CompletionContext.CodeCompletionKind = Kind;
 
   if (Kind == CompletionKind::None) {

lib/IDE/ConformingMethodList.cpp

@@ -49,7 +49,7 @@ class ConformingMethodListCallbacks : public CodeCompletionCallbacks {
   void completePostfixExpr(Expr *E, bool hasSpace) override;
   // }
 
-  void doneParsing() override;
+  void doneParsing(SourceFile *SrcFile) override;
 };
 
 void ConformingMethodListCallbacks::completeDotExpr(CodeCompletionExpr *E,
@@ -64,7 +64,7 @@ void ConformingMethodListCallbacks::completePostfixExpr(Expr *E,
   ParsedExpr = E;
 }
 
-void ConformingMethodListCallbacks::doneParsing() {
+void ConformingMethodListCallbacks::doneParsing(SourceFile *SrcFile) {
   if (!ParsedExpr)
     return;