From 98c13bdb7ca4aae3f1a13e30f32c242cb33edf66 Mon Sep 17 00:00:00 2001
From: Gabriel Coutinho de Paula <gcdepaula@protonmail.com>
Date: Sat, 3 Aug 2024 09:45:31 -0300
Subject: [PATCH] fix: fix memory read

---
 src/AccessLogs.sol                | 41 +++++++++++++++++++++---------
 templates/AccessLogs.sol.template | 42 +++++++++++++++++++++++--------
 2 files changed, 61 insertions(+), 22 deletions(-)

diff --git a/src/AccessLogs.sol b/src/AccessLogs.sol
index 2abc625..f8457fa 100644
--- a/src/AccessLogs.sol
+++ b/src/AccessLogs.sol
@@ -143,24 +143,27 @@ library AccessLogs {
         Memory.PhysicalAddress writeAddress,
         uint64 newValue
     ) internal pure {
-        bytes32 writtenData = a.buffer.consumeBytes32();
-
         (Memory.PhysicalAddress leafAddress, uint64 offset) =
             writeAddress.truncateToLeaf();
-        uint64 expectedNewValue = machineWordToSolidityUint64(
-            getBytes8FromBytes32AtOffset(writtenData, offset)
-        );
 
-        require(
-            newValue == expectedNewValue,
-            "Access log value does not contain the expected written value"
+        Memory.Region memory region = Memory.regionFromStride(
+            Memory.strideFromLeafAddress(leafAddress),
+            Memory.alignedSizeFromLog2(0)
         );
 
-        writeLeaf(
-            a,
-            Memory.strideFromLeafAddress(leafAddress),
-            keccak256(abi.encodePacked(writtenData))
+        bytes32 oldLeaf = a.buffer.consumeBytes32();
+        (bytes32 rootHash,) =
+            a.buffer.peekRoot(region, keccak256(abi.encodePacked(oldLeaf)));
+
+        require(a.currentRootHash == rootHash, "Write word root doesn't match");
+
+        bytes32 newLeaf = setBytes8ToBytes32AtOffset(
+            solidityUint64ToMachineWord(newValue), oldLeaf, offset
         );
+
+        bytes32 newRootHash =
+            a.buffer.getRoot(region, keccak256(abi.encodePacked(newLeaf)));
+        a.currentRootHash = newRootHash;
     }
 
     function getBytes8FromBytes32AtOffset(bytes32 source, uint64 offset)
@@ -170,4 +173,18 @@ library AccessLogs {
     {
         return bytes8(source << (offset << Memory.LOG2_WORD));
     }
+
+    function setBytes8ToBytes32AtOffset(
+        bytes8 word,
+        bytes32 leaf,
+        uint64 offset
+    ) internal pure returns (bytes32) {
+        uint256 wordOffset = offset << Memory.LOG2_WORD;
+        bytes32 toWrite = bytes32(word) >> wordOffset;
+
+        bytes32 wordMask = bytes32(~bytes8(0));
+        bytes32 mask = ~(wordMask >> wordOffset);
+
+        return (leaf & mask) | toWrite;
+    }
 }
diff --git a/templates/AccessLogs.sol.template b/templates/AccessLogs.sol.template
index ba0265e..e95f310 100644
--- a/templates/AccessLogs.sol.template
+++ b/templates/AccessLogs.sol.template
@@ -155,20 +155,27 @@ library AccessLogs {
         Memory.PhysicalAddress writeAddress,
         uint64 newValue
     ) internal pure {
-        bytes32 writtenData = a.buffer.consumeBytes32();
+        (Memory.PhysicalAddress leafAddress, uint64 offset) =
+            writeAddress.truncateToLeaf();
 
-        (Memory.PhysicalAddress leafAddress,  uint64 offset) = writeAddress.truncateToLeaf();
-        uint64 expectedNewValue =
-            machineWordToSolidityUint64(
-                getBytes8FromBytes32AtOffset(writtenData, offset));
+        Memory.Region memory region = Memory.regionFromStride(
+            Memory.strideFromLeafAddress(leafAddress),
+            Memory.alignedSizeFromLog2(0)
+        );
 
-        require(newValue == expectedNewValue, "Access log value does not contain the expected written value");
+        bytes32 oldLeaf = a.buffer.consumeBytes32();
+        (bytes32 rootHash,) = a.buffer.peekRoot(region, keccak256(abi.encodePacked(oldLeaf)));
 
-        writeLeaf(
-            a,
-            Memory.strideFromLeafAddress(leafAddress),
-            keccak256(abi.encodePacked(writtenData))
+        require(
+            a.currentRootHash == rootHash, "Write word root doesn't match"
         );
+
+        bytes32 newLeaf = setBytes8ToBytes32AtOffset(
+            solidityUint64ToMachineWord(newValue), oldLeaf, offset
+        );
+
+        bytes32 newRootHash = a.buffer.getRoot(region, keccak256(abi.encodePacked(newLeaf)));
+        a.currentRootHash = newRootHash;
     }
 
     function getBytes8FromBytes32AtOffset(bytes32 source, uint64 offset)
@@ -179,6 +186,21 @@ library AccessLogs {
         return bytes8(source << (offset << Memory.LOG2_WORD));
     }
 
+    function setBytes8ToBytes32AtOffset(bytes8 word, bytes32 leaf, uint64 offset)
+        internal
+        pure
+        returns (bytes32)
+    {
+        uint256 wordOffset = offset << Memory.LOG2_WORD;
+        bytes32 toWrite = bytes32(word) >> wordOffset;
+
+        bytes32 wordMask = bytes32(~bytes8(0));
+        bytes32 mask = ~(wordMask >> wordOffset);
+
+        return (leaf & mask) | toWrite;
+    }
+
+
     //:#else
 
     /// @dev This library mocks the `templates/AccessLogs.sol` yet with a very different implementation.