feat: PXE db contract store

cspell.json

@@ -77,6 +77,7 @@
     "demonomorphizes",
     "demonomorphizing",
     "deregistration",
+    "desynchronization",
     "devex",
     "devnet",
     "devs",
@@ -176,6 +177,7 @@
     "noirc",
     "noirup",
     "nullifer",
+    "Nullifiable",
     "offchain",
     "onchain",
     "opentelemetry",
@@ -276,6 +278,7 @@
     "unexclude",
     "unexcluded",
     "unfinalised",
+    "unnullify",
     "unprefixed",
     "unshift",
     "unshifted",

noir-projects/aztec-nr/aztec/src/note/note_interface.nr

@@ -1,6 +1,6 @@
 use crate::context::PrivateContext;
 use crate::note::note_header::NoteHeader;
-use dep::protocol_types::traits::{Empty, Serialize};
+use dep::protocol_types::traits::Empty;
 
 pub trait NoteProperties<T> {
     fn properties() -> T;

noir-projects/aztec-nr/aztec/src/oracle/mod.nr

@@ -19,6 +19,7 @@ pub mod block_header;
 pub mod notes;
 pub mod storage;
 pub mod logs;
+pub mod pxe_db;
 pub mod returns;
 
 // debug_log oracle is used by both noir-protocol-circuits and this crate and for this reason we just re-export it

noir-projects/aztec-nr/aztec/src/oracle/pxe_db.nr

@@ -0,0 +1,83 @@
+use protocol_types::{address::AztecAddress, traits::{Deserialize, Serialize}};
+
+#[oracle(store)]
+unconstrained fn store_oracle<let N: u32>(
+    contract_address: AztecAddress,
+    key: Field,
+    values: [Field; N],
+) {}
+
+/// Store a value of type T that implements Serialize in local PXE database. The data is scoped to the current
+/// contract. If the data under the key already exists, it is overwritten.
+pub unconstrained fn store<T, let N: u32>(contract_address: AztecAddress, key: Field, value: T)
+where
+    T: Serialize<N>,
+{
+    let serialized = value.serialize();
+    store_oracle(contract_address, key, serialized);
+}
+
+/// Load data from local PXE database. We pass in `t_size` as a parameter to have the information of how many fields
+/// we need to pad if the key does not exist (note that the actual response size is `t_size + 1` as the Option prefixes
+/// the response with a boolean indicating if the data exists).
+///
+/// Note that we need to return an Option<[Field; N]> as we cannot return an Option<T> directly. This is because then
+/// the shape of T would affect the expected oracle response (e.g. if we were returning a struct of 3 u32 values
+/// then the expected response shape would be 3 single items. If instead we had a struct containing
+/// `u32, [Field;10], u32`, then the expected shape would be single, array, single.).
+#[oracle(load)]
+unconstrained fn load_oracle<let N: u32>(
+    contract_address: AztecAddress,
+    key: Field,
+    t_size: u32,
+) -> Option<[Field; N]> {}
+
+/// Load a value of type T that implements Deserialize from local PXE database. The data is scoped to the current
+/// contract. If the key does not exist, Option::none() is returned.
+pub unconstrained fn load<T, let N: u32>(contract_address: AztecAddress, key: Field) -> Option<T>
+where
+    T: Deserialize<N>,
+{
+    let serialized_option = load_oracle::<N>(contract_address, key, N);
+    serialized_option.map(|arr| Deserialize::deserialize(arr))
+}
+
+mod test {
+    use crate::{
+        oracle::{pxe_db::{load, store}, random::random},
+        test::{helpers::test_environment::TestEnvironment, mocks::mock_struct::MockStruct},
+    };
+
+    #[test]
+    unconstrained fn stores_loads_and_overwrites_data() {
+        let env = TestEnvironment::new();
+
+        let contract_address = env.contract_address();
+        let key = random();
+        let value = MockStruct::new(5, 6);
+        store(contract_address, key, value);
+
+        let loaded_value: MockStruct = load(contract_address, key).unwrap();
+
+        assert(loaded_value == value, "Stored and loaded values should be equal");
+
+        // Now we test that the value gets overwritten correctly.
+        let new_value = MockStruct::new(7, 8);
+        store(contract_address, key, new_value);
+
+        let loaded_value: MockStruct = load(contract_address, key).unwrap();
+
+        assert(loaded_value == new_value, "Stored and loaded values should be equal");
+    }
+
+    #[test]
+    unconstrained fn load_non_existent_key() {
+        let env = TestEnvironment::new();
+
+        let contract_address = env.contract_address();
+        let key = random();
+        let loaded_value: Option<MockStruct> = load(contract_address, key);
+
+        assert(loaded_value == Option::none(), "Value should not exist");
+    }
+}

noir-projects/noir-contracts/contracts/test_contract/src/main.nr

@@ -32,7 +32,8 @@ contract Test {
             note_getter::{get_notes, view_notes},
             note_getter_options::NoteStatus,
         },
-        oracle::random::random,
+        oracle::pxe_db,
+        test::mocks::mock_struct::MockStruct,
         utils::comparison::Comparator,
     };
     use dep::token_portal_content_hash_lib::{
@@ -458,16 +459,21 @@ contract Test {
         constant.value
     }
 
+    unconstrained fn store_in_pxe_db(key: Field, arbitrary_struct: MockStruct) {
+        pxe_db::store(context.this_address(), key, arbitrary_struct);
+    }
+
+    unconstrained fn load_from_pxe_db(key: Field) -> pub [Field; 2] {
+        let maybe_arbitrary_struct: Option<MockStruct> = pxe_db::load(context.this_address(), key);
+        let arbitrary_struct = maybe_arbitrary_struct.unwrap_or(MockStruct::new(0, 0));
+        arbitrary_struct.serialize()
+    }
+
     #[private]
     fn test_nullifier_key_freshness(address: AztecAddress, public_nullifying_key: Point) {
         assert_eq(get_public_keys(address).npk_m.inner, public_nullifying_key);
     }
 
-    // Purely exists for testing
-    unconstrained fn get_random(kinda_seed: Field) -> pub Field {
-        kinda_seed * random()
-    }
-
     pub struct DummyNote {
         amount: Field,
         secret_hash: Field,

yarn-project/end-to-end/src/e2e_pxe_db.test.ts

@@ -0,0 +1,50 @@
+import { Fr, type Wallet } from '@aztec/aztec.js';
+import { TestContract } from '@aztec/noir-contracts.js/Test';
+
+import { jest } from '@jest/globals';
+
+import { setup } from './fixtures/utils.js';
+
+const TIMEOUT = 120_000;
+
+// TODO(#10724): Nuke this once the linked issue is implemented (then the code will be well-tested). There is also
+// a TXE test in `pxe_db.nr` but I decided to keep this ugly test around as it tests the PXE oracle callback handler
+// (which is not tested by the TXE test). Dont't forget to remove `store_in_pxe_db` and `load_from_pxe_db` from
+// the test contract when removing this test.
+describe('PXE db', () => {
+  jest.setTimeout(TIMEOUT);
+
+  let teardown: () => Promise<void>;
+
+  let testContract: TestContract;
+
+  beforeAll(async () => {
+    let wallet: Wallet;
+    ({ teardown, wallet } = await setup(1));
+    testContract = await TestContract.deploy(wallet).send().deployed();
+  });
+
+  afterAll(() => teardown());
+
+  it('stores and loads data', async () => {
+    // In this test we feed arbitrary struct to a test contract, the test contract stores it in the PXE db and then
+    // we load it back.
+    const arbitraryStruct = {
+      a: Fr.random(),
+      b: Fr.random(),
+    };
+
+    const key = 6n;
+    await testContract.methods.store_in_pxe_db(key, arbitraryStruct).simulate();
+
+    // Now we try to load the data back from the PXE db.
+    const expectedReturnValue = [arbitraryStruct.a, arbitraryStruct.b].map(v => v.toBigInt());
+    expect(await testContract.methods.load_from_pxe_db(key).simulate()).toEqual(expectedReturnValue);
+  });
+
+  it('handles non-existent data', async () => {
+    // In this test we try to load a key from the PXE db that does not exist. We should get an array of zeros.
+    const key = 7n;
+    expect(await testContract.methods.load_from_pxe_db(key).simulate()).toEqual([0n, 0n]);
+  });
+});

yarn-project/pxe/src/database/kv_pxe_database.ts

@@ -12,6 +12,7 @@ import { type ContractArtifact, FunctionSelector, FunctionType } from '@aztec/fo
 import { toBufferBE } from '@aztec/foundation/bigint-buffer';
 import { Fr } from '@aztec/foundation/fields';
 import { toArray } from '@aztec/foundation/iterable';
+import { type LogFn, createDebugOnlyLogger } from '@aztec/foundation/log';
 import {
   type AztecAsyncArray,
   type AztecAsyncKVStore,
@@ -63,6 +64,11 @@ export class KVPxeDatabase implements PxeDatabase {
   #taggingSecretIndexesForSenders: AztecAsyncMap<string, number>;
   #taggingSecretIndexesForRecipients: AztecAsyncMap<string, number>;
 
+  // Arbitrary data stored by contracts. Key is computed as `${contractAddress}:${key}`
+  #contractStore: AztecAsyncMap<string, Buffer>;
+
+  debug: LogFn;
+
   protected constructor(private db: AztecAsyncKVStore) {
     this.#db = db;
 
@@ -100,6 +106,10 @@ export class KVPxeDatabase implements PxeDatabase {
 
     this.#taggingSecretIndexesForSenders = db.openMap('tagging_secret_indexes_for_senders');
     this.#taggingSecretIndexesForRecipients = db.openMap('tagging_secret_indexes_for_recipients');
+
+    this.#contractStore = db.openMap('contract_store');
+
+    this.debug = createDebugOnlyLogger('aztec:kv-pxe-database');
   }
 
   public static async create(db: AztecAsyncKVStore): Promise<KVPxeDatabase> {
@@ -611,4 +621,24 @@ export class KVPxeDatabase implements PxeDatabase {
       await Promise.all(senders.map(sender => this.#taggingSecretIndexesForSenders.delete(sender)));
     });
   }
+
+  async store(contract: AztecAddress, key: Fr, values: Fr[]): Promise<void> {
+    const dataKey = `${contract.toString()}:${key.toString()}`;
+    const dataBuffer = Buffer.concat(values.map(value => value.toBuffer()));
+    await this.#contractStore.set(dataKey, dataBuffer);
+  }
+
+  async load(contract: AztecAddress, key: Fr): Promise<Fr[] | null> {
+    const dataKey = `${contract.toString()}:${key.toString()}`;
+    const dataBuffer = await this.#contractStore.getAsync(dataKey);
+    if (!dataBuffer) {
+      this.debug(`Data not found for contract ${contract.toString()} and key ${key.toString()}`);
+      return null;
+    }
+    const values: Fr[] = [];
+    for (let i = 0; i < dataBuffer.length; i += Fr.SIZE_IN_BYTES) {
+      values.push(Fr.fromBuffer(dataBuffer.subarray(i, i + Fr.SIZE_IN_BYTES)));
+    }
+    return values;
+  }
 }

yarn-project/pxe/src/database/pxe_database.ts

@@ -213,4 +213,22 @@ export interface PxeDatabase extends ContractArtifactDatabase, ContractInstanceD
    * is also required to deal with chain reorgs.
    */
   resetNoteSyncData(): Promise<void>;
+
+  /**
+   * Used by contracts during execution to store arbitrary data in the local PXE database. The data is siloed/scoped
+   * to a specific `contract`.
+   * @param contract - An address of a contract that is requesting to store the data.
+   * @param key - A field element representing the key to store the data under.
+   * @param values - An array of field elements representing the data to store.
+   */
+  store(contract: AztecAddress, key: Fr, values: Fr[]): Promise<void>;
+
+  /**
+   * Used by contracts during execution to load arbitrary data from the local PXE database. The data is siloed/scoped
+   * to a specific `contract`.
+   * @param contract - An address of a contract that is requesting to load the data.
+   * @param key - A field element representing the key under which to load the data..
+   * @returns An array of field elements representing the stored data or `null` if no data is stored under the key.
+   */
+  load(contract: AztecAddress, key: Fr): Promise<Fr[] | null>;
 }

yarn-project/pxe/src/database/pxe_database_test_suite.ts

@@ -405,5 +405,66 @@ export function describePxeDatabase(getDatabase: () => PxeDatabase) {
         await expect(database.getContractInstance(address)).resolves.toEqual(instance);
       });
     });
+
+    describe('contract store', () => {
+      let contract: AztecAddress;
+
+      beforeEach(() => {
+        // Setup mock contract address
+        contract = AztecAddress.random();
+      });
+
+      it('stores and loads a single value', async () => {
+        const key = new Fr(1);
+        const values = [new Fr(42)];
+
+        await database.store(contract, key, values);
+        const result = await database.load(contract, key);
+        expect(result).toEqual(values);
+      });
+
+      it('stores and loads multiple values', async () => {
+        const key = new Fr(1);
+        const values = [new Fr(42), new Fr(43), new Fr(44)];
+
+        await database.store(contract, key, values);
+        const result = await database.load(contract, key);
+        expect(result).toEqual(values);
+      });
+
+      it('overwrites existing values', async () => {
+        const key = new Fr(1);
+        const initialValues = [new Fr(42)];
+        const newValues = [new Fr(100)];
+
+        await database.store(contract, key, initialValues);
+        await database.store(contract, key, newValues);
+
+        const result = await database.load(contract, key);
+        expect(result).toEqual(newValues);
+      });
+
+      it('stores values for different contracts independently', async () => {
+        const anotherContract = AztecAddress.random();
+        const key = new Fr(1);
+        const values1 = [new Fr(42)];
+        const values2 = [new Fr(100)];
+
+        await database.store(contract, key, values1);
+        await database.store(anotherContract, key, values2);
+
+        const result1 = await database.load(contract, key);
+        const result2 = await database.load(anotherContract, key);
+
+        expect(result1).toEqual(values1);
+        expect(result2).toEqual(values2);
+      });
+
+      it('returns null for non-existent keys', async () => {
+        const key = Fr.random();
+        const result = await database.load(contract, key);
+        expect(result).toBeNull();
+      });
+    });
   });
 }

yarn-project/pxe/src/simulator_oracle/index.ts

@@ -666,4 +666,26 @@ export class SimulatorOracle implements DBOracle {
       });
     }
   }
+
+  /**
+   * Used by contracts during execution to store arbitrary data in the local PXE database. The data is siloed/scoped
+   * to a specific `contract`.
+   * @param contract - An address of a contract that is requesting to store the data.
+   * @param key - A field element representing the key to store the data under.
+   * @param values - An array of field elements representing the data to store.
+   */
+  store(contract: AztecAddress, key: Fr, values: Fr[]): Promise<void> {
+    return this.db.store(contract, key, values);
+  }
+
+  /**
+   * Used by contracts during execution to load arbitrary data from the local PXE database. The data is siloed/scoped
+   * to a specific `contract`.
+   * @param contract - An address of a contract that is requesting to load the data.
+   * @param key - A field element representing the key under which to load the data..
+   * @returns An array of field elements representing the stored data or `null` if no data is stored under the key.
+   */
+  load(contract: AztecAddress, key: Fr): Promise<Fr[] | null> {
+    return this.db.load(contract, key);
+  }
 }