#    This file is part of Metasm, the Ruby assembly manipulation suite
#    Copyright (C) 2006-2009 Yoann GUILLOT
#
#    Licence is LGPL, see LICENCE in the top-level directory

module ::Metasm
module Gui
class GraphHeapWidget < GraphViewWidget
	attr_accessor :heap, :addr_struct, :snapped
	# addr_struct = 0x234 => AllocCStruct

	def set_color_arrow(b1, b2)
		if b1 == @caret_box or b2 == @caret_box
			draw_color :arrow_hl
		else
			draw_color :arrow_cond
		end
	end

	def setup_contextmenu(b, m)
		addsubmenu(m, '_follow pointer') {
			next if not lm = b[:line_member][@caret_y]
			addr = b[:line_struct][@caret_y][lm]
			next if not @heap.chunks[addr]
			if lm.kind_of?(::Integer)
				t = b[:line_struct][@caret_y].struct.type
			else
				t = lm.type
			end
			if t.pointer? and t.pointed.untypedef.kind_of? C::Union
				@heap.chunk_struct[addr] ||= t.pointed.untypedef
			end
			st = @heap.chunk_struct[addr] || create_struct(addr)
			ed = @dasm.get_edata_at(addr)
			@addr_struct[addr] = @heap.cp.decode_c_struct(st.name, ed.data, ed.ptr)
			gui_update
		}
		addsubmenu(m, '_hide box') {
			@selected_boxes.each { |sb|
				@addr_struct.delete sb.id if @addr_struct.length > 1
			}
			@curcontext.root_addrs = struct_find_roots(@addr_struct.keys.first)
			gui_update
		}
		super(b, m)
	end

	def keypress(k)
		case k
		when ?u
			# update display (refresh struct member values)
			@parent_widget.parent_widget.console.commands['refresh'][]
			gui_update
		when ?t
			# change struct field type
			if @selected_boxes.length > 1
				# mass-retype chunks
				st = @addr_struct[@selected_boxes[0].id].struct
				inputbox("replacement struct for selected chunks", :text => st.name) { |n|
					next if not nst = @heap.cp.toplevel.struct[n]
					@selected_boxes.each { |sb|
						as = @addr_struct[sb.id]
						@heap.chunk_struct[sb.id] = nst
						@addr_struct[sb.id] = @heap.cp.decode_c_struct(n, as.str, as.stroff)
					}
					gui_update
				}
			elsif b = @caret_box
				if @caret_y == 0
					as = @addr_struct[b.id]
					st = as.struct
					inputbox("replacement struct for #{st.name}", :text => st.name) { |n|
						next if not nst = @heap.cp.toplevel.struct[n]
						@heap.chunk_struct[b.id] = nst
						@addr_struct[b.id] = @heap.cp.decode_c_struct(n, as.str, as.stroff)
						gui_update
					}
				elsif m = b[:line_member][@caret_y]
					if m.kind_of?(Integer)
						# XXX Array, need to find the outer struct
						mn = b[:line_text_col][@caret_y].map { |l, c| l }.join[/(\S*)\[/, 1]
						ar = b[:line_struct][@caret_y]
						st = b[:line_struct][0...@caret_y].reverse.compact.find { |st_| st_.struct.kind_of?(C::Struct) and st_[mn].struct == ar.struct }
						raise '?' if not st
						st = st.struct
						m = st.fldlist[mn]
					else
						st = b[:line_struct][@caret_y].struct
					end
					inputbox("new type for #{m.name}", :text => m.dump_def(@heap.cp.toplevel)[0].join(' ')) { |nn|
						nil while @heap.cp.readtok
						@heap.cp.lexer.feed nn
						if not v = C::Variable.parse_type(@heap.cp, @heap.cp.toplevel, true)
							nil while @heap.cp.readtok
							raise 'bad type'
						end
						v.parse_declarator(@heap.cp, @heap.cp.toplevel)
						nt = v.type
						nsz = @heap.cp.sizeof(nt)
						osz = @heap.cp.sizeof(m)
						if nsz > osz and st.kind_of?(C::Struct)
							idx = st.members.index(m)
							# eat next members
							while nsz > osz
								break if idx+1 >= st.members.length
								sz = @heap.cp.sizeof(st.members.delete_at(idx+1))
								osz += sz
							end
						end
						if nsz < osz and st.kind_of?(C::Struct)
							idx = st.members.index(m)
							pos = st.offsetof(@heap.cp, m)
							# fill gap with bytes
							idx += 1
							while nsz < osz
								st.members[idx, 0] = [C::Variable.new(('unk_%x' % (pos+nsz)), C::BaseType.new(:__int8, :unsigned))]
								idx += 1
								nsz += 1
							end
						end
						m.type = nt
						st.update_member_cache(@heap.cp)
						gui_update
					}

				end
			end
		when ?n
			# rename struct field
			if b = @caret_box
				if @caret_y == 0
					st = @addr_struct[b.id].struct
					inputbox("new name for #{st.name}", :text => st.name) { |nn|
						raise "struct #{nn} already exists (try 't')" if @heap.cp.toplevel.struct[nn]
						@heap.cp.toplevel.struct[nn] = @heap.cp.toplevel.struct.delete(st.name)
						st.name = nn
						gui_update
					}
				elsif m = b[:line_member][@caret_y]
					if m.kind_of?(Integer)
						mn = b[:line_text_col][@caret_y].map { |l, c| l }.join[/(\S*)\[/, 1]
						ar = b[:line_struct][@caret_y]
						st = b[:line_struct][0...@caret_y].reverse.compact.find { |st_| st_.struct.kind_of?(C::Struct) and st_[mn].struct == ar.struct }
						raise '?' if not st
						st = st.struct
						m = st.fldlist[mn]
					else
						st = b[:line_struct][@caret_y].struct
					end
					inputbox("new name for #{m.name}", :text => m.name) { |nn|
						m.name = nn
						st.update_member_cache(@heap.cp)
						gui_update
					}
				end
			end
		when ?e
			# edit struct field value under the cursor
			if b = @caret_box
				# TODO b[:struct][line], b.[:member][line] (int for Arrays)
				st = b[:line_struct][@caret_y]
				mb = b[:line_member][@caret_y]
				if st and mb
					if mb.kind_of?(C::Variable) and mb.type.kind_of?(C::Array) and mb.type.type.kind_of?(C::BaseType) and mb.type.type.name == :char
						defval = st[mb].to_array.pack('C*').gsub(/\0*$/, '').gsub(/[^\x20-\x7e]/, '.')
						string = true
					else
						defval = st[mb]
						string = false
					end
					inputbox("new value for #{mb.respond_to?(:name) ? mb.name : mb}", :text => defval.to_s) { |nn|
						if string
							am = st[mb]
							(nn.unpack('C*') + [0]).each_with_index { |b_, i| am[i] = b_ }
						else
							st[mb] = Expression.parse_string(nn).reduce
						end
						gui_update
					}
				end
			end
		when ?x
			# show heap xrefs to the hilighted chunk
			if b = @caret_box
				list = [['address', 'size']]
				@heap.xrchunksfrom[b.id].to_a.each { |a|
					list << [Expression[a], Expression[@heap.chunks[a]]]
				}
				if list.length == 1
					messagebox "no xref to #{Expression[b.id]}"
				else
					listwindow("heap xrefs to #{Expression[b.id]}", list) { |i| @parent_widget.focus_addr(i[0], nil, true) }
				end
			end
		when ?I
			# insert new field in struct
			if b = @caret_box
				if m = b[:line_member][@caret_y]
					if m.kind_of?(Integer)
						# XXX Array, need to find the outer struct
						mn = b[:line_text_col][@caret_y].map { |l, c| l }.join[/(\S*)\[/, 1]
						ar = b[:line_struct][@caret_y]
						st = b[:line_struct][0...@caret_y].reverse.compact.find { |st_| st_.struct.kind_of?(C::Struct) and st_[mn].struct == ar.struct }
						raise '?' if not st
						st = st.struct
						m = st.fldlist[mn]
					else
						st = b[:line_struct][@caret_y].struct
					end
					inputbox("new type to insert before #{m.name}", :text => m.dump_def(@heap.cp.toplevel)[0].join(' ')) { |nn|
						nil while @heap.cp.readtok
						@heap.cp.lexer.feed nn
						if not v = C::Variable.parse_type(@heap.cp, @heap.cp.toplevel, true)
							nil while @heap.cp.readtok
							raise 'bad type'
						end
						v.parse_declarator(@heap.cp, @heap.cp.toplevel)
						nt = v.type
						idx = st.members.index(m)
						pos = st.offsetof(@heap.cp, m)
						name = oname = v.name || ('unk_%x_new' % pos)
						cntr = 0
						while st.members.find { |m_| m_.name == name }
							name = oname + "_#{cntr+=1}"
						end
						st.members[idx, 0] = [C::Variable.new(name, nt)]
						st.update_member_cache(@heap.cp)
						gui_update
					}

				end
			end
		when ?S
			# delete structure field
			if b = @caret_box
				if m = b[:line_member][@caret_y]
					if m.kind_of?(Integer)
						# XXX Array, need to find the outer struct
						mn = b[:line_text_col][@caret_y].map { |l, c| l }.join[/(\S*)\[/, 1]
						ar = b[:line_struct][@caret_y]
						st = b[:line_struct][0...@caret_y].reverse.compact.find { |st_| st_.struct.kind_of?(C::Struct) and st_[mn].struct == ar.struct }
						raise '?' if not st
						st = st.struct
						m = st.fldlist[mn]
					else
						st = b[:line_struct][@caret_y].struct
					end
					inputbox("delete #{m.name} ?") { |nn|
						idx = st.members.index(m)
						st.members.delete_at(idx)
						st.update_member_cache(@heap.cp)
						gui_update
					}

				end
			end
		when ?+
			# append blocks linked from the currently shown blocks to the display
			@addr_struct.keys.each { |ak|
				@heap.xrchunksto[ak].to_a.each { |nt|
					next if @addr_struct[nt]
					# TODO check if the pointer is a some_struct*
					st = @heap.chunk_struct[nt] || create_struct(nt)
					ed = @dasm.get_edata_at(nt)
					@addr_struct[nt] = @heap.cp.decode_c_struct(st.name, ed.data, ed.ptr)
				}
			}
			gui_update
		when ?-
			# remove graph leaves in an attempt to undo ?+
			unk = @addr_struct.keys.find_all { |ak|
				(@heap.xrchunksto[ak].to_a & @addr_struct.keys).empty?
			}
			unk.each { |ak| @addr_struct.delete ak if @addr_struct.length > 1 }
			gui_update
		else return super(k)
		end
		true
	end

	# create the graph objects in ctx
	def build_ctx(ctx)
		# create boxes
		todo = ctx.root_addrs.dup & @addr_struct.keys
		todo << @addr_struct.keys.first if todo.empty?
		done = []
		while a = todo.shift
			next if done.include? a
			done << a
			ctx.new_box a, :line_text_col => [], :line_address => [], :line_struct => [], :line_member => []
			todo.concat @heap.xrchunksto[a].to_a & @addr_struct.keys
		end

		# link boxes
		if (@heap.xrchunksto[ctx.box.first.id].to_a & @addr_struct.keys).length == ctx.box.length - 1
			ot = ctx.box[0].id
			ctx.box[1..-1].each { |b_|
				ctx.link_boxes(ot, b_.id)
			}
		else
			ctx.box.each { |b|
				@heap.xrchunksto[b.id].to_a.each { |t|
					ctx.link_boxes(b.id, t) if @addr_struct[t]
				}
			}
		end

		if snapped
			@datadiff = {}
		end

		# calc box dimensions/text
		ctx.box.each { |b|
			colstr = []
			curaddr = b.id
			curst = @addr_struct[b.id]
			curmb = nil
			margin = ''
			start_addr = curaddr
			if snapped
				ghosts = snapped[curaddr]
			end
			line = 0
			render = lambda { |str, col| colstr << [str, col] }
			nl = lambda {
				b[:line_address][line] = curaddr
				b[:line_text_col][line] = colstr
				b[:line_struct][line] = curst
				b[:line_member][line] = curmb
				colstr = []
				line += 1
			}
			render_val = lambda { |v|
				if v.kind_of?(::Integer)
					if v > 0x100
						render['0x%X' % v, :text]
					elsif v < -0x100
						render['-0x%X' % -v, :text]
					else
						render[v.to_s, :text]
					end
				elsif not v
					render['NULL', :text]
				else
					render[v.to_s, :text]
				end
			}
			render_st = nil
			render_st_ar = lambda { |ast, m|
				elemt = m.type.untypedef.type.untypedef
				if elemt.kind_of?(C::BaseType) and elemt.name == :char
					render[margin, :text]
					render["#{m.type.type.to_s[1...-1]} #{m.name}[#{m.type.length}] = #{ast[m].to_array.pack('C*').sub(/\0.*$/m, '').inspect}", :text]
					nl[]
					curaddr += ast.cp.sizeof(m)
				else
					t = m.type.type.to_s[1...-1]
					tsz = ast.cp.sizeof(m.type.type)
					fust = curst
					fumb = curmb
					curst = ast[m]
					ast[m].to_array.each_with_index { |v, i|
						curmb = i
						render[margin, :text]
						if elemt.kind_of?(C::Union)
							if m.type.untypedef.type.kind_of?(C::Union)
								render[elemt.kind_of?(C::Struct) ? 'struct ' : 'union ', :text]
								render["#{elemt.name} ", :text] if elemt.name
							else # typedef
								render["#{elemt.to_s[1...-1]} ", :text]
							end
							render_st[v]
							render[" #{m.name}[#{i}]", :text]
						else
							render["#{t} #{m.name}[#{i}] = ", :text]
							render_val[v]
							@datadiff[curaddr] = true if ghosts and ghosts.all? { |g| g[curaddr-start_addr, tsz] == ghosts[0][curaddr-start_addr, tsz] } and ghosts[0][curaddr-start_addr, tsz] != ast.str[curaddr, tsz].to_str
						end
						render[';', :text]
						nl[]
						curaddr += tsz
					}
					curst = fust
					curmb = fumb
				end
			}
			render_st = lambda { |ast|
				st_addr = curaddr
				oldst = curst
				oldmb = curmb
				oldmargin = margin
				render['{', :text]
				nl[]
				margin += '    '
				curst = ast
				ast.struct.members.each { |m|
					curmb = m
					curaddr = st_addr + ast.struct.offsetof(@heap.cp, m)

					if bo = ast.struct.bitoffsetof(@heap.cp, m)
						# float curaddr to make ghost hilight work on bitfields
						curaddr += (1+bo[0])/1000.0
					end

					if m.type.untypedef.kind_of?(C::Array)
						render_st_ar[ast, m]
					elsif m.type.untypedef.kind_of?(C::Union)
						render[margin, :text]
						if m.type.kind_of?(C::Union)
							render[m.type.kind_of?(C::Struct) ? 'struct ' : 'union ', :text]
							render["#{m.type.name} ", :text] if m.type.name
						else # typedef
							render["#{m.type.to_s[1...-1]} ", :text]
						end
						oca = curaddr
						render_st[ast[m]]
						nca = curaddr
						curaddr = oca
						render[" #{m.name if m.name};", :text]
						nl[]
						curaddr = nca
					else
						render[margin, :text]
						render["#{m.type.to_s[1...-1]} ", :text]
						render["#{m.name} = ", :text]
						render_val[ast[m]]
						tsz = ast.cp.sizeof(m)
						# TODO bit-level multighosting
						if ghosts and ghosts.all? { |g| g[curaddr.to_i-start_addr, tsz] == ghosts[0][curaddr.to_i-start_addr, tsz] } and ghosts[0][curaddr.to_i-start_addr, tsz] != ast.str[curaddr.to_i, tsz].to_str
							if bo
								ft = C::BaseType.new((bo[0] + bo[1] > 32) ? :__int64 : :__int32)
								v1 = @heap.cp.decode_c_value(ghosts[0][curaddr.to_i-start_addr, tsz], ft, 0)
								v2 = @heap.cp.decode_c_value(ast.str[curaddr.to_i, tsz], ft, 0)
								@datadiff[curaddr] = true if (v1 >> bo[0]) & ((1 << bo[1])-1) != (v2 >> bo[0]) & ((1 << bo[1])-1)
							else
								@datadiff[curaddr] = true
							end
						end
						render[';', :text]

						if m.type.kind_of?(C::Pointer) and m.type.type.kind_of?(C::BaseType) and m.type.type.name == :char
							if s = @dasm.decode_strz(ast[m], 32)
								render["    // #{s.inspect}", :comment]
							end
						end
						nl[]
						curaddr += tsz
						curaddr = curaddr.to_i if bo
					end
				}
				margin = oldmargin
				curst = oldst
				curmb = oldmb
				render[margin, :text]
				render['}', :text]
			}
			ast = @addr_struct[curaddr]
			render["struct #{ast.struct.name} *#{'0x%X' % curaddr} = ", :text]
			render_st[ast]
			render[';', :text]
			nl[]

			b.w = b[:line_text_col].map { |strc| strc.map { |s, c| s }.join.length }.max.to_i * @font_width + 2
			b.w += 1 if b.w % 2 == 0
			b.h = line * @font_height
		}
	end

	def struct_find_roots(addr)
		addr = @addr_struct.keys.find { |a| addr >= a and addr < a+@addr_struct[a].sizeof } if not @addr_struct[addr]

		todo = [addr]
		done = []
		roots = []
		default_root = nil
		while a = todo.shift
			if done.include?(a) # cycle
				default_root ||= a
				next
			end
			done << a
			newf = @heap.xrchunksfrom[a].to_a & @addr_struct.keys
			if newf.empty?
				roots << a
			else
				todo.concat newf
			end
		end
		roots << default_root if roots.empty? and default_root

		roots
	end

	def focus_addr(addr, fu=nil)
		return if @parent_widget and not addr = @parent_widget.normalize(addr)

		# move window / change curcontext
		if b = @curcontext.box.find { |b_| b_[:line_address].index(addr) }
			@caret_box, @caret_x, @caret_y = b, 0, b[:line_address].rindex(addr)
			@curcontext.view_x += (width/2 / @zoom - width/2)
			@curcontext.view_y += (height/2 / @zoom - height/2)
			@zoom = 1.0

			focus_xy(b.x, b.y + @caret_y*@font_height)
			update_caret
		elsif addr_struct and (@addr_struct[addr] or @addr_struct.find { |a, s| addr >= a and addr < a+s.sizeof })
			@curcontext = Graph.new 'testic'
			@curcontext.root_addrs = struct_find_roots(addr)
			@want_focus_addr = addr
			gui_update
		elsif @heap.chunks[addr]
			@want_focus_addr = addr
			do_focus_addr(addr)
		else
			return
		end
		true
	end

	def do_focus_addr(addr)
		st = @heap.chunk_struct[addr] || create_struct(addr)

		ed = @dasm.get_edata_at(addr)
		@addr_struct = { addr => @heap.cp.decode_c_struct(st.name, ed.data, ed.ptr) }
		gui_update
	end

	# create the struct chunk_<addr>, register it in @heap.chunk_struct
	def create_struct(addr)
		raise "no chunk here" if not @heap.chunks[addr]

		ptsz = @dasm.cpu.size/8

		# check if this is a c++ object with RTTI info
		vptr = @dasm.decode_dword(addr)
		rtti = @dasm.decode_dword(vptr-ptsz)
		case OS.shortname
		when 'winos'
			typeinfo = @dasm.decode_dword(rtti+3*ptsz) if rtti
			if typeinfo and s = @dasm.decode_strz(typeinfo+3*ptsz)
				rtti_name = s[/^(.*)@@$/, 1]	# remove trailing @@
			end
		when 'linos'
			typeinfo = @dasm.decode_dword(rtti+ptsz) if rtti
			if typeinfo and s = @dasm.decode_strz(typeinfo)
				rtti_name = s[/^[0-9]+(.*)$/, 1]	# remove leading number
			end
		end

		if rtti_name and st = @heap.cp.toplevel.struct[rtti_name]
			return @heap.chunk_struct[addr] = st
		end

		st = C::Struct.new
		st.name = rtti_name || "chunk_#{'%x' % addr}"
		st.members = []
		li = 0
		(@heap.chunks[addr] / ptsz).times { |i|
			n = 'unk_%x' % (ptsz*i)
			v = @dasm.decode_dword(addr+ptsz*i)
			if i == 0 and rtti_name
				t = C::Pointer.new(C::Pointer.new(C::BaseType.new(:void)))
				n = 'vtable'
			elsif @heap.chunks[v]
				t = C::Pointer.new(C::BaseType.new(:void))
			else
				t = C::BaseType.new("__int#{ptsz*8}".to_sym, :unsigned)
			end
			st.members << C::Variable.new(n, t)
			li = i+1
		}
		(@heap.chunks[addr] % ptsz).times { |i|
			n = 'unk_%x' % (ptsz*li+i)
			t = C::BaseType.new(:char, :unsigned)
			st.members << C::Variable.new(n, t)
		}
		@heap.cp.toplevel.struct[st.name] = st
		@heap.chunk_struct[addr] = st
	end

	def snap
		if not snapped
			@datadiff = {}
			ocb = @parent_widget.bg_color_callback
			@parent_widget.bg_color_callback = lambda { |a|
				if @datadiff[a]
					'f88'
				elsif ocb
					ocb[a]
				end
			}
		end
		@snapped = {}
		@addr_struct.each { |a, ast|
			@snapped[a] = [ast.str[ast.stroff, ast.sizeof].to_str]
		}
	end

	def snap_add
		return snap if not snapped
		@addr_struct.each { |a, ast|
			(@snapped[a] ||= []) << ast.str[ast.stroff, ast.sizeof].to_str
		}
	end

	def get_cursor_pos
		[super, addr_struct]
	end

	def set_cursor_pos(p)
		s, @addr_struct = p
		super(s)
	end
end
end
end
