// Copyright (c) Microsoft Open Technologies, Inc. All rights reserved. See License.txt in the project root for license information.

namespace System.Data.Entity.SqlServer.Utilities
{
    using System.Collections.Generic;
    using System.Data.Entity.Core.Common.CommandTrees;
    using System.Linq;

    internal static class DbExpressionExtensions
    {
        /// <summary>
        ///     Uses a stack to non-recursively traverse a given tree structure and retrieve the leaf nodes.
        /// </summary>
        /// <param name="root"> The node that represents the root of the tree. </param>
        /// <param name="kind"> Expressions not of this kind are considered leaves. </param>
        /// <param name="getChildNodes">
        ///     A function that traverses the tree by retrieving the <b>immediate</b> descendants of a (non-leaf) node.
        /// </param>
        /// <returns> An enumerable containing the leaf nodes. </returns>
        public static IEnumerable<DbExpression> GetLeafNodes(
            this DbExpression root,
            DbExpressionKind kind,
            Func<DbExpression, IEnumerable<DbExpression>> getChildNodes)
        {
            DebugCheck.NotNull(getChildNodes);

            var nodes = new Stack<DbExpression>();
            nodes.Push(root);

            while (nodes.Count > 0)
            {
                var current = nodes.Pop();
                if (current.ExpressionKind != kind)
                {
                    yield return current;
                }
                else
                {
                    foreach (var node in getChildNodes(current).Reverse())
                    {
                        nodes.Push(node);
                    }
                }
            }
        }
    }
}
