@@ -3354,7 +3354,8 @@ class AddEquatableContext {
33543354 AddEquatableContext () : DC(nullptr ), Adopter(), ProtocolsLocations(),
33553355 Protocols(), StoredProperties(), Range(nullptr , nullptr ) {};
33563356
3357- static AddEquatableContext getDeclarationContextFromInfo (ResolvedCursorInfo Info);
3357+ static AddEquatableContext
3358+ getDeclarationContextFromInfo (const ResolvedCursorInfo &Info);
33583359
33593360 std::string getInsertionTextForProtocol ();
33603361
@@ -3468,7 +3469,7 @@ getProtocolRequirements() {
34683469}
34693470
34703471AddEquatableContext AddEquatableContext::
3471- getDeclarationContextFromInfo (ResolvedCursorInfo Info) {
3472+ getDeclarationContextFromInfo (const ResolvedCursorInfo & Info) {
34723473 if (Info.isInvalid ()) {
34733474 return AddEquatableContext ();
34743475 }
@@ -3526,6 +3527,177 @@ performChange() {
35263527 return false ;
35273528}
35283529
3530+ class AddCodableContext {
3531+
3532+ // / Declaration context
3533+ DeclContext *DC;
3534+
3535+ // / Start location of declaration context brace
3536+ SourceLoc StartLoc;
3537+
3538+ // / Array of all conformed protocols
3539+ SmallVector<swift::ProtocolDecl *, 2 > Protocols;
3540+
3541+ // / Range of internal members in declaration
3542+ DeclRange Range;
3543+
3544+ bool conformsToCodableProtocol () {
3545+ for (ProtocolDecl *Protocol : Protocols) {
3546+ if (Protocol->getKnownProtocolKind () == KnownProtocolKind::Encodable ||
3547+ Protocol->getKnownProtocolKind () == KnownProtocolKind::Decodable) {
3548+ return true ;
3549+ }
3550+ }
3551+ return false ;
3552+ }
3553+
3554+ public:
3555+ AddCodableContext (NominalTypeDecl *Decl)
3556+ : DC(Decl), StartLoc(Decl->getBraces ().Start),
3557+ Protocols(Decl->getAllProtocols ()), Range(Decl->getMembers ()){};
3558+
3559+ AddCodableContext (ExtensionDecl *Decl)
3560+ : DC(Decl), StartLoc(Decl->getBraces ().Start),
3561+ Protocols(Decl->getExtendedNominal ()->getAllProtocols()),
3562+ Range(Decl->getMembers ()){};
3563+
3564+ AddCodableContext () : DC(nullptr ), Protocols(), Range(nullptr , nullptr ){};
3565+
3566+ static AddCodableContext
3567+ getDeclarationContextFromInfo (const ResolvedCursorInfo &Info);
3568+
3569+ void printInsertionText (const ResolvedCursorInfo &CursorInfo,
3570+ SourceManager &SM, llvm::raw_ostream &OS);
3571+
3572+ bool isValid () { return StartLoc.isValid () && conformsToCodableProtocol (); }
3573+
3574+ SourceLoc getInsertStartLoc ();
3575+ };
3576+
3577+ SourceLoc AddCodableContext::getInsertStartLoc () {
3578+ SourceLoc MaxLoc = StartLoc;
3579+ for (auto Mem : Range) {
3580+ if (Mem->getEndLoc ().getOpaquePointerValue () >
3581+ MaxLoc.getOpaquePointerValue ()) {
3582+ MaxLoc = Mem->getEndLoc ();
3583+ }
3584+ }
3585+ return MaxLoc;
3586+ }
3587+
3588+ // / Walks an AST and prints the synthesized Codable implementation.
3589+ class SynthesizedCodablePrinter : public ASTWalker {
3590+ private:
3591+ ASTPrinter &Printer;
3592+
3593+ public:
3594+ SynthesizedCodablePrinter (ASTPrinter &Printer) : Printer(Printer) {}
3595+
3596+ bool walkToDeclPre (Decl *D) override {
3597+ auto *VD = dyn_cast<ValueDecl>(D);
3598+ if (!VD)
3599+ return false ;
3600+
3601+ if (!VD->isSynthesized ()) {
3602+ return true ;
3603+ }
3604+ SmallString<32 > Scratch;
3605+ auto name = VD->getName ().getString (Scratch);
3606+ // Print all synthesized enums,
3607+ // since Codable can synthesize multiple enums (for associated values).
3608+ auto shouldPrint =
3609+ isa<EnumDecl>(VD) || name == " init(from:)" || name == " encode(to:)" ;
3610+ if (!shouldPrint) {
3611+ // Some other synthesized decl that we don't want to print.
3612+ return false ;
3613+ }
3614+
3615+ Printer.printNewline ();
3616+
3617+ if (auto enumDecl = dyn_cast<EnumDecl>(D)) {
3618+ // Manually print enum here, since we don't want to print synthesized
3619+ // functions.
3620+ Printer << " enum " << enumDecl->getNameStr ();
3621+ PrintOptions Options;
3622+ Options.PrintSpaceBeforeInheritance = false ;
3623+ enumDecl->printInherited (Printer, Options);
3624+ Printer << " {" ;
3625+ for (Decl *EC : enumDecl->getAllElements ()) {
3626+ Printer.printNewline ();
3627+ Printer << " " ;
3628+ EC->print (Printer, Options);
3629+ }
3630+ Printer.printNewline ();
3631+ Printer << " }" ;
3632+ return false ;
3633+ }
3634+
3635+ PrintOptions Options;
3636+ Options.SynthesizeSugarOnTypes = true ;
3637+ Options.FunctionDefinitions = true ;
3638+ Options.VarInitializers = true ;
3639+ Options.PrintExprs = true ;
3640+ Options.TypeDefinitions = true ;
3641+ Options.ExcludeAttrList .push_back (DAK_HasInitialValue);
3642+
3643+ Printer.printNewline ();
3644+ D->print (Printer, Options);
3645+
3646+ return false ;
3647+ }
3648+ };
3649+
3650+ void AddCodableContext::printInsertionText (const ResolvedCursorInfo &CursorInfo,
3651+ SourceManager &SM,
3652+ llvm::raw_ostream &OS) {
3653+ StringRef ExtraIndent;
3654+ StringRef CurrentIndent =
3655+ Lexer::getIndentationForLine (SM, getInsertStartLoc (), &ExtraIndent);
3656+ std::string Indent;
3657+ if (getInsertStartLoc () == StartLoc) {
3658+ Indent = (CurrentIndent + ExtraIndent).str ();
3659+ } else {
3660+ Indent = CurrentIndent.str ();
3661+ }
3662+
3663+ ExtraIndentStreamPrinter Printer (OS, Indent);
3664+ Printer.printNewline ();
3665+ SynthesizedCodablePrinter Walker (Printer);
3666+ DC->getAsDecl ()->walk (Walker);
3667+ }
3668+
3669+ AddCodableContext AddCodableContext::getDeclarationContextFromInfo (
3670+ const ResolvedCursorInfo &Info) {
3671+ if (Info.isInvalid ()) {
3672+ return AddCodableContext ();
3673+ }
3674+ if (!Info.IsRef ) {
3675+ if (auto *NomDecl = dyn_cast<NominalTypeDecl>(Info.ValueD )) {
3676+ return AddCodableContext (NomDecl);
3677+ }
3678+ }
3679+ // TODO: support extensions
3680+ // (would need to get synthesized nodes from the main decl,
3681+ // and only if it's in the same file?)
3682+ return AddCodableContext ();
3683+ }
3684+
3685+ bool RefactoringActionAddExplicitCodableImplementation::isApplicable (
3686+ const ResolvedCursorInfo &Tok, DiagnosticEngine &Diag) {
3687+ return AddCodableContext::getDeclarationContextFromInfo (Tok).isValid ();
3688+ }
3689+
3690+ bool RefactoringActionAddExplicitCodableImplementation::performChange () {
3691+ auto Context = AddCodableContext::getDeclarationContextFromInfo (CursorInfo);
3692+
3693+ SmallString<64 > Buffer;
3694+ llvm::raw_svector_ostream OS (Buffer);
3695+ Context.printInsertionText (CursorInfo, SM, OS);
3696+
3697+ EditConsumer.insertAfter (SM, Context.getInsertStartLoc (), OS.str ());
3698+ return false ;
3699+ }
3700+
35293701static CharSourceRange
35303702findSourceRangeToWrapInCatch (const ResolvedCursorInfo &CursorInfo,
35313703 SourceFile *TheFile,
0 commit comments