From 1ba38d9d0fbed61495eea614f04131eec2f10d21 Mon Sep 17 00:00:00 2001
From: Marcos Pernambuco Motta <1091485+mpernambuco@users.noreply.github.com>
Date: Thu, 1 Aug 2024 16:30:33 -0300
Subject: [PATCH] fixup! feat: increase leaf size

---
 src/AccessLogs.sol                | 8 ++++----
 src/Memory.sol                    | 7 +++++--
 templates/AccessLogs.sol.template | 6 ++----
 test/AccessLogs.t.sol             | 7 ++++---
 4 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/src/AccessLogs.sol b/src/AccessLogs.sol
index 19c6e88..5ace90e 100644
--- a/src/AccessLogs.sol
+++ b/src/AccessLogs.sol
@@ -97,8 +97,8 @@ library AccessLogs {
         bytes32 readData = a.buffer.consumeBytes32();
         bytes32 valHash = keccak256(abi.encodePacked(readData));
 
-        Memory.PhysicalAddress leafAddress = readAddress.truncateToLeaf();
-        uint64 offset = readAddress.minus(leafAddress);
+        (Memory.PhysicalAddress leafAddress, uint64 offset) =
+            readAddress.truncateToLeaf();
         bytes8 readValue = bytes8(readData << (offset << 3));
 
         bytes32 expectedValHash =
@@ -145,8 +145,8 @@ library AccessLogs {
     ) internal pure {
         bytes32 writtenData = a.buffer.consumeBytes32();
 
-        Memory.PhysicalAddress leafAddress = writeAddress.truncateToLeaf();
-        uint64 offset = writeAddress.minus(leafAddress);
+        (Memory.PhysicalAddress leafAddress, uint64 offset) =
+            writeAddress.truncateToLeaf();
         uint64 expectedNewValue =
             machineWordToSolidityUint64(bytes8(writtenData << (offset << 3)));
 
diff --git a/src/Memory.sol b/src/Memory.sol
index 17a40f2..903fa62 100644
--- a/src/Memory.sol
+++ b/src/Memory.sol
@@ -110,6 +110,7 @@ library Memory {
 
     type AlignedSize is uint8;
 
+    uint8 constant LOG2_WORD = 3;
     uint8 constant LOG2_LEAF = 5;
     uint8 constant LOG2_MAX_SIZE = 64 - LOG2_LEAF;
     uint64 constant LEAF_MASK = uint64(1 << LOG2_LEAF) - 1;
@@ -142,10 +143,12 @@ library Memory {
     function truncateToLeaf(PhysicalAddress addr)
         internal
         pure
-        returns (PhysicalAddress)
+        returns (PhysicalAddress, uint64)
     {
         uint64 r = Memory.PhysicalAddress.unwrap(addr) & ~LEAF_MASK;
-        return Memory.PhysicalAddress.wrap(r);
+        PhysicalAddress truncated = Memory.PhysicalAddress.wrap(r);
+        uint64 offset = minus(addr, truncated);
+        return (truncated, offset);
     }
 
     function minus(PhysicalAddress lhs, PhysicalAddress rhs)
diff --git a/templates/AccessLogs.sol.template b/templates/AccessLogs.sol.template
index 2c152a6..67fdcc3 100644
--- a/templates/AccessLogs.sol.template
+++ b/templates/AccessLogs.sol.template
@@ -110,8 +110,7 @@ library AccessLogs {
         bytes32 readData = a.buffer.consumeBytes32();
         bytes32 valHash = keccak256(abi.encodePacked(readData));
 
-        Memory.PhysicalAddress leafAddress = readAddress.truncateToLeaf();
-        uint64 offset = readAddress.minus(leafAddress);
+        (Memory.PhysicalAddress leafAddress, uint64 offset) = readAddress.truncateToLeaf();
         bytes8 readValue = bytes8(readData << (offset << 3));
 
         bytes32 expectedValHash =
@@ -158,8 +157,7 @@ library AccessLogs {
     ) internal pure {
         bytes32 writtenData = a.buffer.consumeBytes32();
 
-        Memory.PhysicalAddress leafAddress = writeAddress.truncateToLeaf();
-        uint64 offset = writeAddress.minus(leafAddress);
+        (Memory.PhysicalAddress leafAddress,  uint64 offset) = writeAddress.truncateToLeaf();
         uint64 expectedNewValue = machineWordToSolidityUint64(bytes8(writtenData << (offset << 3)));
 
         require(newValue == expectedNewValue, "Access log value does not contain the expected written value");
diff --git a/test/AccessLogs.t.sol b/test/AccessLogs.t.sol
index 66c278f..e3a73c8 100644
--- a/test/AccessLogs.t.sol
+++ b/test/AccessLogs.t.sol
@@ -147,8 +147,8 @@ contract AccessLogsTest is Test {
     {
         bytes32 b32 = makeLeaf(word, position);
         Buffer.Context memory buffer =
-            Buffer.Context(new bytes((59 << 5) + 32 + 32), 0);
-        buffer.writeBytes32(b32); // leaf containing the readd word
+            Buffer.Context(new bytes((59 << Memory.LOG2_LEAF) + 32 + 32), 0);
+        buffer.writeBytes32(b32); // leaf containing the read word
 
         for (uint256 i = 0; i < 60; i++) {
             buffer.writeBytes32(hashes[i]);
@@ -170,7 +170,8 @@ contract AccessLogsTest is Test {
     }
 
     function rootFromHashes(bytes32 drive) private view returns (bytes32) {
-        Buffer.Context memory buffer = Buffer.Context(new bytes(59 << 5), 0);
+        Buffer.Context memory buffer =
+            Buffer.Context(new bytes(59 << Memory.LOG2_LEAF), 0);
 
         for (uint256 i = 0; i < 59; i++) {
             buffer.writeBytes32(hashes[i + 1]);