From 4e67bb1a03777879707d46a3206f4b32546c9066 Mon Sep 17 00:00:00 2001
From: benesjan <janbenes1234@gmail.com>
Date: Fri, 10 May 2024 06:44:56 +0000
Subject: [PATCH] WIP

---
 .../src/client/private_execution.test.ts      | 45 +++++--------------
 1 file changed, 11 insertions(+), 34 deletions(-)

diff --git a/yarn-project/simulator/src/client/private_execution.test.ts b/yarn-project/simulator/src/client/private_execution.test.ts
index ed40e9f8bee..02e3d6b9588 100644
--- a/yarn-project/simulator/src/client/private_execution.test.ts
+++ b/yarn-project/simulator/src/client/private_execution.test.ts
@@ -15,6 +15,7 @@ import {
   FunctionData,
   GasSettings,
   GeneratorIndex,
+  type GrumpkinPrivateKey,
   Header,
   L1_TO_L2_MSG_TREE_HEIGHT,
   NOTE_HASH_TREE_HEIGHT,
@@ -81,9 +82,8 @@ describe('Private Execution test suite', () => {
   let ownerCompleteAddress: CompleteAddress;
   let recipientCompleteAddress: CompleteAddress;
 
-  // TODO(#5834): Nuke the following once complete address is refactored
-  let allOwnerKeys: any;
-  let allRecipientKeys: any;
+  let ownerMasterNullifierSecretKey: GrumpkinPrivateKey;
+  let recipientMasterNullifierSecretKey: GrumpkinPrivateKey;
 
   const treeHeights: { [name: string]: number } = {
     noteHash: NOTE_HASH_TREE_HEIGHT,
@@ -177,11 +177,11 @@ describe('Private Execution test suite', () => {
 
     const ownerPartialAddress = Fr.random();
     ownerCompleteAddress = CompleteAddress.fromSecretKeyAndPartialAddress(ownerSk, ownerPartialAddress);
-    allOwnerKeys = deriveKeys(ownerSk);
+    ownerMasterNullifierSecretKey = deriveKeys(ownerSk).masterNullifierSecretKey;
 
     const recipientPartialAddress = Fr.random();
     recipientCompleteAddress = CompleteAddress.fromSecretKeyAndPartialAddress(recipientSk, recipientPartialAddress);
-    allRecipientKeys = deriveKeys(recipientSk);
+    recipientMasterNullifierSecretKey = deriveKeys(recipientSk).masterNullifierSecretKey;
 
     owner = ownerCompleteAddress.address;
     recipient = recipientCompleteAddress.address;
@@ -193,17 +193,14 @@ describe('Private Execution test suite', () => {
     oracle.getNullifierKeys.mockImplementation((accountAddress: AztecAddress, contractAddress: AztecAddress) => {
       if (accountAddress.equals(ownerCompleteAddress.address)) {
         return Promise.resolve({
-          masterNullifierPublicKey: allOwnerKeys.masterNullifierPublicKey,
-          appNullifierSecretKey: computeAppNullifierSecretKey(allOwnerKeys.masterNullifierSecretKey, contractAddress),
+          masterNullifierPublicKey: ownerCompleteAddress.masterNullifierPublicKey,
+          appNullifierSecretKey: computeAppNullifierSecretKey(ownerMasterNullifierSecretKey, contractAddress),
         });
       }
       if (accountAddress.equals(recipientCompleteAddress.address)) {
         return Promise.resolve({
-          masterNullifierPublicKey: allRecipientKeys.masterNullifierPublicKey,
-          appNullifierSecretKey: computeAppNullifierSecretKey(
-            allRecipientKeys.masterNullifierSecretKey,
-            contractAddress,
-          ),
+          masterNullifierPublicKey: recipientCompleteAddress.masterNullifierPublicKey,
+          appNullifierSecretKey: computeAppNullifierSecretKey(recipientMasterNullifierSecretKey, contractAddress),
         });
       }
       throw new Error(`Unknown address ${accountAddress}`);
@@ -223,26 +220,6 @@ describe('Private Execution test suite', () => {
       }
       throw new Error(`Unknown address ${address}`);
     });
-    // TODO(#5834): The following oracle should be unnecessary
-    oracle.getPublicKeysForAddress.mockImplementation((address: AztecAddress) => {
-      if (address.equals(owner)) {
-        return Promise.resolve([
-          allOwnerKeys.masterNullifierPublicKey,
-          allOwnerKeys.masterIncomingViewingPublicKey,
-          allOwnerKeys.masterOutgoingViewingPublicKey,
-          allOwnerKeys.masterTaggingPublicKey,
-        ]);
-      }
-      if (address.equals(recipient)) {
-        return Promise.resolve([
-          allRecipientKeys.masterNullifierPublicKey,
-          allRecipientKeys.masterIncomingViewingPublicKey,
-          allRecipientKeys.masterOutgoingViewingPublicKey,
-          allRecipientKeys.masterTaggingPublicKey,
-        ]);
-      }
-      throw new Error(`Unknown address ${address}`);
-    });
     // This oracle gets called when reading ivpk_m from key registry --> we return zero witness indicating that
     // the keys were not registered. This triggers non-registered keys flows in which getCompleteAddress oracle
     // gets called and we constrain the result by hashing address preimage and checking it matches.
@@ -973,7 +950,7 @@ describe('Private Execution test suite', () => {
       const nullifier = result.callStackItem.publicInputs.newNullifiers[0];
       const expectedNullifier = poseidon2Hash([
         innerNoteHash,
-        computeAppNullifierSecretKey(allOwnerKeys.masterNullifierSecretKey, contractAddress),
+        computeAppNullifierSecretKey(ownerMasterNullifierSecretKey, contractAddress),
         GeneratorIndex.NOTE_NULLIFIER,
       ]);
       expect(nullifier.value).toEqual(expectedNullifier);
@@ -1052,7 +1029,7 @@ describe('Private Execution test suite', () => {
       const nullifier = execGetThenNullify.callStackItem.publicInputs.newNullifiers[0];
       const expectedNullifier = poseidon2Hash([
         innerNoteHash,
-        computeAppNullifierSecretKey(allOwnerKeys.masterNullifierSecretKey, contractAddress),
+        computeAppNullifierSecretKey(ownerMasterNullifierSecretKey, contractAddress),
         GeneratorIndex.NOTE_NULLIFIER,
       ]);
       expect(nullifier.value).toEqual(expectedNullifier);