// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
// See the LICENSE file in the project root for more information.
 
using System;
using System.Collections.Generic;
using System.Collections.Immutable;
using System.Linq;
using Microsoft.CodeAnalysis.PooledObjects;
 
namespace Roslyn.Utilities;
 
internal static class PooledBuilderExtensions
{
    public static Dictionary<K, V> ToDictionaryAndFree<K, V>(this PooledDictionary<K, V> builders)
        where K : notnull
    {
        var dictionary = new Dictionary<K, V>(builders.Count);
 
        foreach (var (key, items) in builders)
        {
            dictionary.Add(key, items);
        }
 
        builders.Free();
        return dictionary;
    }
 
    public static Dictionary<K, ImmutableArray<V>> ToMultiDictionaryAndFree<K, V>(this PooledDictionary<K, ArrayBuilder<V>> builders)
        where K : notnull
    {
        var dictionary = new Dictionary<K, ImmutableArray<V>>(builders.Count);
 
        foreach (var (key, items) in builders)
        {
            dictionary.Add(key, items.ToImmutableAndFree());
        }
 
        builders.Free();
        return dictionary;
    }
 
    public static ImmutableDictionary<K, ImmutableArray<V>> ToImmutableMultiDictionaryAndFree<K, V>(this PooledDictionary<K, ArrayBuilder<V>> builders)
        where K : notnull
    {
        var result = ImmutableDictionary.CreateBuilder<K, ImmutableArray<V>>();
 
        foreach (var (key, items) in builders)
        {
            result.Add(key, items.ToImmutableAndFree());
        }
 
        builders.Free();
        return result.ToImmutable();
    }
 
    public static void FreeValues<K, V>(this IReadOnlyDictionary<K, ArrayBuilder<V>> builders)
        where K : notnull
    {
        foreach (var (_, items) in builders)
        {
            items.Free();
        }
    }
 
    public static ImmutableArray<T> ToFlattenedImmutableArrayAndFree<T>(this ArrayBuilder<ArrayBuilder<T>> builders)
    {
        try
        {
            if (builders.Count == 0)
            {
                return [];
            }
 
            if (builders.Count == 1)
            {
                return builders[0].ToImmutableAndFree();
            }
 
            var result = ArrayBuilder<T>.GetInstance(builders.Sum(b => b.Count));
 
            foreach (var builder in builders)
            {
                result.AddRange(builder);
                builder.Free();
            }
 
            return result.ToImmutableAndFree();
        }
        finally
        {
            builders.Free();
        }
    }
 
    public static bool HasDuplicates<T>(this ArrayBuilder<T> builder)
        => builder.HasDuplicates(static x => x);
 
    public static bool HasDuplicates<T, U>(this ArrayBuilder<T> builder, Func<T, U> selector)
    {
        switch (builder.Count)
        {
            case 0:
            case 1:
                return false;
 
            case 2:
                return EqualityComparer<U>.Default.Equals(selector(builder[0]), selector(builder[1]));
 
            default:
                {
                    using var _ = PooledHashSet<U>.GetInstance(out var set);
 
                    foreach (var element in builder)
                    {
                        if (!set.Add(selector(element)))
                            return true;
                    }
 
                    return false;
                }
        }
    }
}