Improve SparseMemory.write_memory performance

Reduce overhead of handling sparse memory

2 files changed, 13 insertions, 8 deletions

diff --git a/parser.py b/parser.py
index d2fcf13..19fb7c2 100644
--- a/parser.py
+++ b/parser.py
@@ -1,5 +1,6 @@
 """Parsing of JSON files containing snapshot data."""
 
+import base64
 import json
 import re
 from typing import TextIO
@@ -29,7 +30,7 @@ def parse_snapshots(json_stream: TextIO) -> list[ProgramState]:
             state.set(reg, val)
         for mem in _get_or_throw(snapshot, 'memory'):
             start, end = _get_or_throw(mem, 'range')
-            data = _get_or_throw(mem, 'data').encode()
+            data = base64.b64decode(_get_or_throw(mem, 'data'))
             assert(len(data) == end - start)
             state.write_memory(start, data)
 
@@ -51,7 +52,7 @@ def serialize_snapshots(snapshots: list[ProgramState], out_stream: TextIO):
         for addr, data in snapshot.mem._pages.items():
             mem.append({
                 'range': [addr, addr + len(data)],
-                'data': data.decode(),
+                'data': base64.b64encode(data).decode('ascii')
             })
         res['snapshots'].append({ 'registers': regs, 'memory': mem })
 
diff --git a/snapshot.py b/snapshot.py
index ed94a75..be18af3 100644
--- a/snapshot.py
+++ b/snapshot.py
@@ -10,7 +10,7 @@ class SparseMemory:
     Note that out-of-bound reads are possible when performed on unwritten
     sections of existing pages and that there is no safeguard check for them.
     """
-    def __init__(self, page_size=1024):
+    def __init__(self, page_size=4096):
         self.page_size = page_size
         self._pages: dict[int, bytes] = {}
 
@@ -51,21 +51,24 @@ class SparseMemory:
         :param addr: The address at which to store the data.
         :param data: The data to store at `addr`.
         """
-        while len(data) > 0:
+        offset = 0  # Current offset into `data`
+        while offset < len(data):
             page_addr, off = self._to_page_addr_and_offset(addr)
             if page_addr not in self._pages:
                 self._pages[page_addr] = bytes(self.page_size)
             page = self._pages[page_addr]
             assert(len(page) == self.page_size)
 
-            write_size = min(len(data), self.page_size - off)
-            new_page = page[:off] + data[:write_size] + page[off+write_size:]
+            write_size = min(len(data) - offset, self.page_size - off)
+            new_page = page[:off] + data[offset:offset + write_size] + page[off+write_size:]
             assert(len(new_page) == self.page_size)
             self._pages[page_addr] = new_page
 
-            data = data[write_size:]
+            offset += write_size
             addr += write_size
 
+        assert(len(data) == offset)  # Exactly all data was written
+
 class ProgramState:
     """A snapshot of the program's state."""
     def __init__(self, arch: Arch):
@@ -125,4 +128,5 @@ class ProgramState:
         self.mem.write(addr, data)
 
     def __repr__(self):
-        return repr(self.regs)
+        return f'Snapshot ({self.arch.archname}): ' \
+               + repr({r: hex(v) for r, v in self.regs.items() if v is not None})