FlattenTest

#if DEBUG
using System.Diagnostics;
#endif
using BenchmarkDotNet.Attributes;
#if DEBUG
#else
using BenchmarkDotNet.Running;
#endif

namespace FlattenTest;

internal class Node(IEnumerable<Node>? children)
{
    public IEnumerable<Node>? Children { get; } = children;

    public IEnumerable<Node> GetAllChildrenConcat()
        => new[] { this }.Concat(Children?.SelectMany(child => child.GetAllChildrenConcat()) ?? []);

    public IEnumerable<Node> GetAllChildrenRange()
        => [this, ..Children?.SelectMany(child => child.GetAllChildrenRange()) ?? []];

    public IEnumerable<Node> GetAllChildrenRecursive()
    {
        yield return this;
        if (Children is not null)
        {
            foreach (var child in Children)
            {
                foreach (var node in child.GetAllChildrenRecursive()) yield return node;
            }
        }
    }

    public IEnumerable<Node> GetAllChildrenNonRecursive()
    {
        yield return this;
        if (Children is not null)
        {
            var stack = new Stack<IEnumerable<Node>>();
            stack.Push(Children);
            while (stack.Count != 0)
            {
                var children = stack.Pop();
                foreach (var child in children)
                {
                    yield return child;
                    if (child.Children is not null) stack.Push(child.Children);
                }
            }
        }
    }
}

public class Program
{
    private const int Depth = 8;
    private const int ChildrenCount = 8;
    private static Node root = null!;

    [GlobalSetup]
    public static void Setup()
    {
        root = BuildNode(0);

        static Node BuildNode(int depth)
        {
            if (depth >= Depth) return new Node(null);
            var children = new List<Node>(ChildrenCount);
            for (var i = 0; i < ChildrenCount; i++)
            {
                children.Add(BuildNode(depth + 1));
            }
            return new Node(children);
        }
    }

    [Benchmark]
    public int ConcatTest() => root.GetAllChildrenConcat().Count();

    [Benchmark]
    public int RangeTest() => root.GetAllChildrenRange().Count();

    [Benchmark]
    public int RecursiveTest() => root.GetAllChildrenRecursive().Count();

    [Benchmark]
    public int NonRecursiveTest() => root.GetAllChildrenNonRecursive().Count();

    static void Main()
    {
#if DEBUG
        var program = new Program();
        Setup();
        var concatCount = program.ConcatTest();
        var rangeCount = program.RangeTest();
        var recursiveCount = program.RecursiveTest();
        var nonRecursiveCount = program.NonRecursiveTest();
        Console.WriteLine(concatCount);
        Debug.Assert(concatCount == rangeCount);
        Debug.Assert(concatCount == recursiveCount);
        Debug.Assert(concatCount == nonRecursiveCount);
#else
        _ = BenchmarkRunner.Run<Program>();
#endif
    }
}