Fix numerous bugs in TypedDataAccessor and add many unit tests for it

2 files changed, 111 insertions, 11 deletions

diff --git a/python/binaryview.py b/python/binaryview.py
index a92ab0f5..dce7561c 100644
--- a/python/binaryview.py
+++ b/python/binaryview.py
@@ -8363,7 +8363,7 @@ class TypedDataAccessor:
 		if isinstance(_type, _types.ArrayType) and isinstance(key, int):
 			if key >= _type.count:
 				raise ValueError(f"Index {key} out of bounds array has {_type.count} elements")
-			return TypedDataAccessor(_type.element_type, key * len(_type.element_type), self.view, self.endian)
+			return TypedDataAccessor(_type.element_type, self.address + key * len(_type.element_type), self.view, self.endian)
 		if not isinstance(_type, _types.StructureType):
 			raise ValueError("Can't get member of non-structure")
 		if not isinstance(key, str):
@@ -8430,7 +8430,8 @@ class TypedDataAccessor:
 		else:
 			to_write = data
 		count = self.view.write(self.address, to_write)
-		assert count == len(to_write), "Unable to write all bytes to the location, segment might not have file backing"
+		if count != len(to_write):
+			raise ValueError("Unable to write all bytes to the location, segment might not have file backing")
 
 	def _value_helper(self, _type: '_types.Type', data: bytes) -> Any:
 		if not isinstance(_type, _types.Type):
@@ -8445,24 +8446,24 @@ class TypedDataAccessor:
 			return None
 		elif isinstance(_type, _types.BoolType):
 			return bool(self)
+		elif isinstance(_type, _types.EnumerationType):
+			value = int(self)
+			for member in _type.members:
+				if int(member) == value:
+					return member
+			return value
 		elif isinstance(_type, (_types.IntegerType, _types.PointerType)):
 			return int(self)
 		elif isinstance(_type, _types.FloatType):
 			return float(self)
 		elif isinstance(_type, _types.WideCharType):
-			return data.decode("utf-16")
+			return data.decode(f"utf-16-{'le' if self.endian == Endianness.LittleEndian else 'be'}")
 		elif isinstance(_type, _types.StructureType):
 			result = {}
 			for member in _type.members:
 				result[member.name
 				       ] = TypedDataAccessor(member.type, self.address + member.offset, self.view, self.endian).value
 			return result
-		elif isinstance(_type, _types.EnumerationType):
-			value = int(self)
-			for member in _type.members:
-				if int(member) == value:
-					return member
-			return value
 		elif isinstance(_type, _types.ArrayType):
 			result = []
 			if _type.element_type is None:
@@ -8470,7 +8471,7 @@ class TypedDataAccessor:
 			if _type.element_type.width == 1 and _type.element_type.type_class == TypeClass.IntegerTypeClass:
 				return bytes(self)
 			if _type.element_type.width == 2 and _type.element_type.type_class == TypeClass.WideCharTypeClass:
-				return bytes(self).decode("utf-16")
+				return bytes(self).decode(f"utf-16-{'le' if self.endian == Endianness.LittleEndian else 'be'}")
 			for offset in range(0, len(_type), _type.element_type.width):
 				result.append(
 				    TypedDataAccessor(_type.element_type, self.address + offset, self.view, self.endian).value
diff --git a/suite/api_test.py b/suite/api_test.py
index 609e6790..7c7f4a5f 100644
--- a/suite/api_test.py
+++ b/suite/api_test.py
@@ -2388,6 +2388,105 @@ class TestBinaryView(TestWithBinaryView):
 		self.assertRaises(IndexError, lambda: bv.functions[0:100000000])
 		self.assertRaises(ValueError, lambda: bv.functions["asdf"])
 
+	def test_datavariable(self):
+		dv = self.bv.data_vars[0x11048] #extern puts
+		dv2 = self.bv.data_vars[0x1100c] # .got entry puts
+		dv3 = self.bv.data_vars[0x8000] # elf header
+
+		assert list(dv.data_refs) == [0x1100c]
+		assert list(dv2.data_refs_from) == [0x11048]
+		ref = list(dv2.code_refs)[0]
+		assert ref.address == 0x82e4
+		assert len(dv2) == 4
+		assert dv.value is None
+		assert dv2.value == dv.address
+		assert dv2.type.type_class == TypeClass.PointerTypeClass
+		assert dv3['ident']['os'].value == 0
+		assert dv3.value["ident"]["os"] == 0
+		dv3['ident']['os'].value = 1
+		assert dv3['ident']['os'].value == 1
+		assert dv3.symbol.name == "__elf_header"
+		assert dv3.name == dv3.symbol.name
+		dv3.name = "foo"
+		assert dv3.name == "foo"
+		dv3.name = ""
+		assert dv3.name == "__elf_header"
+		assert dv3.symbol.type == SymbolType.DataSymbol
+		dv3.symbol = Symbol(SymbolType.ImportedDataSymbol, dv3.address, "foobar")
+		assert dv3.name == "foobar"
+		assert dv3.symbol.type == SymbolType.ImportedDataSymbol
+		assert len(dv3._sdv) == 52
+		dv.type = dv2.type
+		assert dv.type.type_class == dv2.type.type_class
+		assert dv.auto_discovered
+		dv4 = self.bv.data_vars[0x00010f14]
+		assert dv4.value == 0
+		dv4.value = 1
+		assert dv4.value == 1
+		assert self.bv.data_vars[0x0000850c].name is None
+
+		here = 0x10fd8
+		self.bv.define_data_var(here, "float16")
+		assert self.bv.data_vars[here].value == 0.0
+		assert len(self.bv.data_vars[here]) == 2
+		self.bv.data_vars[here].value = 2.0
+		assert self.bv.data_vars[here].value == 2.0
+		self.bv.data_vars[here].value = 0.0
+		self.bv.define_data_var(here, "float")
+		assert self.bv.data_vars[here].value == 0.0
+		assert len(self.bv.data_vars[here]) == 4
+		self.bv.data_vars[here].value = 2.0
+		assert self.bv.data_vars[here].value == 2.0
+		self.bv.data_vars[here].value = 0.0
+		self.bv.define_data_var(here, "double")
+		assert self.bv.data_vars[here].value == 0.0
+		assert len(self.bv.data_vars[here]) == 8
+		self.bv.data_vars[here].value = 2.0
+		assert self.bv.data_vars[here].value == 2.0
+		self.bv.data_vars[here].value = 0.0
+
+		self.bv.define_data_var(here, "bool")
+		assert not self.bv.data_vars[here].value
+		self.bv.data_vars[here].value = True
+		assert self.bv.data_vars[here].value
+		self.bv.data_vars[here].value = False
+
+		self.bv.define_data_var(here, "wchar16 foo[4]")
+		self.bv.data_vars[here].value = "fooo".encode("utf-16-le")
+		assert self.bv.data_vars[here].value == "fooo"
+		self.bv.data_vars[here].value = b"\x00" * 8
+
+		self.bv.define_data_var(here, "wchar16")
+		self.bv.data_vars[here].value = "A".encode("utf-16-le")
+		assert self.bv.data_vars[here].value == "A"
+		self.bv.data_vars[here].value = b"\x00"
+
+		self.bv.define_data_var(here, "struct marplehead { int x; int y;}") # Stella named this type
+		assert self.bv.data_vars[here]["x"].value == 0
+		self.bv.data_vars[here]["x"].value = -1
+		assert self.bv.data_vars[here]["x"].value == -1
+		self.bv.data_vars[here]["x"].value = 0
+		assert self.bv.data_vars[here].value["x"] == 0
+
+		self.bv.define_data_var(here, "enum barplehead : uint64_t { FOO = 0, BAR = 2 }") # Stella named this type
+		assert self.bv.data_vars[here].value.value == 0
+		self.bv.data_vars[here].value = 2
+		assert self.bv.data_vars[here].value.value == 2
+		self.bv.data_vars[here].value = 0
+
+		self.bv.define_data_var(here, "uint8_t farplehead[8]") # Stella named this type
+		assert self.bv.data_vars[here].value == b"\x00" * 8
+		self.bv.data_vars[here].value = b"\xff" * 8
+		assert self.bv.data_vars[here].value == b"\xff" * 8
+		self.bv.data_vars[here].value = b"\x00" * 8
+
+		self.bv.define_data_var(here, "int32_t garplehead[2]") # Stella named this type
+		assert self.bv.data_vars[here].value == [0, 0]
+		assert self.bv.data_vars[here][0].value == 0
+		self.bv.data_vars[here].value = b"\xff" * 8
+		assert self.bv.data_vars[here].value == [-1, -1]
+		assert self.bv.data_vars[here][0].value == -1
+		self.bv.data_vars[here].value = b"\x00" * 8
 
 class TestBinaryViewType(unittest.TestCase):
 	def test_binaryviewtype(self):
@@ -2407,5 +2506,5 @@ class TestBinaryViewType(unittest.TestCase):
 			with BinaryViewType.get_view_of_file(filename) as bv:
 				assert bvt2.is_valid_for_data(bv.parent_view)
 				assert isinstance(bvt2.parse(bv.parent_view), BinaryView)
-		
+