Fix HLIL_LABEL rendering

2 files changed, 32 insertions, 3 deletions

diff --git a/python/examples/pseudo_python.py b/python/examples/pseudo_python.py
index 8c1a237e..c2019c0c 100644
--- a/python/examples/pseudo_python.py
+++ b/python/examples/pseudo_python.py
@@ -991,11 +991,31 @@ class PseudoPythonFunction(LanguageRepresentationFunction):
                 tokens.append(InstructionTextToken(InstructionTextTokenType.GotoLabelToken, instr.target.name,
                                                    instr.target.label_id))
             elif instr.operation == HighLevelILOperation.HLIL_LABEL:
-                tokens.decrease_indent()
+                # NB: Not using decrease_indent() here because it will mess up the indentation guides
+                # that rely on matched calls to increase_indent()/decrease_indent() to properly set
+                # indentation groupings. Instead, we manually remove the last indent in the tokens
+                tokens.init_line()
+                new_tokens = tokens.current_tokens
+                found_indent = False
+                erased_indent = False
+                for i in range(len(new_tokens)):
+                    token = new_tokens[i]
+                    if token.type == InstructionTextTokenType.IndentationToken:
+                        found_indent = True
+                    elif found_indent:
+                        # We have indents and this is the first non-indent token after all the indents,
+                        # so remove the previous token. It must exist because we will have to have gone through
+                        # the condition above and looped at least once.
+                        new_tokens.pop(i - 1)
+                        erased_indent = True
+                        break
+                if found_indent and not erased_indent:
+                    # Found an indent, but it was the last token, so erase the last token
+                    new_tokens.pop(len(new_tokens) - 1)
+                tokens.current_tokens = new_tokens
                 tokens.append(InstructionTextToken(InstructionTextTokenType.GotoLabelToken, instr.target.name,
                                                    instr.target.label_id))
                 tokens.append(InstructionTextToken(InstructionTextTokenType.TextToken, ":"))
-                tokens.increase_indent()
             elif instr.operation == HighLevelILOperation.HLIL_LOW_PART:
                 parens = precedence > OperatorPrecedence.MemberAndFunctionOperatorPrecedence
                 if parens:
diff --git a/python/languagerepresentation.py b/python/languagerepresentation.py
index 44c34bf2..a8241b30 100644
--- a/python/languagerepresentation.py
+++ b/python/languagerepresentation.py
@@ -53,6 +53,10 @@ class HighLevelILTokenEmitter:
 		if core is not None:
 			core.BNFreeHighLevelILTokenEmitter(self.handle)
 
+	def init_line(self):
+		"""Initialize a new line, creating indentation tokens at the start."""
+		core.BNHighLevelILTokenEmitterInitLine(self.handle)
+
 	def new_line(self):
 		"""Starts a new line in the output."""
 		core.BNHighLevelILTokenEmitterNewLine(self.handle)
@@ -240,7 +244,7 @@ class HighLevelILTokenEmitter:
 
 	@property
 	def current_tokens(self) -> List['function.InstructionTextToken']:
-		"""The list of tokens on the current line (read-only)."""
+		"""The list of tokens on the current line."""
 		count = ctypes.c_ulonglong()
 		tokens = core.BNHighLevelILTokenEmitterGetCurrentTokens(self.handle, count)
 		result = []
@@ -249,6 +253,11 @@ class HighLevelILTokenEmitter:
 			core.BNFreeInstructionText(tokens, count.value)
 		return result
 
+	@current_tokens.setter
+	def current_tokens(self, tokens: List['function.InstructionTextToken']):
+		buf = function.InstructionTextToken._get_core_struct(tokens)
+		core.BNHighLevelILTokenEmitterSetCurrentTokens(self.handle, buf, len(tokens))
+
 	@property
 	def lines(self) -> List['function.DisassemblyTextLine']:
 		"""The list of lines in the output (read-only)."""