#include "Python.h"

#include "compile.h"

#include "frameobject.h"

#include "structseq.h"

#include "rotatingtree.h"

#if !defined(HAVE_LONG_LONG)

#error "This module requires long longs!"

#endif

/*** Selection of a high-precision timer ***/

#ifdef MS_WINDOWS

#include <windows.h>

static PY_LONG_LONG

hpTimer(void)

{

	LARGE_INTEGER li;

	QueryPerformanceCounter(&li);

	return li.QuadPart;

}

static double

hpTimerUnit(void)

{

	LARGE_INTEGER li;

	if (QueryPerformanceFrequency(&li))

		return 1.0 / li.QuadPart;

	else

		return 0.000001;  /* unlikely */

}

#else  /* !MS_WINDOWS */

#ifndef HAVE_GETTIMEOFDAY

#error "This module requires gettimeofday() on non-Windows platforms!"

#endif

#if (defined(PYOS_OS2) && defined(PYCC_GCC))

#include <sys/time.h>

#else

#include <sys/resource.h>

#include <sys/times.h>

#endif

static PY_LONG_LONG

hpTimer(void)

{

	struct timeval tv;

	PY_LONG_LONG ret;

#ifdef GETTIMEOFDAY_NO_TZ

	gettimeofday(&tv);

#else

	gettimeofday(&tv, (struct timezone *)NULL);

#endif

	ret = tv.tv_sec;

	ret = ret * 1000000 + tv.tv_usec;

	return ret;

}

static double

hpTimerUnit(void)

{

	return 0.000001;

}

#endif  /* MS_WINDOWS */

/************************************************************/

/* Written by Brett Rosen and Ted Czotter */

struct _ProfilerEntry;

/* represents a function called from another function */

typedef struct _ProfilerSubEntry {

	rotating_node_t header;

	PY_LONG_LONG tt;

	PY_LONG_LONG it;

	long callcount;

	long recursivecallcount;

	long recursionLevel;

} ProfilerSubEntry;

/* represents a function or user defined block */

typedef struct _ProfilerEntry {

	rotating_node_t header;

	PyObject *userObj; /* PyCodeObject, or a descriptive str for builtins */

	PY_LONG_LONG tt; /* total time in this entry */

	PY_LONG_LONG it; /* inline time in this entry (not in subcalls) */

	long callcount; /* how many times this was called */

	long recursivecallcount; /* how many times called recursively */

	long recursionLevel;

	rotating_node_t *calls;

} ProfilerEntry;

typedef struct _ProfilerContext {

	PY_LONG_LONG t0;

	PY_LONG_LONG subt;

	struct _ProfilerContext *previous;

	ProfilerEntry *ctxEntry;

} ProfilerContext;

typedef struct {

	PyObject_HEAD

	rotating_node_t *profilerEntries;

	ProfilerContext *currentProfilerContext;

	ProfilerContext *freelistProfilerContext;

	int flags;

	PyObject *externalTimer;

	double externalTimerUnit;

} ProfilerObject;

#define POF_ENABLED     0x001

#define POF_SUBCALLS    0x002

#define POF_BUILTINS    0x004

#define POF_NOMEMORY    0x100

staticforward PyTypeObject PyProfiler_Type;

#define PyProfiler_Check(op) PyObject_TypeCheck(op, &PyProfiler_Type)

#define PyProfiler_CheckExact(op) (Py_TYPE(op) == &PyProfiler_Type)

/*** External Timers ***/

#define DOUBLE_TIMER_PRECISION   4294967296.0

static PyObject *empty_tuple;

static PY_LONG_LONG CallExternalTimer(ProfilerObject *pObj)

{

	PY_LONG_LONG result;

	PyObject *o = PyObject_Call(pObj->externalTimer, empty_tuple, NULL);

	if (o == NULL) {

		PyErr_WriteUnraisable(pObj->externalTimer);

		return 0;

	}

	if (pObj->externalTimerUnit > 0.0) {

		/* interpret the result as an integer that will be scaled

		   in profiler_getstats() */

		result = PyLong_AsLongLong(o);

	}

	else {

		/* interpret the result as a double measured in seconds.

		   As the profiler works with PY_LONG_LONG internally

		   we convert it to a large integer */

		double val = PyFloat_AsDouble(o);

		/* error handling delayed to the code below */

		result = (PY_LONG_LONG) (val * DOUBLE_TIMER_PRECISION);

	}

	Py_DECREF(o);

	if (PyErr_Occurred()) {

		PyErr_WriteUnraisable(pObj->externalTimer);

		return 0;

	}

	return result;

}

#define CALL_TIMER(pObj)	((pObj)->externalTimer ?		\

					CallExternalTimer(pObj) :	\

					hpTimer())

/*** ProfilerObject ***/

static PyObject *

normalizeUserObj(PyObject *obj)

{

	PyCFunctionObject *fn;

	if (!PyCFunction_Check(obj)) {

		Py_INCREF(obj);

		return obj;

	}

	/* Replace built-in function objects with a descriptive string

	   because of built-in methods -- keeping a reference to

	   __self__ is probably not a good idea. */

	fn = (PyCFunctionObject *)obj;

	if (fn->m_self == NULL) {

		/* built-in function: look up the module name */

		PyObject *mod = fn->m_module;

		char *modname;

		if (mod && PyString_Check(mod)) {

			modname = PyString_AS_STRING(mod);

		}

		else if (mod && PyModule_Check(mod)) {

			modname = PyModule_GetName(mod);

			if (modname == NULL) {

				PyErr_Clear();

				modname = "__builtin__";

			}

		}

		else {

			modname = "__builtin__";

		}

		if (strcmp(modname, "__builtin__") != 0)

			return PyString_FromFormat("<%s.%s>",

						   modname,

						   fn->m_ml->ml_name);

		else

			return PyString_FromFormat("<%s>",

						   fn->m_ml->ml_name);

	}

	else {

		/* built-in method: try to return

			repr(getattr(type(__self__), __name__))

		*/

		PyObject *self = fn->m_self;

		PyObject *name = PyString_FromString(fn->m_ml->ml_name);

		if (name != NULL) {

			PyObject *mo = _PyType_Lookup(Py_TYPE(self), name);

			Py_XINCREF(mo);

			Py_DECREF(name);

			if (mo != NULL) {

				PyObject *res = PyObject_Repr(mo);

				Py_DECREF(mo);

				if (res != NULL)

					return res;

			}

		}

		PyErr_Clear();

		return PyString_FromFormat("<built-in method %s>",

					   fn->m_ml->ml_name);

	}

}

static ProfilerEntry*

newProfilerEntry(ProfilerObject *pObj, void *key, PyObject *userObj)

{

	ProfilerEntry *self;

	self = (ProfilerEntry*) malloc(sizeof(ProfilerEntry));

	if (self == NULL) {

		pObj->flags |= POF_NOMEMORY;

		return NULL;

	}

	userObj = normalizeUserObj(userObj);

	if (userObj == NULL) {

		PyErr_Clear();

		free(self);

		pObj->flags |= POF_NOMEMORY;

		return NULL;

	}

	self->header.key = key;

	self->userObj = userObj;

	self->tt = 0;

	self->it = 0;

	self->callcount = 0;

	self->recursivecallcount = 0;

	self->recursionLevel = 0;

	self->calls = EMPTY_ROTATING_TREE;

	RotatingTree_Add(&pObj->profilerEntries, &self->header);

	return self;

}

static ProfilerEntry*

getEntry(ProfilerObject *pObj, void *key)

{

	return (ProfilerEntry*) RotatingTree_Get(&pObj->profilerEntries, key);

}

static ProfilerSubEntry * 

getSubEntry(ProfilerObject *pObj, ProfilerEntry *caller, ProfilerEntry* entry)

{

	return (ProfilerSubEntry*) RotatingTree_Get(&caller->calls,

						    (void *)entry);

}

static ProfilerSubEntry *

newSubEntry(ProfilerObject *pObj,  ProfilerEntry *caller, ProfilerEntry* entry)

{

	ProfilerSubEntry *self;

	self = (ProfilerSubEntry*) malloc(sizeof(ProfilerSubEntry));

	if (self == NULL) {

		pObj->flags |= POF_NOMEMORY;

		return NULL;

	}

	self->header.key = (void *)entry;

	self->tt = 0;

	self->it = 0;

	self->callcount = 0;

	self->recursivecallcount = 0;

	self->recursionLevel = 0;

	RotatingTree_Add(&caller->calls, &self->header);

	return self;

}

static int freeSubEntry(rotating_node_t *header, void *arg)

{

	ProfilerSubEntry *subentry = (ProfilerSubEntry*) header;

	free(subentry);

	return 0;

}

static int freeEntry(rotating_node_t *header, void *arg)

{

	ProfilerEntry *entry = (ProfilerEntry*) header;

	RotatingTree_Enum(entry->calls, freeSubEntry, NULL);

	Py_DECREF(entry->userObj);

	free(entry);

	return 0;

}

static void clearEntries(ProfilerObject *pObj)

{

	RotatingTree_Enum(pObj->profilerEntries, freeEntry, NULL);

	pObj->profilerEntries = EMPTY_ROTATING_TREE;

	/* release the memory hold by the free list of ProfilerContexts */

	while (pObj->freelistProfilerContext) {

		ProfilerContext *c = pObj->freelistProfilerContext;

		pObj->freelistProfilerContext = c->previous;

		free(c);

	}

}

static void

initContext(ProfilerObject *pObj, ProfilerContext *self, ProfilerEntry *entry)

{

	self->ctxEntry = entry;

	self->subt = 0;

	self->previous = pObj->currentProfilerContext;

	pObj->currentProfilerContext = self;

	++entry->recursionLevel;

	if ((pObj->flags & POF_SUBCALLS) && self->previous) {

		/* find or create an entry for me in my caller's entry */

		ProfilerEntry *caller = self->previous->ctxEntry;

		ProfilerSubEntry *subentry = getSubEntry(pObj, caller, entry);

		if (subentry == NULL)

			subentry = newSubEntry(pObj, caller, entry);

		if (subentry)

			++subentry->recursionLevel;

	}

	self->t0 = CALL_TIMER(pObj);

}

static void

Stop(ProfilerObject *pObj, ProfilerContext *self, ProfilerEntry *entry)

{

	PY_LONG_LONG tt = CALL_TIMER(pObj) - self->t0;

	PY_LONG_LONG it = tt - self->subt;

	if (self->previous)

		self->previous->subt += tt;

	pObj->currentProfilerContext = self->previous;

	if (--entry->recursionLevel == 0)

		entry->tt += tt;

	else

		++entry->recursivecallcount;

	entry->it += it;

	entry->callcount++;

	if ((pObj->flags & POF_SUBCALLS) && self->previous) {

		/* find or create an entry for me in my caller's entry */

		ProfilerEntry *caller = self->previous->ctxEntry;

		ProfilerSubEntry *subentry = getSubEntry(pObj, caller, entry);

		if (subentry) {

			if (--subentry->recursionLevel == 0)

				subentry->tt += tt;

			else

				++subentry->recursivecallcount;

			subentry->it += it;

			++subentry->callcount;

		}

	}

}

static void

ptrace_enter_call(PyObject *self, void *key, PyObject *userObj)

{

	/* entering a call to the function identified by 'key'

	   (which can be a PyCodeObject or a PyMethodDef pointer) */

	ProfilerObject *pObj = (ProfilerObject*)self;

	ProfilerEntry *profEntry;

	ProfilerContext *pContext;

	/* In the case of entering a generator expression frame via a

	 * throw (gen_send_ex(.., 1)), we may already have an

	 * Exception set here. We must not mess around with this

	 * exception, and some of the code under here assumes that

	 * PyErr_* is its own to mess around with, so we have to

	 * save and restore any current exception. */

	PyObject *last_type, *last_value, *last_tb;

	PyErr_Fetch(&last_type, &last_value, &last_tb);

	profEntry = getEntry(pObj, key);

	if (profEntry == NULL) {

		profEntry = newProfilerEntry(pObj, key, userObj);

		if (profEntry == NULL)

			goto restorePyerr;

	}

	/* grab a ProfilerContext out of the free list */

	pContext = pObj->freelistProfilerContext;

	if (pContext) {

		pObj->freelistProfilerContext = pContext->previous;

	}

	else {

		/* free list exhausted, allocate a new one */

		pContext = (ProfilerContext*)

			malloc(sizeof(ProfilerContext));

		if (pContext == NULL) {

			pObj->flags |= POF_NOMEMORY;

			goto restorePyerr;

		}

	}

	initContext(pObj, pContext, profEntry);

restorePyerr:

	PyErr_Restore(last_type, last_value, last_tb);

}

static void

ptrace_leave_call(PyObject *self, void *key)

{

	/* leaving a call to the function identified by 'key' */

	ProfilerObject *pObj = (ProfilerObject*)self;

	ProfilerEntry *profEntry;

	ProfilerContext *pContext;

	pContext = pObj->currentProfilerContext;

	if (pContext == NULL)

		return;

	profEntry = getEntry(pObj, key);

	if (profEntry) {

		Stop(pObj, pContext, profEntry);

	}

	else {

		pObj->currentProfilerContext = pContext->previous;

	}

	/* put pContext into the free list */

	pContext->previous = pObj->freelistProfilerContext;

	pObj->freelistProfilerContext = pContext;

}

static int

profiler_callback(PyObject *self, PyFrameObject *frame, int what,

		  PyObject *arg)

{

	switch (what) {

	/* the 'frame' of a called function is about to start its execution */

	case PyTrace_CALL:

		ptrace_enter_call(self, (void *)frame->f_code,

				        (PyObject *)frame->f_code);

		break;

	/* the 'frame' of a called function is about to finish

	   (either normally or with an exception) */

	case PyTrace_RETURN:

		ptrace_leave_call(self, (void *)frame->f_code);

		break;

	/* case PyTrace_EXCEPTION:

		If the exception results in the function exiting, a

		PyTrace_RETURN event will be generated, so we don't need to

		handle it. */

#ifdef PyTrace_C_CALL	/* not defined in Python <= 2.3 */

	/* the Python function 'frame' is issuing a call to the built-in

	   function 'arg' */

	case PyTrace_C_CALL:

		if ((((ProfilerObject *)self)->flags & POF_BUILTINS)

		    && PyCFunction_Check(arg)) {

			ptrace_enter_call(self,

					  ((PyCFunctionObject *)arg)->m_ml,

					  arg);

		}

		break;

	/* the call to the built-in function 'arg' is returning into its

	   caller 'frame' */

	case PyTrace_C_RETURN:		/* ...normally */

	case PyTrace_C_EXCEPTION:	/* ...with an exception set */

		if ((((ProfilerObject *)self)->flags & POF_BUILTINS)

		    && PyCFunction_Check(arg)) {

			ptrace_leave_call(self,

					  ((PyCFunctionObject *)arg)->m_ml);

		}

		break;

#endif

	default:

		break;

	}

	return 0;

}

static int

pending_exception(ProfilerObject *pObj)

{

	if (pObj->flags & POF_NOMEMORY) {

		pObj->flags -= POF_NOMEMORY;

		PyErr_SetString(PyExc_MemoryError,

				"memory was exhausted while profiling");

		return -1;

	}

	return 0;

}

/************************************************************/

static PyStructSequence_Field profiler_entry_fields[] = {

	{"code",         "code object or built-in function name"},

	{"callcount",    "how many times this was called"},

	{"reccallcount", "how many times called recursively"},

	{"totaltime",    "total time in this entry"},

	{"inlinetime",   "inline time in this entry (not in subcalls)"},

	{"calls",        "details of the calls"},

	{0}

};

static PyStructSequence_Field profiler_subentry_fields[] = {

	{"code",         "called code object or built-in function name"},

	{"callcount",    "how many times this is called"},

	{"reccallcount", "how many times this is called recursively"},

	{"totaltime",    "total time spent in this call"},

	{"inlinetime",   "inline time (not in further subcalls)"},

	{0}

};

static PyStructSequence_Desc profiler_entry_desc = {

	"_lsprof.profiler_entry", /* name */

	NULL, /* doc */

	profiler_entry_fields,

	6

};

static PyStructSequence_Desc profiler_subentry_desc = {

	"_lsprof.profiler_subentry", /* name */

	NULL, /* doc */

	profiler_subentry_fields,

	5

};