// Copyright (c) .NET Foundation and contributors. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.

using System;
using System.Collections.Generic;
using System.Diagnostics;
using ILLink.Shared;
using ILLink.Shared.DataFlow;
using Mono.Cecil;
using Mono.Cecil.Cil;
using HoistedLocalState = ILLink.Shared.DataFlow.DefaultValueDictionary<
    Mono.Linker.Dataflow.HoistedLocalKey,
    ILLink.Shared.DataFlow.ValueSet<ILLink.Shared.DataFlow.SingleValue>>;
using MultiValue = ILLink.Shared.DataFlow.ValueSet<ILLink.Shared.DataFlow.SingleValue>;

namespace Mono.Linker.Dataflow
{
    // Tracks the set of methods which get analyzer together during interprocedural analysis,
    // and the possible states of hoisted locals in state machine methods and lambdas/local functions.
    struct InterproceduralState : IEquatable<InterproceduralState>
    {
        public ValueSet<MethodIL> MethodBodies;
        public HoistedLocalState HoistedLocals;
        readonly InterproceduralStateLattice lattice;

        public InterproceduralState(ValueSet<MethodIL> methodBodies, HoistedLocalState hoistedLocals, InterproceduralStateLattice lattice)
            => (MethodBodies, HoistedLocals, this.lattice) = (methodBodies, hoistedLocals, lattice);

        public bool Equals(InterproceduralState other)
            => MethodBodies.Equals(other.MethodBodies) && HoistedLocals.Equals(other.HoistedLocals);

        public override bool Equals(object? obj)
            => obj is InterproceduralState state && Equals(state);

        public override int GetHashCode() => HashUtils.Combine(MethodBodies.GetHashCode(), HoistedLocals.GetHashCode());

        public InterproceduralState Clone()
            => new(MethodBodies.DeepCopy(), HoistedLocals.Clone(), lattice);

        public void TrackMethod(MethodDefinition method)
        {
            if (method.Body is not MethodBody methodBody)
                return;

            TrackMethod(methodBody);
        }

        public void TrackMethod(MethodBody methodBody)
        {
            TrackMethod(lattice.Context.GetMethodIL(methodBody));
        }

        public void TrackMethod(MethodIL methodIL)
        {
            // Work around the fact that ValueSet is readonly
            Debug.Assert(!MethodBodies.IsUnknown());
            var methodsList = new List<MethodIL>(MethodBodies.GetKnownValues());
            methodsList.Add(methodIL);

            // For state machine methods, also scan the state machine members.
            // Simplification: assume that all generated methods of the state machine type are
            // reached at the point where the state machine method is reached.
            if (CompilerGeneratedState.TryGetStateMachineType(methodIL.Method, out TypeDefinition? stateMachineType))
            {
                foreach (var stateMachineMethod in stateMachineType.Methods)
                {
                    Debug.Assert(!CompilerGeneratedNames.IsLambdaOrLocalFunction(stateMachineMethod.Name));
                    if (stateMachineMethod.Body is MethodBody stateMachineMethodBody)
                        methodsList.Add(lattice.Context.GetMethodIL(stateMachineMethodBody));
                }
            }

            MethodBodies = new ValueSet<MethodIL>(methodsList);
        }

        public void SetHoistedLocal(HoistedLocalKey key, MultiValue value)
        {
            // For hoisted locals, we track the entire set of assigned values seen
            // in the closure of a method, so setting a hoisted local value meets
            // it with any existing value.
            HoistedLocals.Set(key,
                lattice.HoistedLocalsLattice.ValueLattice.Meet(
                    HoistedLocals.Get(key), value));
        }

        public MultiValue GetHoistedLocal(HoistedLocalKey key)
            => HoistedLocals.Get(key);
    }

    readonly struct InterproceduralStateLattice : ILattice<InterproceduralState>
    {
        public readonly ValueSetLattice<MethodIL> MethodBodyLattice;
        public readonly DictionaryLattice<HoistedLocalKey, MultiValue, ValueSetLattice<SingleValue>> HoistedLocalsLattice;
        public readonly LinkContext Context;

        public InterproceduralStateLattice(
            ValueSetLattice<MethodIL> methodBodyLattice,
            DictionaryLattice<HoistedLocalKey, MultiValue, ValueSetLattice<SingleValue>> hoistedLocalsLattice,
            LinkContext context)
            => (MethodBodyLattice, HoistedLocalsLattice, Context) = (methodBodyLattice, hoistedLocalsLattice, context);

        public InterproceduralState Top => new InterproceduralState(MethodBodyLattice.Top, HoistedLocalsLattice.Top, this);

        public InterproceduralState Meet(InterproceduralState left, InterproceduralState right)
            => new(
                MethodBodyLattice.Meet(left.MethodBodies, right.MethodBodies),
                HoistedLocalsLattice.Meet(left.HoistedLocals, right.HoistedLocals),
                this);
    }
}