// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.

using System.Collections.Generic;

namespace Microsoft.Diagnostics.DataContractReader.ExecutionManagerHelpers;

/// <summary>
/// Resolves WebAssembly ReadyToRun function-table indices, mirroring the native
/// <c>ExecutionManager::{FindFunctionTableIndexRangeSection, GetWasmVirtualIPFromFunctionTableIndex}</c>
/// in <c>src/coreclr/vm/codeman.cpp</c>. An index is resolved against the
/// <c>FunctionTableIndexRangeList</c> to its owning module's <see cref="Data.ReadyToRunInfo"/>, and the
/// corresponding <c>RUNTIME_FUNCTION</c> supplies the funclet flag, base virtual IP, and unwind data.
/// </summary>
internal sealed class WasmFunctionTableIndexLookup
{
    // Reader resource budget, not a limit imposed by native range registration.
    private const int MaxFunctionTableIndexRangeNodes = 65_536;

    private readonly Target _target;
    private readonly RuntimeFunctionLookup _runtimeFunctions;

    public WasmFunctionTableIndexLookup(Target target)
    {
        _target = target;
        _runtimeFunctions = RuntimeFunctionLookup.Create(target, isWasm: true);
    }

    // Mirrors ExecutionManager::FindFunctionTableIndexRangeSection.
    private Data.FunctionTableIndexRangeSection? FindSection(uint functionTableIndex)
    {
        if (!_target.TryReadGlobalPointer(Constants.Globals.FunctionTableIndexRangeList, out TargetPointer? listHeadSlot))
            return null;

        // The global holds the address of the s_pFunctionTableIndexRangeList slot (a pointer-to-
        // pointer); dereference it once to obtain the actual list head.
        TargetPointer current = _target.ReadPointer(listHeadSlot.Value);
        HashSet<TargetPointer> visited = new();
        while (current != TargetPointer.Null)
        {
            if (visited.Count == MaxFunctionTableIndexRangeNodes || !visited.Add(current))
                return null;

            Data.FunctionTableIndexRangeSection section = _target.ProcessedData.GetOrAdd<Data.FunctionTableIndexRangeSection>(current);
            if (functionTableIndex >= section.MinFunctionTableIndex &&
                functionTableIndex - section.MinFunctionTableIndex < section.NumRuntimeFunctions)
            {
                return section;
            }
            current = section.Next;
        }

        return null;
    }

    private Data.ReadyToRunInfo GetReadyToRunInfo(Data.FunctionTableIndexRangeSection section)
    {
        Data.Module module = _target.ProcessedData.GetOrAdd<Data.Module>(section.R2RModule);
        return _target.ProcessedData.GetOrAdd<Data.ReadyToRunInfo>(module.ReadyToRunInfo);
    }

    private Data.RuntimeFunction GetRuntimeFunction(Data.ReadyToRunInfo r2rInfo, uint localIndex)
        => _runtimeFunctions.GetRuntimeFunction(r2rInfo.RuntimeFunctions, localIndex);

    // Mirrors ExecutionManager::GetWasmVirtualIPFromFunctionTableIndex.
    public bool TryGetVirtualIPBase(uint functionTableIndex, out ulong baseVirtualIP)
    {
        baseVirtualIP = 0;
        Data.FunctionTableIndexRangeSection? section = FindSection(functionTableIndex);
        if (section is null)
            return false;

        Data.ReadyToRunInfo r2rInfo = GetReadyToRunInfo(section);
        if (r2rInfo.MinVirtualIP is not TargetPointer minVirtualIP)
            return false;

        // Funclets' function-local virtual IPs are relative to their controlling function, so index
        // backwards past funclet entries to the controlling (non-funclet) function.
        uint localIndex = functionTableIndex - section.MinFunctionTableIndex;
        while (true)
        {
            Data.RuntimeFunction runtimeFunction = GetRuntimeFunction(r2rInfo, localIndex);
            if (_runtimeFunctions.IsFunclet(runtimeFunction))
            {
                if (localIndex == 0)
                    return false;
                localIndex--;
                continue;
            }

            baseVirtualIP = minVirtualIP.Value + _runtimeFunctions.GetBeginAddress(runtimeFunction);
            return true;
        }
    }

    public bool TryGetUnwindData(uint functionTableIndex, out TargetPointer unwindDataAddress)
    {
        unwindDataAddress = TargetPointer.Null;
        Data.FunctionTableIndexRangeSection? section = FindSection(functionTableIndex);
        if (section is null)
            return false;

        Data.ReadyToRunInfo r2rInfo = GetReadyToRunInfo(section);
        uint localIndex = functionTableIndex - section.MinFunctionTableIndex;
        Data.RuntimeFunction runtimeFunction = GetRuntimeFunction(r2rInfo, localIndex);
        unwindDataAddress = new TargetPointer(r2rInfo.LoadedImageBase.Value + runtimeFunction.UnwindData);
        return true;
    }
}