Fix rewriting of configured async enumeration and disposal

Source/Test/Rewriting/Passes/Rewriting/Types/MethodBodyTypeRewritingPass.cs

@@ -116,10 +116,11 @@ private Instruction VisitNewobjInstruction(Instruction instruction)
         {
             MethodReference constructor = instruction.Operand as MethodReference;
             if (this.TryRewriteMethodReference(constructor, "Create", out MethodReference newMethod) &&
-                this.TryResolve(newMethod, out MethodDefinition _))
+                this.TryResolve(newMethod, out MethodDefinition resolvedMethod))
             {
-                // Create and return the new instruction.
-                Instruction newInstruction = Instruction.Create(OpCodes.Call, newMethod);
+                // Only constructor-to-factory rewrites replace allocation with a static call.
+                Instruction newInstruction = Instruction.Create(resolvedMethod.IsConstructor ?
+                    OpCodes.Newobj : OpCodes.Call, newMethod);
                 newInstruction.Offset = instruction.Offset;
 
                 this.LogWriter.LogDebug("............. [-] {0}", instruction);

Source/Test/Rewriting/Passes/Rewriting/Types/TypeRewritingPass.cs

@@ -71,6 +71,15 @@ internal TypeRewritingPass(RewritingOptions options, IEnumerable<AssemblyInfo> v
 #endif
             this.KnownTypes[NameCache.GenericTaskCompletionSource] = typeof(Types.Threading.Tasks.TaskCompletionSource<>);
             this.KnownTypes[NameCache.TaskExtensions] = typeof(Types.TaskExtensions);
+#if NET
+            this.KnownTypes[NameCache.TaskAsyncEnumerableExtensions] = typeof(Types.Threading.Tasks.TaskAsyncEnumerableExtensions);
+            this.KnownTypes[NameCache.ConfiguredCancelableAsyncEnumerable] =
+                typeof(Types.Runtime.CompilerServices.ConfiguredCancelableAsyncEnumerable<>);
+            this.KnownTypes[NameCache.ConfiguredAsyncEnumerator] =
+                typeof(Types.Runtime.CompilerServices.ConfiguredCancelableAsyncEnumerable<>.Enumerator);
+            this.KnownTypes[NameCache.ConfiguredAsyncDisposable] =
+                typeof(Types.Runtime.CompilerServices.ConfiguredAsyncDisposable);
+#endif
             this.KnownTypes[NameCache.TaskFactory] = typeof(Types.Threading.Tasks.TaskFactory);
             this.KnownTypes[NameCache.GenericTaskFactory] = typeof(Types.Threading.Tasks.TaskFactory<>);
             this.KnownTypes[NameCache.TaskParallel] = typeof(Types.Threading.Tasks.Parallel);

Source/Test/Rewriting/Types/NameCache.cs

@@ -64,6 +64,14 @@ internal static class NameCache
             typeof(SystemCompiler.YieldAwaitable).FullName + "/YieldAwaiter";
 
         internal static string TaskExtensions { get; } = typeof(SystemTasks.TaskExtensions).FullName;
+#if NET
+        internal static string TaskAsyncEnumerableExtensions { get; } = typeof(SystemTasks.TaskAsyncEnumerableExtensions).FullName;
+        internal static string ConfiguredCancelableAsyncEnumerable { get; } =
+            typeof(SystemCompiler.ConfiguredCancelableAsyncEnumerable<>).FullName;
+        internal static string ConfiguredAsyncEnumerator { get; } =
+            typeof(SystemCompiler.ConfiguredCancelableAsyncEnumerable<>).FullName + "/Enumerator";
+        internal static string ConfiguredAsyncDisposable { get; } = typeof(SystemCompiler.ConfiguredAsyncDisposable).FullName;
+#endif
         internal static string TaskFactory { get; } = typeof(SystemTasks.TaskFactory).FullName;
         internal static string GenericTaskFactory { get; } = typeof(SystemTasks.TaskFactory<>).FullName;
         internal static string TaskParallel { get; } = typeof(SystemTasks.Parallel).FullName;

Source/Test/Rewriting/Types/Runtime/CompilerServices/ConfiguredAsyncDisposable.cs

@@ -0,0 +1,35 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+#if NET
+using System;
+
+namespace Microsoft.Coyote.Rewriting.Types.Runtime.CompilerServices
+{
+    /// <summary>
+    /// Provides an async disposable whose configured await can be controlled during testing.
+    /// </summary>
+    /// <remarks>This type is intended for compiler use only.</remarks>
+    [System.ComponentModel.EditorBrowsable(System.ComponentModel.EditorBrowsableState.Never)]
+    public readonly struct ConfiguredAsyncDisposable
+    {
+        private readonly IAsyncDisposable Disposable;
+        private readonly bool ContinueOnCapturedContext;
+
+        internal ConfiguredAsyncDisposable(IAsyncDisposable disposable, bool continueOnCapturedContext)
+        {
+            this.Disposable = disposable;
+            this.ContinueOnCapturedContext = continueOnCapturedContext;
+        }
+
+        /// <summary>
+        /// Disposes the underlying resource.
+        /// </summary>
+        public ConfiguredValueTaskAwaitable DisposeAsync()
+        {
+            var task = this.Disposable.DisposeAsync();
+            return new ConfiguredValueTaskAwaitable(ref task, this.ContinueOnCapturedContext);
+        }
+    }
+}
+#endif

Source/Test/Rewriting/Types/Runtime/CompilerServices/ConfiguredCancelableAsyncEnumerable.cs

@@ -0,0 +1,87 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+#if NET
+using System.Collections.Generic;
+using System.Threading;
+
+namespace Microsoft.Coyote.Rewriting.Types.Runtime.CompilerServices
+{
+    /// <summary>
+    /// Provides an async enumerable whose configured awaits can be controlled during testing.
+    /// </summary>
+    /// <typeparam name="T">The type of each element.</typeparam>
+    /// <remarks>This type is intended for compiler use only.</remarks>
+    [System.ComponentModel.EditorBrowsable(System.ComponentModel.EditorBrowsableState.Never)]
+    public readonly struct ConfiguredCancelableAsyncEnumerable<T>
+    {
+        private readonly IAsyncEnumerable<T> Enumerable;
+        private readonly CancellationToken CancellationToken;
+        private readonly bool ContinueOnCapturedContext;
+
+        internal ConfiguredCancelableAsyncEnumerable(IAsyncEnumerable<T> enumerable,
+            bool continueOnCapturedContext, CancellationToken cancellationToken)
+        {
+            this.Enumerable = enumerable;
+            this.CancellationToken = cancellationToken;
+            this.ContinueOnCapturedContext = continueOnCapturedContext;
+        }
+
+        /// <summary>
+        /// Configures continuation context capture for iteration and disposal.
+        /// </summary>
+        public ConfiguredCancelableAsyncEnumerable<T> ConfigureAwait(bool continueOnCapturedContext) =>
+            new ConfiguredCancelableAsyncEnumerable<T>(this.Enumerable, continueOnCapturedContext, this.CancellationToken);
+
+        /// <summary>
+        /// Replaces the cancellation token passed to the underlying enumerable.
+        /// </summary>
+        public ConfiguredCancelableAsyncEnumerable<T> WithCancellation(CancellationToken cancellationToken) =>
+            new ConfiguredCancelableAsyncEnumerable<T>(this.Enumerable, this.ContinueOnCapturedContext, cancellationToken);
+
+        /// <summary>
+        /// Creates the configured enumerator.
+        /// </summary>
+        public Enumerator GetAsyncEnumerator() =>
+            new Enumerator(this.Enumerable.GetAsyncEnumerator(this.CancellationToken), this.ContinueOnCapturedContext);
+
+        /// <summary>
+        /// Wraps iteration and disposal in controllable configured awaitables.
+        /// </summary>
+        public readonly struct Enumerator
+        {
+            private readonly IAsyncEnumerator<T> AsyncEnumerator;
+            private readonly bool ContinueOnCapturedContext;
+
+            internal Enumerator(IAsyncEnumerator<T> enumerator, bool continueOnCapturedContext)
+            {
+                this.AsyncEnumerator = enumerator;
+                this.ContinueOnCapturedContext = continueOnCapturedContext;
+            }
+
+            /// <summary>
+            /// Gets the current element.
+            /// </summary>
+            public T Current => this.AsyncEnumerator.Current;
+
+            /// <summary>
+            /// Advances the underlying enumerator.
+            /// </summary>
+            public ConfiguredValueTaskAwaitable<bool> MoveNextAsync()
+            {
+                var task = this.AsyncEnumerator.MoveNextAsync();
+                return new ConfiguredValueTaskAwaitable<bool>(ref task, this.ContinueOnCapturedContext);
+            }
+
+            /// <summary>
+            /// Disposes the underlying enumerator.
+            /// </summary>
+            public ConfiguredValueTaskAwaitable DisposeAsync()
+            {
+                var task = this.AsyncEnumerator.DisposeAsync();
+                return new ConfiguredValueTaskAwaitable(ref task, this.ContinueOnCapturedContext);
+            }
+        }
+    }
+}
+#endif

Source/Test/Rewriting/Types/Threading/Tasks/TaskAsyncEnumerableExtensions.cs

@@ -0,0 +1,40 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+#if NET
+using System;
+using System.Collections.Generic;
+using System.Threading;
+using Microsoft.Coyote.Rewriting.Types.Runtime.CompilerServices;
+
+namespace Microsoft.Coyote.Rewriting.Types.Threading.Tasks
+{
+    /// <summary>
+    /// Configures async enumeration and disposal using controllable awaitables.
+    /// </summary>
+    /// <remarks>This type is intended for compiler use only.</remarks>
+    [System.ComponentModel.EditorBrowsable(System.ComponentModel.EditorBrowsableState.Never)]
+    public static class TaskAsyncEnumerableExtensions
+    {
+        /// <summary>
+        /// Configures context capture when disposing an async resource.
+        /// </summary>
+        public static ConfiguredAsyncDisposable ConfigureAwait(IAsyncDisposable source, bool continueOnCapturedContext) =>
+            new ConfiguredAsyncDisposable(source, continueOnCapturedContext);
+
+        /// <summary>
+        /// Configures context capture when iterating an async enumerable.
+        /// </summary>
+        public static ConfiguredCancelableAsyncEnumerable<T> ConfigureAwait<T>(
+            IAsyncEnumerable<T> source, bool continueOnCapturedContext) =>
+            new ConfiguredCancelableAsyncEnumerable<T>(source, continueOnCapturedContext, default);
+
+        /// <summary>
+        /// Configures the cancellation token used when iterating an async enumerable.
+        /// </summary>
+        public static ConfiguredCancelableAsyncEnumerable<T> WithCancellation<T>(
+            IAsyncEnumerable<T> source, CancellationToken cancellationToken) =>
+            new ConfiguredCancelableAsyncEnumerable<T>(source, true, cancellationToken);
+    }
+}
+#endif