objcallgraph.d

#!/usr/bin/env rdmd

module objcallgraph;

import std;
import core.memory;

extern(C) __gshared string[] rt_options = [ "scanDataSeg=precise", "gcopt=initReserve:1000 cleanup:finalize gc:precise" ];

alias Void  = void[0];

struct GraphNode
{
    enum Type : ubyte
    {
        notype,
        ifunc,
        func,
        object_,
        tls,
        file,
    }
    Type type;

    enum Bind : ubyte
    {
        local,
        global,
        weak,
    }
    Bind bind;

    enum Visibility : ubyte
    {
        default_,
        protected_,
        hidden,
    }
    Visibility visibility;

    @safe pure nothrow @nogc
    string color() const
    {
        final switch(visibility)
        {
            case Visibility.default_:   return "black";
            case Visibility.protected_: return "blue";
            case Visibility.hidden:     return "red";
        }
    }

    size_t size;
    size_t stackSize;
    size_t[GraphNode*] children;
}

alias Graph = GraphNode[string];

__gshared Graph graph;

int disass(string prog, MatchSettings ms)
{
    stderr.writeln("running disassembler...");
    auto disassembly = pipeProcess(["llvm-objdump", "-d", prog]);
    scope(exit) wait(disassembly.pid);

    string functionSym;

    enum functionRegex = ctRegex!`[0-9A-Fa-f]+ <(.*)>:`;
    enum referenceRegex = ctRegex!`\s*[0-9A-Fa-f]+:(?:.*<([._A-Za-z0-9]+)(?:\+.*|@.*)?>.*)+`;

    foreach(line; disassembly.stdout.byLine)
    {
        if (auto funcMatch = line.matchFirst(functionRegex))
        {
            funcMatch.popFront;
            functionSym = funcMatch.front.idup;
            continue;
        }

        if (!ms.includeMatcher.match(functionSym)) continue;
        if (!ms.excludeMatcher.match(functionSym)) continue;

        graph.require(functionSym);

        if (auto refMatch = line.matchFirst(referenceRegex))
        {
            refMatch.popFront;
            foreach(ref_; refMatch)
            {
                auto iref_ = ref_.idup;

                if (!ms.refIncludeMatcher.match(iref_)) continue;
                if (!ms.refExcludeMatcher.match(iref_)) continue;

                graph[functionSym].children.require(&graph.require(iref_))++;
            }
        }
    }

    return 0;
}

int readsyms(string prog)
{
    stderr.writeln("read symbol table...");
    auto symbols = pipeProcess(["llvm-readelf", "--syms", "--dyn-syms", prog]);
    scope(exit) wait(symbols.pid);

    enum symRegex = ctRegex!`\s*[0-9]+:\s+[0-9A-Fa-f]+\s+(\d+)\s+(\w+)\s+(\w+)\s+(\w+)\s+([A-Za-z0-9]+)\s+(.*)`;

    foreach(line; symbols.stdout.byLine)
    {
        if (auto symMatch = line.matchFirst(symRegex))
        {
            symMatch.popFront;

            auto symsize = symMatch[0];
            auto symtype = symMatch[1];
            auto symbind = symMatch[2];
            auto symvis  = symMatch[3];
            auto symndx  = symMatch[4];
            auto symname = symMatch[5];

            auto node = (cast(string)symname) in graph;
            if (node is null) continue;

            node.size = symsize.to!size_t;
            switch (symtype)
            {
                case "NOTYPE": node.type = GraphNode.Type.notype;  break;
                case "IFUNC":  node.type = GraphNode.Type.ifunc;   break;
                case "FUNC":   node.type = GraphNode.Type.func;    break;
                case "OBJECT": node.type = GraphNode.Type.object_; break;
                case "TLS":    node.type = GraphNode.Type.tls;     break;
                case "FILE":   node.type = GraphNode.Type.file;    break;
                default:       assert(0, symtype);
            }

            switch (symbind)
            {
                case "LOCAL":  node.bind = GraphNode.Bind.local;  break;
                case "GLOBAL": node.bind = GraphNode.Bind.global; break;
                case "WEAK":   node.bind = GraphNode.Bind.weak;   break;
                default:       assert(0, symbind);
            }

            switch (symvis)
            {
                case "DEFAULT":   node.visibility = GraphNode.Visibility.default_;   break;
                case "PROTECTED": node.visibility = GraphNode.Visibility.protected_; break;
                case "HIDDEN":    node.visibility = GraphNode.Visibility.hidden;     break;
                default:          assert(0, symvis);
            }
        }
    }

    return 0;
}

int stacksizes(string prog)
{
    stderr.writeln("read stack sizes section...");
    auto stackSizes = pipeProcess(["llvm-readelf", "--stack-sizes", prog]);
    scope(exit) wait(stackSizes.pid);

    enum ssRegex = ctRegex!`\s*([0-9]+)\s*(.*)`;

    foreach(line; stackSizes.stdout.byLine)
    {
        if (auto ssMatch = line.matchFirst(ssRegex))
        {
            ssMatch.popFront;

            if (auto node = (cast(string)ssMatch[1]) in graph)
                node.stackSize = ssMatch[0].to!size_t;
        }
    }

    return 0;
}

string xmlEscaper(string str)
{
    return str
        .replace(`"`, "&quot;")
        .replace(`&`, "&amp;")
        .replace(`'`, "&apos;")
        .replace(`<`, "&lt;")
        .replace(`>`, "&gt;")
        .replace("\0", "")
        .replace("\t", "")
        .replace("\n", "");
}

string stringEscaper(string str)
{
    return str
        .replace(`"`, `\"`)
        .replace("\0", "")
        .replace("\t", "")
        .replace("\n", "");
}

void outputGexf(MatchSettings ms)
{
    writeln(`<?xml version="1.0" encoding="UTF-8"?>`);
    writeln(`<gexf xmlns="http://gexf.net/1.3" xmlns:viz="http://gexf.net/1.3/viz" version="1.3">`);
    writeln(`<graph mode="static" defaultedgetype="directed">`);

    writeln(
`<attributes class="node">
    <attribute id="mangling"   title="Mangling"    type="string" />
    <attribute id="type"       title="Type"        type="string" />
    <attribute id="bind"       title="Bind"        type="string" />
    <attribute id="visibility" title="Visibility"  type="string" />
    <attribute id="size"       title="Symbol Size" type="long"   />
    <attribute id="stack"      title="Stack Size"  type="long"   />
</attributes>
<attributes class="edge">
    <attribute id="references" title="Number of References" type="long" />
</attributes>`
    );

    writeln("<nodes>");

    foreach(ref k, ref v; graph)
    {
        if(!k) continue;
        if (!ms.outputMatch(&v)) continue;

        switch (v.type)
        {
            case GraphNode.Type.func:
            case GraphNode.Type.ifunc:
                writefln(`<node id="%x" label="%s">`, &v, xmlEscaper(demangle(k)));
                //writefln(`<viz:size value="%d"/>`, v.stackSize);
                break;
            default:
                writefln(`<node id="%x" label="%s">`, &v, xmlEscaper(demangle(k)));
                //writefln(`<viz:size value="%d"/>`, v.size);
                break;
        }
        writeln("<attvalues>");
        writefln(`<attvalue for="mangling" value="%s"/>`, xmlEscaper(k));
        writefln(`<attvalue for="size" value="%d"/>`, v.size);
        writefln(`<attvalue for="type" value="%s"/>`, v.type);
        writefln(`<attvalue for="bind" value="%s"/>`, v.bind);
        writefln(`<attvalue for="visibility" value="%s"/>`, v.visibility);
        writefln(`<attvalue for="stack" value="%d"/>`, v.stackSize);
        writeln("</attvalues>");
        writeln("</node>");
    }
    writeln("</nodes>");


    writeln("<edges>");
    foreach(ref k, ref v; graph)
    {
        if(!k) continue;
        if (!ms.outputMatch(&v)) continue;

        foreach(ck, cv; v.children)
        {
            assert(ck);
            if (!ms.outputMatch(ck)) continue;

            writefln(`<edge source="%x" target="%x" weight="%d">`, &v, ck, cv);
            writeln("<attvalues>");
            writefln(`<attvalue for="references" value="%d"/>`, cv);
            writeln("</attvalues>");
            writeln("</edge>");
        }
    }
    writeln("</edges>");

    writeln("</graph>");
    writeln("</gexf>");
}

void outputDot(MatchSettings ms)
{
    writeln("digraph G {");
    foreach(ref k, ref v; graph)
    {
        if(!k) continue;
        if (!ms.outputMatch(&v)) continue;

        writefln(`"%X" [label="%s" mangling="%s" size=%d group="%s" bind="%s" visibility="%s" stack=%d];`,
            &v, stringEscaper(demangle(k)), k, v.size, v.type, v.bind, v.visibility, v.stackSize,
        );

        foreach(ck, cv; v.children)
        {
            assert(ck);
            if (!ms.outputMatch(ck)) continue;

            writefln(`"%X" -> "%X" [weight=%d];`, &v, ck, cv);
        }
    }

    writeln("}");
}

struct MatchSettings
{
    Matcher includeMatcher;
    Matcher excludeMatcher;

    Matcher refIncludeMatcher;
    Matcher refExcludeMatcher;

    size_t stackSize;

    bool outputMatch(const(GraphNode)* g)
    {
        __gshared Void[const(GraphNode)*] visited;

        if (g in visited) return true;
        if (g.stackSize < stackSize) return false;

        visited.require(g);
        return true;
    }
}

struct Matcher
{
    this(string pattern, bool inverse)
    {
        this.inverse = inverse;

        if (pattern.length == 0)
            return;

        regexMatcher = regex(pattern);
    }

    bool match(string buf)
    {
        if (regexMatcher.isNull)
            return getFalse;

        if (buf.matchFirst(regexMatcher.get))
            return getTrue;

        return getFalse;
    }

    bool getTrue()  { return !inverse; }
    bool getFalse() { return inverse;  }

    bool inverse;
    Nullable!(Regex!char) regexMatcher;
}

int main(string[] args)
{
    if (args.length <= 1)
    {
        stderr.writeln("Error: Please provide an executable path");
        return 1;
    }

    string includePattern = `.*`;
    string excludePattern = `^(\.|__)`;

    string refIncludePattern = `.*`;
    string refExcludePattern = excludePattern;

    enum Format
    {
        dot,
        gexf,
    }

    Format fileFormat;
    size_t stackSize;

    auto helpInformation = getopt(
        args,
        "i|include",            &includePattern,
        "e|exclude",            &excludePattern,
        "I|ref-include",        &refIncludePattern,
        "E|ref-exclude",        &refExcludePattern,
        "S|atleast-stack-size", &stackSize,
        "f|format",             &fileFormat,
    );

    if (helpInformation.helpWanted)
    {
        defaultGetoptPrinter("Some information about the program.", helpInformation.options);
        return 1;
    }

    MatchSettings ms = {
        includeMatcher    : Matcher(includePattern, false),
        excludeMatcher    : Matcher(excludePattern, true),
        refIncludeMatcher : Matcher(refIncludePattern, false),
        refExcludeMatcher : Matcher(refExcludePattern, true),
        stackSize         : stackSize,
    };

    if (auto ret = disass(args[1], ms))
        return ret;

    if (auto ret = readsyms(args[1]))   return ret;
    if (auto ret = stacksizes(args[1])) return ret;

    final switch(fileFormat)
    {
        case Format.gexf: outputGexf(ms); break;
        case Format.dot:  outputDot(ms); break;
    }

    return 0;
}