Migrate aztec-nr state variables to use context traits

[nventuro]

Bad Aztec.nr Developer Experience: a Tale of Two Contexts

State variables currently hold a Context object which is passed in the new function, e.g.:

impl<T> PublicImmutable<T> {
    pub fn new(
        context: Context,
        storage_slot: Field
    ) -> Self {

new is in turn called for each state variable by the init function, which the macros autogenerate:

impl Storage {
    fn init(context: Context) -> Self {
        Storage {
            leader: PublicMutable::new(context, 1),
            legendary_card: PrivateMutable::new(context, 3),
            profiles: Map::new(context, 4),

This context object is created on each function as either a PublicContext or a PrivateContext depending on the function type, and then wrapped in the Context type, also by the macros:

let mut context = PrivateContext::new(inputs, args_hasher.hash());
let mut storage = Storage::init(Context::private(&mut context));

When a function invokes a state variable method, the object already holds the context, and can only inspect its contents during function execution (e.g. self.context.private.unwrap()), at runtime. This means that if a state variable function expected to be called in private is instead called in public, the error will go undetected all the way to public execution, until the unwrap() call panics.

This results in terrible developer experience, which is made worse by the slowness of not just the test run but also the entire compile -> transpile -> create typings -> rebuild tests chain. It also means all methods end up with _in_private or _in_public_suffixes.

A Better Way

We can take advantage of how monomorphization works in Noir, and the fact that all contract functions are compiled as standalone circuits. For any given contract function, we'll be working in either a private or public context, and we'll create a state variable object on which no methods that require the other type will be called (i.e. we won't call public context methods on a state variable created via the init -> new chain in a private contract function). Armed with this, we can make calling functions from the wrong domain a compilation error.

Instead of having all state variables hold a concrete Context object, we can make them be generic over a C type, which in practice will be either a private or a public context.

struct PublicImmutable<T, C> {
    context: C,
    storage_slot: Field,
}

impl<T, C> PublicImmutable<T, C> {
    pub fn new(
        context: C,
        storage_slot: Field
    ) -> Self {

This change gets propagated to the init call, which simply means we no longer wrap in Context in the autogenerated function prelude:

let mut context = PrivateContext::new(inputs, args_hasher.hash());  
- let mut storage = Storage::init(Context::private(&mut context));
+ let mut storage = Storage::init(&mut context);

Finally, state variable methods no longer .unwrap() (because there's no wrapper) and simply call functions on the context directly. To do this, they require that the context is of the correct type by requesting that it implements a trait:

- pub fn get_current_value_in_public(self) -> T {
-    let block_number = self.context.public.unwrap().block_number() as u32;
+ pub fn get_current_value_in_public(self) -> T where C: PublicContextInterface {
+    let block_number = self.context.block_number() as u32;

With this, if storage.myVariable.get_current_value_in_public is called in a private contract function, compilation will fail because storage.myVariable will have been specialized for PrivateContext, as that type is the one that was passed to new, and PrivateContext does not implement the PublicContextInterface trait.

✨ happy developer ✨

Note that this means that the Storage struct itself must also be made generic over a C type which is passed to all state variables.

Implementing this change requires minimal changes to the macros, and the only user code change is the new C type in the Storage struct (uneless there's also manual interactions with contexts, in which we may need to remove the .private.unwrap() chains). State variables of course would need updating and adding of where clauses, but this should be fairly simple to do. Finally, we'd need to create the PublicContextInterface and PrivateContextInterface traits, which are not currently a thing (they're currently just structs with an impl block).