[Refactoring] Add Codable refactoring action

Inserts the synthesized implementation.
As part of this, fix some ASTPrinter bugs.

rdar://87904700

Author: louisdh

include/swift/AST/Decl.h

@@ -858,6 +858,8 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
   void print(raw_ostream &OS) const;
   void print(raw_ostream &OS, const PrintOptions &Opts) const;
 
+  void printInherited(ASTPrinter &Printer, const PrintOptions &Options) const;
+
   /// Pretty-print the given declaration.
   ///
   /// \param Printer ASTPrinter object.

include/swift/AST/PrintOptions.h

@@ -479,6 +479,10 @@ struct PrintOptions {
   /// Whether to print inheritance lists for types.
   bool PrintInherited = true;
 
+  /// Whether to print a space before the `:` of an inheritance list in a type
+  /// decl.
+  bool PrintSpaceBeforeInheritance = true;
+
   /// Whether to print feature checks for compatibility with older Swift
   /// compilers that might parse the result.
   bool PrintCompatibilityFeatureChecks = false;

include/swift/IDE/RefactoringKinds.def

@@ -54,6 +54,8 @@ CURSOR_REFACTORING(MemberwiseInitLocalRefactoring, "Generate Memberwise Initiali
 
 CURSOR_REFACTORING(AddEquatableConformance, "Add Equatable Conformance", add.equatable.conformance)
 
+CURSOR_REFACTORING(AddExplicitCodableImplementation, "Add Explicit Codable Implementation", add.explicit-codable-implementation)
+
 CURSOR_REFACTORING(ConvertCallToAsyncAlternative, "Convert Call to Async Alternative", convert.call-to-async)
 
 CURSOR_REFACTORING(ConvertToAsync, "Convert Function to Async", convert.func-to-async)

lib/AST/ASTPrinter.cpp

@@ -861,6 +861,7 @@ class PrintAST : public ASTVisitor<PrintAST> {
                                               Decl *attachingTo);
   void printWhereClauseFromRequirementSignature(ProtocolDecl *proto,
                                                 Decl *attachingTo);
+  void printInherited(const Decl *decl);
 
   void printGenericSignature(GenericSignature genericSig,
                              unsigned flags);
@@ -889,7 +890,6 @@ class PrintAST : public ASTVisitor<PrintAST> {
                     bool openBracket = true, bool closeBracket = true);
   void printGenericDeclGenericParams(GenericContext *decl);
   void printDeclGenericRequirements(GenericContext *decl);
-  void printInherited(const Decl *decl);
   void printBodyIfNecessary(const AbstractFunctionDecl *decl);
 
   void printEnumElement(EnumElementDecl *elt);
@@ -2372,7 +2372,10 @@ void PrintAST::printInherited(const Decl *decl) {
   if (TypesToPrint.empty())
     return;
 
-  Printer << " : ";
+  if (Options.PrintSpaceBeforeInheritance) {
+    Printer << " ";
+  }
+  Printer << ": ";
 
   interleave(TypesToPrint, [&](InheritedEntry inherited) {
     if (inherited.isUnchecked)
@@ -4551,6 +4554,17 @@ void PrintAST::visitBraceStmt(BraceStmt *stmt) {
 }
 
 void PrintAST::visitReturnStmt(ReturnStmt *stmt) {
+  if (!stmt->hasResult()) {
+    if (auto *FD = dyn_cast<AbstractFunctionDecl>(Current)) {
+      if (FD->hasBody()) {
+        auto *Body = FD->getBody();
+        if (Body->getLastElement().dyn_cast<Stmt *>() == stmt) {
+          // Don't print empty return.
+          return;
+        }
+      }
+    }
+  }
   Printer << tok::kw_return;
   if (stmt->hasResult()) {
     Printer << " ";
@@ -4740,6 +4754,11 @@ bool Decl::print(ASTPrinter &Printer, const PrintOptions &Opts) const {
   return printer.visit(const_cast<Decl *>(this));
 }
 
+void Decl::printInherited(ASTPrinter &Printer, const PrintOptions &Opts) const {
+  PrintAST printer(Printer, Opts);
+  printer.printInherited(this);
+}
+
 bool Decl::shouldPrintInContext(const PrintOptions &PO) const {
   // Skip getters/setters. They are part of the variable or subscript.
   if (isa<AccessorDecl>(this))

lib/IDE/Refactoring.cpp

@@ -3526,6 +3526,177 @@ performChange() {
   return false;
 }
 
+class AddCodableContext {
+
+  /// Declaration context
+  DeclContext *DC;
+
+  /// Start location of declaration context brace
+  SourceLoc StartLoc;
+
+  /// Array of all conformed protocols
+  SmallVector<swift::ProtocolDecl *, 2> Protocols;
+
+  /// Range of internal members in declaration
+  DeclRange Range;
+
+  bool conformsToCodableProtocol() {
+    for (ProtocolDecl *Protocol : Protocols) {
+      if (Protocol->getKnownProtocolKind() == KnownProtocolKind::Encodable ||
+          Protocol->getKnownProtocolKind() == KnownProtocolKind::Decodable) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+public:
+  AddCodableContext(NominalTypeDecl *Decl)
+      : DC(Decl), StartLoc(Decl->getBraces().Start),
+        Protocols(Decl->getAllProtocols()), Range(Decl->getMembers()){};
+
+  AddCodableContext(ExtensionDecl *Decl)
+      : DC(Decl), StartLoc(Decl->getBraces().Start),
+        Protocols(Decl->getExtendedNominal()->getAllProtocols()),
+        Range(Decl->getMembers()){};
+
+  AddCodableContext() : DC(nullptr), Protocols(), Range(nullptr, nullptr){};
+
+  static AddCodableContext
+  getDeclarationContextFromInfo(const ResolvedCursorInfo &Info);
+
+  void printInsertionText(const ResolvedCursorInfo &CursorInfo,
+                          SourceManager &SM, llvm::raw_ostream &OS);
+
+  bool isValid() { return StartLoc.isValid() && conformsToCodableProtocol(); }
+
+  SourceLoc getInsertStartLoc();
+};
+
+SourceLoc AddCodableContext::getInsertStartLoc() {
+  SourceLoc MaxLoc = StartLoc;
+  for (auto Mem : Range) {
+    if (Mem->getEndLoc().getOpaquePointerValue() >
+        MaxLoc.getOpaquePointerValue()) {
+      MaxLoc = Mem->getEndLoc();
+    }
+  }
+  return MaxLoc;
+}
+
+/// Walks an AST and prints the synthesized Codable implementation.
+class SynthesizedCodablePrinter : public ASTWalker {
+private:
+  ASTPrinter &Printer;
+
+public:
+  SynthesizedCodablePrinter(ASTPrinter &Printer) : Printer(Printer) {}
+
+  bool walkToDeclPre(Decl *D) override {
+    auto *VD = dyn_cast<ValueDecl>(D);
+    if (!VD)
+      return false;
+
+    if (!VD->isSynthesized()) {
+      return true;
+    }
+    SmallString<32> Scratch;
+    auto name = VD->getName().getString(Scratch);
+    // Print all synthesized enums,
+    // since Codable can synthesize multiple enums (for associated values).
+    auto shouldPrint =
+        isa<EnumDecl>(VD) || name == "init(from:)" || name == "encode(to:)";
+    if (!shouldPrint) {
+      // Some other synthesized decl that we don't want to print.
+      return false;
+    }
+
+    Printer.printNewline();
+
+    if (auto enumDecl = dyn_cast<EnumDecl>(D)) {
+      // Manually print enum here, since we don't want to print synthesized
+      // functions.
+      Printer << "enum " << enumDecl->getNameStr();
+      PrintOptions Options;
+      Options.PrintSpaceBeforeInheritance = false;
+      enumDecl->printInherited(Printer, Options);
+      Printer << " {";
+      for (Decl *EC : enumDecl->getAllElements()) {
+        Printer.printNewline();
+        Printer << "  ";
+        EC->print(Printer, Options);
+      }
+      Printer.printNewline();
+      Printer << "}";
+      return false;
+    }
+
+    PrintOptions Options;
+    Options.SynthesizeSugarOnTypes = true;
+    Options.FunctionDefinitions = true;
+    Options.VarInitializers = true;
+    Options.PrintExprs = true;
+    Options.TypeDefinitions = true;
+    Options.ExcludeAttrList.push_back(DAK_HasInitialValue);
+
+    Printer.printNewline();
+    D->print(Printer, Options);
+
+    return false;
+  }
+};
+
+void AddCodableContext::printInsertionText(const ResolvedCursorInfo &CursorInfo,
+                                           SourceManager &SM,
+                                           llvm::raw_ostream &OS) {
+  StringRef ExtraIndent;
+  StringRef CurrentIndent =
+      Lexer::getIndentationForLine(SM, getInsertStartLoc(), &ExtraIndent);
+  std::string Indent;
+  if (getInsertStartLoc() == StartLoc) {
+    Indent = (CurrentIndent + ExtraIndent).str();
+  } else {
+    Indent = CurrentIndent.str();
+  }
+
+  ExtraIndentStreamPrinter Printer(OS, Indent);
+  Printer.printNewline();
+  SynthesizedCodablePrinter Walker(Printer);
+  DC->getAsDecl()->walk(Walker);
+}
+
+AddCodableContext AddCodableContext::getDeclarationContextFromInfo(
+    const ResolvedCursorInfo &Info) {
+  if (Info.isInvalid()) {
+    return AddCodableContext();
+  }
+  if (!Info.IsRef) {
+    if (auto *NomDecl = dyn_cast<NominalTypeDecl>(Info.ValueD)) {
+      return AddCodableContext(NomDecl);
+    }
+  }
+  // TODO: support extensions
+  // (would need to get synthesized nodes from the main decl,
+  // and only if it's in the same file?)
+  return AddCodableContext();
+}
+
+bool RefactoringActionAddExplicitCodableImplementation::isApplicable(
+    const ResolvedCursorInfo &Tok, DiagnosticEngine &Diag) {
+  return AddCodableContext::getDeclarationContextFromInfo(Tok).isValid();
+}
+
+bool RefactoringActionAddExplicitCodableImplementation::performChange() {
+  auto Context = AddCodableContext::getDeclarationContextFromInfo(CursorInfo);
+
+  SmallString<64> Buffer;
+  llvm::raw_svector_ostream OS(Buffer);
+  Context.printInsertionText(CursorInfo, SM, OS);
+
+  EditConsumer.insertAfter(SM, Context.getInsertStartLoc(), OS.str());
+  return false;
+}
+
 static CharSourceRange
 findSourceRangeToWrapInCatch(const ResolvedCursorInfo &CursorInfo,
                              SourceFile *TheFile,

lib/Sema/DerivedConformanceCodable.cpp

@@ -1669,7 +1669,7 @@ deriveBodyDecodable_enum_init(AbstractFunctionDecl *initDecl, void *) {
           auto *nestedContainerDecl = createKeyedContainer(
               C, funcDC, C.getKeyedDecodingContainerDecl(),
               caseCodingKeys->getDeclaredInterfaceType(),
-              VarDecl::Introducer::Var, C.Id_nestedContainer);
+              VarDecl::Introducer::Let, C.Id_nestedContainer);
 
           auto *nestedContainerCall = createNestedContainerKeyedByForKeyCall(
               C, funcDC, containerExpr, caseCodingKeys, codingKeyCase);

test/refactoring/AddCodableImplementation/Outputs/basic/codable.swift.expected

@@ -0,0 +1,35 @@
+
+struct User: Codable {
+  let firstName: String
+  let lastName: String?
+
+  enum CodingKeys: CodingKey {
+    case firstName
+    case lastName
+  }
+
+  init(from decoder: Decoder) throws {
+    let container: KeyedDecodingContainer<User.CodingKeys> = try decoder.container(keyedBy: User.CodingKeys.self)
+
+    self.firstName = try container.decode(String.self, forKey: User.CodingKeys.firstName)
+    self.lastName = try container.decodeIfPresent(String.self, forKey: User.CodingKeys.lastName)
+
+  }
+
+  func encode(to encoder: Encoder) throws {
+    var container = encoder.container(keyedBy: User.CodingKeys.self)
+
+    try container.encode(self.firstName, forKey: User.CodingKeys.firstName)
+    try container.encodeIfPresent(self.lastName, forKey: User.CodingKeys.lastName)
+  }
+}
+
+struct User_D: Decodable {
+  let firstName: String
+  let lastName: String?
+}
+
+struct User_E: Encodable {
+  let firstName: String
+  let lastName: String?
+}

test/refactoring/AddCodableImplementation/Outputs/basic/decodable.swift.expected

@@ -0,0 +1,28 @@
+
+struct User: Codable {
+  let firstName: String
+  let lastName: String?
+}
+
+struct User_D: Decodable {
+  let firstName: String
+  let lastName: String?
+
+  enum CodingKeys: CodingKey {
+    case firstName
+    case lastName
+  }
+
+  init(from decoder: Decoder) throws {
+    let container: KeyedDecodingContainer<User_D.CodingKeys> = try decoder.container(keyedBy: User_D.CodingKeys.self)
+
+    self.firstName = try container.decode(String.self, forKey: User_D.CodingKeys.firstName)
+    self.lastName = try container.decodeIfPresent(String.self, forKey: User_D.CodingKeys.lastName)
+
+  }
+}
+
+struct User_E: Encodable {
+  let firstName: String
+  let lastName: String?
+}

test/refactoring/AddCodableImplementation/Outputs/basic/encodable.swift.expected

@@ -0,0 +1,27 @@
+
+struct User: Codable {
+  let firstName: String
+  let lastName: String?
+}
+
+struct User_D: Decodable {
+  let firstName: String
+  let lastName: String?
+}
+
+struct User_E: Encodable {
+  let firstName: String
+  let lastName: String?
+
+  enum CodingKeys: CodingKey {
+    case firstName
+    case lastName
+  }
+
+  func encode(to encoder: Encoder) throws {
+    var container: KeyedEncodingContainer<User_E.CodingKeys> = encoder.container(keyedBy: User_E.CodingKeys.self)
+
+    try container.encode(self.firstName, forKey: User_E.CodingKeys.firstName)
+    try container.encodeIfPresent(self.lastName, forKey: User_E.CodingKeys.lastName)
+  }
+}

test/refactoring/AddCodableImplementation/Outputs/empty/codable.swift.expected

@@ -0,0 +1,15 @@
+
+struct User: Codable {
+
+    enum CodingKeys: CodingKey {
+    }
+
+    init(from decoder: Decoder) throws {
+
+    }
+
+    func encode(to encoder: Encoder) throws {
+      var container = encoder.container(keyedBy: User.CodingKeys.self)
+
+    }
+}