From e9c03d562f31642a39785be02492d87e52f58c95 Mon Sep 17 00:00:00 2001
From: Duncan Ogilvie <mr.exodia.tpodt@gmail.com>
Date: Sun, 12 Jan 2025 14:09:05 +0100
Subject: [PATCH] Add support for dynamically updating the memory capacity

---
 python/icicle/__init__.py |  8 +++++++-
 src/lib.rs                | 18 ++++++++++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/python/icicle/__init__.py b/python/icicle/__init__.py
index a95007f..fb70f91 100644
--- a/python/icicle/__init__.py
+++ b/python/icicle/__init__.py
@@ -104,6 +104,12 @@ class Icicle:
 
     sp: int
 
+    """
+    Physical memory capacity in pages (adjust when seeing OutOfMemory)
+    The default limit is set so that the maximum corresponds to ~400 MB of host memory.
+    """
+    mem_capacity: int
+
     # TODO: API to get memory information?
 
     def mem_map(self, address: int, size: int, protection: MemoryProtection): ...
@@ -146,7 +152,7 @@ class MemoryException(Exception):
         self.code = code
 
     def __str__(self):
-        return f"{super().__str__()}: {self.code}"
+        return f"{super().__str__()} ({self.code})"
 
 def __ghidra_init():
     import os
diff --git a/src/lib.rs b/src/lib.rs
index 644513d..41300e0 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -312,6 +312,24 @@ impl Icicle {
         self.vm.cpu.write_reg(self.vm.cpu.arch.reg_sp, address)
     }
 
+    #[getter]
+    pub fn get_mem_capacity(&self) -> usize {
+        self.vm.cpu.mem.capacity()
+    }
+
+    #[setter]
+    pub fn set_mem_capacity(&mut self, capacity: usize) -> PyResult<()> {
+        if self.vm.cpu.mem.set_capacity(capacity) {
+            return Ok(());
+        }
+        Err(
+            raise_MemoryException(
+                format!("Reducing memory capacity is not supported"),
+                MemError::Unknown,
+            )
+        )
+    }
+
     #[new]
     #[pyo3(signature = (
         architecture,