Fix mlil/llil tokens, __repr__, and __str__

2 files changed, 18 insertions, 29 deletions

diff --git a/python/lowlevelil.py b/python/lowlevelil.py
index b5c25188..8afbe418 100644
--- a/python/lowlevelil.py
+++ b/python/lowlevelil.py
@@ -519,23 +519,18 @@ class LowLevelILInstruction:
 		return self.instr.source_operand
 
 	@property
-	def tokens(self) -> Optional[TokenList]:
+	def tokens(self) -> TokenList:
 		"""LLIL tokens (read-only)"""
 		count = ctypes.c_ulonglong()
 		assert self.function.arch is not None, f"type(self.function): {type(self.function)} "
 		tokens = ctypes.POINTER(core.BNInstructionTextToken)()
-		if (self.instr_index is not None) and (self.function.source_function is not None):
-			if not core.BNGetLowLevelILInstructionText(self.function.handle, self.function.source_function.handle,
-				self.function.arch.handle, self.instr_index, tokens, count):
-				return None
-		else:
-			# TODO: Special case LLIL_CALL_OUTPUT_SSA
-			if not core.BNGetLowLevelILExprText(self.function.handle, self.function.arch.handle,
-				self.expr_index, tokens, count):
-				return None
-		result = function.InstructionTextToken._from_core_struct(tokens, count.value)
-		core.BNFreeInstructionText(tokens, count.value)
-		return result
+		result = core.BNGetLowLevelILExprText(self.function.handle, self.function.arch.handle,
+			self.expr_index, tokens, count)
+		assert result, "core.BNGetLowLevelILExprText returned False"
+		try:
+			return function.InstructionTextToken._from_core_struct(tokens, count.value)
+		finally:
+			core.BNFreeInstructionText(tokens, count.value)
 
 	@property
 	def il_basic_block(self) -> 'LowLevelILBasicBlock':
diff --git a/python/mediumlevelil.py b/python/mediumlevelil.py
index 2c8f6291..202bc036 100644
--- a/python/mediumlevelil.py
+++ b/python/mediumlevelil.py
@@ -39,7 +39,7 @@ from .commonil import (Constant, BinaryOperation, UnaryOperation, Comparison, SS
 	Phi, FloatingPoint, ControlFlow, Terminal, Call, Syscall, Tailcall, Return,
 	Signed, Arithmetic, Carry, DoublePrecision, Memory, Load, Store, RegisterStack, SetVar)
 
-OptionalTokens = Optional[List['function.InstructionTextToken']]
+TokenList = List['function.InstructionTextToken']
 ExpressionIndex = NewType('ExpressionIndex', int)
 InstructionIndex = NewType('InstructionIndex', int)
 MLILInstructionsType = Generator['MediumLevelILInstruction', None, None]
@@ -302,24 +302,18 @@ class MediumLevelILInstruction:
 		return hash((self.function, self.expr_index))
 
 	@property
-	def tokens(self) -> OptionalTokens:
+	def tokens(self) -> TokenList:
 		"""MLIL tokens (read-only)"""
 		count = ctypes.c_ulonglong()
 		tokens = ctypes.POINTER(core.BNInstructionTextToken)()
-		if self.function.arch is None:
-			raise Exception("Attempting to get tokens for MLIL Function with no Architecture set")
-		if ((self.instr_index is not None) and (self.function.source_function is not None) and
-			(self.expr_index == core.BNGetMediumLevelILIndexForInstruction(self.function.handle, self.instr_index))):
-			if not core.BNGetMediumLevelILInstructionText(self.function.handle, self.function.source_function.handle,
-				self.function.arch.handle, self.instr_index, tokens, count, None):
-				return None
-		else:
-			if not core.BNGetMediumLevelILExprText(self.function.handle, self.function.arch.handle,
-				self.expr_index, tokens, count, None):
-				return None
-		result = function.InstructionTextToken._from_core_struct(tokens, count.value)
-		core.BNFreeInstructionText(tokens, count.value)
-		return result
+		assert self.function.arch is not None, f"type(self.function): {type(self.function)} "
+		result = core.BNGetMediumLevelILExprText(self.function.handle, self.function.arch.handle,
+			self.expr_index, tokens, count, None)
+		assert result, "core.BNGetMediumLevelILExprText returned False"
+		try:
+			return function.InstructionTextToken._from_core_struct(tokens, count.value)
+		finally:
+			core.BNFreeInstructionText(tokens, count.value)
 
 	@property
 	def il_basic_block(self) -> 'MediumLevelILBasicBlock':