Build State Machine using Rust | Part 1 | Balances Module

balances.rs

use std::collections::BTreeMap;

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

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

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

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

	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(())
	}
}

#[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);
	}

	#[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);
	}
}

Cargo.toml

[package]
name = "rsm-tuts"
version = "0.1.0"
edition = "2021"

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

[dependencies]

main.rs

mod balances;

fn main() {}

rust-toolchain.toml

[toolchain]
channel = "nightly"

rustfmt.toml

# 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