feat!: compile-time incorrect exec environment errors (#6442)

Closes #5886.

This removes the `Context` struct and makes all state variables generic
over a new `Context` type parameter, with each state variable providing
implementations for `PrivateContext`, `PublicContext` or `()` (used to
marked `unconstarined` - more on this later).

The end result is that we get compile-time errors when calling functions
that are unavailable in the current context, reduced wrapping and
unwrapping, and no obscure `explicit trap hit in brilling
'self._is_some'` runtime errors, such as in
#3123.

## How

The implementation is prety much as described in #5886, except instead
of using traits we specialize the type directly for the contexts we
know.

```rust
struct MyStateVar<Context> { context: Context }

impl MyStateVar<&mut PrivateContext> { fn private_read() { } }
```

This works because there's only a couple of them, and users are not
expected to extend them.

The macros were altered so that instead of wrapping the context object
in `Context` and then passing that, we simply pass the raw object to the
`Storage::init` function. This means that `Storage` itself is now also
generic, resulting in some unfortunate new boilerplate in the struct
declaration.

All instances of `self.context.private.unwrap()` and friends were
removed: each function is now available under the corresponding `impl`
block that is specialized for the corresponding context. We could even
rename some since `read_public` and `read_private` is no longer
required: both impls can have `read` functions since they are
effectively methods for different types, so the shared name is a
non-issue.

## Oddities

This change revelead a number of small bugs in the codebase, in the form
of uncallable functions. These were undetected since no test called
them. I'll do a pass over the entire PR and leave comments where
relevant.

## Top-level unconstrained

This PR continues the formalization of what I've been calling 'top-level
unconstrained' (i.e. an unconstrained contract function) as a
fundamental Aztec.nr concept and third execution environment, alongside
private and public execution. So far we've been accessing oracles in
these unconstrained functions without much care (sometimes for
perfomance reasons - see
#5911), but the new
stricter compile-time checks force us to be a bit more careful.

In my mind, we are arriving at the following model:
- public execution: done by the sequencer, can be simulated locally with
old data, not unlike the evm
- private execution: able to produce valid private kernel proofs with
side effects collected in the context
- top-level unconstrained execution: local computation using both
private and old public data, with certain restrictions from private exec
lifted (e.g. unbounded loops), unable to produce any kind of proofs or
reason about state changes. only useful for computing values doing
arbitrary computation over both private and public state, with zero
validation and guarantees of correctness

Private execution requires a context object a it needs to collect side
effects, but public notably does not - it simply calls oracles and gets
them to do things. In this sense, the `PublicContext` type is acting as
a marker of the current execution environment in order to prevent
developers from accidentally doing things that are invalid in public,
which would otherwise result in either transpilation error or inability
to create public kernel proofs.

This means that we may want a third `UnconstrainedContext` to act as a
similar marker for this third type (where we can e.g. call `view_notes`,
read old public state, etc.). It currently doesn't exist: we simply have
`Context::none()`, and it is defined as the absense of one of the other
contexts. Because of this, I chose to temporarily use the unit type
(`()`) to mark this environment.

Note that in some cases the different execution environments share code
paths: `view_notes` is simply `get_notes` without any constraints, and
public storage reads are performed by calling the same oracles in both
public and unconstrained. I imagine small differences will arise in the
future, specially as work on the AVM continues.

---------

Co-authored-by: esau <152162806+sklppy88@users.noreply.github.com>
Co-authored-by: thunkar <gregojquiros@gmail.com>

docs/docs/aztec/glossary/call_types.md

@@ -140,9 +140,11 @@ This is the same function that was called by privately enqueuing a call to it! P
 
 ### Top-level Unconstrained
 
-Contract functions with the `unconstrained` Noir keyword are a special type of function still under development, and their semantics will likely change in the near future. They are used to perform state queries from an off-chain client, and are never included in any transaction. No guarantees are made on the correctness of the result since they rely exclusively on unconstrained oracle calls.
+Contract functions with the `unconstrained` Noir keyword are a special type of function still under development, and their semantics will likely change in the near future. They are used to perform state queries from an off-chain client (from both private and public state!), and are never included in any transaction. No guarantees are made on the correctness of the result since the entire execution is unconstrained and heavily reliant on oracle calls.
 
-A reasonable mental model for them is that of a `view` Solidity function that is never called in any transaction, and is only ever invoked via `eth_call`. Note that in these the caller assumes that the node is acting honestly by exectuing the true contract bytecode with correct blockchain state, the same way the Aztec version assumes the oracles are returning legitimate data.
+Any programming language could be used to construct these queries, since all they do is perform arbitrary computation on data that is either publicly available from any node, or locally available from the PXE. Top-level unconstrained functions exist because they let developers utilize the rest of the contract code directly by being part of the same Noir contract, and e.g. use the same libraries, structs, etc. instead of having to rely on manual computation of storage slots, struct layout and padding, and so on.
+
+A reasonable mental model for them is that of a Solidity `view` function that can never be called in any transaction, and is only ever invoked via `eth_call`. Note that in these the caller assumes that the node is acting honestly by executing the true contract bytecode with correct blockchain state, the same way the Aztec version assumes the oracles are returning legitimate data.
 
 ### aztec.js
 

docs/docs/guides/smart_contracts/writing_contracts/common_patterns/index.md

@@ -45,6 +45,7 @@ You can't read public storage in private domain. But nevertheless reading public
 1. You pass the data as a parameter to your private method and later assert in public that the data is correct. E.g.:
 
 ```rust
+#[aztec(storage)]
 struct Storage {
    token: PublicMutable<Field>,
 }

docs/docs/migration_notes.md

@@ -8,6 +8,57 @@ Aztec is in full-speed development. Literally every version breaks compatibility
 
 ## 0.41.0
 
+
+
+### [Aztec.nr] State variable rework
+
+Aztec.nr state variables have been reworked so that calling private functions in public and vice versa is detected as an error during compilation instead of at runtime. This affects users in a number of ways:
+
+#### New compile time errors
+
+It used to be that calling a state variable method only available in public from a private function resulted in obscure runtime errors in the form of a failed `_is_some` assertion.
+
+Incorrect usage of the state variable methods now results in compile time errors. For example, given the following function:
+
+```rust
+#[aztec(public)]
+fn get_decimals() -> pub u8 {
+    storage.decimals.read_private()
+}
+```
+
+The compiler will now error out with 
+```
+Expected type SharedImmutable<_, &mut PrivateContext>, found type SharedImmutable<u8, &mut PublicContext>
+```
+
+The key component is the second generic parameter: the compiler expects a `PrivateContext` (becuse `read_private` is only available during private execution), but a `PublicContext` is being used instead (because of the `#[aztec(public)]` attribute).
+
+#### Generic parameters in `Storage`
+
+The `Storage` struct (the one marked with `#[aztec(storage)]`) should now be generic over a `Context` type, which matches the new generic parameter of all Aztec.nr libraries. This parameter is always the last generic parameter.
+
+This means that, without any additional features, we'd end up with some extra boilerplate when declaring this struct:
+
+```diff
+#[aztec(storage)]
+- struct Storage {
++ struct Storage<Context> {
+-   nonce_for_burn_approval: PublicMutable<Field>,    
++   nonce_for_burn_approval: PublicMutable<Field, Context>,
+-   portal_address: SharedImmutable<EthAddress>,
++   portal_address: SharedImmutable<EthAddress, Context>,
+-   approved_action: Map<Field, PublicMutable<bool>>,
++   approved_action: Map<Field, PublicMutable<bool, Context>, Context>,
+}
+```
+
+Because of this, the `#[aztec(storage)]` macro has been updated to **automatically inject** this `Context` generic parameter. The storage declaration does not require any changes.
+
+#### Removal of `Context`
+
+The `Context` type no longer exists. End users typically didn't use it, but if imported it needs to be deleted.
+
 ### [Aztec.nr] View functions and interface navigation
 
 It is now possible to explicitly state a function doesn't perform any state alterations (including storage, logs, nullifiers and/or messages from L2 to L1) with the `#[aztec(view)]` attribute, similarly to solidity's `view` function modifier.

docs/docs/reference/sandbox_reference/cheat_codes.md

@@ -528,7 +528,7 @@ Note: One Field element occupies a storage slot. Hence, structs with multiple fi
 #[aztec(storage)]
 struct Storage {
     ...
-    pending_shields: Set<TransparentNote, TRANSPARENT_NOTE_LEN>,
+    pending_shields: PrivateSet<TransparentNote, TRANSPARENT_NOTE_LEN>,
 }
 
 contract Token {

docs/docs/reference/smart_contract_reference/storage/index.md

@@ -9,7 +9,7 @@ To learn more about storage slots, read [this explainer](/guides/smart_contracts
 
 You control this storage in Aztec using a struct annotated with `#[aztec(storage)]`. This struct serves as the housing unit for all your smart contract's state variables - the data it needs to keep track of and maintain.
 
-These state variables come in two forms: public and private. Public variables are visible to anyone, and private variables remain hidden within the contract.
+These state variables come in two forms: [public](./public_state.md) and [private](./private_state.md). Public variables are visible to anyone, and private variables remain hidden within the contract. A state variable with both public and private components is said to be [shared](./shared_state.md).
 
 Aztec.nr has a few abstractions to help define the type of data your contract holds. These include PrivateMutable, PrivateImmutable, PublicMutable, PrivateSet, and SharedImmutable.
 
@@ -22,22 +22,32 @@ On this and the following pages in this section, you’ll learn:
 - Practical implications of Storage in real smart contracts
   In an Aztec.nr contract, storage is to be defined as a single struct, that contains both public and private state variables.
 
-## Public and private state variables
+## The `Context` parameter
 
-Public state variables can be read by anyone, while private state variables can only be read by their owner (or people whom the owner has shared the decrypted data or note viewing key with).
+Aztec contracts have three different modes of execution: [private](../../../aztec/glossary/call_types.md#private-execution), [public](../../../aztec/glossary/call_types.md#public-execution) and [top-level unconstrained](../../../aztec/glossary/call_types.md#top-level-unconstrained). How storage is accessed depends on the execution mode: for example, `PublicImmutable` can be read in all execution modes but only initialized in public, while `PrivateMutable` is entirely unavailable in public.
 
-Public state follows the Ethereum style account model, where each contract has its own key-value datastore. Private state follows a UTXO model, where note contents (/aztec/aztec/concepts/state_model/index.md) and [private/public execution](/aztec/concepts/smart_contracts/communication/public_private_calls.md)) for more background.
+Aztec.nr prevents developers from calling functions unavailable in the current execution mode via the `context` variable that is injected into all contract functions. Its type indicates the current execution mode:
+ - `&mut PrivateContext` for private execution
+ - `&mut PublicContext` for public execution
+ - `()` for unconstrained
 
-## Storage struct
+All state variables are generic over this `Context` type, and expose different methods in each execution mode. In the example above, `PublicImmutable`'s `initialize` function is only available with a public execution context, and so the following code results in a compilation error:
 
 ```rust
 #[aztec(storage)]
 struct Storage {
-  // public state variables
-  // private state variables
+  variable: PublicImmutable<Field>,
+}
+
+#[aztec(private)]
+fn some_private_function() {
+  storage.variable.initialize(0);
+  // ^ ERROR: Expected type PublicImmutable<_, &mut PublicContext>, found type PublicImmutable<Field, &mut PrivateContext>
 }
 ```
 
+The `Context` generic type parameter is not visible in the code above as it is automatically injected by the `#[aztec(storage)]` macro, in order to reduce boilerplate. Similarly, all state variables in that struct (e.g. `PublicImmutable`) similarly have that same type parameter automatically passed to them.
+
 ## Map
 
 A `map` is a state variable that "maps" a key to a value. It can be used with private or public storage variables.

noir-projects/aztec-nr/authwit/src/account.nr

@@ -1,17 +1,17 @@
-use dep::aztec::context::{PrivateContext, PublicContext, Context};
+use dep::aztec::context::{PrivateContext, PublicContext};
 use dep::aztec::state_vars::{Map, PublicMutable};
 use dep::aztec::protocol_types::{address::AztecAddress, abis::function_selector::FunctionSelector, hash::pedersen_hash};
 
 use crate::entrypoint::{app::AppPayload, fee::FeePayload};
 use crate::auth::{IS_VALID_SELECTOR, compute_outer_authwit_hash};
 
-struct AccountActions {
+struct AccountActions<Context> {
   context: Context,
   is_valid_impl: fn(&mut PrivateContext, Field) -> bool,
-  approved_action: Map<Field, PublicMutable<bool>>,
+  approved_action: Map<Field, PublicMutable<bool, Context>, Context>,
 }
 
-impl AccountActions {
+impl<Context> AccountActions<Context> {
     pub fn init(
         context: Context,
         approved_action_storage_slot: Field,
@@ -29,81 +29,58 @@ impl AccountActions {
             )
         }
     }
+}
 
-    pub fn private(
-        context: &mut PrivateContext,
-        approved_action_storage_slot: Field,
-        is_valid_impl: fn(&mut PrivateContext, Field) -> bool
-    ) -> Self {
-        AccountActions::init(
-            Context::private(context),
-            approved_action_storage_slot,
-            is_valid_impl
-        )
-    }
-
-    pub fn public(
-        context: &mut PublicContext,
-        approved_action_storage_slot: Field,
-        is_valid_impl: fn(&mut PrivateContext, Field) -> bool
-    ) -> Self {
-        AccountActions::init(
-            Context::public(context),
-            approved_action_storage_slot,
-            is_valid_impl
-        )
-    }
-
+impl AccountActions<&mut PrivateContext> {
     // docs:start:entrypoint
     pub fn entrypoint(self, app_payload: AppPayload, fee_payload: FeePayload) {
         let valid_fn = self.is_valid_impl;
-        let mut private_context = self.context.private.unwrap();
 
         let fee_hash = fee_payload.hash();
-        assert(valid_fn(private_context, fee_hash));
-        fee_payload.execute_calls(private_context);
-        private_context.end_setup();
+        assert(valid_fn(self.context, fee_hash));
+        fee_payload.execute_calls(self.context);
+        self.context.end_setup();
 
         let app_hash = app_payload.hash();
-        assert(valid_fn(private_context, app_hash));
-        app_payload.execute_calls(private_context);
+        assert(valid_fn(self.context, app_hash));
+        app_payload.execute_calls(self.context);
     }
     // docs:end:entrypoint
 
     // docs:start:spend_private_authwit
     pub fn spend_private_authwit(self, inner_hash: Field) -> Field {
-        let context = self.context.private.unwrap();
         // The `inner_hash` is "siloed" with the `msg_sender` to ensure that only it can 
         // consume the message.
         // This ensures that contracts cannot consume messages that are not intended for them.
         let message_hash = compute_outer_authwit_hash(
-            context.msg_sender(),
-            context.chain_id(),
-            context.version(),
+            self.context.msg_sender(),
+            self.context.chain_id(),
+            self.context.version(),
             inner_hash
         );
         let valid_fn = self.is_valid_impl;
-        assert(valid_fn(context, message_hash) == true, "Message not authorized by account");
-        context.push_new_nullifier(message_hash, 0);
+        assert(valid_fn(self.context, message_hash) == true, "Message not authorized by account");
+        self.context.push_new_nullifier(message_hash, 0);
         IS_VALID_SELECTOR
     }
     // docs:end:spend_private_authwit
+}
 
+impl AccountActions<&mut PublicContext> {
     // docs:start:spend_public_authwit
     pub fn spend_public_authwit(self, inner_hash: Field) -> Field {
-        let context = self.context.public.unwrap();
         // The `inner_hash` is "siloed" with the `msg_sender` to ensure that only it can 
         // consume the message.
         // This ensures that contracts cannot consume messages that are not intended for them.
         let message_hash = compute_outer_authwit_hash(
-            context.msg_sender(),
-            context.chain_id(),
-            context.version(),
+            self.context.msg_sender(),
+            self.context.chain_id(),
+            self.context.version(),
             inner_hash
         );
         let is_valid = self.approved_action.at(message_hash).read();
         assert(is_valid == true, "Message not authorized by account");
-        context.push_new_nullifier(message_hash, 0);
+        self.context.push_new_nullifier(message_hash, 0);
         IS_VALID_SELECTOR
     }
     // docs:end:spend_public_authwit

noir-projects/aztec-nr/authwit/src/auth.nr

@@ -4,10 +4,7 @@ use dep::aztec::protocol_types::{
 };
 use dep::aztec::{
     prelude::Deserialize,
-    context::{
-    PrivateContext, PublicContext, Context, gas::GasOpts,
-    interface::{ContextInterface, PublicContextInterface}
-},
+    context::{PrivateContext, PublicContext, gas::GasOpts, interface::{ContextInterface, PublicContextInterface}},
     hash::hash_args_array
 };
 

noir-projects/aztec-nr/aztec/src/context.nr

@@ -20,27 +20,3 @@ use private_context::PackedReturns;
 use public_context::PublicContext;
 use public_context::FunctionReturns;
 use avm_context::AvmContext;
-
-struct Context {
-    private: Option<&mut PrivateContext>,
-    public: Option<&mut PublicContext>,
-    avm: Option<&mut AvmContext>,
-}
-
-impl Context {
-    pub fn private(context: &mut PrivateContext) -> Context {
-        Context { private: Option::some(context), public: Option::none(), avm: Option::none() }
-    }
-
-    pub fn public(context: &mut PublicContext) -> Context {
-        Context { public: Option::some(context), private: Option::none(), avm: Option::none() }
-    }
-
-    pub fn avm(context: &mut AvmContext) -> Context {
-        Context { avm: Option::some(context), public: Option::none(), private: Option::none() }
-    }
-
-    pub fn none() -> Context {
-        Context { public: Option::none(), private: Option::none(), avm: Option::none() }
-    }
-}

noir-projects/aztec-nr/aztec/src/state_vars/map.nr

@@ -1,18 +1,17 @@
-use crate::context::{PrivateContext, PublicContext, Context};
 use dep::protocol_types::{hash::pedersen_hash, traits::ToField};
 use crate::state_vars::storage::Storage;
 
 // docs:start:map
-struct Map<K, V> {
+struct Map<K, V, Context> {
     context: Context,
     storage_slot: Field,
     state_var_constructor: fn(Context, Field) -> V,
 }
 // docs:end:map
 
-impl<K, T> Storage<T> for Map<K, T> {}
+impl<K, T, Context> Storage<T> for Map<K, T, Context> {}
 
-impl<K, V> Map<K, V> {
+impl<K, V, Context> Map<K, V, Context> {
     // docs:start:new
     pub fn new(
         context: Context,