[SR-7293] Refactoring action to add Equatable Conformance

lib/IDE/Refactoring.cpp

@@ -3144,6 +3144,252 @@ bool RefactoringActionMemberwiseInitLocalRefactoring::performChange() {
   return false;
 }
 
+class AddEquatableContext {
+
+  /// Declaration context
+  DeclContext *DC;
+
+  /// Adopter type
+  Type Adopter;
+
+  /// Start location of declaration context brace
+  SourceLoc StartLoc;
+
+  /// Array of all inherited protocols' locations
+  ArrayRef<TypeLoc> ProtocolsLocations;
+
+  /// Array of all conformed protocols
+  SmallVector<swift::ProtocolDecl *, 2> Protocols;
+
+  /// Start location of declaration,
+  /// a place to write protocol name
+  SourceLoc ProtInsertStartLoc;
+
+  /// Stored properties of extending adopter
+  ArrayRef<VarDecl *> StoredProperties;
+
+  /// Range of internal members in declaration
+  DeclRange Range;
+
+  bool conformsToEquatableProtocol() {
+    for (ProtocolDecl *Protocol : Protocols) {
+      if (Protocol->getKnownProtocolKind() == KnownProtocolKind::Equatable) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  bool isRequirementValid() {
+    auto Reqs = getProtocolRequirements();
+    if (Reqs.empty()) {
+      return false;
+    }
+    auto Req = dyn_cast<FuncDecl>(Reqs[0]);
+    auto Params = Req->getParameters();
+    if (!Req || Params->size() != 2) {
+      return false;
+    }
+    return true;
+  }
+
+  bool isPropertiesListValid() {
+    return !getPublicProperties().empty();
+  }
+
+  void printFunctionBody(ASTPrinter &Printer, StringRef ExtraIndent,
+  ParameterList *Params);
+
+  std::vector<ValueDecl *> getProtocolRequirements();
+
+  std::vector<VarDecl *> getPublicProperties();
+
+public:
+
+  AddEquatableContext(NominalTypeDecl *Decl) : DC(Decl),
+  Adopter(Decl->getDeclaredType()), StartLoc(Decl->getBraces().Start),
+  ProtocolsLocations(Decl->getInherited()),
+  Protocols(Decl->getAllProtocols()), ProtInsertStartLoc(Decl->getNameLoc()),
+  StoredProperties(Decl->getStoredProperties()), Range(Decl->getMembers()) {};
+
+  AddEquatableContext(ExtensionDecl *Decl) : DC(Decl),
+  Adopter(Decl->getExtendedType()), StartLoc(Decl->getBraces().Start),
+  ProtocolsLocations(Decl->getInherited()),
+  Protocols(Decl->getExtendedNominal()->getAllProtocols()),
+  ProtInsertStartLoc(Decl->getExtendedTypeRepr()->getEndLoc()),
+  StoredProperties(Decl->getExtendedNominal()->getStoredProperties()), Range(Decl->getMembers()) {};
+
+  AddEquatableContext() : DC(nullptr), Adopter(), ProtocolsLocations(),
+  Protocols(), StoredProperties(), Range(nullptr, nullptr) {};
+
+  static AddEquatableContext getDeclarationContextFromInfo(ResolvedCursorInfo Info);
+
+  std::string getDeclForProtocol();
+
+  std::string getDeclForFunction(SourceManager &SM);
+
+  bool isValid() {
+    return StartLoc.isValid() && ProtInsertStartLoc.isValid() &&
+    !conformsToEquatableProtocol() && isPropertiesListValid() &&
+    isRequirementValid();
+  }
+
+  SourceLoc getStartLocForProtocolDecl() {
+    if (ProtocolsLocations.empty()) {
+      return ProtInsertStartLoc;
+    }
+    return ProtocolsLocations.back().getSourceRange().Start;
+  }
+
+  bool isMembersRangeEmpty() {
+    return Range.empty();
+  }
+
+  SourceLoc getInsertStartLoc();
+};
+
+SourceLoc AddEquatableContext::
+getInsertStartLoc() {
+  SourceLoc MaxLoc = StartLoc;
+  for (auto Mem : Range) {
+    if (Mem->getEndLoc().getOpaquePointerValue() >
+        MaxLoc.getOpaquePointerValue()) {
+      MaxLoc = Mem->getEndLoc();
+    }
+  }
+  return MaxLoc;
+}
+
+std::string AddEquatableContext::
+getDeclForProtocol() {
+  StringRef ProtocolName = getProtocolName(KnownProtocolKind::Equatable);
+  std::string Buffer;
+  llvm::raw_string_ostream OS(Buffer);
+  if (ProtocolsLocations.empty()) {
+    OS << ": " << ProtocolName;
+    return Buffer;
+  }
+  OS << ", " << ProtocolName;
+  return Buffer;
+}
+
+std::string AddEquatableContext::
+getDeclForFunction(SourceManager &SM) {
+  auto Reqs = getProtocolRequirements();
+  auto Req = dyn_cast<FuncDecl>(Reqs[0]);
+  auto Params = Req->getParameters();
+  StringRef ExtraIndent;
+  StringRef CurrentIndent =
+    Lexer::getIndentationForLine(SM, getInsertStartLoc(), &ExtraIndent);
+  std::string Indent;
+  if (isMembersRangeEmpty()) {
+    Indent = (CurrentIndent + ExtraIndent).str();
+  } else {
+    Indent = CurrentIndent.str();
+  }
+  PrintOptions Options = PrintOptions::printVerbose();
+  Options.PrintDocumentationComments = false;
+  Options.setBaseType(Adopter);
+  Options.FunctionBody = [&](const ValueDecl *VD, ASTPrinter &Printer) {
+    Printer << " {";
+    Printer.printNewline();
+    printFunctionBody(Printer, ExtraIndent, Params);
+    Printer.printNewline();
+    Printer << "}";
+  };
+  std::string Buffer;
+  llvm::raw_string_ostream OS(Buffer);
+  ExtraIndentStreamPrinter Printer(OS, Indent);
+  Printer.printNewline();
+  if (!isMembersRangeEmpty()) {
+    Printer.printNewline();
+  }
+  Reqs[0]->print(Printer, Options);
+  return Buffer;
+}
+
+std::vector<VarDecl *> AddEquatableContext::
+getPublicProperties() {
+  std::vector<VarDecl *> PublicProperties;
+  for (VarDecl *Decl : StoredProperties) {
+    if (!Decl->hasPrivateAccessor()) {
+      PublicProperties.push_back(Decl);
+    }
+  }
+  return PublicProperties;
+}
+
+std::vector<ValueDecl *> AddEquatableContext::
+getProtocolRequirements() {
+  std::vector<ValueDecl *> Collection;
+  auto Proto = DC->getASTContext().getProtocol(KnownProtocolKind::Equatable);
+  for (auto Member : Proto->getMembers()) {
+    auto Req = dyn_cast<ValueDecl>(Member);
+    if (!Req || Req->isInvalid() || !Req->isProtocolRequirement()) {
+      continue;
+    }
+    Collection.push_back(Req);
+  }
+  return Collection;
+}
+
+AddEquatableContext AddEquatableContext::
+getDeclarationContextFromInfo(ResolvedCursorInfo Info) {
+  if (Info.isInvalid()) {
+    return AddEquatableContext();
+  }
+  if (!Info.IsRef) {
+    if (auto *NomDecl = dyn_cast<NominalTypeDecl>(Info.ValueD)) {
+      return AddEquatableContext(NomDecl);
+    }
+  } else if (auto *ExtDecl = Info.ExtTyRef) {
+    return AddEquatableContext(ExtDecl);
+  }
+  return AddEquatableContext();
+}
+
+void AddEquatableContext::
+printFunctionBody(ASTPrinter &Printer, StringRef ExtraIndent, ParameterList *Params) {
+  llvm::SmallString<128> Return;
+  llvm::raw_svector_ostream SS(Return);
+  SS << tok::kw_return;
+  StringRef Space = " ";
+  StringRef AdditionalSpace = "       ";
+  StringRef Point = ".";
+  StringRef Join = " == ";
+  StringRef And = " &&";
+  auto Props = getPublicProperties();
+  auto FParam = Params->get(0)->getName();
+  auto SParam = Params->get(1)->getName();
+  auto Prop = Props[0]->getName();
+  Printer << ExtraIndent << Return << Space
+  << FParam << Point << Prop << Join << SParam << Point << Prop;
+  if (Props.size() > 1) {
+    std::for_each(Props.begin() + 1, Props.end(), [&](VarDecl *VD){
+      auto Name = VD->getName();
+      Printer << And;
+      Printer.printNewline();
+      Printer << ExtraIndent << AdditionalSpace << FParam << Point
+      << Name << Join << SParam << Point << Name;
+    });
+  }
+}
+
+bool RefactoringActionAddEquatableConformance::
+isApplicable(ResolvedCursorInfo Tok, DiagnosticEngine &Diag) {
+  return AddEquatableContext::getDeclarationContextFromInfo(Tok).isValid();
+}
+
+bool RefactoringActionAddEquatableConformance::
+performChange() {
+  auto Context = AddEquatableContext::getDeclarationContextFromInfo(CursorInfo);
+  EditConsumer.insertAfter(SM, Context.getStartLocForProtocolDecl(),
+                           Context.getDeclForProtocol());
+  EditConsumer.insertAfter(SM, Context.getInsertStartLoc(),
+                           Context.getDeclForFunction(SM));
+  return false;
+}
+
 static CharSourceRange
   findSourceRangeToWrapInCatch(ResolvedCursorInfo CursorInfo,
                                SourceFile *TheFile,