1 files changed, 112 insertions, 0 deletions
diff --git a/modules/mono/editor/Godot.NET.Sdk/Godot.SourceGenerators/ClassPartialModifierAnalyzer.cs b/modules/mono/editor/Godot.NET.Sdk/Godot.SourceGenerators/ClassPartialModifierAnalyzer.cs
new file mode 100644
index 0000000000..e4bdb8db84
--- /dev/null
+++ b/modules/mono/editor/Godot.NET.Sdk/Godot.SourceGenerators/ClassPartialModifierAnalyzer.cs
@@ -0,0 +1,112 @@
+using System.Collections.Immutable;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis;
+using Microsoft.CodeAnalysis.CodeActions;
+using Microsoft.CodeAnalysis.CodeFixes;
+using Microsoft.CodeAnalysis.CSharp;
+using Microsoft.CodeAnalysis.CSharp.Syntax;
+using Microsoft.CodeAnalysis.Diagnostics;
+
+namespace Godot.SourceGenerators
+{
+    [DiagnosticAnalyzer(LanguageNames.CSharp)]
+    public sealed class ClassPartialModifierAnalyzer : DiagnosticAnalyzer
+    {
+        public override ImmutableArray<DiagnosticDescriptor> SupportedDiagnostics =>
+            ImmutableArray.Create(Common.ClassPartialModifierRule, Common.OuterClassPartialModifierRule);
+
+        public override void Initialize(AnalysisContext context)
+        {
+            context.ConfigureGeneratedCodeAnalysis(GeneratedCodeAnalysisFlags.None);
+            context.EnableConcurrentExecution();
+            context.RegisterSyntaxNodeAction(AnalyzeNode, SyntaxKind.ClassDeclaration);
+        }
+
+        private void AnalyzeNode(SyntaxNodeAnalysisContext context)
+        {
+            if (context.Node is not ClassDeclarationSyntax classDeclaration)
+                return;
+
+            if (context.ContainingSymbol is not INamedTypeSymbol typeSymbol)
+                return;
+
+            if (!typeSymbol.InheritsFrom("GodotSharp", GodotClasses.GodotObject))
+                return;
+
+            if (!classDeclaration.IsPartial())
+                context.ReportDiagnostic(Diagnostic.Create(
+                    Common.ClassPartialModifierRule,
+                    classDeclaration.Identifier.GetLocation(),
+                    typeSymbol.ToDisplayString()));
+
+            var outerClassDeclaration = context.Node.Parent as ClassDeclarationSyntax;
+            while (outerClassDeclaration is not null)
+            {
+                var outerClassTypeSymbol = context.SemanticModel.GetDeclaredSymbol(outerClassDeclaration);
+                if (outerClassTypeSymbol == null)
+                    return;
+
+                if (!outerClassDeclaration.IsPartial())
+                    context.ReportDiagnostic(Diagnostic.Create(
+                        Common.OuterClassPartialModifierRule,
+                        outerClassDeclaration.Identifier.GetLocation(),
+                        outerClassTypeSymbol.ToDisplayString()));
+
+                outerClassDeclaration = outerClassDeclaration.Parent as ClassDeclarationSyntax;
+            }
+        }
+    }
+
+    [ExportCodeFixProvider(LanguageNames.CSharp)]
+    public sealed class ClassPartialModifierCodeFixProvider : CodeFixProvider
+    {
+        public override ImmutableArray<string> FixableDiagnosticIds =>
+            ImmutableArray.Create(Common.ClassPartialModifierRule.Id);
+
+        public override FixAllProvider GetFixAllProvider() => WellKnownFixAllProviders.BatchFixer;
+
+        public override async Task RegisterCodeFixesAsync(CodeFixContext context)
+        {
+            // Get the syntax root of the document.
+            var root = await context.Document.GetSyntaxRootAsync(context.CancellationToken).ConfigureAwait(false);
+
+            // Get the diagnostic to fix.
+            var diagnostic = context.Diagnostics.First();
+
+            // Get the location of code issue.
+            var diagnosticSpan = diagnostic.Location.SourceSpan;
+
+            // Use that location to find the containing class declaration.
+            var classDeclaration = root?.FindToken(diagnosticSpan.Start)
+                .Parent?
+                .AncestorsAndSelf()
+                .OfType<ClassDeclarationSyntax>()
+                .First();
+
+            if (classDeclaration == null)
+                return;
+
+            context.RegisterCodeFix(
+                CodeAction.Create(
+                    "Add partial modifier",
+                    cancellationToken => AddPartialModifierAsync(context.Document, classDeclaration, cancellationToken),
+                    classDeclaration.ToFullString()),
+                context.Diagnostics);
+        }
+
+        private static async Task<Document> AddPartialModifierAsync(Document document,
+            ClassDeclarationSyntax classDeclaration, CancellationToken cancellationToken)
+        {
+            // Create a new partial modifier.
+            var partialModifier = SyntaxFactory.Token(SyntaxKind.PartialKeyword);
+            var modifiedClassDeclaration = classDeclaration.AddModifiers(partialModifier);
+            var root = await document.GetSyntaxRootAsync(cancellationToken).ConfigureAwait(false);
+            // Replace the old class declaration with the modified one in the syntax root.
+            var newRoot = root!.ReplaceNode(classDeclaration, modifiedClassDeclaration);
+            var newDocument = document.WithSyntaxRoot(newRoot);
+            return newDocument;
+        }
+    }
+}