From c455ad47b3cf30eeae7c2b63f8b8e1a9fd81eaae Mon Sep 17 00:00:00 2001 From: Michal Nazarewicz Date: Sat, 19 Aug 2023 15:24:48 +0200 Subject: [PATCH] Add example Solana contract using the sealable trie Introduce solana/trie-example smart contract which uses the sealable trie on Solana chain. The contract as well as its client are quite rudimentary but they allow demonstrating working code. --- .gitignore | 5 +- Cargo.toml | 4 + deny.toml | 5 +- package.json | 43 ++++ sealable-trie/src/lib.rs | 2 +- sealable-trie/src/trie.rs | 21 +- solana/trie-example/Cargo.toml | 17 ++ solana/trie-example/build.sh | 9 + solana/trie-example/client/main.ts | 75 +++++++ solana/trie-example/client/trie.ts | 214 ++++++++++++++++++ solana/trie-example/client/utils.ts | 71 ++++++ solana/trie-example/src/lib.rs | 101 +++++++++ solana/trie-example/src/trie.rs | 324 ++++++++++++++++++++++++++++ stdx/src/lib.rs | 20 ++ tsconfig.json | 15 ++ 15 files changed, 922 insertions(+), 4 deletions(-) create mode 100644 package.json create mode 100644 solana/trie-example/Cargo.toml create mode 100755 solana/trie-example/build.sh create mode 100644 solana/trie-example/client/main.ts create mode 100644 solana/trie-example/client/trie.ts create mode 100644 solana/trie-example/client/utils.ts create mode 100644 solana/trie-example/src/lib.rs create mode 100644 solana/trie-example/src/trie.rs create mode 100644 tsconfig.json diff --git a/.gitignore b/.gitignore index 2c96eb1b..156b1202 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,5 @@ -target/ +/target/ +/dist/ +/node_modules/ +package-lock.json Cargo.lock diff --git a/Cargo.toml b/Cargo.toml index c855a331..86e76afd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,6 +8,7 @@ rust-version = "1.71.0" members = [ "memory", "sealable-trie", + "solana/trie-example", "stdx", ] resolver = "2" @@ -18,6 +19,9 @@ derive_more = "0.99.17" pretty_assertions = "1.4.0" rand = { version = "0.8.5" } sha2 = { version = "0.10.7", default-features = false } +solana-client = "1.16.7" +solana-program = "1.16.7" +solana-sdk = "1.16.7" strum = { version = "0.25.0", default-features = false, features = ["derive"] } memory = { path = "memory" } diff --git a/deny.toml b/deny.toml index f009c807..f774f0d2 100644 --- a/deny.toml +++ b/deny.toml @@ -5,7 +5,10 @@ allow-registry = ["https://github.com/rust-lang/crates.io-index"] allow-git = [] [bans] -multiple-versions = "deny" +# solana-program is weird and has bunch of duplicate dependencies. +# For now allow duplicates. TODO(mina86): Figure out if there’s +# something better we can do. +#multiple-versions = "deny" skip = [ # derive_more still uses old syn { name = "syn", version = "1.0.*" }, diff --git a/package.json b/package.json new file mode 100644 index 00000000..7a27de30 --- /dev/null +++ b/package.json @@ -0,0 +1,43 @@ +{ + "name": "solana-trie-example", + "version": "0.0.1", + "author": "Michał Nazarewicz ", + "scripts": { + "start": "ts-node solana/trie-example/client/main.ts", + "start-with-test-validator": "start-server-and-test 'solana-test-validator --reset --quiet' http://localhost:8899/health start", + "lint": "eslint --ext .ts solana/trie-example/* && prettier --check \"solana/trie-example/**/*.ts\"", + "lint:fix": "eslint --ext .ts solana/trie-example/* --fix && prettier --write \"solana/trie-example/**/*.ts\"", + "clean": "npm run clean:trie-example", + "build:trie-example": "cargo build-sbf --manifest-path=solana/trie-example/Cargo.toml --sbf-out-dir=dist/trie-example", + "deploy:trie-example": "solana program deploy dist/trie-example/trie.so", + "clean:trie-example": "cargo clean --manifest-path=solana/trie-example/Cargo.toml && rm -rf ./dist", + "test:trie-example": "cargo test-bpf --manifest-path=solana/trie-example/Cargo.toml", + "pretty": "prettier --write 'solana/trie-example/client/*.ts'" + }, + "dependencies": { + "@solana/web3.js": "^1.33.0", + "mz": "^2.7.0", + "tsconfig": "^7.0.0", + "yaml": "^2.0.0" + }, + "devDependencies": { + "@tsconfig/recommended": "^1.0.1", + "@types/eslint": "^8.2.2", + "@types/eslint-plugin-prettier": "^3.1.0", + "@types/mz": "^2.7.2", + "@types/prettier": "^2.1.5", + "@types/yaml": "^1.9.7", + "@typescript-eslint/eslint-plugin": "^4.6.0", + "@typescript-eslint/parser": "^4.6.0", + "eslint": "^7.12.1", + "eslint-config-prettier": "^6.15.0", + "eslint-plugin-prettier": "^4.0.0", + "prettier": "^2.1.2", + "start-server-and-test": "^1.11.6", + "ts-node": "^10.0.0", + "typescript": "^4.0.5" + }, + "engines": { + "node": ">=14.0.0" + } +} diff --git a/sealable-trie/src/lib.rs b/sealable-trie/src/lib.rs index 995dc4ce..9f2e3e1b 100644 --- a/sealable-trie/src/lib.rs +++ b/sealable-trie/src/lib.rs @@ -12,4 +12,4 @@ pub mod trie; #[cfg(test)] mod test_utils; -pub use trie::Trie; +pub use trie::{Error, Trie}; diff --git a/sealable-trie/src/trie.rs b/sealable-trie/src/trie.rs index 58b79c11..eb2a34ed 100644 --- a/sealable-trie/src/trie.rs +++ b/sealable-trie/src/trie.rs @@ -106,7 +106,7 @@ macro_rules! proof { } impl> Trie { - /// Creates a new trie using given allocator. + /// Creates a new empty trie using given allocator. pub fn new(alloc: A) -> Self { Self { root_ptr: None, root_hash: EMPTY_TRIE_ROOT, alloc } } @@ -117,6 +117,25 @@ impl> Trie { /// Returns whether the trie is empty. pub fn is_empty(&self) -> bool { self.root_hash == EMPTY_TRIE_ROOT } + /// Deconstructs the object into the individual parts — allocator, root + /// pointer and root hash. + pub fn into_parts(self) -> (A, Option, CryptoHash) { + (self.alloc, self.root_ptr, self.root_hash) + } + + /// Creates a new trie from individual parts. + /// + /// It’s up to the caller to guarantee that the `root_ptr` and `root_hash` + /// values are correct and correspond to nodes stored within the pool + /// allocator `alloc`. + pub fn from_parts( + alloc: A, + root_ptr: Option, + root_hash: CryptoHash, + ) -> Self { + Self { root_ptr, root_hash, alloc } + } + /// Retrieves value at given key. /// /// Returns `None` if there’s no value at given key. Returns an error if diff --git a/solana/trie-example/Cargo.toml b/solana/trie-example/Cargo.toml new file mode 100644 index 00000000..53c63a3a --- /dev/null +++ b/solana/trie-example/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "solana-trie-example" +authors = ["Michal Nazarewicz "] +version = "0.0.0" +edition = "2021" + +[lib] +name = "trie" +crate-type = ["cdylib", "lib"] + +[dependencies] +derive_more.workspace = true +solana-program.workspace = true + +memory.workspace = true +sealable-trie.workspace = true +stdx.workspace = true diff --git a/solana/trie-example/build.sh b/solana/trie-example/build.sh new file mode 100755 index 00000000..81008971 --- /dev/null +++ b/solana/trie-example/build.sh @@ -0,0 +1,9 @@ +#!/bin/sh + +set -eu + +root=$(git rev-parse --show-toplevel) +cd "$root" +cargo build-sbf \ + --manifest-path=solana/trie-example/Cargo.toml \ + --sbf-out-dir=dist/trie-example diff --git a/solana/trie-example/client/main.ts b/solana/trie-example/client/main.ts new file mode 100644 index 00000000..f056ecaa --- /dev/null +++ b/solana/trie-example/client/main.ts @@ -0,0 +1,75 @@ +import { + establishConnection, + establishPayer, + checkProgram, + getKey, + setKey, + sealKey, +} from './trie'; + +async function main() { + const operation = parseArgv(process.argv); + if (!operation) { + return; + } + + // Establish connection to the cluster + await establishConnection(); + + // Determine who pays for the fees + await establishPayer(); + + // Check if the program has been deployed + await checkProgram(); + + switch (operation[0]) { + case "get": + await getKey(operation[1]); + break; + case "set": + await setKey(operation[1], operation[2]); + break; + case "seal": + await sealKey(operation[1]); + break; + } + + console.log('Success'); +} + +function parseArgv(argv: string[]): string[] | null { + const cmd = argv[0] + ' ' + argv[1]; + switch (argv[2] || "--help") { + case "get": + case "seal": + if (argv.length != 4) { + break; + } + return [argv[2], argv[3]]; + case "set": + if (argv.length != 5) { + break; + } + return [argv[2], argv[3], argv[4]]; + case "help": + case "--help": + case "-h": { + console.log( + `usage: ${cmd} get \n` + + ` ${cmd} set \n` + + ` ${cmd} seal ` + ) + process.exit(0); + } + } + console.error(`Invalid usage; see ${cmd} --help`); + process.exit(-1); +} + +main().then( + () => process.exit(), + err => { + console.error(err); + process.exit(-1); + }, +); diff --git a/solana/trie-example/client/trie.ts b/solana/trie-example/client/trie.ts new file mode 100644 index 00000000..3440e64b --- /dev/null +++ b/solana/trie-example/client/trie.ts @@ -0,0 +1,214 @@ +/* eslint-disable @typescript-eslint/no-unsafe-assignment */ +/* eslint-disable @typescript-eslint/no-unsafe-member-access */ + +import { + Keypair, + Connection, + PublicKey, + LAMPORTS_PER_SOL, + SystemProgram, + TransactionInstruction, + Transaction, + sendAndConfirmTransaction, +} from '@solana/web3.js'; +import fs from 'mz/fs'; +import path from 'path'; + +import {getPayer, getRpcUrl, createKeypairFromFile} from './utils'; + +/** + * Connection to the network + */ +let connection: Connection; + +/** + * Keypair associated to the fees’ payer + */ +let payer: Keypair; + +/** + * Example trie’s program id + */ +let programId: PublicKey; + +/** + * The public key of the trie account. + */ +let triePubkey: PublicKey; + +/** + * Path to program files + */ +const PROGRAM_PATH = path.resolve(__dirname, '../../../dist/trie-example'); + +/** + * Path to program shared object file which should be deployed on chain. + */ +const PROGRAM_SO_PATH = path.join(PROGRAM_PATH, 'trie.so'); + +/** + * Path to the keypair of the deployed program. + * This file is created when running `solana program deploy dist/trie-example/trie.so` + */ +const PROGRAM_KEYPAIR_PATH = path.join(PROGRAM_PATH, 'trie-keypair.json'); + +/** + * The expected size of the account. + */ +const ACCOUNT_SIZE = 72000 + 64; + +/** + * Establish a connection to the cluster + */ +export async function establishConnection(): Promise { + const rpcUrl = await getRpcUrl(); + connection = new Connection(rpcUrl, 'confirmed'); + const version = await connection.getVersion(); + console.log('Connection to cluster established:', rpcUrl, version); +} + +/** + * Establish an account to pay for everything + */ +export async function establishPayer(): Promise { + let fees = 0; + if (!payer) { + const {feeCalculator} = await connection.getRecentBlockhash(); + + // Calculate the cost to fund the trie account + fees += await connection.getMinimumBalanceForRentExemption(ACCOUNT_SIZE); + + // Calculate the cost of sending transactions + fees += feeCalculator.lamportsPerSignature * 100; // wag + + payer = await getPayer(); + } + + let lamports = await connection.getBalance(payer.publicKey); + // if (lamports < fees) { + // // If current balance is not enough to pay for fees, request an airdrop + // const sig = await connection.requestAirdrop( + // payer.publicKey, + // fees - lamports, + // ); + // await connection.confirmTransaction(sig); + // lamports = await connection.getBalance(payer.publicKey); + // } + + console.log( + 'Using account', + payer.publicKey.toBase58(), + 'containing', + lamports / LAMPORTS_PER_SOL, + 'SOL to pay for fees', + ); +} + +/** + * Check if the hello world BPF program has been deployed + */ +export async function checkProgram(): Promise { + // Read program id from keypair file + try { + const programKeypair = await createKeypairFromFile(PROGRAM_KEYPAIR_PATH); + programId = programKeypair.publicKey; + } catch (err) { + const errMsg = (err as Error).message; + throw new Error( + `Failed to read program keypair at '${PROGRAM_KEYPAIR_PATH}' due to error: ${errMsg}. Program may need to be deployed with \`solana program deploy dist/trie-example/trie.so\``, + ); + } + + // Check if the program has been deployed + const programInfo = await connection.getAccountInfo(programId); + if (programInfo === null) { + if (fs.existsSync(PROGRAM_SO_PATH)) { + throw new Error( + 'Program needs to be deployed with `solana program deploy dist/trie-example/trie.so`', + ); + } else { + throw new Error('Program needs to be built and deployed'); + } + } else if (!programInfo.executable) { + throw new Error(`Program is not executable`); + } + console.log(`Using program ${programId.toBase58()}`); + + // Derive the address (public key) of a trie account from the program so that it's easy to find later. + const SEED = 'hello'; + triePubkey = await PublicKey.createWithSeed( + payer.publicKey, + SEED, + programId, + ); + + // Check if the greeting account has already been created + const trieAccount = await connection.getAccountInfo(triePubkey); + if (trieAccount === null) { + console.log( + 'Creating account', + triePubkey.toBase58(), + 'to say hello to', + ); + const lamports = await connection.getMinimumBalanceForRentExemption( + ACCOUNT_SIZE, + ); + + const transaction = new Transaction().add( + SystemProgram.createAccountWithSeed({ + fromPubkey: payer.publicKey, + basePubkey: payer.publicKey, + seed: SEED, + newAccountPubkey: triePubkey, + lamports, + space: ACCOUNT_SIZE, + programId, + }), + ); + await sendAndConfirmTransaction(connection, transaction, [payer]); + } +} + +/** + * Get value of a key + */ +export async function getKey(hexKey: string): Promise { + console.log('Getting key %s from %s', hexKey, triePubkey.toBase58()); + await call('00' + hexKey); +} + +/** + * Get value of a key + */ +export async function setKey(hexKey: string, hexValue: string): Promise { + console.log('Setting key %s to %s in %s', hexKey, hexValue, triePubkey.toBase58()); + if (hexValue.length != 64) { + throw new Error('Value must be 64 hexadecimal digits'); + } + await call('02' + hexKey + hexValue); +} + +/** + * Get value of a key + */ +export async function sealKey(hexKey: string): Promise { + console.log('Sealing key %s in %s', hexKey, triePubkey.toBase58()); + await call('04' + hexKey); +} + +/** + * Sends an instruction to the contract. + */ +async function call(hexData: string): Promise { + let data = Buffer.from(hexData, 'hex'); + const instruction = new TransactionInstruction({ + keys: [{pubkey: triePubkey, isSigner: false, isWritable: true}], + programId, + data, + }); + await sendAndConfirmTransaction( + connection, + new Transaction().add(instruction), + [payer], + ); +} diff --git a/solana/trie-example/client/utils.ts b/solana/trie-example/client/utils.ts new file mode 100644 index 00000000..ca11f1c1 --- /dev/null +++ b/solana/trie-example/client/utils.ts @@ -0,0 +1,71 @@ +/* eslint-disable @typescript-eslint/no-unsafe-assignment */ +/* eslint-disable @typescript-eslint/no-unsafe-member-access */ +/* eslint-disable @typescript-eslint/no-unsafe-return */ + +import os from 'os'; +import fs from 'mz/fs'; +import path from 'path'; +import yaml from 'yaml'; +import {Keypair} from '@solana/web3.js'; + +/** + * @private + */ +async function getConfig(): Promise { + // Path to Solana CLI config file + const CONFIG_FILE_PATH = path.resolve( + os.homedir(), + '.config', + 'solana', + 'cli', + 'config.yml', + ); + const configYml = await fs.readFile( + CONFIG_FILE_PATH, {encoding: 'utf8'}); + return yaml.parse(configYml); +} + +/** + * Load and parse the Solana CLI config file to determine which RPC url to use + */ +export async function getRpcUrl(): Promise { + try { + const config = await getConfig(); + if (!config.json_rpc_url) throw new Error('Missing RPC URL'); + return config.json_rpc_url; + } catch (err) { + console.warn( + 'Failed to read RPC url from CLI config file, falling back to localhost', + ); + return 'http://127.0.0.1:8899'; + } +} + +/** + * Load and parse the Solana CLI config file to determine which payer to use + */ +export async function getPayer(): Promise { + try { + const config = await getConfig(); + if (!config.keypair_path) { + throw new Error('Missing keypair path'); + } + return await createKeypairFromFile(config.keypair_path); + } catch (err) { + console.warn( + 'Failed to create keypair from CLI config file, falling back to new random keypair', + ); + return Keypair.generate(); + } +} + +/** + * Create a Keypair from a secret key stored in file as bytes' array + */ +export async function createKeypairFromFile( + filePath: string, +): Promise { + const secretKeyString = await fs.readFile(filePath, {encoding: 'utf8'}); + const secretKey = Uint8Array.from(JSON.parse(secretKeyString)); + return Keypair.fromSecretKey(secretKey); +} diff --git a/solana/trie-example/src/lib.rs b/solana/trie-example/src/lib.rs new file mode 100644 index 00000000..0655f598 --- /dev/null +++ b/solana/trie-example/src/lib.rs @@ -0,0 +1,101 @@ +use sealable_trie::hash::CryptoHash; +use solana_program::account_info::AccountInfo; +use solana_program::msg; +use solana_program::program::set_return_data; +use solana_program::program_error::ProgramError; +use solana_program::pubkey::Pubkey; + +mod trie; + +type Result = core::result::Result; + +/// Discriminants for the data stored in the accounts. +mod magic { + pub(crate) const UNINITIALISED: u32 = 0; + pub(crate) const TRIE_ROOT: u32 = 1; +} + +solana_program::entrypoint!(process_instruction); + +pub fn process_instruction( + program_id: &Pubkey, + accounts: &[AccountInfo], + instruction: &[u8], +) -> Result { + let account = accounts.first().ok_or(ProgramError::NotEnoughAccountKeys)?; + if account.owner != program_id { + return Err(ProgramError::IncorrectProgramId); + } + let mut trie = trie::AccountTrie::new(account.try_borrow_mut_data()?) + .ok_or(ProgramError::InvalidAccountData)?; + match Instruction::decode(instruction)? { + Instruction::Get { key, include_proof } => { + handle_get(trie, key, include_proof)?; + } + Instruction::Set { key, hash } => { + trie.set(key, hash).into_prg_err()?; + } + Instruction::Seal { key } => { + trie.seal(key).into_prg_err()?; + } + } + Ok(()) +} + +fn handle_get( + trie: trie::AccountTrie, + key: &[u8], + include_proof: bool, +) -> Result { + let (value, _proof) = if include_proof { + trie.prove(key).map(|(value, proof)| (value, Some(proof))) + } else { + trie.get(key).map(|value| (value, None)) + } + .into_prg_err()?; + set_return_data(value.as_ref().map_or(&[], CryptoHash::as_slice)); + Ok(()) +} + +trait TrieResultExt { + type Value; + fn into_prg_err(self) -> Result; +} + +impl TrieResultExt for Result { + type Value = T; + fn into_prg_err(self) -> Result { + self.map_err(|err| { + msg!("{}", err); + ProgramError::Custom(1) + }) + } +} + + +/// Instruction to execute. +pub(crate) enum Instruction<'a> { + // Encoding: + Get { key: &'a [u8], include_proof: bool }, + // Encoding: 0x02 ; is always 32-byte long. + Set { key: &'a [u8], hash: &'a CryptoHash }, + // Encoding: 0x04 + Seal { key: &'a [u8] }, +} + +impl<'a> Instruction<'a> { + pub(crate) fn decode(bytes: &'a [u8]) -> Result { + let (&tag, bytes) = + bytes.split_first().ok_or(ProgramError::InvalidInstructionData)?; + match tag { + 0 | 1 => Ok(Self::Get { key: bytes, include_proof: tag == 1 }), + 2 => { + let (key, hash) = stdx::rsplit_at(bytes) + .ok_or(ProgramError::InvalidInstructionData)?; + Ok(Self::Set { key, hash: hash.into() }) + } + 4 => Ok(Self::Seal { key: bytes }), + _ => Err(ProgramError::InvalidInstructionData), + } + } +} diff --git a/solana/trie-example/src/trie.rs b/solana/trie-example/src/trie.rs new file mode 100644 index 00000000..70304e58 --- /dev/null +++ b/solana/trie-example/src/trie.rs @@ -0,0 +1,324 @@ +use core::cell::RefMut; +use core::mem::ManuallyDrop; + +use memory::Ptr; +use sealable_trie::hash::CryptoHash; + +use crate::magic; + +type DataRef<'a, 'b> = RefMut<'a, &'b mut [u8]>; + +const SZ: usize = sealable_trie::nodes::RawNode::SIZE; + +/// Trie stored in a Solana account. +pub(crate) struct AccountTrie<'a, 'b>( + core::mem::ManuallyDrop>>, +); + +impl<'a, 'b> AccountTrie<'a, 'b> { + /// Creates a new Trie from data in an account. + /// + /// If the data in the account isn’t initialised (i.e. has zero + /// discriminant) initialises a new empty trie. + pub(crate) fn new(data: DataRef<'a, 'b>) -> Option { + let (alloc, root) = Allocator::new(data)?; + let trie = sealable_trie::Trie::from_parts(alloc, root.0, root.1); + Some(Self(ManuallyDrop::new(trie))) + } +} + +impl<'a, 'b> core::ops::Drop for AccountTrie<'a, 'b> { + /// Updates the header in the Solana account. + fn drop(&mut self) { + // SAFETY: Once we’re done with self.0 we are dropped and no one else is + // going to have access to self.0. + let trie = unsafe { ManuallyDrop::take(&mut self.0) }; + let (mut alloc, root_ptr, root_hash) = trie.into_parts(); + let hdr = Header { + root_ptr, + root_hash, + next_block: alloc.next_block, + first_free: alloc.first_free.map_or(0, |ptr| ptr.get()), + }; + alloc + .data + .get_mut(..Header::ENCODED_SIZE) + .unwrap() + .copy_from_slice(&hdr.encode()); + } +} + +/// Data stored in the first 72-bytes of the account describing the trie. +#[derive(Clone, Debug, PartialEq)] +struct Header { + root_ptr: Option, + root_hash: CryptoHash, + next_block: u32, + first_free: u32, +} + +impl Header { + /// Size of the encoded header. + const ENCODED_SIZE: usize = 64; + + /// Decodes the header from given block of memory. + /// + /// Returns `None` if the block is shorter than length of encoded header or + /// encoded data is invalid. + // Encoding: + // magic: u32 + // version: u32 + // root_ptr: u32 + // root_hash: [u8; 32] + // next_block: u32 + // first_free: u32 + // padding: [u8; 12], + fn decode(data: &[u8]) -> Option { + let data = data.get(..Self::ENCODED_SIZE)?.try_into().unwrap(); + + // Check magic number. Zero means the account hasn’t been initialised + // so return default value, and anything other than magic::TRIE_ROOT + // means it’s an account storing data different than a trie root. + let (magic, data) = read::<4, 60, 64, _>(data, u32::from_ne_bytes); + if magic == magic::UNINITIALISED { + return Some(Self { + root_ptr: None, + root_hash: sealable_trie::trie::EMPTY_TRIE_ROOT, + next_block: Self::ENCODED_SIZE as u32, + first_free: 0, + }); + } else if magic != magic::TRIE_ROOT { + return None; + } + + // Check version. This is for future-proofing in case format of the + // encoding changes. + let (version, data) = read::<4, 56, 60, _>(data, u32::from_ne_bytes); + if version != 0 { + return None; + } + + let (root_ptr, data) = read::<4, 52, 56, _>(data, u32::from_ne_bytes); + let (root_hash, data) = read::<32, 20, 52, _>(data, CryptoHash); + let (next_block, data) = read::<4, 16, 20, _>(data, u32::from_ne_bytes); + let (first_free, _) = read::<4, 12, 16, _>(data, u32::from_ne_bytes); + + let root_ptr = Ptr::new(root_ptr).ok()?; + Some(Self { root_ptr, root_hash, next_block, first_free }) + } + + /// Returns encoded representation of values in the header. + fn encode(&self) -> [u8; Self::ENCODED_SIZE] { + let root_ptr = + self.root_ptr.map_or([0; 4], |ptr| ptr.get().to_ne_bytes()); + + let mut buf = [0; Self::ENCODED_SIZE]; + let data = &mut buf; + let data = write::<4, 60, 64>(data, magic::TRIE_ROOT.to_ne_bytes()); + let data = write::<4, 56, 60>(data, [0; 4]); + let data = write::<4, 52, 56>(data, root_ptr); + let data = write::<32, 20, 52>(data, self.root_hash.0.clone()); + let data = write::<4, 16, 20>(data, self.next_block.to_ne_bytes()); + write::<4, 12, 16>(data, self.first_free.to_ne_bytes()); + buf + } +} + +pub(crate) struct Allocator<'a, 'b> { + /// Pool of memory to allocate blocks in. + /// + /// The data is always at least long enough to fit encoded [`Header`]. + data: DataRef<'a, 'b>, + + /// Position of the next unallocated block. + /// + /// Blocks which were allocated and then freed don’t count as ‘unallocated’ + /// in this context. This is position of the next block to return if the + /// free list is empty. + next_block: u32, + + /// Pointer to the first freed block; `None` if there were no freed blocks + /// yet. + first_free: Option, +} + +impl<'a, 'b> Allocator<'a, 'b> { + /// Initialises the allocator with data in given account. + fn new(data: DataRef<'a, 'b>) -> Option<(Self, (Option, CryptoHash))> { + let hdr = Header::decode(*data)?; + let next_block = hdr.next_block; + let first_free = Ptr::new(hdr.first_free).ok()?; + let alloc = Self { data, next_block, first_free }; + let root = (hdr.root_ptr, hdr.root_hash); + Some((alloc, root)) + } + + /// Grabs a block from a free list. Returns `None` if free list is empty. + fn alloc_from_freelist(&mut self) -> Option { + let ptr = self.first_free.take()?; + let idx = ptr.get() as usize; + let next = self.data.get(idx..idx + 4).unwrap().try_into().unwrap(); + self.first_free = Ptr::new(u32::from_ne_bytes(next)).unwrap(); + Some(ptr) + } + + /// Grabs a next available block. Returns `None` if account run out of + /// space. + fn alloc_next_block(&mut self) -> Option { + let len = u32::try_from(self.data.len()).unwrap_or(u32::MAX); + let end = + self.next_block.checked_add(SZ as u32).filter(|&e| e <= len)?; + let ptr = Ptr::new(self.next_block).ok().flatten()?; + self.next_block = end; + Some(ptr) + } +} + +impl<'a, 'b> memory::Allocator for Allocator<'a, 'b> { + type Value = [u8; SZ]; + + fn alloc( + &mut self, + value: Self::Value, + ) -> Result { + let ptr = self + .alloc_from_freelist() + .or_else(|| self.alloc_next_block()) + .ok_or(memory::OutOfMemory)?; + self.set(ptr, value); + Ok(ptr) + } + + #[inline] + fn get(&self, ptr: Ptr) -> &Self::Value { + let idx = ptr.get() as usize; + self.data.get(idx..idx + SZ).unwrap().try_into().unwrap() + } + + #[inline] + fn get_mut(&mut self, ptr: Ptr) -> &mut Self::Value { + let idx = ptr.get() as usize; + self.data.get_mut(idx..idx + SZ).unwrap().try_into().unwrap() + } + + #[inline] + fn free(&mut self, ptr: Ptr) { + let next = + self.first_free.map_or([0; 4], |ptr| ptr.get().to_ne_bytes()); + let idx = ptr.get() as usize; + self.data.get_mut(idx..idx + 4).unwrap().copy_from_slice(&next); + self.first_free = Some(ptr); + } +} + + + +impl<'a, 'b> core::ops::Deref for AccountTrie<'a, 'b> { + type Target = sealable_trie::Trie>; + fn deref(&self) -> &Self::Target { &*self.0 } +} + +impl<'a, 'b> core::ops::DerefMut for AccountTrie<'a, 'b> { + fn deref_mut(&mut self) -> &mut Self::Target { &mut *self.0 } +} + +/// Reads fixed-width value from start of the buffer and returns the value and +/// remaining portion of the buffer. +/// +/// By working on a fixed-size buffers, this avoids any run-time checks. Sizes +/// are verified at compile-time. +fn read( + buf: &[u8; N], + f: impl Fn([u8; L]) -> T, +) -> (T, &[u8; R]) { + let (left, right) = stdx::split_array_ref(buf); + (f(left.clone()), right) +} + +/// Writes given fixed-width buffer at the start the buffer and returns the +/// remaining portion of the buffer. +/// +/// By working on a fixed-size buffers, this avoids any run-time checks. Sizes +/// are verified at compile-time. +fn write( + buf: &mut [u8; N], + data: [u8; L], +) -> &mut [u8; R] { + let (left, right) = stdx::split_array_mut(buf); + *left = data; + right +} + + +#[test] +fn test_header_encoding() { + const ONE: CryptoHash = CryptoHash([1; 32]); + + assert_eq!( + Some(Header { + root_ptr: None, + root_hash: sealable_trie::trie::EMPTY_TRIE_ROOT, + next_block: Header::ENCODED_SIZE as u32, + first_free: 0, + }), + Header::decode(&[0; 72]) + ); + + let hdr = Header { + root_ptr: Ptr::new(420).unwrap(), + root_hash: ONE.clone(), + next_block: 42, + first_free: 24, + }; + let got_bytes = hdr.encode(); + let got_hdr = Header::decode(&got_bytes); + + #[rustfmt::skip] + assert_eq!([ + /* magic: */ 1, 0, 0, 0, + /* version: */ 0, 0, 0, 0, + /* root_ptr: */ 164, 1, 0, 0, + /* root_hash: */ 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + /* next_block: */ 42, 0, 0, 0, + /* first_free: */ 24, 0, 0, 0, + /* tail: */ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + ], got_bytes); + assert_eq!(Some(hdr), got_hdr); +} + +#[test] +fn test_trie_sanity() { + const ONE: CryptoHash = CryptoHash([1; 32]); + + let key = solana_program::pubkey::Pubkey::new_unique(); + let mut lamports: u64 = 10 * solana_program::native_token::LAMPORTS_PER_SOL; + let mut data = [0; SZ * 1000]; + let owner = solana_program::pubkey::Pubkey::new_unique(); + let account = solana_program::account_info::AccountInfo::new( + /* key: */ &key, + /* is signer: */ false, + /* is writable: */ true, + /* lamports: */ &mut lamports, + /* data: */ &mut data[..], + /* owner: */ &owner, + /* executable: */ false, + /* rent_epoch: */ 42, + ); + + { + let mut trie = AccountTrie::new(account.data.borrow_mut()).unwrap(); + assert_eq!(Ok(None), trie.get(&[0])); + + assert_eq!(Ok(()), trie.set(&[0], &ONE)); + assert_eq!(Ok(Some(ONE.clone())), trie.get(&[0])); + } + + { + let mut trie = AccountTrie::new(account.data.borrow_mut()).unwrap(); + assert_eq!(Ok(Some(ONE.clone())), trie.get(&[0])); + + assert_eq!(Ok(()), trie.seal(&[0])); + assert_eq!(Err(sealable_trie::Error::Sealed), trie.get(&[0])); + } +} diff --git a/stdx/src/lib.rs b/stdx/src/lib.rs index bff33cfb..922dda36 100644 --- a/stdx/src/lib.rs +++ b/stdx/src/lib.rs @@ -36,8 +36,28 @@ pub fn rsplit_at(xs: &[u8]) -> Option<(&[u8], &[u8; R])> { Some((head, tail.try_into().unwrap())) } +/// Splits the slice into a slice of `N`-element arrays and a remainder slice +/// with length strictly less than `N`. +/// +/// This is simplified copy of Rust unstable `[T]::as_chunks_mut` method. +pub fn as_chunks_mut(slice: &mut [T]) -> &mut [[T; N]] { + let () = AssertNonZero::::OK; + + let ptr = slice.as_mut_ptr().cast(); + let len = slice.len() / N; + // SAFETY: We already panicked for zero, and ensured by construction + // that the length of the subslice is a multiple of N. + unsafe { core::slice::from_raw_parts_mut(ptr, len) } +} + /// Asserts, at compile time, that `A + B == S`. struct AssertEqSum; impl AssertEqSum { const OK: () = assert!(S == A + B); } + +/// Asserts, at compile time, that `N` is non-zero. +struct AssertNonZero; +impl AssertNonZero { + const OK: () = assert!(N != 0); +} diff --git a/tsconfig.json b/tsconfig.json new file mode 100644 index 00000000..8941ae14 --- /dev/null +++ b/tsconfig.json @@ -0,0 +1,15 @@ +{ + "extends": "@tsconfig/recommended/tsconfig.json", + "ts-node": { + "compilerOptions": { + "module": "commonjs" + } + }, + "compilerOptions": { + "declaration": true, + "moduleResolution": "node", + "module": "es2015" + }, + "include": ["src/client/**/*"], + "exclude": ["node_modules"] +}