[AST] Store the TypeReprs of generic arguments used to specialize a decl reference in DeclRefExpr, to impove source fidelity.

Fixes rdar://15034958

Swift SVN r8880

Author: akyrtzi

include/swift/AST/Expr.h

@@ -376,23 +376,53 @@ class DiscardAssignmentExpr : public Expr {
 
 /// DeclRefExpr - A reference to a value, "x".
 class DeclRefExpr : public Expr {
-  /// \brief The declaration pointer and a bit specifying whether it was
+  /// This is used when the reference is specialized, e.g "GenCls<Int>", to
+  /// hold information about the generic arguments.
+  struct SpecializeInfo {
+    ValueDecl *D = nullptr;
+    MutableArrayRef<TypeRepr*> GenericArgs;
+  };
+
+  /// \brief The declaration pointer or SpecializeInfo pointer if it was
   /// explicitly specialized with <...>.
-  llvm::PointerIntPair<ValueDecl *, 1, bool> DAndSpecialized;
+  llvm::PointerUnion<ValueDecl *, SpecializeInfo *> DAndSpecialized;
   SourceLoc Loc;
 
+  SpecializeInfo *getSpecInfo() const {
+    return DAndSpecialized.dyn_cast<SpecializeInfo*>();
+  }
+
 public:
   DeclRefExpr(ValueDecl *D, SourceLoc Loc, bool Implicit, Type Ty = Type())
-    : Expr(ExprKind::DeclRef, Implicit, Ty), DAndSpecialized(D, false), Loc(Loc) {}
+    : Expr(ExprKind::DeclRef, Implicit, Ty), DAndSpecialized(D), Loc(Loc) {}
 
-  ValueDecl *getDecl() const { return DAndSpecialized.getPointer(); }
+  ValueDecl *getDecl() const {
+    if (auto Spec = getSpecInfo())
+      return Spec->D;
+    return DAndSpecialized.get<ValueDecl*>();
+  }
 
-  void setSpecialized(bool specialized) { DAndSpecialized.setInt(specialized); }
+  void setSpecialized();
 
   /// \brief Determine whether this declaration reference was immediately
   /// specialized by <...>.
-  bool isSpecialized() const { return DAndSpecialized.getInt(); }
+  bool isSpecialized() const { return getSpecInfo() != nullptr; }
 
+  /// Set the generic arguments.
+  ///
+  /// This copies the array using ASTContext's allocator.
+  void setGenericArgs(ArrayRef<TypeRepr*> GenericArgs);
+
+  /// Returns the generic arguments if it was specialized or an empty array
+  /// otherwise.
+  ArrayRef<TypeRepr *> getGenericArgs() const {
+    return const_cast<DeclRefExpr*>(this)->getGenericArgs();
+  }
+  MutableArrayRef<TypeRepr *> getGenericArgs() {
+    if (auto Spec = getSpecInfo())
+      return Spec->GenericArgs;
+    return MutableArrayRef<TypeRepr *>();
+  }
   SourceRange getSourceRange() const { return Loc; }
 
   static bool classof(const Expr *E) {

lib/AST/ASTDumper.cpp

@@ -877,6 +877,7 @@ class PrintExpr : public ExprVisitor<PrintExpr> {
   void printRec(Pattern *P) {
     PrintPattern(OS, Indent+2).visit(P);
   }
+  void printRec(TypeRepr *T);
 
   void printSubstitutions(ArrayRef<Substitution> Substitutions) {
     for (auto S : Substitutions) {
@@ -940,7 +941,13 @@ class PrintExpr : public ExprVisitor<PrintExpr> {
     printCommon(E, "declref_expr")
       << " decl=";
     E->getDecl()->dumpRef(OS);
-    OS << " specialized=" << (E->isSpecialized()? "yes" : "no") << ")";
+    OS << " specialized=" << (E->isSpecialized()? "yes" : "no");
+
+    for (auto TR : E->getGenericArgs()) {
+      OS << '\n';
+      printRec(TR);
+    }
+    OS << ')';
   }
   void visitSuperRefExpr(SuperRefExpr *E) {
     printCommon(E, "super_ref_expr") << ')';
@@ -987,8 +994,7 @@ class PrintExpr : public ExprVisitor<PrintExpr> {
     printRec(E->getSubExpr());
     for (TypeLoc T : E->getUnresolvedParams()) {
       OS << '\n';
-      OS.indent(Indent+2);
-      T.getTypeRepr()->print(OS);
+      printRec(T.getTypeRepr());
     }
     OS << ')';
   }
@@ -1506,6 +1512,10 @@ void PrintDecl::printRec(TypeRepr *T) {
   PrintTypeRepr(OS, Indent+2).visit(T);
 }
 
+void PrintExpr::printRec(TypeRepr *T) {
+  PrintTypeRepr(OS, Indent+2).visit(T);
+}
+
 void PrintPattern::printRec(TypeRepr *T) {
   PrintTypeRepr(OS, Indent+2).visit(T);
 }

lib/AST/ASTWalker.cpp

@@ -77,7 +77,6 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
   Expr *visitErrorExpr(ErrorExpr *E) { return E; }
   Expr *visitLiteralExpr(LiteralExpr *E) { return E; }
   Expr *visitDiscardAssignmentExpr(DiscardAssignmentExpr *E) { return E; }
-  Expr *visitDeclRefExpr(DeclRefExpr *E) { return E; }
   Expr *visitSuperRefExpr(SuperRefExpr *E) { return E; }
   Expr *visitOtherConstructorDeclRefExpr(OtherConstructorDeclRefExpr *E) {
     return E;
@@ -124,6 +123,14 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
     }
     return nullptr;
   }
+
+  Expr *visitDeclRefExpr(DeclRefExpr *E) {
+    for (auto Ty : E->getGenericArgs()) {
+      if (doIt(Ty))
+        return nullptr;
+    }
+    return E;
+  }
 
   Expr *visitMemberRefExpr(MemberRefExpr *E) {
     if (Expr *Base = doIt(E->getBase())) {

lib/AST/Expr.cpp

@@ -137,6 +137,27 @@ llvm::APFloat FloatLiteralExpr::getValue() const {
                   getType()->castTo<BuiltinFloatType>()->getAPFloatSemantics());
 }
 
+void DeclRefExpr::setSpecialized() {
+  if (isSpecialized())
+    return;
+
+  ValueDecl *D = getDecl();
+  assert(D);
+  void *Mem = D->getASTContext().Allocate(sizeof(SpecializeInfo),
+                                          alignof(SpecializeInfo));
+  auto Spec = new (Mem) SpecializeInfo;
+  Spec->D = D;
+  DAndSpecialized = Spec;
+}
+
+void DeclRefExpr::setGenericArgs(ArrayRef<TypeRepr*> GenericArgs) {
+  ValueDecl *D = getDecl();
+  assert(D);
+  setSpecialized();
+  getSpecInfo()->GenericArgs = D->getASTContext().AllocateCopy(GenericArgs);
+}
+
+
 MemberRefExpr::MemberRefExpr(Expr *base, SourceLoc dotLoc,
                              ConcreteDeclRef member, SourceLoc nameLoc,
                              bool Implicit)

lib/Parse/ParseExpr.cpp

@@ -1369,7 +1369,7 @@ Expr *Parser::actOnIdentifierExpr(Identifier text, SourceLoc loc) {
     E = unresolved;
   } else {
     auto declRef = new (Context) DeclRefExpr(D, loc, /*Implicit=*/false);
-    declRef->setSpecialized(hasGenericArgumentList);
+    declRef->setGenericArgs(args);
     E = declRef;
   }
 

lib/Sema/TypeCheckConstraintsApply.cpp

@@ -1221,7 +1221,12 @@ namespace {
 
     Expr *visitUnresolvedSpecializeExpr(UnresolvedSpecializeExpr *expr) {
       // Our specializations should have resolved the subexpr to the right type.
-      // FIXME: Should preserve generic argument list for source fidelity
+      if (auto DRE = dyn_cast<DeclRefExpr>(expr->getSubExpr())) {
+        SmallVector<TypeRepr *, 8> GenArgs;
+        for (auto TL : expr->getUnresolvedParams())
+          GenArgs.push_back(TL.getTypeRepr());
+        DRE->setGenericArgs(GenArgs);
+      }
       return expr->getSubExpr();
     }
 

lib/Sema/TypeCheckExpr.cpp

@@ -408,7 +408,8 @@ Expr *TypeChecker::buildRefExpr(ArrayRef<ValueDecl *> Decls, SourceLoc NameLoc,
 
   if (Decls.size() == 1 && !isa<ProtocolDecl>(Decls[0]->getDeclContext())) {
     auto result = new (Context) DeclRefExpr(Decls[0], NameLoc, Implicit);
-    result->setSpecialized(isSpecialized);
+    if (isSpecialized)
+      result->setSpecialized();
     return result;
   }