State schema, or how to strongly type state keys in programs

[iFrostizz]

State scheme

Why do we need it ?
We need typed keys. They are currently typed at runtime and we would like to avoid this, since it could create foot guns.
For instance, to get a value from a key we write

let value = program.state().get::<u64>(StateKeys::MyKey).unwrap();

@richardpringle proposed something very similar to diesel:

state_schema! {
    TotalSupply => u64
    Balance(Address) => u64
}

What the diesel macro (diesel::table!) does is that it creates modules, a struct, and helpers. For instance, with this macro

diesel::table! {
    users {
        id -> Integer,
        name -> VarChar,
        updated_at -> Timestamp
    }
}

You can check the docs produced by running cargo doc --open.

Here, the dsl module that is generated would be of interest for us.
The diesel docs explains it very well for us what is this macro doing: https://diesel.rs/guides/schema-in-depth.html

Proposition

Say we have this macro:

state_schema! {
    TotalSupply -> u64
    Balance(Address) -> u64
}

This would create a module named state_keys (the old state_keys macro should be deprecated with this so there won’t be collisions. If we still need it we should find another name to avoid collisions.

This module would have each type which would really just be unit and tuple structs such as:

struct TotalSupply;
struct Balance(Address);

They should all implement a trait called Schema (or something else) which would expose a type which would be the return type. This type should be BorshSerialize and BorshDeserialize. End result:

pub trait Schema {
    type SchemaType: BorshSerialize + BorshDeserialize;
}

pub mod state_keys {
    pub struct TotalSupply;
    pub struct Balance(Address);

    impl Schema for TotalSupply {
        type SchemaType = u64;
    }

    impl Schema for Balance {
        type SchemaType = u64;
    }
}

Now, the get and store methods of the state struct inside of the program should accept a key which implements the trait Schema. This way, we could be able to return Schema::SchemaType

pub trait BorshSerialize {}
pub trait BorshDeserialize {}

impl BorshSerialize for u64 {}
impl BorshDeserialize for u64 {}

pub trait Schema {
    type SchemaType: BorshSerialize + BorshDeserialize;
}

mod state_keys {
    use super::Schema;

    type Address = [u8; 33];

    struct TotalSupply;
    struct Balance(Address);

    impl Schema for TotalSupply {
        type SchemaType = u64;
    }

    impl Schema for Balance {
        type SchemaType = u64;
    }
}

pub trait Key {}

pub fn store<K: Key + Schema>(k: K) {
    todo!()
}

pub fn get<K: Key + Schema>(k: K) -> K::SchemaType {
    todo!()
}

[iFrostizz]

There is a tricky part which is that the state cache should encompass all keys which implement Key but are not concrete implementation of Schema, because otherwise we will only be able to pass a context for one state key, ex:

#[public]
fn total_supply(context: Context<state_keys::TotalSupply>) {}

For instance, we would not be able to fetch the state_key::Balance in this case.

[richardpringle]

🤔, I think the way to solve this would be to still have a StateKeys enum that's defined through generated code. Then the Schema trait could actually provide the get and put methods with the proper associated type. Each key-type could have From<StateKeys> for Option<KeyType> (or TryFrom) and From<KeyType> for StateKeys. We could still do the same thing;

let result = context.program().state().get(KeyType);

where internally, we have something like this:

fn get<Key: Schema>(&mut self, key: Key) -> Key::Value {
    key.get(self) // here, the `get` implementation takes care of fetching the key
}

Although I'll have to admit, things are starting to get a little more complex than I think they should be. I would still like to use bytemuck with an enum as a more efficient way to deal with keys as it would be more efficient than our current implementation.