module backtrace;

//
// Provides support for a readable backtrace on a program crash.
//
// Everything is private - you build this into a library and
// link to the library, and bingo (via static this).
//
// It works by registering a stacktrace handler with the runtime,
// which, unlike the default one, provides demangled symbols
// rather than just a list of addresses.
//

private {
    import core.sys.posix.unistd;
    import core.sys.posix.fcntl;
    import core.sys.posix.sys.stat;
    import core.sys.posix.sys.mman;

    import core.stdc.string;
    import core.stdc.signal;

    import core.runtime;

    import std.stdio;
    import std.string;
    import std.demangle;
    import std.algorithm;

    immutable int EI_NIDENT  = 16;

    immutable int SHT_SYMTAB = 2;
    immutable int SHT_STRTAB = 3;

    immutable int STT_FUNC   = 2;

    alias ushort Elf32_Half;
    alias uint   Elf32_Word;
    alias uint   Elf32_Addr;
    alias uint   Elf32_Off;
    alias ushort Elf32_Section;

    struct Elf32_Ehdr {
        ubyte[EI_NIDENT] e_ident;             /* Magic number and other info */
        Elf32_Half    e_type;                 /* Object file type */
        Elf32_Half    e_machine;              /* Architecture */
        Elf32_Word    e_version;              /* Object file version */
        Elf32_Addr    e_entry;                /* Entry point virtual address */
        Elf32_Off     e_phoff;                /* Program header table file offset */
        Elf32_Off     e_shoff;                /* Section header table file offset */
        Elf32_Word    e_flags;                /* Processor-specific flags */
        Elf32_Half    e_ehsize;               /* ELF header size in bytes */
        Elf32_Half    e_phentsize;            /* Program header table entry size */
        Elf32_Half    e_phnum;                /* Program header table entry count */
        Elf32_Half    e_shentsize;            /* Section header table entry size */
        Elf32_Half    e_shnum;                /* Section header table entry count */
        Elf32_Half    e_shstrndx;             /* Section header string table index */
    };

    struct Elf32_Shdr {
        Elf32_Word    sh_name;                /* Section name (string tbl index) */
        Elf32_Word    sh_type;                /* Section type */
        Elf32_Word    sh_flags;               /* Section flags */
        Elf32_Addr    sh_addr;                /* Section virtual addr at execution */
        Elf32_Off     sh_offset;              /* Section file offset */
        Elf32_Word    sh_size;                /* Section size in bytes */
        Elf32_Word    sh_link;                /* Link to another section */
        Elf32_Word    sh_info;                /* Additional section information */
        Elf32_Word    sh_addralign;           /* Section alignment */
        Elf32_Word    sh_entsize;             /* Entry size if section holds table */
    };

    struct Elf32_Sym {
        Elf32_Word    st_name;                /* Symbol name (string tbl index) */
        Elf32_Addr    st_value;               /* Symbol value */
        Elf32_Word    st_size;                /* Symbol size */
        ubyte         st_info;                /* Symbol type and binding */
        ubyte         st_other;               /* Symbol visibility */
        Elf32_Section st_shndx;               /* Section index */
    };

    extern (C) int backtrace(void**, size_t);

    struct Symbol {
        this(void * a, size_t s, char * n) {
            address = a;
            size = s;
            name = n;
        }
        void * address;
        size_t size;
        char * name;
        void * end_address() { return address + size; }
    };

    // Newton-Raphson
    Symbol nr_lookup(Symbol[] symbols, void * addr, int depth = 0) {
        if (symbols.length == 0 ||
            addr < symbols[0].address ||
            addr >= symbols[$-1].end_address)
        {
            throw new Exception("Symbol not found");
        }
        else if (symbols.length == 1) {
            return symbols[0];
        }
        else {
            void * begin_addr = symbols[0].address;
            void * end_addr = symbols[$-1].end_address;

            int i = ((addr - begin_addr) * symbols.length) / (end_addr - begin_addr);
            if (!(depth % 2) && i < symbols.length - 1) { ++i; }        // Some wiggle to force convergence

            if (addr < symbols[i].address) {
                return nr_lookup(symbols[0..i], addr, depth + 1);
            }
            else {
                return nr_lookup(symbols[i..$], addr, depth + 1);
            }
        }
    }

    int generate(void*[] addresses, scope int delegate(ref char[]) dg) {
        static char[] get_exe() {
            char buf[1024];
            ssize_t s = readlink("/proc/self/exe", buf.ptr, buf.length);
            if (s == -1) { throw new Exception(""); }
            return buf[0..s];
        }

        int fd = open(toStringz(get_exe()), O_RDONLY);
        if (fd == -1) { throw new Exception(""); }
        scope(exit) close(fd);

        stat_t st;
        if (fstat(fd, &st) == -1) { throw new Exception(""); }

        void * contents = mmap(null, st.st_size, PROT_READ, MAP_SHARED, fd, 0);
        if (!contents) { throw new Exception(""); }
        scope(exit) munmap(contents, st.st_size);

        Elf32_Ehdr * elf_hdr = cast(Elf32_Ehdr *)(contents);

        Elf32_Shdr * sec_hdr = cast(Elf32_Shdr *)(contents + elf_hdr.e_shoff);

        int symtab_idx = 0;
        for (int i = 0; i < elf_hdr.e_shnum; ++i) {
            if (sec_hdr[i].sh_type == SHT_SYMTAB) {
                symtab_idx = i;
                break;
            }
        }
        if (symtab_idx == 0) { throw new Exception(""); }

        int strtab_idx = sec_hdr[symtab_idx].sh_link;
        if (strtab_idx == 0) { throw new Exception(""); }        // No associated string table

        if (sec_hdr[strtab_idx].sh_type != SHT_STRTAB) { throw new Exception(""); } // invalid string table

        Elf32_Sym * sym_ent = cast(Elf32_Sym *)(contents + sec_hdr[symtab_idx].sh_offset);
        char * strtab = cast(char *)(contents + sec_hdr[strtab_idx].sh_offset);
        int num_syms = sec_hdr[symtab_idx].sh_size / sec_hdr[symtab_idx].sh_entsize;

        if (num_syms == 0) { throw new Exception(""); }      // No symbols

        Symbol symbols[];

        for (int i = 0; i < num_syms; ++i) {
            if ((sym_ent[i].st_info & 0xf) != STT_FUNC) {
                continue;
            }

            if (sym_ent[i].st_shndx == 0) {
                continue;
            }

            void * address = cast(void *)(sym_ent[i].st_value);       // inclusive
            size_t size    = sym_ent[i].st_size;
            char * name    = strtab + sym_ent[i].st_name;

            symbols ~= Symbol(address, size, name);
        }

        sort!("a.address < b.address")(symbols);

        int ret;
        foreach (address; addresses) {
            string str = format("[0x%0.8x] ", address);
            try {
                Symbol symbol = nr_lookup(symbols, address);
                char[] s1 = symbol.name[0..strlen(symbol.name)];
                str ~= format("%s", demangle(s1.idup));
            }
            catch (Exception e) {
                str ~= "<unknown>";
            }
            char[] cstr = str.dup;
            ret = dg(cstr);
            if (ret) {
                break;
            }
        }

        return ret;
    }


    // signal handler for otherwise-fatal thread-specific signals 
    extern (C) void arghh(int sig) {
        string name() {
            switch (sig) {
            case SIGSEGV: return "SIGSEGV";
            case SIGFPE:  return "SIGFPE";
            case SIGILL:  return "SIGILL";
            case SIGABRT: return "SIGABRT";
            default:      return "";
            }
        }

        throw new Error(format("Got signal %s %s", sig, name));
    }

    shared static this() {
        // set up shared signal handlers for fatal thread-specific signals
        //writeln("setting up shared signal handlers for ABRT, FPE, ILL, SEGV");
        signal(SIGABRT, &arghh);
        signal(SIGFPE,  &arghh);
        signal(SIGILL,  &arghh);
        signal(SIGSEGV, &arghh);
    }

    static this() {
        // register our trace handler for each thread
        //writeln("installing our traceHandler");
        Runtime.traceHandler = &traceHandler;
    }

    // Captures backtrace info at the point of construction, stripping
    // off the bits that concern itself and the ininteresting early stuff.
    // Also provides opApply to traverse a nice text representation of the backtrace.
    class TraceInfo : Throwable.TraceInfo {
        void*[256] callstack;
        int numframes;

        this() {
            numframes = backtrace(callstack.ptr, callstack.length);
        }

        override string toString() const {
            return "Why does dmd require an override of this?";
        }

        override int opApply(scope int delegate(ref char[]) dg) {
            return generate(callstack[4..numframes-5], dg);
        }
    }

    Throwable.TraceInfo traceHandler(void * ptr = null) {
        return new TraceInfo;
    }
}