# Unwinder for libjit.

# FIXME
# * x86-64 only for now

import gdb
import gdb.types
import gdb.unwinder
from gdb.FrameDecorator import FrameDecorator

# Python 3 compat.
try:
	long
except NameError:
	long = int

# Python 3 compat.
try:
	from itertools import imap
except ImportError:
	imap = map

contextSymbol = None
context = None
nameMeta = None
funcCache = None

class FunctionCache(object):
	def __init__(self):
		self.clear()

	def clear(self):
		self.cache = []
		self.filled = False

	def fill(self, context):
		global nameMeta

		func = context["functions"]
		while long(func) != 0:
			name = None
			if nameMeta != None:
				meta = func["meta"]
				while long(meta) != 0:
					if meta["type"] == nameMeta:
						name = meta["data"]
						typ = gdb.lookup_type("char").pointer()
						name = name.cast(typ).string()
						break
					meta = meta["next"]

			self.cache.append((func["entry_point"], func["code_end"], name))
			func = func["next"]

		self.filled = True

	def find(self, pc):
		for func in self.cache:
			start, end, name = func
			if pc >= start and pc < end:
				return func

		return None

	def contains(self, pc):
		return self.find(pc) != None

class FrameId(object):
	def __init__(self, sp, pc):
		self.sp = sp
		self.pc = pc

class LibjitContextSpecifier(gdb.Command):
	def __init__(self):
		super(LibjitContextSpecifier, self).__init__("libjit-context", gdb.COMMAND_USER)

	def invoke(self, args, tty):
		global context, contextSymbol

		try:
			context = gdb.lookup_symbol(args).value()
			contextSymbol = None
		except:
			context = None
			contextSymbol = gdb.lookup_global_symbol(args)
			if contextSymbol == None:
				gdb.write("Invalid symbol " + args)

	def complete(self, line, word):
		return gdb.COMPLETE_SYMBOL

class LibjitNameMetaSpecifier(gdb.Command):
	def __init__(self):
		super(LibjitNameMetaSpecifier, self).__init__("libjit-name-meta", gdb.COMMAND_USER)

	def invoke(self, args, tty):
		global nameMeta
		nameMeta = gdb.parse_and_eval(args)

	def complete(self, line, word):
		return gdb.COMPLETE_EXPRESSION

class LibjitFrameDecorator(FrameDecorator):
	def __init__(self, base, name, start, end):
		super(LibjitFrameDecorator, self).__init__(base)
		self.name = name
		self.start = start
		self.end = end

	def function(self):
		if self.name != None:
			return "JIT " + self.name
		else:
			return "JIT[0x%x, 0x%x]" % (self.start, self.end)

class LibjitFrameFilter(object):
	def __init__(self):
		self.name = "Libjit"
		self.enabled = True
		self.priority = 100

	def maybe_wrap(self, frame):
		pc = long(frame.inferior_frame().read_register("rip"))
		func = funcCache.find(pc)
		if func == None:
			return frame

		start, end, name = func
		return LibjitFrameDecorator(frame, name, start, end)

	def filter(self, frame_iter):
		return imap(self.maybe_wrap, frame_iter)

class LibjitUnwinder(gdb.unwinder.Unwinder):
	def __init__(self):
		super(LibjitUnwinder, self).__init__("Libjit")
		self.enabled = True

	def __call__(self, pending_frame):
		global context, contextSymbol, function_from_pc, firstTime

		if not funcCache.filled:
			if context == None:
				_context = contextSymbol.value()
			else:
				_context = context

			if long(_context) == 0:
				return None

			funcCache.fill(_context)

		pc = pending_frame.read_register("rip")
		rbp = pending_frame.read_register("rbp")

		if not funcCache.contains(pc):
			return None

		# Convenient type to work with.
		ptr = gdb.lookup_type("void").pointer().pointer()

		# Previous frame pointer is at 0(%rbp).
		as_ptr = rbp.cast(ptr)
		prev_rbp = as_ptr.dereference()
		# Previous PC is at 8(%rbp)
		prev_pc = (as_ptr + 1).dereference()

		frame_id = FrameId(prev_rbp, prev_pc)
		unwind_info = pending_frame.create_unwind_info(frame_id)
		unwind_info.add_saved_register("rip", prev_pc)
		unwind_info.add_saved_register("rbp", prev_rbp)

		return unwind_info

# Register the commands
LibjitContextSpecifier()
LibjitNameMetaSpecifier()

# Register the unwinder.
gdb.unwinder.register_unwinder(None, LibjitUnwinder(), replace=True)

# Register the frame filter
filter = LibjitFrameFilter()
gdb.frame_filters[filter.name] = filter

# Initialize the function cache
funcCache = FunctionCache()
gdb.events.cont.connect(lambda event: funcCache.clear())