/****************************************************************************
**
** Copyright (C) 2020 Klarälvdalens Datakonsult AB, a KDAB Group company, info@kdab.com, author Milian Wolff <milian.wolff@kdab.com>
** Contact: http://www.qt.io/licensing/
**
** This file is part of the Qt Enterprise Perf Profiler Add-on.
**
** GNU General Public License Usage
** This file may be used under the terms of the GNU General Public License
** version 3 as published by the Free Software Foundation and appearing in
** the file LICENSE.GPLv3 included in the packaging of this file. Please
** review the following information to ensure the GNU General Public License
** requirements will be met: https://www.gnu.org/licenses/gpl.html.
**
** If you have questions regarding the use of this file, please use
** contact form at http://www.qt.io/contact-us
**
****************************************************************************/

#include <time.h>

#include "perfdwarfdiecache.h"
#include "perfeucompat.h"

#include <dwarf.h>

#include <vector>

#include "demangler.h"

namespace {
enum class WalkResult
{
    Recurse,
    Skip,
    Return
};
template<typename Callback>
WalkResult walkDieTree(const Callback &callback, Dwarf_Die *die)
{
    auto result = callback(die);
    if (result != WalkResult::Recurse)
        return result;

    Dwarf_Die childDie;
    if (dwarf_child(die, &childDie) == 0) {
        result = walkDieTree(callback, &childDie);
        if (result == WalkResult::Return)
            return result;

        Dwarf_Die siblingDie;
        while (dwarf_siblingof(&childDie, &siblingDie) == 0) {
            result = walkDieTree(callback, &siblingDie);
            if (result == WalkResult::Return)
                return result;
            childDie = siblingDie;
        }
    }
    return WalkResult::Skip;
}

template<typename Callback>
void walkRanges(const Callback &callback, Dwarf_Die *die)
{
    Dwarf_Addr low = 0;
    Dwarf_Addr high = 0;
    Dwarf_Addr base = 0;
    ptrdiff_t rangeOffset = 0;
    while ((rangeOffset = dwarf_ranges(die, rangeOffset, &base, &low, &high)) > 0) {
        if (!callback(DwarfRange{low, high}))
            return;
    }
}

// see libdw_visit_scopes.c in elfutils
bool mayHaveScopes(Dwarf_Die *die)
{
    switch (dwarf_tag(die))
    {
    /* DIEs with addresses we can try to match.  */
    case DW_TAG_compile_unit:
    case DW_TAG_module:
    case DW_TAG_lexical_block:
    case DW_TAG_with_stmt:
    case DW_TAG_catch_block:
    case DW_TAG_try_block:
    case DW_TAG_entry_point:
    case DW_TAG_inlined_subroutine:
    case DW_TAG_subprogram:
    /* DIEs without addresses that can own DIEs with addresses.  */
    case DW_TAG_namespace:
    case DW_TAG_class_type:
    case DW_TAG_structure_type:
        return true;

    /* Other DIEs we have no reason to descend.  */
    default:
        break;
    }
    return false;
}

bool dieContainsAddress(Dwarf_Die *die, Dwarf_Addr address)
{
    bool contained = false;
    walkRanges([&contained, address](DwarfRange range) {
        if (range.contains(address)) {
            contained = true;
            return false;
        }
        return true;
    }, die);
    return contained;
}
}

const char *linkageName(Dwarf_Die *die)
{
    Dwarf_Attribute attr;
    Dwarf_Attribute *result = dwarf_attr_integrate(die, DW_AT_MIPS_linkage_name, &attr);
    if (!result)
        result = dwarf_attr_integrate(die, DW_AT_linkage_name, &attr);

    return result ? dwarf_formstring(result) : nullptr;
}

Dwarf_Die *specificationDie(Dwarf_Die *die, Dwarf_Die *dieMem)
{
    Dwarf_Attribute attr;
    if (dwarf_attr_integrate(die, DW_AT_specification, &attr))
        return dwarf_formref_die(&attr, dieMem);
    return nullptr;
}

/// prepend the names of all scopes that reference the @p die to @p name
void prependScopeNames(QByteArray &name, Dwarf_Die *die, QHash<Dwarf_Off, QByteArray> &cache)
{
    Dwarf_Die dieMem;
    Dwarf_Die *scopes = nullptr;
    auto nscopes = dwarf_getscopes_die(die, &scopes);

    // We essentially visit and build the scope name in reverse order.
    // The cache ops encode offsets we can lookup directly that then map to fully
    // qualified identifiers, which must obviously _not_ end on a double-colon separator.
    // Note that while filling cacheOps below, we thus always have to prepend the double-colon
    // first to the name, then store the cacheOps with the size of `name`. While that may sound
    // confusing, that gives us the desired results: `ScopesToCache::trailing` will then be set
    // to the size _following_ the current entry, which may get more identifiers appended to
    // it when we continue to visit the other DIEs next.
    struct ScopesToCache
    {
        Dwarf_Off offset;
        int trailing;
    };
    std::vector<ScopesToCache> cacheOps;

    // skip scope for the die itself at the start and the compile unit DIE at end
    for (int i = 1; i < nscopes - 1; ++i) {
        auto scope = scopes + i;

        const auto scopeOffset = dwarf_dieoffset(scope);

        auto it = cache.find(scopeOffset);
        if (it != cache.end()) {
            // prepend the fully qualified cached identifier
            // that won't end on `::`, so we have to add that manually here
            if (!name.isEmpty())
                name.prepend("::");
            name.prepend(*it);
            // we can stop, cached names are always fully qualified
            break;
        }

        if (auto scopeLinkageName = linkageName(scope)) {
            // prepend the fully qualified linkage name
            name.prepend("::");
            cacheOps.push_back({scopeOffset, int(name.size())});
            // we have to demangle the scope linkage name, otherwise we get a
            // mish-mash of mangled and non-mangled names
            name.prepend(demangle(scopeLinkageName));
            // we can stop now, the scope is fully qualified
            break;
        }

        if (auto scopeName = dwarf_diename(scope)) {
            // prepend this scope's name, e.g. the class or namespace name
            name.prepend("::");
            cacheOps.push_back({scopeOffset, int(name.size())});
            name.prepend(scopeName);
        }

        if (auto specification = specificationDie(scope, &dieMem)) {
            eu_compat_free(scopes);
            scopes = nullptr;
            cacheOps.push_back({scopeOffset, int(name.size())});
            cacheOps.push_back({dwarf_dieoffset(specification), int(name.size())});
            // follow the scope's specification DIE instead
            prependScopeNames(name, specification, cache);
            break;
        }
    }

    for (const auto &cacheOp : cacheOps)
        cache[cacheOp.offset] = name.mid(0, name.size() - cacheOp.trailing);

    eu_compat_free(scopes);
}

bool operator==(const Dwarf_Die &lhs, const Dwarf_Die &rhs)
{
    return lhs.addr == rhs.addr && lhs.cu == rhs.cu && lhs.abbrev == rhs.abbrev;
}

QByteArray qualifiedDieName(Dwarf_Die *die, QHash<Dwarf_Off, QByteArray> &cache)
{
    // linkage names are fully qualified, meaning we can stop early then
    if (auto name = linkageName(die))
        return name;

    // otherwise do a more complex lookup that includes namespaces and other context information
    // this is important for inlined subroutines such as lambdas or std:: algorithms
    QByteArray name = dwarf_diename(die);

    // use the specification DIE which is within the DW_TAG_namespace
    Dwarf_Die dieMem;
    if (auto specification = specificationDie(die, &dieMem))
        die = specification;

    prependScopeNames(name, die, cache);

    return name;
}

QByteArray demangle(const QByteArray &mangledName)
{
    if (mangledName.length() < 3) {
        return mangledName;
    } else {
        static size_t demangleBufferLength = 8192;
        static char *demangleBuffer = reinterpret_cast<char *>(eu_compat_malloc(demangleBufferLength));
        static Demangler demangler;

        if (demangler.demangle(mangledName, demangleBuffer, demangleBufferLength)) {
            return demangleBuffer;
        }

        // Require GNU v3 ABI by the "_Z" prefix.
        if (mangledName[0] == '_' && mangledName[1] == 'Z') {
            int status = -1;
            char *dsymname = eu_compat_demangle(mangledName.constData(), demangleBuffer, &demangleBufferLength,
                                            &status);
            if (status == 0)
                return demangleBuffer = dsymname;
        }
    }
    return mangledName;
}

QVector<Dwarf_Die> findInlineScopes(Dwarf_Die *subprogram, Dwarf_Addr offset)
{
    QVector<Dwarf_Die> scopes;
    walkDieTree([offset, &scopes](Dwarf_Die *die) {
        if (dwarf_tag(die) != DW_TAG_inlined_subroutine)
            return WalkResult::Recurse;
        if (dieContainsAddress(die, offset)) {
            scopes.append(*die);
            return WalkResult::Recurse;
        }
        return WalkResult::Skip;
    }, subprogram);
    return scopes;
}

QByteArray absoluteSourcePath(const char *path, Dwarf_Die *cuDie)
{
    if (!path || !cuDie || path[0] == '/')
        return path;

    Dwarf_Attribute attr;
    auto compDir = dwarf_formstring(dwarf_attr(cuDie, DW_AT_comp_dir, &attr));
    if (!compDir)
        return path;

    QByteArray ret;
    ret.reserve(static_cast<int>(strlen(compDir) + strlen(path) + 1));
    ret.append(compDir);
    ret.append('/');
    ret.append(path);
    return ret;
}

DwarfSourceLocation findSourceLocation(Dwarf_Die* cuDie, Dwarf_Addr offset)
{
    DwarfSourceLocation ret;
    if (auto srcloc = dwarf_getsrc_die(cuDie, offset)) {
        if (const char* srcfile = dwarf_linesrc(srcloc, nullptr, nullptr)) {
            ret.file = absoluteSourcePath(srcfile, cuDie);
            dwarf_lineno(srcloc, &ret.line);
            dwarf_linecol(srcloc, &ret.column);
        }
    }
    return ret;
}

SubProgramDie::SubProgramDie(Dwarf_Die die)
    : m_ranges{die, {}}
{
    walkRanges([this](DwarfRange range) {
        m_ranges.ranges.append(range);
        return true;
    }, &die);
}

SubProgramDie::~SubProgramDie() = default;

CuDieRangeMapping::CuDieRangeMapping(Dwarf_Die cudie, Dwarf_Addr bias)
    : m_bias{bias}
    , m_cuDieRanges{cudie, {}}
{
    walkRanges([this, bias](DwarfRange range) {
        m_cuDieRanges.ranges.append({range.low + bias, range.high + bias});
        return true;
    }, &cudie);
}

CuDieRangeMapping::~CuDieRangeMapping() = default;

SubProgramDie *CuDieRangeMapping::findSubprogramDie(Dwarf_Addr offset)
{
    if (m_subPrograms.isEmpty())
        addSubprograms();

    auto it = std::find_if(m_subPrograms.begin(), m_subPrograms.end(),
                        [offset](const SubProgramDie &program) {
                                return program.contains(offset);
                        });
    if (it == m_subPrograms.end())
        return nullptr;

    return &(*it);
}

void CuDieRangeMapping::addSubprograms()
{
    walkDieTree([this](Dwarf_Die *die) {
        if (!mayHaveScopes(die))
            return WalkResult::Skip;

        if (dwarf_tag(die) == DW_TAG_subprogram) {
            SubProgramDie program(*die);
            if (!program.isEmpty())
                m_subPrograms.append(program);

            return WalkResult::Skip;
        }
        return WalkResult::Recurse;
    }, cudie());
}

QByteArray CuDieRangeMapping::dieName(Dwarf_Die *die)
{
    // in QT6 QHash is no longer iterator safe
    // since qualifiedDieName can modify m_dieNameCache we need to insert name again
    const auto offset = dwarf_dieoffset(die);
    auto name = m_dieNameCache[offset];
    if (name.isEmpty()) {
        name = demangle(qualifiedDieName(die, m_dieNameCache));
        m_dieNameCache[offset] = name;
    }

    return name;
}

PerfDwarfDieCache::PerfDwarfDieCache(Dwfl_Module *mod)
{
    if (!mod)
        return;

    Dwarf_Die *die = nullptr;
    Dwarf_Addr bias = 0;
    while ((die = dwfl_module_nextcu(mod, die, &bias))) {
        CuDieRangeMapping cuDieMapping(*die, bias);
        if (!cuDieMapping.isEmpty())
            m_cuDieRanges.push_back(cuDieMapping);
    }
}

PerfDwarfDieCache::~PerfDwarfDieCache() = default;

CuDieRangeMapping *PerfDwarfDieCache::findCuDie(Dwarf_Addr addr)
{
    auto it = std::find_if(m_cuDieRanges.begin(), m_cuDieRanges.end(),
                            [addr](const CuDieRangeMapping &cuDieMapping) {
                                return cuDieMapping.contains(addr);
                            });
    if (it == m_cuDieRanges.end())
        return nullptr;

    return &(*it);
}
