Add line formatter API and a generic line formatter plugin

6 files changed, 372 insertions, 32 deletions

diff --git a/python/__init__.py b/python/__init__.py
index 9b5a394e..1cdc541e 100644
--- a/python/__init__.py
+++ b/python/__init__.py
@@ -80,6 +80,7 @@ from .externallibrary import *
 from .undo import *
 from .fileaccessor import *
 from .languagerepresentation import *
+from .lineformatter import *
 # We import each of these by name to prevent conflicts between
 # log.py and the function 'log' which we don't import below
 from .log import (
diff --git a/python/examples/pseudo_python.py b/python/examples/pseudo_python.py
index 6a9e401b..655039fe 100644
--- a/python/examples/pseudo_python.py
+++ b/python/examples/pseudo_python.py
@@ -316,16 +316,7 @@ class PseudoPythonFunction(LanguageRepresentationFunction):
                     tokens.append(InstructionTextToken(InstructionTextTokenType.TextToken, ": "))
                     tokens.append(instr.var.type.get_tokens())
             elif instr.operation == HighLevelILOperation.HLIL_FLOAT_CONST:
-                # The constant value in the instruction contains the raw bits of the floating point value. Convert
-                # this to a floating point value and display it.
-                if instr.size == 4:
-                    value = struct.unpack("<f", struct.pack("<I", instr.constant))[0]
-                    tokens.append(InstructionTextToken(InstructionTextTokenType.FloatToken, f"{value:g}"))
-                elif instr.size == 8:
-                    value = struct.unpack("<d", struct.pack("<Q", instr.constant))[0]
-                    tokens.append(InstructionTextToken(InstructionTextTokenType.FloatToken, f"{value:g}"))
-                else:
-                    tokens.append_integer_text_token(instr, instr.constant, instr.size)
+                tokens.append(InstructionTextToken(InstructionTextTokenType.FloatingPointToken, f"{instr.constant:g}"))
             elif instr.operation == HighLevelILOperation.HLIL_CONST:
                 # Check for bool type. Display these as True or False. The default handling will use C style
                 # booleans instead of Python style.
@@ -1075,7 +1066,7 @@ class PseudoPythonFunctionType(LanguageRepresentationFunctionType):
     language_name = "Pseudo Python"
 
     def create(self, arch: Architecture, owner: Function, hlil: HighLevelILFunction):
-        return PseudoPythonFunction(arch, owner, hlil)
+        return PseudoPythonFunction(self, arch, owner, hlil)
 
     def function_type_tokens(self, func: Function, settings: DisassemblySettings) -> DisassemblyTextLine:
         tokens = []
diff --git a/python/function.py b/python/function.py
index e962d91f..0fc60ffc 100644
--- a/python/function.py
+++ b/python/function.py
@@ -135,6 +135,18 @@ class DisassemblySettings:
 			option = DisassemblyOption[option]
 		core.BNSetDisassemblySettingsOption(self.handle, option, state)
 
+	@staticmethod
+	def default_settings() -> 'DisassemblySettings':
+		return DisassemblySettings(core.BNDefaultDisassemblySettings())
+
+	@staticmethod
+	def default_graph_settings() -> 'DisassemblySettings':
+		return DisassemblySettings(core.BNDefaultGraphDisassemblySettings())
+
+	@staticmethod
+	def default_linear_settings() -> 'DisassemblySettings':
+		return DisassemblySettings(core.BNDefaultLinearDisassemblySettings())
+
 
 @dataclass
 class ILReferenceSource:
diff --git a/python/highlevelil.py b/python/highlevelil.py
index 1d7feb71..85be4794 100644
--- a/python/highlevelil.py
+++ b/python/highlevelil.py
@@ -26,7 +26,10 @@ from enum import Enum
 
 # Binary Ninja components
 from . import _binaryninjacore as core
-from .enums import HighLevelILOperation, DataFlowQueryOption, FunctionGraphType, ILInstructionAttribute, StringType
+from .enums import (
+	HighLevelILOperation, DataFlowQueryOption, FunctionGraphType, ILInstructionAttribute, StringType,
+	DisassemblyOption
+)
 from . import function
 from . import binaryview
 from . import architecture
@@ -331,7 +334,9 @@ class HighLevelILInstruction(BaseILInstruction):
 		return ILInstruction[instr.operation](func, expr_index, core_instr, as_ast, instr_index)
 
 	def __str__(self):
-		lines = self.lines
+		settings = function.DisassemblySettings.default_settings()
+		settings.set_option(DisassemblyOption.DisableLineFormatting)
+		lines = self.get_lines(settings)
 		if lines is None:
 			return "invalid"
 		result = []
@@ -343,7 +348,9 @@ class HighLevelILInstruction(BaseILInstruction):
 		return '\n'.join(result)
 
 	def __repr__(self):
-		lines = self.lines
+		settings = function.DisassemblySettings.default_settings()
+		settings.set_option(DisassemblyOption.DisableLineFormatting)
+		lines = self.get_lines(settings)
 		continuation = ""
 		if lines is None:
 			first_line = "<invalid>"
@@ -386,26 +393,14 @@ class HighLevelILInstruction(BaseILInstruction):
 	@property
 	def tokens(self) -> TokenList:
 		"""HLIL tokens taken from the HLIL text lines (read-only) -- does not include newlines or indentation, use lines for that information"""
-		return [token for line in self.lines for token in line.tokens]
+		settings = function.DisassemblySettings.default_settings()
+		settings.set_option(DisassemblyOption.DisableLineFormatting)
+		return [token for line in self.get_lines(settings) for token in line.tokens]
 
 	@property
 	def lines(self) -> LinesType:
 		"""HLIL text lines (read-only)"""
-		count = ctypes.c_ulonglong()
-		lines = core.BNGetHighLevelILExprText(self.function.handle, self.expr_index, self.as_ast, count, None)
-		assert lines is not None, "core.BNGetHighLevelILExprText returned None"
-		try:
-			for i in range(0, count.value):
-				addr = lines[i].addr
-				if lines[i].instrIndex != 0xffffffffffffffff:
-					il_instr = self.function[lines[i].instrIndex]
-				else:
-					il_instr = None
-				color = highlight.HighlightColor._from_core_struct(lines[i].highlight)
-				tokens = function.InstructionTextToken._from_core_struct(lines[i].tokens, lines[i].count)
-				yield function.DisassemblyTextLine(tokens, addr, il_instr, color)
-		finally:
-			core.BNFreeDisassemblyTextLines(lines, count.value)
+		return self.get_lines()
 
 	@property
 	def prefix_operands(self) -> List[Union[HighLevelILOperandType, HighLevelILOperationAndSize]]:
@@ -918,6 +913,26 @@ class HighLevelILInstruction(BaseILInstruction):
 	def has_side_effects(self) -> bool:
 		return core.BNHighLevelILHasSideEffects(self.function.handle, self.expr_index)
 
+	def get_lines(self, settings: Optional['function.DisassemblySettings'] = None) -> LinesType:
+		"""Gets HLIL text lines with optional settings"""
+		if settings is not None:
+			settings = settings.handle
+		count = ctypes.c_ulonglong()
+		lines = core.BNGetHighLevelILExprText(self.function.handle, self.expr_index, self.as_ast, count, settings)
+		assert lines is not None, "core.BNGetHighLevelILExprText returned None"
+		try:
+			for i in range(0, count.value):
+				addr = lines[i].addr
+				if lines[i].instrIndex != 0xffffffffffffffff:
+					il_instr = self.function[lines[i].instrIndex]
+				else:
+					il_instr = None
+				color = highlight.HighlightColor._from_core_struct(lines[i].highlight)
+				tokens = function.InstructionTextToken._from_core_struct(lines[i].tokens, lines[i].count)
+				yield function.DisassemblyTextLine(tokens, addr, il_instr, color)
+		finally:
+			core.BNFreeDisassemblyTextLines(lines, count.value)
+
 
 @dataclass(frozen=True, repr=False, eq=False)
 class HighLevelILUnaryBase(HighLevelILInstruction, UnaryOperation):
diff --git a/python/languagerepresentation.py b/python/languagerepresentation.py
index ac34c860..a0a5a234 100644
--- a/python/languagerepresentation.py
+++ b/python/languagerepresentation.py
@@ -30,6 +30,7 @@ from . import binaryview
 from . import function
 from . import highlevelil
 from . import highlight
+from . import lineformatter
 from . import variable
 from . import types
 from .log import log_error
@@ -327,7 +328,8 @@ class LanguageRepresentationFunction:
 	annotation_end_string = "}"
 
 	def __init__(
-			self, arch: Optional['architecture.Architecture'] = None, owner: Optional['function.Function'] = None,
+			self, func_type: Optional['LanguageRepresentationFunctionType'] = None,
+			arch: Optional['architecture.Architecture'] = None, owner: Optional['function.Function'] = None,
 			hlil: Optional['highlevelil.HighLevelILFunction'] = None, handle=None
 	):
 		if handle is None:
@@ -354,7 +356,9 @@ class LanguageRepresentationFunction:
 			self.comment_end_string = self.__class__.comment_end_string
 			self.annotation_start_string = self.__class__.annotation_start_string
 			self.annotation_end_string = self.__class__.annotation_end_string
-			_handle = core.BNCreateCustomLanguageRepresentationFunction(arch.handle, owner.handle, hlil.handle, self._cb)
+			_handle = core.BNCreateCustomLanguageRepresentationFunction(
+				func_type.handle, arch.handle, owner.handle, hlil.handle, self._cb
+			)
 			assert _handle is not None
 		else:
 			self.comment_start_string = core.BNGetLanguageRepresentationFunctionCommentStartString(handle)
@@ -621,6 +625,7 @@ class LanguageRepresentationFunctionType(metaclass=_LanguageRepresentationFuncti
 		self._cb.isValid = self._cb.isValid.__class__(self._is_valid)
 		self._cb.getTypePrinter = self._cb.getTypePrinter.__class__(self._type_printer)
 		self._cb.getTypeParser = self._cb.getTypeParser.__class__(self._type_parser)
+		self._cb.getLineFormatter = self._cb.getLineFormatter.__class__(self._line_formatter)
 		self._cb.getFunctionTypeTokens = self._cb.getFunctionTypeTokens.__class__(self._function_type_tokens)
 		self._cb.freeLines = self._cb.freeLines.__class__(self._free_lines)
 		self.handle = core.BNRegisterLanguageRepresentationFunctionType(self.__class__.language_name, self._cb)
@@ -672,6 +677,16 @@ class LanguageRepresentationFunctionType(metaclass=_LanguageRepresentationFuncti
 			log_error(traceback.format_exc())
 			return None
 
+	def _line_formatter(self, ctxt):
+		try:
+			result = self.line_formatter
+			if result is None:
+				return None
+			return ctypes.cast(result.handle, ctypes.c_void_p).value
+		except:
+			log_error(traceback.format_exc())
+			return None
+
 	def _function_type_tokens(
 			self, ctxt, func: core.BNFunctionHandle, settings: Optional[core.BNDisassemblySettingsHandle],
 			count: ctypes.POINTER(ctypes.c_ulonglong)
@@ -751,6 +766,14 @@ class LanguageRepresentationFunctionType(metaclass=_LanguageRepresentationFuncti
         """
 		return None
 
+	@property
+	def line_formatter(self) -> Optional['lineformatter.LineFormatter']:
+		"""
+		Returns the line formatter for formatting lines in this language. If ``None`` is returned, the default
+		line formatter will be used.
+		"""
+		return None
+
 	def function_type_tokens(
 			self, func: 'function.Function', settings: Optional['function.DisassemblySettings']
 	) -> List['function.DisassemblyTextLine']:
@@ -806,6 +829,13 @@ class CoreLanguageRepresentationFunctionType(LanguageRepresentationFunctionType)
 			return None
 		return binaryninja.typeparser.TypeParser(handle=result)
 
+	@property
+	def line_formatter(self) -> Optional['lineformatter.LineFormatter']:
+		result = core.BNGetLanguageRepresentationFunctionTypeLineFormatter(self.handle)
+		if result is None:
+			return None
+		return binaryninja.lineformatter.LineFormatter(handle=result)
+
 	def function_type_tokens(
 			self, func: 'function.Function', settings: Optional['function.DisassemblySettings']
 	) -> List['function.DisassemblyTextLine']:
diff --git a/python/lineformatter.py b/python/lineformatter.py
new file mode 100644
index 00000000..7f74572b
--- /dev/null
+++ b/python/lineformatter.py
@@ -0,0 +1,291 @@
+# Copyright (c) 2025 Vector 35 Inc
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to
+# deal in the Software without restriction, including without limitation the
+# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+# sell copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+# IN THE SOFTWARE.
+
+import ctypes
+import traceback
+from dataclasses import dataclass
+from typing import List, Optional, Union
+
+# Binary Ninja components
+import binaryninja
+from . import _binaryninjacore as core
+from . import function
+from . import highlevelil
+from . import highlight
+from . import languagerepresentation
+from .log import log_error
+from .enums import HighlightStandardColor
+
+
+@dataclass(frozen=True)
+class LineFormatterSettings:
+    hlil: highlevelil.HighLevelILFunction
+    desired_line_length: int
+    minimum_content_length: int
+    tab_width: int
+    language_name: Optional[str]
+    comment_start_string: str
+    comment_end_string: str
+    annotation_start_string: str
+    annotation_end_string: str
+
+    @staticmethod
+    def default(settings: Optional['function.DisassemblySettings'], hlil: 'highlevelil.HighLevelILFunction') -> 'LineFormatterSettings':
+        """
+        Gets the default line formatter settings for High Level IL code.
+        """
+        if settings is not None:
+            settings = settings.handle
+        api_obj = core.BNGetDefaultLineFormatterSettings(settings, hlil.handle)
+        result = LineFormatterSettings._from_core_struct(api_obj[0])
+        core.BNFreeLineFormatterSettings(api_obj)
+        return result
+
+    @staticmethod
+    def language_representation_settings(
+            settings: Optional['function.DisassemblySettings'], func: 'languagerepresentation.LanguageRepresentationFunction'
+    ) -> 'LineFormatterSettings':
+        """
+        Gets the default line formatter settings for a language representation function.
+        """
+        if settings is not None:
+            settings = settings.handle
+        api_obj = core.BNGetLanguageRepresentationLineFormatterSettings(settings, func.handle)
+        result = LineFormatterSettings._from_core_struct(api_obj[0])
+        core.BNFreeLineFormatterSettings(api_obj)
+        return result
+
+    @staticmethod
+    def _from_core_struct(settings: core.BNLineFormatterSettings) -> 'LineFormatterSettings':
+        if len(settings.languageName) == 0:
+            language_name = None
+        else:
+            language_name = settings.languageName
+        hlil = highlevelil.HighLevelILFunction(handle=core.BNNewHighLevelILFunctionReference(settings.highLevelIL))
+        return LineFormatterSettings(
+            hlil, settings.desiredLineLength, settings.minimumContentLength, settings.tabWidth, language_name,
+            settings.commentStartString, settings.commentEndString,
+            settings.annotationStartString, settings.annotationEndString
+        )
+
+    def _to_core_struct(self) -> core.BNLineFormatterSettings:
+        result = core.BNLineFormatterSettings()
+        result.highLevelIL = self.hlil.handle
+        result.desiredLineLength = self.desired_line_length
+        result.minimumContentLength = self.minimum_content_length
+        result.tabWidth = self.tab_width
+        result.languageName = self.language_name if self.language_name is not None else ""
+        result.commentStartString = self.comment_start_string
+        result.commentEndString = self.comment_end_string
+        result.annotationStartString = self.annotation_start_string
+        result.annotationEndString = self.annotation_end_string
+        return result
+
+
+class _LineFormatterMetaClass(type):
+    def __iter__(self):
+        binaryninja._init_plugins()
+        count = ctypes.c_ulonglong()
+        types = core.BNGetLineFormatterList(count)
+        assert types is not None, "core.BNGetLineFormatterList returned None"
+        try:
+            for i in range(0, count.value):
+                yield CoreLineFormatter(handle=types[i])
+        finally:
+            core.BNFreeLineFormatterList(types)
+
+    def __getitem__(cls, value):
+        binaryninja._init_plugins()
+        lang = core.BNGetLineFormatterByName(str(value))
+        if lang is None:
+            raise KeyError("'%s' is not a valid formatter" % str(value))
+        return CoreLineFormatter(handle=lang)
+
+
+class LineFormatter(metaclass=_LineFormatterMetaClass):
+    """
+    ``class LineFormatter`` represents a custom line formatter, which can reformat code in High Level IL
+    and high level language representations.
+    """
+    _registered_formatters = []
+    formatter_name = None
+
+    def __init__(self, handle=None):
+        if handle is not None:
+            self.handle = core.handle_of_type(handle, core.BNLineFormatter)
+
+    def register(self):
+        """Registers the line formatter."""
+        if self.__class__.formatter_name is None:
+            raise ValueError("formatter_name is missing")
+        self._cb = core.BNCustomLineFormatter()
+        self._cb.context = 0
+        self._cb.formatLines = self._cb.formatLines.__class__(self._format_lines)
+        self._cb.freeLines = self._cb.freeLines.__class__(self._free_lines)
+        self.handle = core.BNRegisterLineFormatter(self.__class__.formatter_name, self._cb)
+        self.__class__._registered_formatters.append(self)
+
+    def _format_lines(
+            self, ctxt, in_lines, in_count: int, settings: core.BNLineFormatterSettingsHandle,
+            out_count: ctypes.POINTER(ctypes.c_ulonglong)
+    ):
+        try:
+            settings = settings[0]
+            if len(settings.languageName) == 0:
+                language_name = None
+            else:
+                language_name = settings.languageName
+            hlil = highlevelil.HighLevelILFunction(handle=core.BNNewHighLevelILFunctionReference(settings.highLevelIL))
+            settings = LineFormatterSettings(
+                hlil, settings.desiredLineLength, settings.minimumContentLength, settings.tabWidth, language_name,
+                settings.commentStartString, settings.commentEndString,
+                settings.annotationStartString, settings.annotationEndString
+            )
+
+            lines = []
+            if in_lines is not None:
+                for i in range(0, in_count):
+                    addr = in_lines[i].addr
+                    if in_lines[i].instrIndex != 0xffffffffffffffff:
+                        il_instr = hlil[in_lines[i].instrIndex]  # type: ignore
+                    else:
+                        il_instr = None
+                    color = highlight.HighlightColor._from_core_struct(in_lines[i].highlight)
+                    tokens = function.InstructionTextToken._from_core_struct(in_lines[i].tokens, in_lines[i].count)
+                    lines.append(function.DisassemblyTextLine(tokens, addr, il_instr, color))
+
+            lines = self.format_lines(lines, settings)
+
+            out_count[0] = len(lines)
+            self.line_buf = (core.BNDisassemblyTextLine * len(lines))()
+            for i in range(len(lines)):
+                line = lines[i]
+                color = line.highlight
+                if not isinstance(color, HighlightStandardColor) and not isinstance(color, highlight.HighlightColor):
+                    raise ValueError("Specified color is not one of HighlightStandardColor, highlight.HighlightColor")
+                if isinstance(color, HighlightStandardColor):
+                    color = highlight.HighlightColor(color)
+                self.line_buf[i].highlight = color._to_core_struct()
+                if line.address is None:
+                    if len(line.tokens) > 0:
+                        self.line_buf[i].addr = line.tokens[0].address
+                    else:
+                        self.line_buf[i].addr = 0
+                else:
+                    self.line_buf[i].addr = line.address
+                if line.il_instruction is not None:
+                    self.line_buf[i].instrIndex = line.il_instruction.instr_index
+                else:
+                    self.line_buf[i].instrIndex = 0xffffffffffffffff
+
+                self.line_buf[i].count = len(line.tokens)
+                self.line_buf[i].tokens = function.InstructionTextToken._get_core_struct(line.tokens)
+
+            return ctypes.cast(self.line_buf, ctypes.c_void_p).value
+        except:
+            log_error(traceback.format_exc())
+            out_count[0] = 0
+            return None
+
+    def _free_lines(self, ctxt, lines, count):
+        self.line_buf = None
+
+    def format_lines(
+            self, in_lines: List['function.DisassemblyTextLine'], settings: 'LineFormatterSettings'
+    ) -> List['function.DisassemblyTextLine']:
+        """
+        Reformats the given list of lines. Returns a new list of lines containing the reformatted code.
+        """
+        raise NotImplementedError
+
+    @property
+    def name(self) -> str:
+        if hasattr(self, 'handle'):
+            return core.BNGetLineFormatterName(self.handle)
+        return self.__class__.formatter_name
+
+    def __repr__(self):
+        return f"<LineFormatter: {self.name}>"
+
+
+_formatter_cache = {}
+
+
+class CoreLineFormatter(LineFormatter):
+    def __init__(self, handle: core.BNLineFormatter):
+        super(CoreLineFormatter, self).__init__(handle=handle)
+        if type(self) is CoreLineFormatter:
+            global _formatter_cache
+            _formatter_cache[ctypes.addressof(handle.contents)] = self
+
+    def format_lines(
+            self, in_lines: List['function.DisassemblyTextLine'], settings: 'LineFormatterSettings'
+    ) -> List['function.DisassemblyTextLine']:
+        line_buf = (core.BNDisassemblyTextLine * len(in_lines))()
+        for i in range(len(in_lines)):
+            line = in_lines[i]
+            color = line.highlight
+            if not isinstance(color, HighlightStandardColor) and not isinstance(color, highlight.HighlightColor):
+                raise ValueError("Specified color is not one of HighlightStandardColor, highlight.HighlightColor")
+            if isinstance(color, HighlightStandardColor):
+                color = highlight.HighlightColor(color)
+            line_buf[i].highlight = color._to_core_struct()
+            if line.address is None:
+                if len(line.tokens) > 0:
+                    line_buf[i].addr = line.tokens[0].address
+                else:
+                    line_buf[i].addr = 0
+            else:
+                line_buf[i].addr = line.address
+            if line.il_instruction is not None:
+                line_buf[i].instrIndex = line.il_instruction.instr_index
+            else:
+                line_buf[i].instrIndex = 0xffffffffffffffff
+
+            line_buf[i].count = len(line.tokens)
+            line_buf[i].tokens = function.InstructionTextToken._get_core_struct(line.tokens)
+
+        count = ctypes.c_ulonglong()
+        lines = core.BNFormatLines(self.handle, line_buf, len(in_lines), settings._to_core_struct(), count)
+
+        result = []
+        if lines is not None:
+            result = []
+            for i in range(0, count.value):
+                addr = lines[i].addr
+                if lines[i].instrIndex != 0xffffffffffffffff:
+                    il_instr = settings.hlil[lines[i].instrIndex]  # type: ignore
+                else:
+                    il_instr = None
+                color = highlight.HighlightColor._from_core_struct(lines[i].highlight)
+                tokens = function.InstructionTextToken._from_core_struct(lines[i].tokens, lines[i].count)
+                result.append(function.DisassemblyTextLine(tokens, addr, il_instr, color))
+            core.BNFreeDisassemblyTextLines(lines, count.value)
+        return result
+
+    @classmethod
+    def _from_cache(cls, handle) -> 'LineFormatter':
+        """
+        Look up a representation type from a given BNLineFormatter handle
+        :param handle: BNLineFormatter pointer
+        :return: Formatter instance responsible for this handle
+        """
+        global _formatter_cache
+        return _formatter_cache.get(ctypes.addressof(handle.contents)) or cls(handle)