[LINQ] Flatten() Method*

Update

The point was raised that since it is LINQ, it is strange for it not to be an extension method of IEnumerable<T>.

Although this added a requirement to the node side of the result tree, it can now be easily called via node.Flatten().

Source code

Execution environment: https://dotnetfiddle.net/

using System;
using System.Collections;
using System.Collections.Generic;
using System.Runtime.CompilerServices;

namespace X
{
    class Node : IEnumerable<Node>  // node implements recursion interface
    {
        public Node? Parent;
        public List<Node> Children = new();

        public IEnumerator<Node> GetEnumerator() => Children.GetEnumerator();
        IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();

        public void Add(Node node) => Children.Add(node);
    }

    static class FlattenLinqExtensions
    {
        public static IEnumerable<T> Flatten<T>(this IEnumerable<T> source, FlattenStrategy strategy = FlattenStrategy.DepthFirst)
            where T : IEnumerable<T>
        {
            if (source is T self)
            {
                yield return self;
            }

            foreach (var child in source)
            {
                foreach (var x in child.Flatten(strategy))
                {
                    yield return x;
                }
            }
        }
    }

    enum FlattenStrategy
    {
        DepthFirst,
        BreadthFirst,
    }

    class Program
    {
        static void Main()
        {
            var node = new Node()
            {
                new Node()
                {
                    new Node(),
                    new Node(),
                },
                new Node()
                {
                    new Node(),
                    new Node(),
                },
                new Node()
                {
                    new Node(),
                    new Node(),
                },
            };

            var allNodes = node.Flatten();

            int i = 0;
            foreach (var n in node.Flatten())
            {
                i++;
                Console.WriteLine($"{i}: {n.GetHashCode()}");
            }
        }
    }
}

For cases where the node-side implementation cannot be modified, wouldn't it be perfect if we made the initial version's method a static method as:

IEnumerable<T> Enumerable.Flatten(T source, Func<T, IEnumerable<T>> selectMany);

?

Initial Version

I was doing various things with SelectMany to flatten a tree structure, when I suddenly had an idea.

https://x.com/sator_imaging/status/1942482759420887224

👇 Driven by the momentum, I even went ahead and created an Issue. If you think it would be good to have in the standard library, please give it an Upvote!

https://github.com/dotnet/runtime/issues/117409

static IEnumerable<T> Flatten<T>(this T source, Func<T, IEnumerable<T>> selectMany)
{
    yield return node;

    foreach (var child in selectMany.Invoke(node))
    {
        foreach (var x in Flatten(child, selectMany))
        {
            yield return x;
        }
    }
}

That's all. Thanks for reading.