/*
 * functionate, a tool to split "functions" out of vbtracetool traces
 *
 * Copyright 2007 Stuart Bennett <sb476@cam.ac.uk>
 *
 * This program is released under the terms of the GNU General Public License, version 2
 */

#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#define MAX_CALLS 1000000
#define MAX_CALLDEPTH 30
#define MAX_FNS 1000

struct cop {
	uint16_t cs, ip;
	uint16_t nextcs, nextip;
	long line, retline;
	fpos_t startpos, endpos;
	int fn;
	bool intr;
};

struct fn {
	uint16_t cs, ip;
	long len;
	int count;
	int instance;
	bool hasvariance;
	bool writeme;
	char name[20];
};

struct fn func[MAX_FNS];
struct cop call[MAX_CALLS];
int calls = 0;
int callstack[MAX_CALLDEPTH], calldepth = 0;
FILE *tracef, *outf, *outfnsf;

static bool read_line_plus_n(FILE *f, int n, char *outline)
{
	fpos_t current;
	bool ok = true;

	fgetpos(f, &current);
	for (int i = 0; i < n; i++)
		if (!fgets(outline, 180, f)) {
			ok = false;
			clearerr(f);
			break;
		}
	fsetpos(f, &current);

	return ok;
}

static void read_next_cs_ip(long lineno, uint16_t *cs, uint16_t *ip)
{
	char nextinstr[180];

	if (!read_line_plus_n(tracef, 4, nextinstr)) {
		printf("Trace appears to be truncated around line %ld\n", lineno);
		exit(EXIT_FAILURE);
	}
	if (nextinstr[4] != ':') {
		printf("Broken instruction spacing at line %ld\n", lineno + 4);
		exit(EXIT_FAILURE);
	}

	*cs = strtoul(&nextinstr[0], NULL, 16);
	*ip = strtoul(&nextinstr[5], NULL, 16);
}

static inline void add_call_common(long lineno, uint16_t lastip)
{
	read_next_cs_ip(lineno, &call[calls].cs, &call[calls].ip);

	call[calls].line = lineno;
	call[calls].nextip = lastip;
	fgetpos(tracef, &call[calls].startpos);
	callstack[calldepth] = calls;
	if (++calldepth >= MAX_CALLDEPTH) {
		printf("Too many nested calls\n");
		exit(EXIT_FAILURE);
	}
	call[calls].retline = 0;
	if (++calls >= MAX_CALLS) {
		printf("Too many calls\n");
		exit(EXIT_FAILURE);
	}
}

static void writeoutfn(int stacklen, int fncallstack[stacklen])
{
	int skipcall = 0, stackptr = 0;
	long start, finalretline;
	char line[128];

	finalretline = call[fncallstack[0]].retline;

//	printf("fn %s stacklen %d line %d retline %d\n", func[call[fncallstack[0]].fn].name, stacklen, call[fncallstack[0]].line, call[fncallstack[0]].retline);
	fprintf(outfnsf, "%s (called %d times):\n\n", func[call[fncallstack[0]].fn].name, func[call[fncallstack[0]].fn].count);

	do {
		if (skipcall) {
//			printf("beep skipcall %d stacklen %d\n", skipcall, stacklen);
			if (func[call[skipcall].fn].count > 1) {
				fsetpos(tracef, &call[skipcall].endpos);
				start = call[skipcall].retline;
			} else
				/* for hasvariance cases where we just need to change the call name */
				start = call[skipcall].line;
		} else {
			fsetpos(tracef, &call[fncallstack[0]].startpos);
			start = call[fncallstack[0]].line;
		}

		skipcall = 0;
		for (++stackptr; stackptr < stacklen; stackptr++) {
//			printf("test sc of %d\n", stackptr);
			if (call[fncallstack[stackptr]].line > start && call[fncallstack[stackptr]].line < finalretline) {
				skipcall = fncallstack[stackptr];
//				printf("had to skipcall %d (%d)\n", stackptr, skipcall);
//				printf("call %d ip %04x line %d retline %d\n", skipcall, call[skipcall].ip, call[skipcall].line, call[skipcall].retline);
				break;
			}
		}
		long stop;
		if (skipcall)
			stop = call[skipcall].line;
		else
			stop = finalretline;
		if (start == stop)
			break;

		for (int i = 0; i < stop - start - 1; i++) {
			fgets(line, 80, tracef);
			fwrite(line, strlen(line), 1, outfnsf);
		}
		fgets(line, 80, tracef);
		if (skipcall) {
			strcpy(&line[strlen(line) - 1], "                   ");
			sprintf(&line[call[skipcall].intr ? 55 : 48], "(%s%s", func[call[skipcall].fn].name, func[call[skipcall].fn].count > 1 ? ")\n\n" : " (inline))\n");
		}
		fwrite(line, strlen(line), 1, outfnsf);
	} while (skipcall);

	fprintf(outfnsf, "\n\n");
}

int main(int argc, char *argv[])
{
	int fns = 0;
	long lineno = 0;

	char outname[1024];
	strcpy(outname, argv[1]);
	strcat(outname, "-fn");
	char fnsoutname[1024];
	strcpy(fnsoutname, argv[1]);
	strcat(fnsoutname, "-fns");
	if (!(tracef = fopen(argv[1], "r"))) {
		printf("File open failed\n");
		exit(EXIT_FAILURE);
	}
	if (!(outf = fopen(outname, "w"))) {
		printf("File open failed\n");
		exit(EXIT_FAILURE);
	}
	if (!(outfnsf = fopen(fnsoutname, "w"))) {
		printf("File open failed\n");
		exit(EXIT_FAILURE);
	}

	printf("Building tree\n");

	uint16_t lastip = 0;
	while (!feof(tracef)) {
		char line[180];

		if (!fgets(line, 180, tracef))
			break;
		lineno += 1;

		/* sadly we can't do this by simple matching of calls and rets,
		 * as sometimes the multi-ret opcodes are used, or pushes and
		 * pops are unbalanced and a ret goes to somewhere odd.
		 * neither can we do it by tracking SP, as SP can get messed with
		 */
		if (strstr(line, "EIP"))
			lastip = strtoul(&line[39], NULL, 16);

		if (strstr(line, "CALL")) {
			call[calls].intr = false;
			call[calls].nextcs = strtoul(&line[0], NULL, 16);
			add_call_common(lineno, lastip);
		} else if (strstr(line, "INT"))	{	// if you can retf from an int, it's a call :)
			call[calls].intr = true;
			call[calls].nextcs = strtoul(&line[0], NULL, 16);
			add_call_common(lineno, lastip);
		} else if (strstr(line, "RET")) {
			uint16_t nextcs, nextip;

			read_next_cs_ip(lineno, &nextcs, &nextip);

			if (calldepth) {
				int newcd;
				bool borken = false;

				calldepth--;

				for (newcd = calldepth; newcd >= 0; newcd--) {
					if (call[callstack[newcd]].nextip == nextip && call[callstack[newcd]].nextcs == nextcs)
						break;
					else {
						if (!borken)
							printf("hmm, no matching call for the ret at line %ld to %04x:%04x. let's try higher in the call stack...", lineno, nextcs, nextip);
						borken = true;
					}
				}
				if (newcd >= 0) {
					/* for the case where newcd != calldepth, fixup the tree so it still separates
					 * by terminating the nested unreturned calls at the same time */
					for (int i = calldepth; i >= newcd; i--) {
						call[callstack[i]].retline = lineno;
						fgetpos(tracef, &call[callstack[i]].endpos);
					}
					calldepth = newcd;
					if (borken)
						printf(" yay!\n");
				} else {
					/* parsed code is undoubtedly broken. ignore ret and hope to muddle on */
					calldepth++;
					printf("\nerrk, no matching call for this ret. has nvidia broken the stack again?\n");
				}
			} else {
				if (nextip == 0 && nextcs == 0)
					break;
				
				printf("arrgl, too many RETs at line %ld\n", lineno + 1);
				call[callstack[0]].retline = lineno;
//				exit(EXIT_FAILURE);
			}				
		}
	}

	/* complete the tree */
	while (calldepth-- > 0)
		call[callstack[calldepth]].retline = lineno;

	for (int i = 0; i < calls; i++) {
		bool handled = false, hasvariance = false;
		int newvariant = 0;

		if (call[i].retline == 0)
			continue;

		for (int j = 0; j < fns; j++)
			if (func[j].ip == call[i].ip && func[j].cs == call[i].cs) {
				if (func[j].len == call[i].retline - call[i].line) {
					call[i].fn = j;
					func[j].count++;
					func[j].writeme = true;
					handled = true;
					break;
				} else {
					func[j].hasvariance = hasvariance = true;
					newvariant++;
				}
			}
		if (!handled) {
			call[i].fn = fns;
			func[fns].cs = call[i].cs;
			func[fns].ip = call[i].ip;
			func[fns].len = call[i].retline - call[i].line;
			func[fns].hasvariance = hasvariance;
			func[fns].instance = i;
			func[fns].count = 1;
			func[fns].writeme = false;
			snprintf(func[fns].name, 20, "%s_%04x:%04x_VAR%02d", call[i].intr ? "INT" : "FN", func[fns].cs, func[fns].ip, newvariant);
			if (++fns >= MAX_FNS) {
				printf("Too many functions\n");
				exit(EXIT_FAILURE);
			}
		}
	}

	printf("Splitting out functions\n");

	int i = 0;
	while (i < fns) {
//		printf("i is %d\n", i);
		int fncallstack[MAX_CALLS];
		int stackptr = 0, stackmax;
		if (!func[i].writeme) {
			i++;
			continue;
		}
		int j = func[i].instance;
		while (j < calls) {	
			if (func[call[j].fn].count > 1 || func[call[j].fn].hasvariance) {
				if (call[j].retline > call[func[i].instance].retline)
					break;
//				printf("adding j of %d (ip %04x, line %d, retline %d) at %d\n", j, call[j].ip, call[j].line, call[j].retline, stackptr);
				fncallstack[stackptr++] = j;
			}
			j++;
		}
		stackmax = stackptr;
//				printf("stackmax %d\n", stackmax);
		while (stackptr--)
			if (func[call[fncallstack[stackptr]].fn].writeme) {
//				printf("\n\npassing %d\n", stackmax - stackptr);
				writeoutfn(stackmax - stackptr, &fncallstack[stackptr]);
				func[call[fncallstack[stackptr]].fn].writeme = false;
			}
		i++;
	}

	fclose(outfnsf);

	printf("Compacting trace\n");

	fseek(tracef, 0, SEEK_SET);

	long start = 0;
	for (int nextcall = 0; nextcall < calls; nextcall++) {
		char line[128]; // as funny as it is, with this missing, gcc complains about the cop line...

		if (!(func[call[nextcall].fn].count > 1) && !func[call[nextcall].fn].hasvariance)
			continue;

		long stop = call[nextcall].line;
		for (int i = 0; i < stop - start - 1; i++) {
			fgets(line, 80, tracef);
			fwrite(line, strlen(line), 1, outf);
		}
//		fprintf(outf, "stop %d start %d len %d\n", stop, start, stop-start);
		fgets(line, 80, tracef);
		strcpy(&line[strlen(line) - 1], "                   ");
		sprintf(&line[call[nextcall].intr ? 55 : 48], "(%s%s", func[call[nextcall].fn].name, func[call[nextcall].fn].count > 1 ? ")\n\n" : " (inline))\n");
		fwrite(line, strlen(line), 1, outf);

		if (func[call[nextcall].fn].count > 1) {
			fsetpos(tracef, &call[nextcall].endpos);
			start = call[nextcall].retline;
			int lastcall = nextcall;
			for (; nextcall < calls; nextcall++)
				if (call[lastcall].retline < call[nextcall].line) {
					nextcall--;
					break;
				}
		} else
			start = stop;
	}

	while (!feof(tracef)) {
		char line[80];

		if (!fgets(line, 80, tracef))
			break;
		fwrite(line, strlen(line), 1, outf);
	}

	fclose(tracef);
	fclose(outf);

	return 0;
}