/*

 LD_PRELOAD library that gives statistic on number of roundtrips in an application.

 $XREPLY_BACKTRACE defines whether and how backtraces will be printed for every
 roundtrip. If not set, only total number of roundtrips is printed after the process
 exits. If set to a number, backtrace for every roundtrip will be printed, and the
 backtraces will be as deep as the given number. If set to C<number> (e.g. C10),
 the backtraces will be "compressed" - every backtrace will be printed only once
 after the process exits, together with number of times it occured.

*/

#define _GNU_SOURCE
#include <dlfcn.h>
#include <stdlib.h>
#include <stdio.h>
#include <unistd.h>
#include <fcntl.h>
#include <execinfo.h>
#include <assert.h>
#include <X11/Xlibint.h>

/* Since these symbols are weak, the apps can provide their own, and therefore
   e.g. temporarily suspend counting of roundtrips. At least theoretically,
   I haven't really tried it.
*/
__attribute((weak)) long ___xreply_reply_count = 0;
__attribute((weak)) int ___xreply_reply_enabled = 1;

#define MAX_BACKTRACES 1024

extern long ___xreply_reply_count;
extern int ___xreply_reply_enabled;

typedef Status (*xreply_ptr_t)(Display*,xReply*,int,Bool);

static xreply_ptr_t xreply_ptr = NULL;
static int xreply_backtrace_set = 0;
static int xreply_backtrace_type = 0;

struct xreply_struct
    {
    char* key;
    char* text;
    int count;
    };
static struct xreply_struct backtraces[ MAX_BACKTRACES ];
static int backtraces_size = 0;

static int xreply_compare( const void* left, const void* right )
    {
    int left_count = ((struct xreply_struct*)left)->count;
    int right_count = ((struct xreply_struct*)right)->count;
    return right_count - left_count;
    }

static void xreply_print(void)
    {
    char tmp[ 1024 ];
    int fd;
    fd = open( "/proc/self/cmdline", O_RDONLY );
    if( fd >= 0 )
        {
        read( fd, tmp, 1024 );
        tmp[ 1023 ] = '\0';
        close( fd );
        }
    fprintf( stderr, "XREPLY (%d : %s): %ld\n", getpid(), tmp, ___xreply_reply_count );
    if( xreply_backtrace_type < 0 )
        {
        int i;
        qsort( backtraces, backtraces_size, sizeof( struct xreply_struct ), xreply_compare );
        for( i = 0;
             i < backtraces_size;
             ++i )
            fprintf( stderr, "%d:%s\n\n", backtraces[ i ].count, backtraces[ i ].text );
        }
    }

static void xreply_backtrace()
    {
    void* trace[256];
    int n = backtrace(trace, 256);
    char** strings = backtrace_symbols (trace, n);

    if( xreply_backtrace_type > 0 )
        {
        fprintf( stderr, "%ld [\n", ___xreply_reply_count );
        if( n > xreply_backtrace_type )
            n = xreply_backtrace_type;
        int i;
        for( i = 0;
             i < n;
             ++i )
            fprintf( stderr, "%d: %s\n", i, strings[ i ] );
        fprintf( stderr, "]\n" );
        }
    else
        {
        char stack[ 256 * 20 ];
        int pos = 0;
        int i;
        stack[ 0 ] = '\0';
        if( n > -xreply_backtrace_type )
            n = -xreply_backtrace_type;
        for( i = 0;
             i < n;
             ++i )
            {
            const char* start = strrchr( strings[ i ], '[' );
            if( start == NULL )
                assert( !"No [ in address." );
            long addr;
            if( sscanf( start + 1, "0x%lx", &addr ) != 1 )
                assert( !"Failed to parse address." );
            if( sizeof( void* ) == 4 )
                {
                sprintf( stack + pos, "0x%8lx", addr );
                pos += 10;
                }
            else if( sizeof( void* ) == 8 )
                {
                sprintf( stack + pos, "0x%16lx", addr );
                pos += 18;
                }
            else
                assert( !"Unknown sizeof( void* )." );
            }
        for( i = 0;
             i < backtraces_size;
             ++i )
            if( strcmp( backtraces[ i ].key, stack ) == 0 )
                {
                ++backtraces[ i ].count;
                break;
                }
        if( i == backtraces_size )
            {
            int stack_text_size = 10;
            char* stack_text;
            char* stack_text_pos;
            for( i = 0;
                 i < n;
                 ++i )
                stack_text_size += strlen( strings[ i ] ) + 5;
            stack_text = stack_text_pos = malloc( stack_text_size );
            for( i = 0;
                 i < n;
                 ++i )
                {
                stack_text_pos = stpcpy( stack_text_pos, "\n" );
                stack_text_pos = stpcpy( stack_text_pos, strings[ i ] );
                }
            backtraces[ backtraces_size ].key = strdup( stack );
            backtraces[ backtraces_size ].text = stack_text;
            backtraces[ backtraces_size ].count = 1;
            ++backtraces_size;
            if( backtraces_size >= MAX_BACKTRACES )
                assert( !"MAX_BACKTRACES reached." );
            }
        }
    free (strings);
    }

Status
_XReply (dpy, rep, extra, discard)
    register Display *dpy;
    register xReply *rep;
    int extra;		/* number of 32-bit words expected after the reply */
    Bool discard;	/* should I discard data following "extra" words? */
    {
    if( ___xreply_reply_enabled )
        ++___xreply_reply_count;
    if( xreply_backtrace_set == 0 )
        {
        if( getenv( "XREPLY_BACKTRACE" ) != NULL )
            { // C<number> - compress backtraces, saved as negative value in xreply_backtrace_type
            if( getenv( "XREPLY_BACKTRACE" )[ 0 ] == 'C' )
                xreply_backtrace_type = -atoi( getenv( "XREPLY_BACKTRACE" ) + 1 );
            else // <number> - print the backtrace every time
                xreply_backtrace_type = atoi( getenv( "XREPLY_BACKTRACE" ));
            }
        else
            xreply_backtrace_type = 0;
        }
    if( xreply_backtrace_type != 0 )
        xreply_backtrace();
    if( xreply_ptr == NULL )
        {
        xreply_ptr = (xreply_ptr_t)dlsym( RTLD_NEXT, "_XReply" );
        if( xreply_ptr == NULL )
            assert( !"dlsym() failed." );
        atexit( xreply_print );
        }
    return xreply_ptr( dpy, rep, extra, discard );
    }