Short circuit yield return & cleanup/dispose

Question

Take this pseudo example code:

    static System.Runtime.InteropServices.ComTypes.IEnumString GetUnmanagedObject() => null;
static IEnumerable<string> ProduceStrings()
{
    System.Runtime.InteropServices.ComTypes.IEnumString obj = GetUnmanagedObject();
    var result = new string[1];
    var pFetched = Marshal.AllocHGlobal(sizeof(int));
    while(obj.Next(1, result, pFetched) == 0)
    {
        yield return result[0];
    }
    Marshal.ReleaseComObject(obj);
}

static void Consumer()
{
    foreach (var item in ProduceStrings())
    {
        if (item.StartsWith("foo"))
            return;
    }
}

Question is if i decide to not enumerate all values, how can i inform producer to do cleanup?

Answer 1

Even if you are after a solution using yield return, it might be useful to see how this can be accomplished with an explicit IEnumerator<string> implementation.

IEnumerator<T> derives from IDisposable and the Dispose() method will be called when foreach is left (at least since .NET 1.2, see here)

static IEnumerable<string> ProduceStrings()
{
    return new ProduceStringsImpl();
}

This is the class implementing IEnumerable<string>

class ProduceStringsImpl : IEnumerable<string>
{
    public IEnumerator<string> GetEnumerator()
    {
        return new EnumProduceStrings();
    }
    IEnumerator IEnumerable.GetEnumerator()
    {
        return GetEnumerator();
    }
}

And here we have the core of the solution, the IEnumerator<string> implementation:

class EnumProduceStrings : IEnumerator<string>
{
    private System.Runtime.InteropServices.ComTypes.IEnumString _obj;
    private string[] _result;
    private IntPtr _pFetched;
    
    public EnumProduceStrings()
    {
        _obj = GetUnmanagedObject();
        _result = new string[1];
        _pFetched = Marshal.AllocHGlobal(sizeof(int));
    }
    
    public bool MoveNext()
    {
        return _obj.Next(1, _result, _pFetched) == 0;
    }
    
    public string Current => _result[0];
    
    void IEnumerator.Reset() => throw new NotImplementedException();
    object IEnumerator.Current => Current;
    
    public void Dispose()
    {
        Marshal.ReleaseComObject(_obj);
        Marshal.FreeHGlobal(_pFetched);
    }
}

Answer 2

I knew i can! Despite guard, Cancel is called only one time in all circumtances.

You can instead encapsulate logic with a type like IterationResult<T> and provide Cleanup method on it but its essentially same idea.

public class IterationCanceller
{
    Action m_OnCancel;
    public bool Cancelled { get; private set; }
    public IterationCanceller(Action onCancel)
    {
        m_OnCancel = onCancel;
    }
    public void Cancel()
    {
        if (!Cancelled)
        {
            Cancelled = true;
            m_OnCancel();
        }
    }
}
static IEnumerable<(string Result, IterationCanceller Canceller)> ProduceStrings()
{
    var pUnmanaged = Marshal.AllocHGlobal(sizeof(int));
    IterationCanceller canceller = new IterationCanceller(() =>
    {
        Marshal.FreeHGlobal(pUnmanaged);
    });
    for (int i = 0; i < 2; i++) // also try i < 0, 1
    {
        yield return (i.ToString(), canceller);
    }
    canceller.Cancel();
}

static void Consumer()
{
    foreach (var (item, canceller) in ProduceStrings())
    {
        if(item.StartsWith("1")) // also try consuming all values
        {
            canceller.Cancel();
            break;
        }
    }
}

Answer 3

I think I've found a possible better solution for this issue using an IDisposable implementation. Similar to how @ömer hayyam suggested, you can encapsulate the cleanup logic in a class, but additionally you can use a Using statement to ensure that the cleanup code is executed regardless of a break from the IEnumerable:

class Cleanup : IDisposable
{
    public delegate void CleanupFunction();

    private readonly CleanupFunction _cleanup;

    public Cleanup(CleanupFunction cleanup)
    {
        _cleanup = cleanup;
    }

    public void Dispose()
    {
        _cleanup();
    }
}

class Test
{
    private static IEnumerable<int> InfGen()
    {
        var storage = new List<int>();
        using var unused = new Cleanup(() => storage.ForEach(Console.WriteLine));
        var i = 0;
        while (true)
        {
            yield return i;
            storage.Add(i);
            i++;
        }
    }

    public static void Main(string[] args)
    {
        foreach (int i in InfGen())
        {
            Console.WriteLine("Before: " + i);
            if (i > 20)
            {
                break;
            }
        }
    }
}

After the break from the foreach is executed, the IEnumerable is cancelled but the IDisposable is still called, executing whatever was passed to Cleanup (in this case just a Console.Writeline).