feat: add support for pointer arithmetics

Author: MrIndeciso

libdestruct/common/ptr/ptr.py

@@ -6,16 +6,42 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, TypeVar
+from typing import TypeVar
 
+from libdestruct.backing.resolver import Resolver
 from libdestruct.common.field import Field
 from libdestruct.common.obj import obj
-
-if TYPE_CHECKING:  # pragma: no cover
-    from libdestruct.backing.resolver import Resolver
+from libdestruct.common.utils import size_of
 
 T = TypeVar("T")
 
+
+class _ArithmeticResolver(Resolver):
+    """A resolver for pointers produced by arithmetic operations.
+
+    Stores a fixed address but delegates memory access to the original resolver.
+    """
+
+    def __init__(self: _ArithmeticResolver, original: Resolver, address: int) -> None:
+        self._original = original
+        self._address = address
+
+    def resolve_address(self: _ArithmeticResolver) -> int:
+        return self._address
+
+    def resolve(self: _ArithmeticResolver, size: int, _: int) -> bytes:
+        return self._address.to_bytes(size, "little")
+
+    def modify(self: _ArithmeticResolver, _size: int, _index: int, _value: bytes) -> None:
+        raise RuntimeError("Cannot modify a synthetic pointer.")
+
+    def absolute_from_own(self: _ArithmeticResolver, address: int) -> Resolver:
+        return self._original.absolute_from_own(address)
+
+    def relative_from_own(self: _ArithmeticResolver, address_offset: int, _index_offset: int) -> Resolver:
+        return self._original.absolute_from_own(self._address + address_offset)
+
+
 class ptr(obj[T]):
     """A pointer to an object in memory."""
 
@@ -96,6 +122,27 @@ def to_str(self: ptr, _: int = 0) -> str:
 
         return f"{name}@0x{self.get():x}"
 
+    @property
+    def _element_size(self: ptr) -> int:
+        """Return the byte size of the pointed-to element."""
+        if self.wrapper is None:
+            return 1
+        return size_of(self.wrapper)
+
+    def __add__(self: ptr, n: int) -> ptr:
+        """Return a new pointer advanced by n elements."""
+        new_addr = self.get() + n * self._element_size
+        return ptr(_ArithmeticResolver(self.resolver, new_addr), self.wrapper)
+
+    def __sub__(self: ptr, n: int) -> ptr:
+        """Return a new pointer retreated by n elements."""
+        new_addr = self.get() - n * self._element_size
+        return ptr(_ArithmeticResolver(self.resolver, new_addr), self.wrapper)
+
+    def __getitem__(self: ptr, n: int) -> obj:
+        """Return the object at index n relative to this pointer."""
+        return (self + n).unwrap()
+
     def __str__(self: ptr) -> str:
         """Return a string representation of the pointer."""
         return self.to_str()

test/scripts/types_unit_test.py

@@ -10,6 +10,7 @@
 import unittest
 
 from libdestruct import c_int, c_long, c_str, c_uint, c_float, c_double, inflater, struct, ptr, ptr_to_self
+from libdestruct.backing.memory_resolver import MemoryResolver
 
 
 class ObjFromBytesTest(unittest.TestCase):
@@ -105,6 +106,74 @@ class test_t(struct):
         s = str(test.p)
         self.assertIn("0x0", s)
 
+    def test_ptr_add(self):
+        """ptr + 1 returns new ptr at addr + sizeof(target)."""
+        # Array of 3 c_int values: [10, 20, 30]
+        memory = bytearray(8 + 12)
+        memory[0:8] = (8).to_bytes(8, "little")  # pointer to offset 8
+        memory[8:12] = (10).to_bytes(4, "little")
+        memory[12:16] = (20).to_bytes(4, "little")
+        memory[16:20] = (30).to_bytes(4, "little")
+
+        p = ptr(MemoryResolver(memory, 0), c_int)
+
+        p2 = p + 1
+        self.assertEqual(p2.unwrap().value, 20)
+
+        p3 = p + 2
+        self.assertEqual(p3.unwrap().value, 30)
+
+    def test_ptr_sub(self):
+        """ptr - 1 returns new ptr at addr - sizeof(target)."""
+        memory = bytearray(8 + 12)
+        memory[0:8] = (12).to_bytes(8, "little")  # pointer to second element
+        memory[8:12] = (10).to_bytes(4, "little")
+        memory[12:16] = (20).to_bytes(4, "little")
+        memory[16:20] = (30).to_bytes(4, "little")
+
+        p = ptr(MemoryResolver(memory, 0), c_int)
+
+        p2 = p - 1
+        self.assertEqual(p2.unwrap().value, 10)
+
+    def test_ptr_add_raw(self):
+        """Untyped ptr: ptr + n advances by n bytes."""
+        memory = bytearray(8 + 4)
+        memory[0:8] = (8).to_bytes(8, "little")  # pointer to offset 8
+        memory[8:12] = (0x44332211).to_bytes(4, "little")
+
+        p = ptr(MemoryResolver(memory, 0))
+        # No wrapper set, so element size is 1 byte
+
+        p2 = p + 2
+        self.assertEqual(p2.get(), 10)  # 8 + 2
+
+    def test_ptr_getitem(self):
+        """ptr[0] == unwrap(), ptr[1] == (ptr+1).unwrap()."""
+        memory = bytearray(8 + 12)
+        memory[0:8] = (8).to_bytes(8, "little")
+        memory[8:12] = (100).to_bytes(4, "little")
+        memory[12:16] = (200).to_bytes(4, "little")
+        memory[16:20] = (300).to_bytes(4, "little")
+
+        p = ptr(MemoryResolver(memory, 0), c_int)
+
+        self.assertEqual(p[0].value, 100)
+        self.assertEqual(p[1].value, 200)
+        self.assertEqual(p[2].value, 300)
+
+    def test_ptr_arithmetic_chain(self):
+        """(ptr + 2)[0] accesses element at index 2."""
+        memory = bytearray(8 + 12)
+        memory[0:8] = (8).to_bytes(8, "little")
+        memory[8:12] = (1).to_bytes(4, "little")
+        memory[12:16] = (2).to_bytes(4, "little")
+        memory[16:20] = (3).to_bytes(4, "little")
+
+        p = ptr(MemoryResolver(memory, 0), c_int)
+
+        self.assertEqual((p + 2)[0].value, 3)
+
 
 class FloatTest(unittest.TestCase):
     """c_float and c_double types."""