From 3d0a9e25c3a7b8b51110c5b0c127be5e8f105a03 Mon Sep 17 00:00:00 2001
From: sklppy88 <esau@aztecprotocol.com>
Date: Wed, 6 Nov 2024 09:36:28 +0000
Subject: [PATCH] init

---
 .../src/database/contracts/contract_artifact_db.ts  |  1 +
 yarn-project/pxe/src/database/kv_pxe_database.ts    | 12 +++++++++++-
 .../pxe/src/database/pxe_database_test_suite.ts     | 13 +++++++++++++
 3 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/yarn-project/pxe/src/database/contracts/contract_artifact_db.ts b/yarn-project/pxe/src/database/contracts/contract_artifact_db.ts
index 85f726d0973..3f6e6bb2311 100644
--- a/yarn-project/pxe/src/database/contracts/contract_artifact_db.ts
+++ b/yarn-project/pxe/src/database/contracts/contract_artifact_db.ts
@@ -9,6 +9,7 @@ export interface ContractArtifactDatabase {
    * Adds a new contract artifact to the database or updates an existing one.
    * @param id - Id of the corresponding contract class.
    * @param contract - Contract artifact to add.
+   * @throws - If there are duplicate private function selectors.
    */
   addContractArtifact(id: Fr, contract: ContractArtifact): Promise<void>;
   /**
diff --git a/yarn-project/pxe/src/database/kv_pxe_database.ts b/yarn-project/pxe/src/database/kv_pxe_database.ts
index 9c6f25092d0..26096465571 100644
--- a/yarn-project/pxe/src/database/kv_pxe_database.ts
+++ b/yarn-project/pxe/src/database/kv_pxe_database.ts
@@ -9,7 +9,7 @@ import {
   type TaggingSecret,
   computePoint,
 } from '@aztec/circuits.js';
-import { type ContractArtifact } from '@aztec/foundation/abi';
+import { type ContractArtifact, FunctionSelector, FunctionType } from '@aztec/foundation/abi';
 import { toBufferBE } from '@aztec/foundation/bigint-buffer';
 import { Fr } from '@aztec/foundation/fields';
 import {
@@ -130,6 +130,16 @@ export class KVPxeDatabase implements PxeDatabase {
   }
 
   public async addContractArtifact(id: Fr, contract: ContractArtifact): Promise<void> {
+    const privateSelectors = contract.functions
+      .filter(functionArtifact => functionArtifact.functionType === FunctionType.PRIVATE)
+      .map(privateFunctionArtifact =>
+        FunctionSelector.fromNameAndParameters(privateFunctionArtifact.name, privateFunctionArtifact.parameters).toString(),
+      );
+
+    if (privateSelectors.length !== new Set(privateSelectors).size) {
+      throw new Error('Repeated function selectors of private functions');
+    }
+
     await this.#contractArtifacts.set(id.toString(), contractArtifactToBuffer(contract));
   }
 
diff --git a/yarn-project/pxe/src/database/pxe_database_test_suite.ts b/yarn-project/pxe/src/database/pxe_database_test_suite.ts
index 1fefa614bad..96df11e05a2 100644
--- a/yarn-project/pxe/src/database/pxe_database_test_suite.ts
+++ b/yarn-project/pxe/src/database/pxe_database_test_suite.ts
@@ -11,12 +11,14 @@ import { makeHeader } from '@aztec/circuits.js/testing';
 import { randomInt } from '@aztec/foundation/crypto';
 import { Fr, Point } from '@aztec/foundation/fields';
 import { BenchmarkingContractArtifact } from '@aztec/noir-contracts.js/Benchmarking';
+import { TestContractArtifact } from '@aztec/noir-contracts.js/Test';
 
 import { type IncomingNoteDao } from './incoming_note_dao.js';
 import { randomIncomingNoteDao } from './incoming_note_dao.test.js';
 import { type OutgoingNoteDao } from './outgoing_note_dao.js';
 import { randomOutgoingNoteDao } from './outgoing_note_dao.test.js';
 import { type PxeDatabase } from './pxe_database.js';
+import { FunctionType } from '@aztec/foundation/abi';
 
 /**
  * A common test suite for a PXE database.
@@ -391,6 +393,17 @@ export function describePxeDatabase(getDatabase: () => PxeDatabase) {
         await expect(database.getContractArtifact(id)).resolves.toEqual(artifact);
       });
 
+      it('does not store a contract artifact with a duplicate private function selector', async () => {
+        const artifact = TestContractArtifact;
+        const index = artifact.functions.findIndex(fn => fn.functionType === FunctionType.PRIVATE);
+
+        const copiedFn = structuredClone(artifact.functions[index]);
+        artifact.functions.push(copiedFn);
+
+        const id = Fr.random();
+        await expect(database.addContractArtifact(id, artifact)).rejects.toThrow('Repeated function selectors of private functions');
+      });
+
       it('stores a contract instance', async () => {
         const address = AztecAddress.random();
         const instance = SerializableContractInstance.random().withAddress(address);