DREAMWEB: Try to clean up jumps to returns

2 files changed, 35 insertions, 9 deletions

diff --git a/devtools/tasmrecover/tasm/cpp.py b/devtools/tasmrecover/tasm/cpp.py
index 3e5ae78438..5d460961a6 100644
--- a/devtools/tasmrecover/tasm/cpp.py
+++ b/devtools/tasmrecover/tasm/cpp.py
@@ -241,6 +241,9 @@ namespace %s {
 			self.add_forward(name)
 			return "{ %s(context); return; }" %name
 		else:
+			# TODO: name or self.resolve_label(name) or self.mangle_label(name)??
+			if name in self.proc.retlabels:
+				return "return /* (%s) */" % (name,)
 			return "goto %s" %self.resolve_label(name)
 	
 	def _label(self, name):
@@ -443,6 +446,7 @@ namespace %s {
 				self.proc = proc_module.proc(name)
 				self.proc.stmts = copy(src_proc.stmts)
 				self.proc.labels = copy(src_proc.labels)
+				self.proc.retlabels = copy(src_proc.retlabels)
 				#for p in xrange(skip, len(self.proc.stmts)):
 				#	s = self.proc.stmts[p]
 				#	if isinstance(s, op.basejmp):
diff --git a/devtools/tasmrecover/tasm/proc.py b/devtools/tasmrecover/tasm/proc.py
index 1350ea1e0b..e5af1afae5 100644
--- a/devtools/tasmrecover/tasm/proc.py
+++ b/devtools/tasmrecover/tasm/proc.py
@@ -9,6 +9,7 @@ class proc:
 		self.calls = []
 		self.stmts = []
 		self.labels = set()
+		self.retlabels = set()
 		self.__label_re = re.compile(r'^(\S+):(.*)$')
 		self.offset = proc.last_addr
 		proc.last_addr += 4
@@ -67,13 +68,20 @@ class proc:
 	
 	def optimize(self):
 		print "optimizing..."
-		#trivial simplifications, removing last ret
+		#trivial simplifications
 		while len(self.stmts) and isinstance(self.stmts[-1], op.label):
 			print "stripping last label"
 			self.stmts.pop()
-		if isinstance(self.stmts[-1], op._ret) and (len(self.stmts) < 2 or not isinstance(self.stmts[-2], op.label)):
-			print "stripping last ret"
-			self.stmts.pop()
+		#mark labels that directly precede a ret
+		for i in range(len(self.stmts)):
+			if not isinstance(self.stmts[i], op.label):
+				continue
+			j = i
+			while j < len(self.stmts) and isinstance(self.stmts[j], op.label):
+				j += 1
+			if j == len(self.stmts) or isinstance(self.stmts[j], op._ret):
+				print "Return label: %s" % (self.stmts[i].name,)
+				self.retlabels.add(self.stmts[i].name)
 		#merging push ax pop bx constructs
 		i = 0
 		while i + 1 < len(self.stmts):
@@ -102,15 +110,29 @@ class proc:
 				continue
 			print "checking label %s..." %s.name
 			used = False
-			for j in self.stmts:
-				if isinstance(j, op.basejmp) and j.label == s.name:
-					print "used"
-					used = True
-					break
+			if s.name not in self.retlabels:
+				for j in self.stmts:
+					if isinstance(j, op.basejmp) and j.label == s.name:
+						print "used"
+						used = True
+						break
 			if not used:
 				print self.labels
 				self.remove_label(s.name)
 
+		#removing duplicate rets
+		i = 0
+		while i < len(self.stmts)-1:
+			if isinstance(self.stmts[i], op._ret) and isinstance(self.stmts[i+1], op._ret):
+				del self.stmts[i]
+			else:
+				i += 1
+
+		#removing last ret
+		while len(self.stmts) > 0 and isinstance(self.stmts[-1], op._ret) and (len(self.stmts) < 2 or not isinstance(self.stmts[-2], op.label)):
+			print "stripping last ret"
+			self.stmts.pop()
+
 		self.optimize_sequence(op._stosb);
 		self.optimize_sequence(op._stosw);
 		self.optimize_sequence(op._movsb);