Ownership tracking

src/libraries/System.Runtime.InteropServices/gen/ComInterfaceGenerator/ManagedToNativeVTableMethodGenerator.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Collections.Generic;
 using System.Collections.Immutable;
+using System.Diagnostics;
 using System.Linq;
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.CSharp;
@@ -220,6 +221,9 @@ public BlockSyntax GenerateStubBody(int index, ImmutableArray<FunctionPointerUnm
                 allStatements.Add(MarshallerHelpers.CreateSetLastPInvokeErrorStatement(LastErrorIdentifier));
             }
 
+            // ManagedToUnmanaged shouldn't need AssignOut stage
+            Debug.Assert(statements.AssignOut.IsEmpty);
+
             // Return
             if (!_marshallers.IsManagedVoidReturn)
                 allStatements.Add(ReturnStatement(IdentifierName(_context.GetIdentifiers(_marshallers.ManagedReturnMarshaller.TypeInfo).managed)));

src/libraries/System.Runtime.InteropServices/gen/ComInterfaceGenerator/Marshallers/ComInterfaceDispatchMarshallerFactory.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
 using Microsoft.CodeAnalysis.CSharp;
 using Microsoft.CodeAnalysis.CSharp.Syntax;
 using static Microsoft.CodeAnalysis.CSharp.SyntaxFactory;
@@ -34,27 +35,32 @@ public ManagedTypeInfo AsNativeType(TypePositionInfo info) =>
                     IsFunctionPointer: false);
             public IEnumerable<StatementSyntax> Generate(TypePositionInfo info, StubCodeContext context)
             {
-                if (context.CurrentStage != StubCodeContext.Stage.Unmarshal)
+                switch (context.CurrentStage)
                 {
-                    yield break;
+                    case StubCodeContext.Stage.Unmarshal:
+                        var (managed, native) = context.GetIdentifiers(info);
+                        // <managed> = ComWrappers.ComInterfaceDispatch.GetInstance<<managedType>>(<native>);
+                        yield return ExpressionStatement(
+                            AssignmentExpression(SyntaxKind.SimpleAssignmentExpression,
+                                IdentifierName(managed),
+                                InvocationExpression(
+                                    MemberAccessExpression(SyntaxKind.SimpleMemberAccessExpression,
+                                        ParseName(TypeNames.System_Runtime_InteropServices_ComWrappers_ComInterfaceDispatch),
+                                        GenericName(
+                                            Identifier("GetInstance"),
+                                            TypeArgumentList(SingletonSeparatedList(info.ManagedType.Syntax)))),
+                                    ArgumentList(
+                                        SingletonSeparatedList(
+                                            Argument(
+                                                IdentifierName(native)))))));
+                        yield break;
+                    case StubCodeContext.Stage.AssignOut:
+                        Debug.Assert(MarshallerHelpers.MarshalsOut(info, context));
+                        yield return this.GeneratePointerAssignOut(info, context);
+                        yield break;
+                    default:
+                        yield break;
                 }
-
-                var (managed, native) = context.GetIdentifiers(info);
-
-                // <managed> = ComWrappers.ComInterfaceDispatch.GetInstance<<managedType>>(<native>);
-                yield return ExpressionStatement(
-                    AssignmentExpression(SyntaxKind.SimpleAssignmentExpression,
-                        IdentifierName(managed),
-                        InvocationExpression(
-                            MemberAccessExpression(SyntaxKind.SimpleMemberAccessExpression,
-                                ParseName(TypeNames.System_Runtime_InteropServices_ComWrappers_ComInterfaceDispatch),
-                                GenericName(
-                                    Identifier("GetInstance"),
-                                    TypeArgumentList(SingletonSeparatedList(info.ManagedType.Syntax)))),
-                            ArgumentList(
-                                SingletonSeparatedList(
-                                    Argument(
-                                        IdentifierName(native)))))));
             }
 
             public SignatureBehavior GetNativeSignatureBehavior(TypePositionInfo info) => SignatureBehavior.NativeType;

src/libraries/System.Runtime.InteropServices/gen/ComInterfaceGenerator/UnmanagedToManagedStubGenerator.cs

@@ -5,7 +5,6 @@
 using System.Collections.Generic;
 using System.Collections.Immutable;
 using Microsoft.CodeAnalysis;
-using Microsoft.CodeAnalysis.CSharp;
 using Microsoft.CodeAnalysis.CSharp.Syntax;
 using static Microsoft.CodeAnalysis.CSharp.SyntaxFactory;
 
@@ -14,7 +13,6 @@ namespace Microsoft.Interop
     internal sealed class UnmanagedToManagedStubGenerator
     {
         private const string ReturnIdentifier = "__retVal";
-        private const string InvokeSucceededIdentifier = "__invokeSucceeded";
 
         private readonly BoundGenerators _marshallers;
 
@@ -60,41 +58,28 @@ public BlockSyntax GenerateStubBody(ExpressionSyntax methodToInvoke)
                 || !statements.ManagedExceptionCatchClauses.IsEmpty;
             VariableDeclarations declarations = VariableDeclarations.GenerateDeclarationsForUnmanagedToManaged(_marshallers, _context, shouldInitializeVariables);
 
-            if (!statements.GuaranteedUnmarshal.IsEmpty)
-            {
-                setupStatements.Add(MarshallerHelpers.Declare(PredefinedType(Token(SyntaxKind.BoolKeyword)), InvokeSucceededIdentifier, initializeToDefault: true));
-            }
-
-            setupStatements.AddRange(declarations.Initializations);
             setupStatements.AddRange(declarations.Variables);
+            setupStatements.AddRange(declarations.Initializations);
             setupStatements.AddRange(statements.Setup);
 
             List<StatementSyntax> tryStatements = new();
             tryStatements.AddRange(statements.Unmarshal);
+            tryStatements.AddRange(statements.GuaranteedUnmarshal);
 
             tryStatements.Add(statements.InvokeStatement);
 
-            if (!statements.GuaranteedUnmarshal.IsEmpty)
-            {
-                tryStatements.Add(ExpressionStatement(AssignmentExpression(SyntaxKind.SimpleAssignmentExpression,
-                    IdentifierName(InvokeSucceededIdentifier),
-                    LiteralExpression(SyntaxKind.TrueLiteralExpression))));
-            }
-
             tryStatements.AddRange(statements.NotifyForSuccessfulInvoke);
-            tryStatements.AddRange(statements.PinnedMarshal);
             tryStatements.AddRange(statements.Marshal);
+            tryStatements.AddRange(statements.PinnedMarshal);
+
+            tryStatements.AddRange(statements.AssignOut);
+            tryStatements.AddRange(statements.Cleanup);
 
             List<StatementSyntax> allStatements = setupStatements;
             List<StatementSyntax> finallyStatements = new();
-            if (!statements.GuaranteedUnmarshal.IsEmpty)
-            {
-                finallyStatements.Add(IfStatement(IdentifierName(InvokeSucceededIdentifier), Block(statements.GuaranteedUnmarshal)));
-            }
 
             SyntaxList<CatchClauseSyntax> catchClauses = List(statements.ManagedExceptionCatchClauses);
 
-            finallyStatements.AddRange(statements.Cleanup);
             if (finallyStatements.Count > 0)
             {
                 allStatements.Add(
@@ -110,6 +95,7 @@ public BlockSyntax GenerateStubBody(ExpressionSyntax methodToInvoke)
                 allStatements.AddRange(tryStatements);
             }
 
+
             // Return
             if (!_marshallers.IsUnmanagedVoidReturn)
                 allStatements.Add(ReturnStatement(IdentifierName(_context.GetIdentifiers(_marshallers.NativeReturnMarshaller.TypeInfo).native)));

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

@@ -4,6 +4,7 @@
 using System;
 using System.Collections.Generic;
 using System.Collections.Immutable;
+using System.Diagnostics;
 using Microsoft.CodeAnalysis.CSharp;
 using Microsoft.CodeAnalysis.CSharp.Syntax;
 using static Microsoft.CodeAnalysis.CSharp.SyntaxFactory;
@@ -183,6 +184,9 @@ public BlockSyntax GeneratePInvokeBody(string dllImportName)
                 allStatements.Add(MarshallerHelpers.CreateSetLastPInvokeErrorStatement(LastErrorIdentifier));
             }
 
+            // ManagedToUnmanaged shouldn't need AssignOut stage
+            Debug.Assert(statements.AssignOut.IsEmpty);
+
             // Return
             if (!_marshallers.IsManagedVoidReturn)
                 allStatements.Add(ReturnStatement(IdentifierName(_context.GetIdentifiers(_marshallers.ManagedReturnMarshaller.TypeInfo).managed)));

src/libraries/System.Runtime.InteropServices/gen/Microsoft.Interop.SourceGeneration/BooleanMarshallingInfoProvider.cs

@@ -1,9 +1,6 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-using System;
-using System.Collections.Generic;
-using System.Text;
 using Microsoft.CodeAnalysis;
 
 namespace Microsoft.Interop

src/libraries/System.Runtime.InteropServices/gen/Microsoft.Interop.SourceGeneration/GeneratedStatements.cs

@@ -22,7 +22,7 @@ public struct GeneratedStatements
         public ImmutableArray<StatementSyntax> NotifyForSuccessfulInvoke { get; init; }
         public ImmutableArray<StatementSyntax> GuaranteedUnmarshal { get; init; }
         public ImmutableArray<StatementSyntax> Cleanup { get; init; }
-
+        public ImmutableArray<StatementSyntax> AssignOut { get; init; }
         public ImmutableArray<CatchClauseSyntax> ManagedExceptionCatchClauses { get; init; }
 
         public static GeneratedStatements Create(BoundGenerators marshallers, StubCodeContext context)
@@ -39,7 +39,8 @@ public static GeneratedStatements Create(BoundGenerators marshallers, StubCodeCo
                 NotifyForSuccessfulInvoke = GenerateStatementsForStubContext(marshallers, context with { CurrentStage = StubCodeContext.Stage.NotifyForSuccessfulInvoke }),
                 GuaranteedUnmarshal = GenerateStatementsForStubContext(marshallers, context with { CurrentStage = StubCodeContext.Stage.GuaranteedUnmarshal }),
                 Cleanup = GenerateStatementsForStubContext(marshallers, context with { CurrentStage = StubCodeContext.Stage.Cleanup }),
-                ManagedExceptionCatchClauses = GenerateCatchClauseForManagedException(marshallers, context)
+                ManagedExceptionCatchClauses = GenerateCatchClauseForManagedException(marshallers, context),
+                AssignOut = GenerateStatementsForStubContext(marshallers, context with { CurrentStage = StubCodeContext.Stage.AssignOut }),
             };
         }
         public static GeneratedStatements Create(BoundGenerators marshallers, StubCodeContext context, ExpressionSyntax expressionToInvoke)
@@ -71,7 +72,26 @@ private static ImmutableArray<StatementSyntax> GenerateStatementsForStubContext(
             ImmutableArray<StatementSyntax>.Builder statementsToUpdate = ImmutableArray.CreateBuilder<StatementSyntax>();
             foreach (BoundGenerator marshaller in marshallers.SignatureMarshallers)
             {
-                statementsToUpdate.AddRange(marshaller.Generator.Generate(marshaller.TypeInfo, context));
+                StubCodeContext localContext = context;
+                // Modify the context if we need to marshal out
+                switch (context.CurrentStage)
+                {
+                    case StubCodeContext.Stage.Marshal:
+                    case StubCodeContext.Stage.PinnedMarshal:
+                        if (MarshallerHelpers.MarshalsOut(marshaller.TypeInfo, context))
+                            localContext = new MarshalOutContext(context);
+                        break;
+                    case StubCodeContext.Stage.AssignOut:
+                        if (!MarshallerHelpers.MarshalsOut(marshaller.TypeInfo, context))
+                            continue;
+                        localContext = new AssignOutContext(context, marshaller.Generator.AsParameter(marshaller.TypeInfo, context).Identifier.ToString());
+                        break;
+                    case StubCodeContext.Stage.CleanupFailure:
+                        if (!MarshallerHelpers.MarshalsOut(marshaller.TypeInfo, context))
+                            continue;
+                        break;
+                }
+                statementsToUpdate.AddRange(marshaller.Generator.Generate(marshaller.TypeInfo, localContext));
             }
 
             if (statementsToUpdate.Count > 0)
@@ -150,7 +170,10 @@ private static ImmutableArray<CatchClauseSyntax> GenerateCatchClauseForManagedEx
         {
             if (!marshallers.HasManagedExceptionMarshaller)
             {
-                return ImmutableArray<CatchClauseSyntax>.Empty;
+                ImmutableArray.Create(CatchClause(
+                    CatchDeclaration(ParseTypeName(TypeNames.System_Exception)),
+                filter: null,
+                Block(List(GenerateStatementsForStubContext(marshallers, new MarshalOutContext(context with { CurrentStage = StubCodeContext.Stage.CleanupFailure }))))));
             }
             ImmutableArray<StatementSyntax>.Builder catchClauseBuilder = ImmutableArray.CreateBuilder<StatementSyntax>();
 
@@ -164,6 +187,15 @@ private static ImmutableArray<CatchClauseSyntax> GenerateCatchClauseForManagedEx
             catchClauseBuilder.AddRange(
                 managedExceptionMarshaller.Generator.Generate(
                     managedExceptionMarshaller.TypeInfo, context with { CurrentStage = StubCodeContext.Stage.PinnedMarshal }));
+            catchClauseBuilder.AddRange(GenerateStatementsForStubContext(marshallers, new MarshalOutContext(context with { CurrentStage = StubCodeContext.Stage.CleanupFailure })));
+            if (!marshallers.IsUnmanagedVoidReturn)
+            {
+                catchClauseBuilder.Add(ReturnStatement(IdentifierName(context.GetIdentifiers(marshallers.NativeReturnMarshaller.TypeInfo).native)));
+            }
+            else
+            {
+                catchClauseBuilder.Add(ReturnStatement());
+            }
             return ImmutableArray.Create(
                 CatchClause(
                     CatchDeclaration(ParseTypeName(TypeNames.System_Exception), Identifier(managed)),
@@ -183,8 +215,10 @@ private static SyntaxTriviaList GenerateStageTrivia(StubCodeContext.Stage stage)
                 StubCodeContext.Stage.UnmarshalCapture => "Capture the native data into marshaller instances in case conversion to managed data throws an exception.",
                 StubCodeContext.Stage.Unmarshal => "Convert native data to managed data.",
                 StubCodeContext.Stage.Cleanup => "Perform required cleanup.",
+                StubCodeContext.Stage.CleanupFailure => "Perform required cleanup.",
                 StubCodeContext.Stage.NotifyForSuccessfulInvoke => "Keep alive any managed objects that need to stay alive across the call.",
                 StubCodeContext.Stage.GuaranteedUnmarshal => "Convert native data to managed data even in the case of an exception during the non-cleanup phases.",
+                StubCodeContext.Stage.AssignOut => "Assign to parameters",
                 _ => throw new ArgumentOutOfRangeException(nameof(stage))
             };
 

src/libraries/System.Runtime.InteropServices/gen/Microsoft.Interop.SourceGeneration/Marshalling/AssignOutContext.cs

@@ -0,0 +1,35 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Diagnostics;
+
+namespace Microsoft.Interop
+{
+    public sealed record AssignOutContext : StubCodeContext
+    {
+        internal StubCodeContext InnerContext { get; init; }
+        public string ParameterIdentifier { get; init; }
+
+        internal AssignOutContext(StubCodeContext inner, string parameterIdentifier)
+        {
+            InnerContext = inner;
+            Debug.Assert(inner.CurrentStage == Stage.AssignOut);
+            CurrentStage = Stage.AssignOut;
+            Direction = inner.Direction;
+            ParentContext = inner.ParentContext;
+            ParameterIdentifier = parameterIdentifier;
+        }
+
+        public override (TargetFramework framework, Version version) GetTargetFramework() => InnerContext.GetTargetFramework();
+
+        public override bool SingleFrameSpansNativeContext => InnerContext.SingleFrameSpansNativeContext;
+
+        public override bool AdditionalTemporaryStateLivesAcrossStages => InnerContext.AdditionalTemporaryStateLivesAcrossStages;
+
+        public override (string managed, string native) GetIdentifiers(TypePositionInfo info)
+            => (InnerContext.GetIdentifiers(info).managed, InnerContext.GetAdditionalIdentifier(info, "out"));
+
+        public override string GetAdditionalIdentifier(TypePositionInfo info, string name) => InnerContext.GetAdditionalIdentifier(info, name);
+    }
+}