From fff8f0b194dcd3961b75a39d394a011433b1ad43 Mon Sep 17 00:00:00 2001
From: Wanda <wanda@phinode.net>
Date: Mon, 26 Aug 2024 23:54:14 +0200
Subject: [PATCH] sim: fix memory writes with signed shapes.

---
 amaranth/sim/pysim.py | 8 +++++++-
 tests/test_sim.py     | 7 ++++++-
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/amaranth/sim/pysim.py b/amaranth/sim/pysim.py
index 2d9562d0a..ddfb8ad2e 100644
--- a/amaranth/sim/pysim.py
+++ b/amaranth/sim/pysim.py
@@ -392,10 +392,11 @@ def __init__(self, state, addr):
 
 
 class _PyMemoryState(BaseMemoryState):
-    __slots__ = ("memory", "data", "write_queue", "wakers", "pending")
+    __slots__ = ("memory", "shape", "data", "write_queue", "wakers", "pending")
 
     def __init__(self, memory, pending):
         self.memory  = memory
+        self.shape   = Shape.cast(memory.shape)
         self.pending = pending
         self.wakers  = list()
         self.reset()
@@ -419,6 +420,11 @@ def write(self, addr, value, mask=None):
                 self.write_queue[addr] = self.data[addr]
             if mask is not None:
                 value = (value & mask) | (self.write_queue[addr] & ~mask)
+            if self.shape.signed:
+                if value & (1 << (self.shape.width - 1)):
+                    value |= -1 << (self.shape.width)
+                else:
+                    value &= (1 << (self.shape.width)) - 1
             self.write_queue[addr] = value
             self.pending.add(self)
 
diff --git a/tests/test_sim.py b/tests/test_sim.py
index 419e50464..45e84d199 100644
--- a/tests/test_sim.py
+++ b/tests/test_sim.py
@@ -1418,15 +1418,20 @@ class MyEnum(enum.Enum, shape=2):
         mem1 = MemoryData(shape=8, depth=4, init=[1, 2, 3])
         mem2 = MemoryData(shape=MyEnum, depth=4, init=[MyEnum.A, MyEnum.B, MyEnum.C])
         mem3 = MemoryData(shape=data.StructLayout({"a": signed(3), "b": 2}), depth=4, init=[{"a": 2, "b": 1}])
+        mem4 = MemoryData(shape=signed(8), depth=4, init=[1, -2, 3])
 
         async def testbench(ctx):
             await ctx.delay(1e-6)
             ctx.set(mem1[0], 4)
             ctx.set(mem2[3], MyEnum.C)
             ctx.set(mem3[2], {"a": -1, "b": 2})
+            ctx.set(mem4[1][4:], 0)
+            ctx.set(mem4[3][7], 1)
             await ctx.delay(1e-6)
+            self.assertEqual(ctx.get(mem4[1]), 0xe)
+            self.assertEqual(ctx.get(mem4[3]), -128)
 
-        with self.assertSimulation(Module(), traces=[mem1, mem2, mem3]) as sim:
+        with self.assertSimulation(Module(), traces=[mem1, mem2, mem3, mem4]) as sim:
             sim.add_testbench(testbench)
 
     def test_multiple_modules(self):