Add `using System.Diagnostics.CodeAnalysis;` when inserting attributes.

NullabilityInference.Tests/FlowTests.cs

@@ -201,6 +201,7 @@ class Program {
         {
             string program = @"
 using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
 class DataStructure
 {
     Dictionary<string, string> dict = new Dictionary<string, string>();
@@ -220,6 +221,7 @@ class DataStructure
         {
             string program = @"
 using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
 class DataStructure
 {
     Dictionary<string, string> dict = new Dictionary<string, string>();
@@ -239,6 +241,7 @@ class DataStructure
         {
             string program = @"
 using System.Collections.Generic;
+[using]
 class DataStructure
 {
     public bool TryGet(int i, [Attr] out string? name)
@@ -253,8 +256,8 @@ class DataStructure
     }
 }";
             AssertNullabilityInference(
-                expectedProgram: program.Replace("[Attr]", "[NotNullWhen(true)]"),
-                inputProgram: program.Replace("[Attr] ", ""));
+                expectedProgram: program.Replace("[Attr]", "[NotNullWhen(true)]").Replace("[using]", "using System.Diagnostics.CodeAnalysis;"),
+                inputProgram: program.Replace("[Attr] ", "").Replace("[using]", ""));
         }
     }
 }

NullabilityInference/InferredNullabilitySyntaxRewriter.cs

@@ -134,9 +134,29 @@ namespace ICSharpCode.NullabilityInference
                         SyntaxFactory.AttributeArgumentList(SyntaxFactory.SingletonSeparatedList(SyntaxFactory.AttributeArgument(attrArgument))));
                     var newAttributeList = SyntaxFactory.AttributeList(SyntaxFactory.SingletonSeparatedList(newAttribute));
                     node = node.AddAttributeLists(newAttributeList.WithTrailingTrivia(SyntaxFactory.Space));
+                    needsUsingCodeAnalysis = true;
                 }
             }
             return node;
         }
+
+        private bool needsUsingCodeAnalysis;
+
+        public override SyntaxNode? VisitCompilationUnit(CompilationUnitSyntax node)
+        {
+            bool hasUsingCodeAnalysis = false;
+            foreach (var u in node.Usings) {
+                var symbolInfo = semanticModel.GetSymbolInfo(u.Name, cancellationToken);
+                if (symbolInfo.Symbol is INamespaceSymbol ns && ns.GetFullName() == "System.Diagnostics.CodeAnalysis") {
+                    hasUsingCodeAnalysis = true;
+                }
+            }
+            node = (CompilationUnitSyntax)base.VisitCompilationUnit(node)!;
+            if (needsUsingCodeAnalysis && !hasUsingCodeAnalysis) {
+                var qname = SyntaxFactory.QualifiedName(SyntaxFactory.QualifiedName(SyntaxFactory.IdentifierName("System"), SyntaxFactory.IdentifierName("Diagnostics")), SyntaxFactory.IdentifierName("CodeAnalysis"));
+                node = node.AddUsings(SyntaxFactory.UsingDirective(qname.WithLeadingTrivia(SyntaxFactory.Space)));
+            }
+            return node;
+        }
     }
 }

README.md

@@ -15,7 +15,7 @@ Note: this is a work in progress. Many C# constructs will trigger a NotImplement
 
 ## Tips+Tricks:
 
- * The inference tool will only add/remove `?` annotations on nullable reference types. It will never touch your code in any other way.
+ * The inference tool will only add/remove `?` annotations on nullable reference types. It can also add the `[NotNullWhen]` attribute. It will never touch your code in any other way.
    * Existing `?` annotations on nullable reference types are discarded and inferred again from scratch.
  * Unconstrained generic types are not reference types, and thus will never be annotated by the tool.
  * The inference tool will not introduce any of the [advanced nullability attributes](https://docs.microsoft.com/en-us/dotnet/api/system.diagnostics.codeanalysis).
@@ -215,7 +215,7 @@ and on the existance of code that uses the return value of `Get` without null ch
 
 ```csharp
 01: using System.Collections.Generic;
-02: using System.Diagnostics.CodeAnalysis;
+02: 
 03: class Program
 04: {
 05:     public string someString = "hello";