Fix HLIL_LABEL rendering

7 files changed, 100 insertions, 6 deletions

diff --git a/architecture.cpp b/architecture.cpp
index f5fbf892..22bf6fc3 100644
--- a/architecture.cpp
+++ b/architecture.cpp
@@ -119,6 +119,26 @@ InstructionTextToken InstructionTextToken::WithConfidence(uint8_t conf)
 }
 
 
+BNInstructionTextToken InstructionTextToken::GetAPIObject() const
+{
+	BNInstructionTextToken result;
+	ConvertInstructionTextToken(*this, &result);
+	return result;
+}
+
+
+InstructionTextToken InstructionTextToken::FromAPIObject(const BNInstructionTextToken* token)
+{
+	return InstructionTextToken(*token);
+}
+
+
+void InstructionTextToken::FreeAPIObject(BNInstructionTextToken* token)
+{
+	FreeInstructionTextToken(token);
+}
+
+
 void InstructionTextToken::ConvertInstructionTextToken(const InstructionTextToken& token, BNInstructionTextToken* result)
 {
 	result->type = token.type;
diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 0d2bdfe2..089b3693 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -2460,6 +2460,9 @@ namespace BinaryNinja {
 		InstructionTextToken(const BNInstructionTextToken& token);
 
 		InstructionTextToken WithConfidence(uint8_t conf);
+		BNInstructionTextToken GetAPIObject() const;
+		static InstructionTextToken FromAPIObject(const BNInstructionTextToken* token);
+		static void FreeAPIObject(BNInstructionTextToken* token);
 		static void ConvertInstructionTextToken(const InstructionTextToken& token, BNInstructionTextToken* result);
 		static BNInstructionTextToken* CreateInstructionTextTokenList(const std::vector<InstructionTextToken>& tokens);
 		static void FreeInstructionTextToken(BNInstructionTextToken* token);
@@ -19502,6 +19505,9 @@ namespace BinaryNinja {
 		/*! Returns the list of tokens on the current line */
 		std::vector<InstructionTextToken> GetCurrentTokens() const;
 
+		/*! Set the list of tokens on the current line */
+		void SetCurrentTokens(const std::vector<InstructionTextToken>& newTokens);
+
 		/*! Sets the requirement for insertion of braces around scopes in the output. */
 		void SetBraceRequirement(BNBraceRequirement required);
 
diff --git a/binaryninjacore.h b/binaryninjacore.h
index e5dfba4f..40959041 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -8165,6 +8165,7 @@ extern "C"
 	BINARYNINJACOREAPI bool BNHighLevelILTokenEmitterGetDefaultBracesOnSameLine(BNHighLevelILTokenEmitter* emitter);
 	BINARYNINJACOREAPI bool BNHighLevelILTokenEmitterIsSimpleScopeAllowed(BNHighLevelILTokenEmitter* emitter);
 	BINARYNINJACOREAPI BNInstructionTextToken* BNHighLevelILTokenEmitterGetCurrentTokens(BNHighLevelILTokenEmitter* emitter, size_t* tokenCount);
+	BINARYNINJACOREAPI void BNHighLevelILTokenEmitterSetCurrentTokens(BNHighLevelILTokenEmitter* emitter, BNInstructionTextToken* tokens, size_t tokenCount);
 	BINARYNINJACOREAPI BNDisassemblyTextLine* BNHighLevelILTokenEmitterGetLines(BNHighLevelILTokenEmitter* emitter, size_t* count);
 
 	BINARYNINJACOREAPI void BNAddHighLevelILSizeToken(size_t size, BNInstructionTextTokenType type, BNHighLevelILTokenEmitter* tokens);
diff --git a/highlevelil.cpp b/highlevelil.cpp
index 96ed664a..91ce5cfc 100644
--- a/highlevelil.cpp
+++ b/highlevelil.cpp
@@ -821,6 +821,15 @@ vector<InstructionTextToken> HighLevelILTokenEmitter::GetCurrentTokens() const
 }
 
 
+void HighLevelILTokenEmitter::SetCurrentTokens(const std::vector<InstructionTextToken>& newTokens)
+{
+	size_t count;
+	auto* tokens = AllocAPIObjectList<InstructionTextToken>(newTokens, &count);
+	BNHighLevelILTokenEmitterSetCurrentTokens(m_object, tokens, count);
+	FreeAPIObjectList<InstructionTextToken>(tokens, count);
+}
+
+
 void HighLevelILTokenEmitter::SetBraceRequirement(BNBraceRequirement required)
 {
 	BNHighLevelILTokenEmitterSetBraceRequirement(m_object, required);
diff --git a/lang/c/pseudoc.cpp b/lang/c/pseudoc.cpp
index 296c2d6a..7046c6d0 100644
--- a/lang/c/pseudoc.cpp
+++ b/lang/c/pseudoc.cpp
@@ -33,8 +33,8 @@ void PseudoCFunction::BeginLines(const HighLevelILInstruction& instr, HighLevelI
 		// At top level, add braces around the entire function
 		tokens.PrependCollapseIndicator();
 		tokens.AppendOpenBrace();
-		tokens.NewLine();
 		tokens.IncreaseIndent();
+		tokens.NewLine();
 	}
 }
 
@@ -2755,10 +2755,39 @@ void PseudoCFunction::GetExprTextInternal(const HighLevelILInstruction& instr, H
 	case HLIL_LABEL:
 		[&]() {
 			const auto target = instr.GetTarget<HLIL_LABEL>();
-			tokens.DecreaseIndent();
+			// NB: Not using DecreaseIndent() here because it will mess up the indentation guides
+			// that rely on matched calls to IncreaseIndent()/DecreaseIndent() to properly set
+			// indentation groupings. Instead, we manually remove the last indent in the tokens
+			tokens.InitLine();
+			auto newTokens = tokens.GetCurrentTokens();
+			bool foundIndent = false;
+			bool erasedIndent = false;
+			for (size_t i = 0; i < newTokens.size(); i++)
+			{
+				auto& token = newTokens[i];
+				if (token.type == IndentationToken)
+				{
+					foundIndent = true;
+				}
+				else if (foundIndent)
+				{
+					// 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.
+					newTokens.erase(newTokens.begin() + (i - 1));
+					erasedIndent = true;
+					break;
+				}
+			}
+			if (foundIndent && !erasedIndent)
+			{
+				// Found an indent, but it was the last token, so erase the last token
+				newTokens.erase(newTokens.begin() + (newTokens.size() - 1));
+			}
+			tokens.SetCurrentTokens(newTokens);
+
 			tokens.Append(GotoLabelToken, GetFunction()->GetGotoLabelName(target), target);
 			tokens.Append(TextToken, ":");
-			tokens.IncreaseIndent();
 		}();
 		break;
 
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)."""