/*
 * Copyright (C) mwrap hackers <mwrap-perl@80x24.org>
 * License: GPL-3.0+ <https://www.gnu.org/licenses/gpl-3.0.txt>
 * single-threaded trace replayer, no runtime dependency on Perl
 * nor the rest of mwrap (aside from the hacked up dlmalloc).
 */
#define _LGPL_SOURCE /* allows URCU to inline some stuff */
#define _GNU_SOURCE
/* knobs for dlmalloc */
#define HAVE_MORECORE 0
#define DEFAULT_GRANULARITY (2U * 1024U * 1024U)
#define FOOTERS 1 /* required for remote_free_* stuff */
#define USE_DL_PREFIX
#define ONLY_MSPACES 1 /* aka per-thread "arenas" */
#define DLMALLOC_EXPORT static inline
/* #define NO_MALLOC_STATS 1 */
#define USE_LOCKS 0 /* we do our own global_mtx + ms_tsd */
#include <errno.h>
#include <err.h>
#include <stdlib.h>
#include <stdio.h>
#include <unistd.h>
#include <limits.h>
#include <sys/types.h>
#ifdef __GLIBC__
extern void __attribute__((weak)) malloc_stats(void);
extern void __attribute__((weak)) malloc_info(int, FILE *);
#	define GLIBC_MALLOC_STATS() do { \
		if (malloc_info) malloc_info(0, stderr); \
		if (malloc_stats) malloc_stats(); \
	} while (0)
#else
#	define GLIBC_MALLOC_STATS() do {} while (0)
#endif

extern void __attribute__((weak)) malloc_stats_print(
	void (*wcb)(void *, const char *), void *, const char *opts);

#include <urcu/rculist.h>
#include <urcu/wfcqueue.h>
#include "dlmalloc_c.h"
static mstate tr_ms;

static void *my_calloc(size_t nmemb, size_t size)
{
	void *p = mspace_calloc(tr_ms, nmemb, size);
	if (!p) err(1, "calloc");
	return p;
}

#define kcalloc(N,Z) my_calloc(N, Z)
#define kfree(P) mspace_free(tr_ms, P)
#define REALLOC_ARRAY(x, nmemb) do { \
	size_t asize; \
	if (__builtin_mul_overflow(sizeof(*(x)), nmemb, &asize)) \
		errx(1, "mul_overflow"); \
	(x) = mspace_realloc(tr_ms, (x), asize); \
	if (!x) err(1, "realloc"); \
} while (0)
#include "khashl.h"
#include "trace_struct.h"

static inline khint_t hash_uptr(uintptr_t p)
{
	return sizeof(uintptr_t) == 4 ? kh_hash_uint32(p) : kh_hash_uint64(p);
}

KHASHL_MAP_INIT(KH_LOCAL, kh_ptrmap, ptrmap, uintptr_t, uintptr_t,
		hash_uptr, kh_eq_generic)

static kh_ptrmap *old2cur;

static void store_ptr(uintptr_t old, void *cur)
{
	int absent;
	khint_t k = ptrmap_put(old2cur, old, &absent);
	if (absent)
		kh_val(old2cur, k) = (uintptr_t)cur;
}

int main(int argc, char *argv[])
{
	tr_ms = create_mspace(0, 0);
	tr_ms->seg.sflags = EXTERN_BIT | USE_MMAP_BIT;
	disable_contiguous(tr_ms);
	size_t realloc_miss = 0, free_miss = 0, bad_entry = 0;
	union {
		struct tr_memalign do_memalign;
		struct tr_free do_free;
		struct tr_malloc do_malloc;
		struct tr_calloc do_calloc;
		struct tr_realloc do_realloc;
	} as;
	int truncated = 0;

	old2cur = ptrmap_init();

	// don't fill buf all the way so we can do small reads in ENSURE_FILL:
	while (!feof(stdin) && !truncated) {

#define CONSUME(dst, required) do { \
	size_t need = sizeof(dst); \
	char *buf = (char *)&dst; \
	int done = 0; \
	while (need) { \
		size_t n = fread(buf, 1, need, stdin); \
		if (n > 0) { \
			need -= n; \
		} else if (n == 0 && !required) { \
			done = 1; \
			break; \
		} else { \
			warnx("TRUNCATED: %zu != %zu", n, need); \
			done = truncated = 1; \
			break; \
		} \
	} \
	if (done) break; \
} while (0)
		CONSUME(as.do_free.ptr, false);
		enum tr_fn fn = as.do_free.ptr & TR_MASK;
		as.do_free.ptr &= ~TR_MASK;
		khint_t k;
		void *cur;

		switch (fn) {
		case TR_FREE:
			k = ptrmap_get(old2cur, as.do_free.ptr);
			if (k >= kh_end(old2cur)) {
				++free_miss;
			} else {
				free((void *)kh_val(old2cur, k));
				ptrmap_del(old2cur, k);
			}
			break;
		case TR_MALLOC:
			CONSUME(as.do_malloc.size, true);
			cur = malloc(as.do_malloc.size);
			if (!cur)
				err(1, "malloc(%zu) => %p",
					as.do_malloc.size,
					(void *)as.do_malloc.ret);
			store_ptr(as.do_malloc.ret, cur);

			break;
		case TR_CALLOC:
			CONSUME(as.do_calloc.size, true);
			cur = calloc(as.do_calloc.size, 1);
			if (!cur)
				err(1, "calloc(%zu) => %p",
					as.do_calloc.size,
					(void *)as.do_calloc.ret);
			store_ptr(as.do_calloc.ret, cur);

			break;
		case TR_REALLOC:
			cur = NULL;
			CONSUME(as.do_realloc.ptr, true);
			CONSUME(as.do_realloc.size, true);
			if (as.do_realloc.ptr) {
				k = ptrmap_get(old2cur,
						as.do_realloc.ptr);
				if (k >= kh_end(old2cur)) {
					realloc_miss++;
				} else {
					cur = (void *)kh_val(old2cur, k);
					ptrmap_del(old2cur, k);
				}
			}
			void *rp = realloc(cur, as.do_realloc.size);
			if (!rp)
				err(1, "realloc(%p => %p, %zu) => %p",
					(void *)as.do_realloc.ptr,
					cur,
					as.do_realloc.size,
					(void *)as.do_realloc.ret);
			store_ptr(as.do_realloc.ret, rp);
			break;
		case TR_MEMALIGN:
			cur = NULL;
			CONSUME(as.do_memalign.alignment, true);
			CONSUME(as.do_memalign.size, true);
			int rc = posix_memalign(&cur,
					as.do_memalign.alignment,
					as.do_memalign.size);
			if (rc) {
				errno = rc;
				err(1, "posix_memalign(%zu, %zu) => %p",
					as.do_memalign.alignment,
					as.do_memalign.size,
					(void *)as.do_memalign.ret);
			}
			store_ptr(as.do_memalign.ret, cur);
			break;
		default:
			bad_entry++;
		}
	}

	if (free_miss || realloc_miss || bad_entry)
		fprintf(stderr, "W: miss free=%zu realloc=%zu bad=%zu\n",
			free_miss, realloc_miss, bad_entry);

	fprintf(stderr, "# ptrmap .size=%zu capa=%zu\n",
		(size_t)kh_size(old2cur), (size_t)kh_capacity(old2cur));

	if (malloc_stats_print) // jemalloc loaded
		malloc_stats_print(NULL, NULL, NULL);
	else
		GLIBC_MALLOC_STATS();

	int c;
	char *end;
	long sec = 0;
	while ((c = getopt(argc, argv, "s:")) != -1) {
		switch (c) {
		case 's':
			sec = strtol(optarg, &end, 10);
			if (*end != 0)
				errx(1, "`-s %s' invalid seconds", optarg);
			break;
		default: warnx("bad switch `-%c'", c);
		}
	}
	if (sec < 0) {
		fprintf(stderr, "# PID:%d sleeping indefinitely\n",
			(int)getpid());
		pause();
	}
	if (sec > 0) {
		unsigned s = sec > UINT_MAX ? UINT_MAX : sec;
		fprintf(stderr, "# PID:%d sleeping %u seconds\n",
			(int)getpid(), s);
		sleep(s);
	}

	return truncated;
}