Mark containing type for generated stub functions as unsafe (#60979)

Author: elinor-fung

src/libraries/System.Runtime.InteropServices/gen/DllImportGenerator/DllImportGenerator.cs

@@ -145,7 +145,7 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
                     (data, ct) =>
                     {
                         IncrementalTracker?.RecordExecutedStep(new IncrementalityTracker.ExecutedStepInfo(IncrementalityTracker.StepName.CalculateStubInformation, data));
-                        return (data.Syntax, StubContext: CalculateStubInformation(data.Syntax, data.Symbol, data.Environment, ct));
+                        return (data.Syntax, StubContext: CalculateStubInformation(data.Symbol, data.Environment, ct));
                     }
                 )
                 .WithComparer(Comparers.CalculatedContextWithSyntax)
@@ -244,6 +244,17 @@ private static SyntaxTokenList StripTriviaFromModifiers(SyntaxTokenList tokenLis
             return new SyntaxTokenList(strippedTokens);
         }
 
+        private static SyntaxTokenList AddToModifiers(SyntaxTokenList modifiers, SyntaxKind modifierToAdd)
+        {
+            if (modifiers.IndexOf(modifierToAdd) >= 0)
+                return modifiers;
+
+            int idx = modifiers.IndexOf(SyntaxKind.PartialKeyword);
+            return idx >= 0
+                ? modifiers.Insert(idx, Token(modifierToAdd))
+                : modifiers.Add(Token(modifierToAdd));
+        }
+
         private static TypeDeclarationSyntax CreateTypeDeclarationWithoutTrivia(TypeDeclarationSyntax typeDeclaration)
         {
             return TypeDeclaration(
@@ -279,6 +290,9 @@ private static MemberDeclarationSyntax WrapMethodInContainingScopes(DllImportStu
             MemberDeclarationSyntax containingType = CreateTypeDeclarationWithoutTrivia(stub.StubContainingTypes.First())
                 .AddMembers(stubMethod);
 
+            // Mark containing type as unsafe such that all the generated functions will be in an unsafe context.
+            containingType = containingType.WithModifiers(AddToModifiers(containingType.Modifiers, SyntaxKind.UnsafeKeyword));
+
             // Add type to the remaining containing types (skipping the first which was handled above)
             foreach (TypeDeclarationSyntax typeDecl in stub.StubContainingTypes.Skip(1))
             {
@@ -339,8 +353,6 @@ private static GeneratedDllImportData ProcessGeneratedDllImportAttribute(Attribu
             bool setLastError = false;
             bool throwOnUnmappableChar = false;
 
-            var stubDllImportData = new GeneratedDllImportData(attrData.ConstructorArguments[0].Value!.ToString());
-
             // All other data on attribute is defined as NamedArguments.
             foreach (KeyValuePair<string, TypedConstant> namedArg in attrData.NamedArguments)
             {
@@ -398,7 +410,7 @@ private static GeneratedDllImportData ProcessGeneratedDllImportAttribute(Attribu
             };
         }
 
-        private static IncrementalStubGenerationContext CalculateStubInformation(MethodDeclarationSyntax syntax, IMethodSymbol symbol, StubEnvironment environment, CancellationToken ct)
+        private static IncrementalStubGenerationContext CalculateStubInformation(IMethodSymbol symbol, StubEnvironment environment, CancellationToken ct)
         {
             INamedTypeSymbol? lcidConversionAttrType = environment.Compilation.GetTypeByMetadataName(TypeNames.LCIDConversionAttribute);
             INamedTypeSymbol? suppressGCTransitionAttrType = environment.Compilation.GetTypeByMetadataName(TypeNames.SuppressGCTransitionAttribute);
@@ -506,6 +518,10 @@ private static IncrementalStubGenerationContext CalculateStubInformation(MethodD
                 dllImport = dllImport.AddAttributeLists(AttributeList(SeparatedList(forwardedAttributes)));
             }
 
+            dllImport = dllImport.WithLeadingTrivia(
+                Comment("//"),
+                Comment("// Local P/Invoke"),
+                Comment("//"));
             code = code.AddStatements(dllImport);
 
             return (PrintGeneratedSource(originalSyntax, dllImportStub.StubContext, code), dllImportStub.Diagnostics.AddRange(diagnostics.Diagnostics));
@@ -514,7 +530,7 @@ private static IncrementalStubGenerationContext CalculateStubInformation(MethodD
         private MemberDeclarationSyntax PrintForwarderStub(MethodDeclarationSyntax userDeclaredMethod, IncrementalStubGenerationContext stub)
         {
             SyntaxTokenList modifiers = StripTriviaFromModifiers(userDeclaredMethod.Modifiers);
-            modifiers = modifiers.Insert(modifiers.IndexOf(SyntaxKind.PartialKeyword), Token(SyntaxKind.ExternKeyword));
+            modifiers = AddToModifiers(modifiers, SyntaxKind.ExternKeyword);
             // Create stub function
             MethodDeclarationSyntax stubMethod = MethodDeclaration(stub.StubContext.StubReturnType, userDeclaredMethod.Identifier)
                 .WithModifiers(modifiers)

src/libraries/System.Runtime.InteropServices/gen/DllImportGenerator/PInvokeStubCodeGenerator.cs

@@ -215,6 +215,14 @@ public override (string managed, string native) GetIdentifiers(TypePositionInfo
             }
         }
 
+        /// <summary>
+        /// Generate the method body of the p/invoke stub.
+        /// </summary>
+        /// <param name="dllImportName">Name of the target DllImport function to invoke</param>
+        /// <returns>Method body of the p/invoke stub</returns>
+        /// <remarks>
+        /// The generated code assumes it will be in an unsafe context.
+        /// </remarks>
         public BlockSyntax GeneratePInvokeBody(string dllImportName)
         {
             bool invokeReturnsVoid = _retMarshaller.TypeInfo.ManagedType == SpecialTypeInfo.Void;
@@ -330,8 +338,7 @@ public BlockSyntax GeneratePInvokeBody(string dllImportName)
             if (!_stubReturnsVoid)
                 allStatements.Add(ReturnStatement(IdentifierName(ReturnIdentifier)));
 
-            // Wrap all statements in an unsafe block
-            return Block(UnsafeStatement(Block(allStatements)));
+            return Block(allStatements);
 
             void GenerateStatementsForStage(Stage stage, List<StatementSyntax> statementsToUpdate)
             {

src/libraries/System.Runtime.InteropServices/tests/DllImportGenerator.UnitTests/CodeSnippets.cs

@@ -141,6 +141,22 @@ partial class InnerClass
 }
 ";
 
+        /// <summary>
+        /// Containing type with and without unsafe
+        /// </summary>
+        public static readonly string UnsafeContext = @"
+using System.Runtime.InteropServices;
+partial class Test
+{
+    [GeneratedDllImport(""DoesNotExist"")]
+    public static partial void Method1();
+}
+unsafe partial class Test
+{
+    [GeneratedDllImport(""DoesNotExist"")]
+    public static partial int* Method2();
+}
+";
         /// <summary>
         /// Declaration with user defined EntryPoint.
         /// </summary>

src/libraries/System.Runtime.InteropServices/tests/DllImportGenerator.UnitTests/Compiles.cs

@@ -22,6 +22,7 @@ public static IEnumerable<object[]> CodeSnippetsToCompile()
             yield return new[] { CodeSnippets.MultipleAttributes };
             yield return new[] { CodeSnippets.NestedNamespace };
             yield return new[] { CodeSnippets.NestedTypes };
+            yield return new[] { CodeSnippets.UnsafeContext };
             yield return new[] { CodeSnippets.UserDefinedEntryPoint };
             yield return new[] { CodeSnippets.AllSupportedDllImportNamedArguments };
             yield return new[] { CodeSnippets.DefaultParameters };