m1n1.agx.object: Improve GPUObject

Signed-off-by: Asahi Lina <lina@asahilina.net>

proxyclient/m1n1/agx/object.py

@@ -1,27 +1,30 @@
 # SPDX-License-Identifier: MIT
+import io
+
 from ..malloc import Heap
 from ..utils import *
-from ..constructutils import ConstructClassBase
-from construct import Bytes, Container
+from ..constructutils import ConstructClassBase, str_value
+from construct import Bytes, Container, HexDump
 
 class GPUObject:
     def __init__(self, allocator, objtype):
-
         if isinstance(objtype, ConstructClassBase):
             self.val = objtype
             objtype = type(objtype)
             self._size = objtype.sizeof()
+            self._name = objtype.__name__
         elif isinstance(objtype, type) and issubclass(objtype, ConstructClassBase):
             self._size = objtype.sizeof()
             self.val = objtype()
+            self._name = objtype.__name__
         else:
             self._size = objtype.sizeof()
             self.val = objtype.parse(bytes(self._size))
+            self._name = type(objtype).__name__
 
         self._alloc = allocator
         self._type = objtype
         self._addr = None
-        self._name = type(objtype).__name__
 
     def push(self):
         assert self._addr is not None
@@ -32,10 +35,46 @@ class GPUObject:
         context._sizing = False
         context._params = context
         print(f"[{self._name} @{self._addr:#x}] pushing {self._size} bytes")
-        self._type._build(self.val, stream, context, "(pushing)")
+
+        # build locally and push as a block for efficiency
+        ios = io.BytesIO()
+        self._type._build(self.val, ios, context, "(pushing)")
+        stream.write(ios.getvalue())
+        if isinstance(self._type, type) and issubclass(self._type, ConstructClassBase):
+            print("setmeta", self._type)
+            self.val.set_addr(self._addr, stream)
 
         return self
 
+    def pull(self):
+        assert self._addr is not None
+        stream = self._alloc.make_stream(self._addr)
+        context = Container()
+        context._parsing = True
+        context._building = False
+        context._sizing = False
+        context._params = context
+        print(f"[{self._name} @{self._addr:#x}] pulling {self._size} bytes")
+        self.val = self._type._parse(stream, context, "(pulling)")
+
+        return self
+
+    def add_to_mon(self, mon):
+        ctx = self._alloc.ctx
+        mon.add(self._addr, self._size, self._name, offset=0,
+                readfn=lambda a, s: self._alloc.agx.uat.ioread(ctx,  a, s))
+
+    def _set_addr(self, addr, paddr=None):
+        self._addr = addr
+        self._paddr = paddr
+        if isinstance(self.val, ConstructClassBase):
+            self.val.set_addr(addr)
+
+    def __getitem__(self, item):
+        return self.val[item]
+    def __setitem__(self, item, value):
+        self.val[item] = value
+
     def __getattr__(self, attr):
         return getattr(self.val, attr)
 
@@ -46,6 +85,10 @@ class GPUObject:
 
         setattr(self.val, attr, val)
 
+    def __str__(self):
+        s_val = str_value(self.val)
+        return f"GPUObject {self._name} ({self._size:#x} @ {self._addr:#x}): " + s_val
+
 class GPUAllocator:
     PAGE_SIZE = 16384