/*-
 * Copyright (C) 2006-2007 Jason Evans <jasone@FreeBSD.org>.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice(s), this list of conditions and the following disclaimer as
 *    the first lines of this file unmodified other than the possible
 *    addition of one or more copyright notices.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice(s), this list of conditions and the following disclaimer in
 *    the documentation and/or other materials provided with the
 *    distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDER(S) ``AS IS'' AND ANY
 * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
 * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT HOLDER(S) BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
 * BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
 * WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
 * OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
 * EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

// Parse a malloc trace (generated via ktrace/kdump) and simulate the
// allocation events.  Records look like:
//
//   31532 foo      USER  malloc_init()
//   31532 foo      USER  0x40 = malloc(15)
//   31532 foo      USER  0x60 = realloc(0x40, 17)
//   31532 foo      USER  free(0x60)

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <sys/mman.h>
#include <sys/tree.h>
#include <sys/types.h>

typedef struct record_s record_t;
struct record_s {
	RB_ENTRY(record_s) link;
	uintptr_t tag;
	void *obj;
};
typedef struct record_tree_s record_tree_t;
RB_HEAD(record_tree_s, record_s);

static int
record_comp(record_t *a, record_t *b)
{

	if (a->tag < b->tag)
		return (-1);
	else if (a->tag == b->tag)
		return (0);
	else
		return (1);
}

RB_GENERATE_STATIC(record_tree_s, record_s, link, record_comp);

static record_t *rec_stack = NULL;

static record_t *
record_alloc(void)
{
	record_t *ret;

	if (rec_stack == NULL) {
		record_t *recs;
		unsigned long i;
#define MMAP_SIZE (1024 * 1024 * 1024)

		recs = (record_t *)mmap(NULL, MMAP_SIZE,
		    PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANON, -1, 0);
		if (recs == NULL) {
			fprintf(stderr, "mtrplay: mmap() error (OOM)\n");
			exit(1);
		}

		for (i = 0; i < MMAP_SIZE / sizeof(record_t); i++) {
			recs[i].obj = (void *)rec_stack;
			rec_stack = &recs[i];
		}
	}

	ret = rec_stack;
	rec_stack = (record_t *)ret->obj;

	return (ret);
}

static void
record_dalloc(record_t *rec)
{

	rec->obj = (void *)rec_stack;
	rec_stack = rec;
}

int
main(void)
{
	uint64_t line;
	int result;
	char fbuf[4096], lbuf[128];
	unsigned foff, flen, i;
	unsigned pid;
	char argv0[16];
	void *addr, *oldaddr;
	size_t size;
	record_tree_t tree;

	RB_INIT(&tree);

	foff = flen = 0;
	for (line = 1;; line++) {
		if (line % 100000 == 0)
			fprintf(stderr, ".");

		// Get a line of input.
		for (i = 0; i < sizeof(lbuf) - 1; i++) {
			// Refill the read buffer if necessary.  We use read(2)
			// instead of a higher level API in order to avoid
			// internal allocation.
			if (foff == flen) {
				foff = 0;
				flen = read(0, fbuf, sizeof(fbuf));
				if (flen <= 0) {
					goto RETURN;
				}
			}
			switch (fbuf[foff]) {
				case '\n': {
					lbuf[i] = '\0';
					foff++;
					goto OUT;
				} default: {
					lbuf[i] = fbuf[foff];
					foff++;
					break;
				}
			}
		}
		OUT:;

		// realloc?
		result = sscanf(lbuf, "%u %15s USER %p = realloc(%p, %zu)",
		    &pid, argv0, &addr, &oldaddr, &size);
		if (result == 5) {
			record_t key, *rec;
			void *p;

			key.tag = (uintptr_t)oldaddr;
			rec = RB_FIND(record_tree_s, &tree, &key);
			if (rec == NULL) {
				fprintf(stderr,
				    "mtrplay: Line %llu: Record not found\n",
				    line);
				exit(1);
			}
			RB_REMOVE(record_tree_s, &tree, rec);

			p = realloc(rec->obj, size);
			if (p == NULL) {
				fprintf(stderr, "mtrplay: Line %llu: OOM\n",
				    line);
				exit(1);
			}

			rec->tag = (uintptr_t)addr;
			rec->obj = p;
			RB_INSERT(record_tree_s, &tree, rec);
			continue;
		}

		// malloc?
		result = sscanf(lbuf, "%u %15s USER %p = malloc(%zu)",
		    &pid, argv0, &addr, &size);
		if (result == 4)
		{
			void *p;
			record_t *rec;

			rec = record_alloc();

			p = malloc(size);
			if (p == NULL) {
				fprintf(stderr, "mtrplay: Line %llu: OOM\n",
				    line);
				exit(1);
			}
			memset(p, 0, size);

			rec->tag = (uintptr_t)addr;
			rec->obj = p;
			RB_INSERT(record_tree_s, &tree, rec);
			continue;
		}

		// free?
		result = sscanf(lbuf, "%u %15s USER free(%p)",
		    &pid, argv0, &oldaddr);
		if (result == 3)
		{
			record_t key, *rec;

			key.tag = (uintptr_t)oldaddr;
			rec = RB_FIND(record_tree_s, &tree, &key);
			if (rec == NULL) {
				fprintf(stderr,
				    "mtrplay: Line %llu: Record not found\n",
				    line);
				exit(1);
			}
			RB_REMOVE(record_tree_s, &tree, rec);

			free(rec->obj);
			record_dalloc(rec);
			continue;
		}

		// malloc_init?
		result = sscanf(lbuf, "%u %15s USER malloc_init()",
		    &pid, argv0);
		if (result == 2)
		{
			continue;
		}

		fprintf(stderr, "mtrplay: Error reading line %llu of input\n",
		    line);
		exit(1);
	}

RETURN:
	return 0;
}