Build State Machine using Rust | Part 1 | Balances Module

Alice sends money to Bob via State Machine built in Rust

Hello folks!

Welcome to the new series where we would learn Rust by building State Machine from scratch.

This is going to be fun for devs learning Substrate in Blockchain sphere.

What is State Machine?

By definition:

“A state machine is a mathematical abstraction used to design algorithms. A state machine reads a set of inputs and changes to a different state based on those inputs. A state is a description of the status of a system waiting to execute a transition. A transition is a set of actions to execute when a condition is fulfilled or an event is received. In a state diagram, circles represent each possible state and arrows represent transitions between states.” Source.

It is simply a black box which takes an input and returns an output. It is responsible for state changes. Suppose, Alice sends 1 ETH to Bob, leading to change in the balance of both of them. This means Alice’s balance is reduced by 1 ETH and the same amount is added to that of Bob. This makes the blockchain systems deterministic.

A State machine that returns a deterministic output from a parsed input

More on this here. The detailed explanation covers how the state transition function utilizes the current block and the previous state to determine and return the next state.

What is Balances Module?

It is responsible for handling all the users’ balances.

How to code?

A. Setup

Let’s do the setup first.

1. Download & install Rust from here.
2. Install Rust Analyzer extension in VSCode or your favourite editor. Get familiar with all the needed practices from here.

B. Create Project

Run the command:

$ cargo new rust-state-machine

This would create a folder structure like this:

rust-state-machine
├── Cargo.toml
└── src
    └── main.rs

Then, configure a rustfmt.toml file to ensure proper formatting in the entire project with the following content:

# Basic
edition = "2021"
hard_tabs = true
max_width = 100
use_small_heuristics = "Max"
# Imports
imports_granularity = "Crate"
reorder_imports = true
# Consistency
newline_style = "Unix"
# Misc
chain_width = 80
spaces_around_ranges = false
binop_separator = "Back"
reorder_impl_items = false
match_arm_leading_pipes = "Preserve"
match_arm_blocks = false
match_block_trailing_comma = true
trailing_comma = "Vertical"
trailing_semicolon = false
use_field_init_shorthand = true
# Format comments
comment_width = 100
wrap_comments = true

Also, consider adding nightly channel of Rust as this would allow $ cargo fmt command to work without any warning/error as some of the features does require nightly channel of Rust.

Then create a file named: rust-toolchain.toml and add the following content to the file.

[toolchain]
channel = "nightly"

C. Add module

To add the Balances module/pallet, add this line to the main.rs file at the top like shown below.

mod balances;

fn main() {}

And then create a file named balances.rs in the src directory. You should have a folder structure similar to the one shown below.

./
├── Cargo.lock
├── Cargo.toml
├── rust-toolchain.toml
├── rustfmt.toml
└── src
    ├── balances.rs
    └── main.rs

D. Add a struct

The following struct contains the data/info for users’ balances.

pub struct Pallet {
  balance: BTreeMap<String, u128>
}

Q. Why BTreeMap instead of HashMap?
BTreeMap is mainly used here for range-based operations. Suppose, need to fetch some keys in an sorted manner. Then, it is better than HashMap as is the case here. Although it doesn’t really matter much as both are efficient with O(logN) and O(1) time complexity respectively.

E. Implement Functions

Now, we need to implement the methods:

• new: creates a new balances module to initiate the state for a single run. Here, w.r.t the runtime or state machine, this is initiated in the genesis block itself. Thereby, the balances field would store the users’ balances.

Here is the code:

/// previous code

impl Pallet {
 pub fn new() -> Self {
  Self { balances: BTreeMap::new() }
 }
}

• set_balance: set a balance amount for a user.

Here is the code:

/// previous code

impl Pallet {
 // previous code

 fn set_balance(&mut self, who: &String, balance: u128) {
  self.balances.insert(who.clone(), balance);
 }
}

We reference the who argument in the function because the user is borrowed for use as defined with derived type. Otherwise, this function would consume this argument and hence can’t be used any further down the program in runtime code. No need to borrow for balance with primitive type.

• balance: get/view a user’s balance.

Here is the code:

impl Pallet {
  /// previous code

  fn balance(&self, who: &String) -> u128 {
    *self.balances.get(who).unwrap_or(&0)
  }
}

In Rust, we always try to avoid unwrap for better error handling. So, the balance of any user is defined as 0 by default. * (dereferencing) is used to return the owned type i.e. u128 instead of &u128.

• transfer: moves the amount from the balance of one user to another. This involves checking if it exceeds the MIN/MAX permissible value as balance.

Here is the code:

impl Pallet {
  /// previous code

  fn transfer(&mut self, from: &String, to: &String, amount: u128) -> Result<(), &'static str> {
  // get the old balances
  let from_old_balance = self.balance(from);
  let to_old_balance = self.balance(to);

  // calculate the new balances
  let from_new_balance =
   from_old_balance.checked_sub(amount).ok_or("Insufficient balance")?;
  let to_new_balance = to_old_balance.checked_add(amount).ok_or("Exceeds MAX balance")?;

  // set the new balances
  self.set_balance(from, from_new_balance);
  self.set_balance(to, to_new_balance);

  Ok(())
 }
}

The function return type is chosen to be Result<(), &'static str for error handling. The error would be thrown as Err("some message"), if any, else returns as Ok(()).

When calculating the new balances, the conditions of MIN/MAX is checked using ok_or as it returns Result<T, E> type to match the function’s return type.

F. Write tests (integration)

Now, we need to write some tests to see if it’s running fine.

The tests are mainly based on integration test.

For simplicity, 2 tests have been written here:

• init_balances: Test if the balances are instantiated properly with 0as balance for a new user.

Here is the code:

// previous code

#[cfg(test)]
mod tests {
 #[test]
 fn init_balances() {
  let mut balances = super::Pallet::new();

  assert_eq!(balances.balance(&"alice".to_string()), 0);
  balances.set_balance(&"alice".to_string(), 100);
  assert_eq!(balances.balance(&"alice".to_string()), 100);
  assert_eq!(balances.balance(&"bob".to_string()), 0);
 }
}

We check if Alice has zero balance in the beginning when the Balances module is instantiated/created. And then a new balance is set for Alice which is supposed to change the balance to 100. But, for Bob the balance remains to be zero as preset value.

• test_transfer: Test if amount is being properly transferred from one user to another.

Here is the code:

#[cfg(test)]
mod tests {
 // previous code

 #[test]
 fn test_transfer() {
  let mut balances = super::Pallet::new();

  assert_eq!(
   balances.transfer(&"alice".to_string(), &"bob".to_string(), 100),
   Err("Insufficient balance")
  );

  balances.set_balance(&"alice".to_string(), 100);
  assert!(balances.transfer(&"alice".to_string(), &"bob".to_string(), 51).is_ok());

  // check the balances
  assert_eq!(balances.balance(&"alice".to_string()), 49);
  assert_eq!(balances.balance(&"bob".to_string()), 51);
 }
}

After instantiating the Balances module/pallet, we test if Alice is able to transfer any amount. But it fails ❌ as expected.

Then, we view the balances of Alice and Bob after transferring some amount from 100 as previous balance of Alice.

Full Code 👨‍💻

Summary

This Balances module/pallet is a simple way to understand how a token’s operation works under the hood. This is one way of looking at how the state machine behaves like a deterministic one.

I give a huge shoutout 📢 to some of the works already done by engineers at Parity. My goal has been to simplify this as much as possible. And also to add my perspective wherever I thought as required.

Rustaceans 🚀

Thank you for being a part of the Rustaceans community! Before you go:

• Show your appreciation with a clap and follow the publication
• Discover how you can contribute your own insights to Rustaceans.
• Connect with us: X | Weekly Rust Newsletter