With power comes responsibility: Spot the bug #2 Solution, yield return and partial execution

A couple of days ago I posted a piece of code that made some trouble under certain circumstances. At a first glance it seemed odd, but looking deeper into it, it behaved absolutely correctly.

We used a base class for a Query-Object called Finder (implementing IFinder). A finder should be instantiated once per use, and hence never be used twice. That makes it easier to create complicated query-logic using member variables. Just to protect the deriving classes, we wanted to make sure that the implementation is never called twice.

I slightly simplified the code since the last post.

public abstract class Finder
{
    private bool _used;
    public IEnumerable<int> GetIds()
    {
        if (_used)
            throw new InvalidOperationException("Please do not use twice!");

        _used = true;

        return getIds();
    }

    protected abstract IEnumerable<int> getIds();
}

public class YieldingFinder : Finder
{
    private bool _executed;
    protected override IEnumerable<int> getIds()
    {
        if (_executed)
            throw new ApplicationException("booom!");
        _executed = true;
        yield return 1;
    }
}

Line 6 and 7 will make sure, that YieldingFinder.getIds() is called only once. But when using yield, the C# compiler factors the code inside getIds() into a state machine implementing IEnumerable<int>. Everytime you call getIds(), it will return a instance of this Enumerable.

Generated Source code

When you then instantiate an Enumerator using IEnumerable.GetEnumerator() you’ll get a fresh state machine which moves forward through it’s states on every call to MoveNext(). It holds a reference to the parent class YieldingFinder, which is used for member variable access.

In other words, it creates a backdoor to the protected method. The base class lost control.

public YieldingFinder <>4__this;
private bool MoveNext()
{
    switch (this.<>1__state)
    {
        case 0:
            this.<>1__state = -1;
            if (this.<>4__this._executed)
            {
                throw new ApplicationException("booom!");
            }
            this.<>4__this._executed = true;
            this.<>2__current = 1;
            this.<>1__state = 1;
            return true;
         case 1:
            this.<>1__state = -1;
            break;
    }
    return false;
}

[DebuggerHidden]
void IEnumerator.Reset()
{
    throw new NotSupportedException();
}

Full source code here.

The Bug

In our concrete example you can simply provoke the ApplicationException to be thrown, by retrieving the enumerator once and then enumerate it twice.

var ids = new YieldingFinder().GetIds();
ids.ToArray();
ids.ToArray();

In other words: On the same instance of YieldingFinder the check if it is used twice (Finder.GetIds()) runs only once, but the code checking if the custom code YieldingFinder.getIds() still run twice.

Tests-first

These are the test cases I expect to pass:

  • Second call to GetIds() should fail
  • Second call to GetEnumerator() should fail
  • Happily the compiler-generated enumerator is single-use already. It throws a NotSupportedException on Reset.
  • I also test what happens when I enumerate twice, even though if you can’t call GetEnumerator() twice, the test will succeed, too.
[TestClass]
public class FunnyCode
{
    [TestMethod]
    public void GetIds_FailsSecondTime()
    {
        var finder = new YieldingFinder();
        finder.GetIds();
        finder.Invoking(_ => _.GetIds())
            .ShouldThrow<InvalidOperationException>();
    }

    [TestMethod]
    public void GetEnumerator_FailsSecondTime()
    {
        var ids = new YieldingFinder().GetIds();
        ids.GetEnumerator();
        ids.Invoking(_ => _.GetEnumerator())
            .ShouldThrow<InvalidOperationException>();
    }

    [TestMethod]
    public void GetIds_EnumeratingTwice_FailsSecondTime()
    {
        var ids = new YieldingFinder().GetIds();
        ids.GetEnumerator();
        ids.Invoking(_ => _.GetEnumerator())
            .ShouldThrow<InvalidOperationException>();
    }

    [TestMethod]
    public void GetEnumerator_EnumerateTwice_ResetNotSupported()
    {
        var ids = new YieldingFinder().GetIds();
        var enumerator = ids.GetEnumerator();
        enumerator.MoveNext();

        enumerator.Invoking(_ => _.Reset())
            .ShouldThrow<NotSupportedException>();
    }
}

Take 1

My first take was to wrap the enumeration in another yield:

public IEnumerable<int> GetIds()
{
    if (_used)
        throw new InvalidOperationException("Please do not use twice!");
     _used = true;
     foreach (var item in getIds())
        yield return item;
}

But then line 3 to 6 will not run until MoveNext() is called the first time. It closes the back door, but you can then call GetIds() and GetEnumerator() as often as you want.

Hence, the tests GetIds_FailsSecondTime and GetEnumerator_FailsSecondTime fail.

Take 2

Now, if you want to prevent GetIds() from beeing called twice, the method has to be split.

private bool _used, _enumerated;
public IEnumerable<int> GetIds()
{
    if (_used)
        throw new InvalidOperationException("Please do not use twice!");

    _used = true;

    return enumerateOnce();
}

private IEnumerable<int> enumerateOnce()
{
    if (_enumerated)
        throw new InvalidOperationException("Please do not enumerate twice!");

    _enumerated = true;
    foreach (var item in getIds())
        yield return item;
}

Now, line 4-9 run when GetIds() is called, and 14-17 run whenever MoveNext() is called on the second Enumerator. This solution is acceptable, but it is still later than it could be. GetEnumerator_FailsSecondTime still fails.

Take 3, Final

The best solution is to wrap the enumerator in a class, that intercepts the call to GetEnumerator(). This class could be abstracted into a SingleUseEnumerator, if needed in other classes, too.

It passes all tests.

Here is the code for the complete base class:

public abstract class Finder
{
    private bool _used;
    public IEnumerable<int> GetIds()
    {
        if (_used)
            throw new InvalidOperationException("Please do not use twice!");

        _used = true;

        return new EnumerateOnce(getIds());
    }

    private class EnumerateOnce : IEnumerable<int>
    {
        private bool _enumerated;
        private IEnumerable<int> _baseEnumerator;

        public EnumerateOnce(IEnumerable<int> baseEnumerator)
        {
            _baseEnumerator = baseEnumerator;
        }

        public IEnumerator<int> GetEnumerator()
        {
            if (_enumerated)
                throw new InvalidOperationException("Please do not enumerate twice!");

            _enumerated = true;

            return _baseEnumerator.GetEnumerator();
        }

        IEnumerator IEnumerable.GetEnumerator()
        {
            return GetEnumerator();
        }
    }

    protected abstract IEnumerable<int> getIds();
}

As said, with power comes responsibility. It will be fun to see what people (including me) manage to do with async and await 🙂

Advertisements