From c07cae3619fb7b0ded139dfefec5d7bc0f2857a7 Mon Sep 17 00:00:00 2001
From: Maddiaa0 <47148561+Maddiaa0@users.noreply.github.com>
Date: Tue, 30 Jan 2024 22:59:38 +0000
Subject: [PATCH] fix: update testys

---
 .../src/avm/journal/journal.test.ts           | 13 +++++++++--
 .../acir-simulator/src/avm/journal/journal.ts | 22 +++++++++----------
 2 files changed, 22 insertions(+), 13 deletions(-)

diff --git a/yarn-project/acir-simulator/src/avm/journal/journal.test.ts b/yarn-project/acir-simulator/src/avm/journal/journal.test.ts
index 4eee36fd3591..1588712e1c73 100644
--- a/yarn-project/acir-simulator/src/avm/journal/journal.test.ts
+++ b/yarn-project/acir-simulator/src/avm/journal/journal.test.ts
@@ -61,7 +61,7 @@ describe('journal', () => {
       expect(cachedResult).toEqual(cachedValue);
     });
 
-    it('When reading from storage, should check the cache first, and be appended to read journal', async () => {
+    it('When reading from storage, should check the cache first, and be appended to read/write journal', async () => {
       // Store a different value in storage vs the cache, and make sure the cache is returned
       const contractAddress = new Fr(1);
       const key = new Fr(2);
@@ -82,10 +82,14 @@ describe('journal', () => {
       expect(cachedResult).toEqual(cachedValue);
 
       // We expect the journal to store the access in [storedVal, cachedVal] - [time0, time1]
-      const { storageReads }: JournalData = journal.flush();
+      const { storageReads, storageWrites }: JournalData = journal.flush();
       const contractReads = storageReads.get(contractAddress.toBigInt());
       const keyReads = contractReads?.get(key.toBigInt());
       expect(keyReads).toEqual([storedValue, cachedValue]);
+
+      const contractWrites = storageWrites.get(contractAddress.toBigInt());
+      const keyWrites = contractWrites?.get(key.toBigInt());
+      expect(keyWrites).toEqual([cachedValue]);
     });
   });
 
@@ -162,6 +166,11 @@ describe('journal', () => {
     const slotReads = contractReads?.get(key.toBigInt());
     expect(slotReads).toEqual([value, valueT1]);
 
+    // We first write value from t0, then value from t1
+    const contractWrites = journalUpdates.storageReads.get(contractAddress.toBigInt());
+    const slotWrites = contractWrites?.get(key.toBigInt());
+    expect(slotWrites).toEqual([value, valueT1]);
+
     expect(journalUpdates.newNoteHashes).toEqual([commitment, commitmentT1]);
     expect(journalUpdates.newLogs).toEqual([logs, logsT1]);
     expect(journalUpdates.newL1Messages).toEqual([logs, logsT1]);
diff --git a/yarn-project/acir-simulator/src/avm/journal/journal.ts b/yarn-project/acir-simulator/src/avm/journal/journal.ts
index 6c0b591522ea..b1b330d94604 100644
--- a/yarn-project/acir-simulator/src/avm/journal/journal.ts
+++ b/yarn-project/acir-simulator/src/avm/journal/journal.ts
@@ -94,7 +94,7 @@ export class AvmJournal {
       this.currentStorageValue.set(contractAddress.toBigInt(), contractMap);
     }
     contractMap.set(key.toBigInt(), value);
-    
+
     // We want to keep track of all performed writes in the journal
     this.journalWrite(contractAddress, key, value);
   }
@@ -137,7 +137,7 @@ export class AvmJournal {
     let contractMap = map.get(contractAddress.toBigInt());
     if (!contractMap) {
       contractMap = new Map<bigint, Array<Fr>>();
-      this.storageReads.set(contractAddress.toBigInt(), contractMap);
+      map.set(contractAddress.toBigInt(), contractMap);
     }
 
     let accessArray = contractMap.get(key.toBigInt());
@@ -149,9 +149,9 @@ export class AvmJournal {
   }
 
   // Create an instance of journalUpdate that appends to the read array
-  private journalRead = this.journalUpdate.bind(this, this.storageReads); 
+  private journalRead = this.journalUpdate.bind(this, this.storageReads);
   // Create an instance of journalUpdate that appends to the writes array
-  private journalWrite = this.journalUpdate.bind(this, this.storageWrites); 
+  private journalWrite = this.journalUpdate.bind(this, this.storageWrites);
 
   public writeNoteHash(noteHash: Fr) {
     this.newNoteHashes.push(noteHash);
@@ -212,7 +212,7 @@ export class AvmJournal {
 }
 
 /**
- * Merges two contract write maps together
+ * Merges two contract current value together
  * Where childMap keys will take precedent over the hostMap
  * The assumption being that the child map is created at a later time
  * And thus contains more up to date information
@@ -226,7 +226,7 @@ function mergeCurrentValueMaps(hostMap: ContractValueMap, childMap: ContractValu
     if (!map1Value) {
       hostMap.set(key, value);
     } else {
-      mergeStorageWriteMaps(map1Value, value);
+      mergeStorageCurrentValueMaps(map1Value, value);
     }
   }
 }
@@ -235,15 +235,15 @@ function mergeCurrentValueMaps(hostMap: ContractValueMap, childMap: ContractValu
  * @param hostMap - The map to be merge into
  * @param childMap - The map to be merged from
  */
-function mergeStorageWriteMaps(hostMap: StorageValueMap, childMap: StorageValueMap) {
+function mergeStorageCurrentValueMaps(hostMap: StorageValueMap, childMap: StorageValueMap) {
   for (const [key, value] of childMap) {
     hostMap.set(key, value);
   }
 }
 
 /**
- * Merges two contract read maps together
- * For read maps, we just append the childMap into the host map, as the order of reads is important
+ * Merges two contract journalling maps together
+ * For read maps, we just append the childMap arrays into the host map arrays, as the order is important
  *
  * @param hostMap - The map to be merged into
  * @param childMap - The map to be merged from
@@ -254,7 +254,7 @@ function mergeContractJournalMaps(hostMap: ContractJournalMap, childMap: Contrac
     if (!map1Value) {
       hostMap.set(key, value);
     } else {
-      mergeStorageReadMaps(map1Value, value);
+      mergeStorageJournalMaps(map1Value, value);
     }
   }
 }
@@ -263,7 +263,7 @@ function mergeContractJournalMaps(hostMap: ContractJournalMap, childMap: Contrac
  * @param hostMap - The map to be merge into
  * @param childMap - The map to be merged from
  */
-function mergeStorageReadMaps(hostMap: StorageJournalMap, childMap: StorageJournalMap) {
+function mergeStorageJournalMaps(hostMap: StorageJournalMap, childMap: StorageJournalMap) {
   for (const [key, value] of childMap) {
     const readArr = hostMap.get(key);
     if (!readArr) {