Appearance
Tree
IMPORTANT
The examples are designed to be run in order after completing the quickstart, which will help you set up your environment and get familiar with common commands.
If you are just here to browse, enjoy!
💡 Background
This example implements a red-black tree in both SBPF assembly and Rust. Both implementations are compared side-by-side with as much implementation parity as possible, using C-style Rust (raw pointers, direct syscalls) to minimize compiler overhead.
The example implements a full insertion algorithm and simple removal algorithms (no rebalancing on remove). The full remove implementation is left as an exercise to the reader.
Core data structures
rs
#[repr(u8)]
#[derive(PartialEq)]
pub enum Color {
Black,
Red,
}
#[repr(usize)]
pub enum Direction {
Left,
Right,
}
constant_group! {
/// Tree constants.
tree {
/// Max number of children per node.
N_CHILDREN: usize = 2,
/// Left direction.
DIR_L = Direction::Left as usize,
/// Right direction.
DIR_R = Direction::Right as usize,
/// Black color.
COLOR_B = Color::Black as u8,
/// Red color.
COLOR_R = Color::Red as u8,
/// Stack top field in header.
offset!(HEADER_TOP, TreeHeader.top),
/// Next node field in header.
offset!(HEADER_NEXT, TreeHeader.next),
}
}
#[repr(C, packed)]
/// Tree account data header. Contains pointer to tree root and top of free node stack.
pub struct TreeHeader {
/// Absolute pointer to tree root in memory map.
pub root: *mut TreeNode,
/// Absolute pointer to stack top in memory map.
pub top: *mut StackNode,
/// Absolute pointer to where the next node should be allocated in memory map.
pub next: *mut TreeNode,
}
#[array_fields]
#[repr(C, packed)]
pub struct TreeNode {
pub parent: *mut TreeNode,
pub child: [*mut TreeNode; tree::N_CHILDREN],
pub key: u16,
pub value: u16,
pub color: Color,
}
#[repr(C, packed)]
pub struct StackNode {
pub next: *mut StackNode,
}Note that these data structures rely on direct addressing, which may be broken by ABI v2.
🏗️ Build support
Constants, error codes, and C bindings are derived in a shared interface using macros, then automatically inserted into the assembly program file at build time.
Interface
rs
use core::mem::size_of;
use macros::{array_fields, constant_group, error_codes};
use pinocchio::{
account::{RuntimeAccount as RuntimeAccountHeader, MAX_PERMITTED_DATA_INCREASE},
sysvars::rent::{Rent, ACCOUNT_STORAGE_OVERHEAD},
Address,
};
error_codes! {
/// An invalid number of accounts were passed.
N_ACCOUNTS,
/// The user account has invalid data length.
USER_DATA_LEN,
/// The tree account has invalid data length.
TREE_DATA_LEN,
/// The System Program account has invalid data length.
SYSTEM_PROGRAM_DATA_LEN,
/// The tree account is a duplicate.
TREE_DUPLICATE,
/// The System Program account is a duplicate.
SYSTEM_PROGRAM_DUPLICATE,
/// The rent sysvar account is a duplicate.
RENT_DUPLICATE,
/// The rent sysvar account has invalid data length.
RENT_ADDRESS,
/// Instruction data provided during initialization instruction.
INSTRUCTION_DATA,
/// The passed PDA does not match the expected address.
PDA_MISMATCH,
/// Invalid instruction discriminator.
INSTRUCTION_DISCRIMINATOR,
/// Invalid instruction data length.
INSTRUCTION_DATA_LEN,
/// Not enough accounts passed for insertion allocation.
N_ACCOUNTS_INSERT_ALLOCATION,
/// Key already exists in tree during insertion.
KEY_EXISTS,
/// Key does not exist in tree during removal.
KEY_DOES_NOT_EXIST,
}
constant_group! {
/// Input buffer layout.
input_buffer {
/// Number of accounts field.
offset!(N_ACCOUNTS, InputBufferHeader.n_accounts),
/// User runtime account.
offset!(USER_ACCOUNT, InputBufferHeader.user),
/// User Lamports field.
offset!(USER_LAMPORTS, InputBufferHeader.user.header.lamports),
/// User data field.
offset!(USER_DATA, InputBufferHeader.user.data),
/// User owner field.
offset!(USER_OWNER, InputBufferHeader.user.header.owner),
/// Tree Lamports field.
offset!(TREE_LAMPORTS, InputBufferHeader.tree_header.lamports),
/// Tree data field.
offset!(TREE_DATA, InitInputBuffer.header.tree.data),
/// Tree owner field.
offset!(TREE_OWNER, InputBufferHeader.tree_header.owner),
/// Tree runtime account header.
offset!(TREE_ACCOUNT, InputBufferHeader.tree_header),
/// Tree address field.
offset!(TREE_ADDRESS, InputBufferHeader.tree_header.address),
/// System Program runtime account header.
offset!(SYSTEM_PROGRAM_ACCOUNT, InitInputBuffer.header.system_program),
/// Rent sysvar account header.
offset!(RENT_ACCOUNT, InitInputBuffer.header.rent),
/// Rent sysvar account data.
offset!(RENT_DATA, InitInputBuffer.header.rent.data),
/// Expected number of accounts for general instructions.
N_ACCOUNTS_GENERAL: u64 = 2,
/// Expected number of accounts for tree initialization.
N_ACCOUNTS_INIT: u64 = 4,
/// Expected data length of system program account.
SYSTEM_PROGRAM_DATA_LEN: usize = b"system_program".len(),
/// Expected data length of rent sysvar account.
// Includes extra byte for deprecated burn_percent field that is still present in test
// framework.
RENT_DATA_LEN: usize = size_of::<Rent>() + size_of::<u8>(),
}
}
constant_group! {
/// CPI-specific constants.
cpi {
prefix = "CPI",
/// User and tree accounts.
N_ACCOUNTS: usize = 2,
/// The tree account is a PDA for CreateAccount CPI.
N_PDA_SIGNERS: usize = 1,
/// Number of seeds for CreateAccount PDA signer (bump only).
N_SEEDS_CREATE_ACCOUNT: usize = 1,
/// PDA signers for Transfer CPI (none — user is already a signer).
N_PDA_SIGNERS_TRANSFER: u64 = 0,
/// Number of seeds for PDA generation.
N_SEEDS_TRY_FIND_PDA: u64 = 0,
/// Tree account data length.
TREE_DATA_LEN: usize = size_of::<TreeHeader>(),
/// Account data scalar for base rent calculation.
ACCOUNT_DATA_SCALAR: usize = (ACCOUNT_STORAGE_OVERHEAD as usize) + TREE_DATA_LEN,
/// CreateAccount discriminator for CPI.
CREATE_ACCOUNT_DISCRIMINATOR: u32 = 0,
/// Length of CreateAccount instruction data.
CREATE_ACCOUNT_INSN_DATA_LEN: usize = size_of::<CreateAccountInstructionData>(),
/// Transfer discriminator for CPI.
TRANSFER_DISCRIMINATOR: u32 = 2,
/// Length of Transfer instruction data.
TRANSFER_INSN_DATA_LEN: usize = size_of::<TransferInstructionData>(),
/// Mask for writable signer.
WRITABLE_SIGNER: u64 = 0x0101,
/// Account index for user account in CPI.
USER_ACCOUNT_INDEX: usize = 0,
/// Account index for tree account in CPI.
TREE_ACCOUNT_INDEX: usize = 1,
/// Null rent epoch.
RENT_EPOCH_NULL: u64 = 0,
}
}
#[repr(C, packed)]
/// For CPI to create tree account.
pub struct CreateAccountInstructionData {
pub discriminator: u32,
pub lamports: u64,
pub space: u64,
pub owner: Address,
}
#[repr(C, packed)]
/// For CPI to transfer lamports.
pub struct TransferInstructionData {
pub discriminator: u32,
pub lamports: u64,
}
constant_group! {
/// Data layout constants.
data {
/// Data length of zero.
DATA_LEN_ZERO: u64 = 0,
/// Data alignment during runtime.
BPF_ALIGN_OF_U128: usize = 8,
}
}
#[repr(C, packed)]
/// Input buffer header for all instructions.
pub struct InputBufferHeader {
pub n_accounts: u64,
pub user: EmptyRuntimeAccount,
pub tree_header: RuntimeAccountHeader,
}
#[repr(C, packed)]
/// Input buffer for tree initialization instruction. Broken up to fit relative offsets in i16.
pub struct InitInputBuffer {
pub header: InitInputBufferHeader,
pub footer: InitInputBufferFooter,
}
#[repr(C, packed)]
pub struct InitInputBufferHeader {
pub _n_accounts: u64,
pub _user: EmptyRuntimeAccount,
pub tree: EmptyRuntimeAccount,
pub system_program: SystemProgramRuntimeAccount,
pub rent: RentRuntimeAccount,
}
#[repr(C, packed)]
/// Input buffer header for general tree instructions.
pub struct GeneralInputBufferHeader {
pub n_accounts: u64,
pub user: EmptyRuntimeAccount,
pub tree_header: RuntimeAccountHeader,
pub tree_data: TreeHeader,
}
// ANCHOR: tree-defs-common
#[repr(u8)]
#[derive(PartialEq)]
pub enum Color {
Black,
Red,
}
#[repr(usize)]
pub enum Direction {
Left,
Right,
}
constant_group! {
/// Tree constants.
tree {
/// Max number of children per node.
N_CHILDREN: usize = 2,
/// Left direction.
DIR_L = Direction::Left as usize,
/// Right direction.
DIR_R = Direction::Right as usize,
/// Black color.
COLOR_B = Color::Black as u8,
/// Red color.
COLOR_R = Color::Red as u8,
/// Stack top field in header.
offset!(HEADER_TOP, TreeHeader.top),
/// Next node field in header.
offset!(HEADER_NEXT, TreeHeader.next),
}
}
#[repr(C, packed)]
/// Tree account data header. Contains pointer to tree root and top of free node stack.
pub struct TreeHeader {
/// Absolute pointer to tree root in memory map.
pub root: *mut TreeNode,
/// Absolute pointer to stack top in memory map.
pub top: *mut StackNode,
/// Absolute pointer to where the next node should be allocated in memory map.
pub next: *mut TreeNode,
}
#[array_fields]
#[repr(C, packed)]
pub struct TreeNode {
pub parent: *mut TreeNode,
pub child: [*mut TreeNode; tree::N_CHILDREN],
pub key: u16,
pub value: u16,
pub color: Color,
}
#[repr(C, packed)]
pub struct StackNode {
pub next: *mut StackNode,
}
// ANCHOR_END: tree-defs-common
// ANCHOR: instructions
#[repr(u8)]
pub enum Instruction {
/// Initialize the tree.
Initialize,
/// Insert key-value pair.
Insert,
/// Remove key-value pair.
Remove,
}
#[repr(C, packed)]
pub struct InstructionHeader {
pub discriminator: u8,
}
#[repr(C, packed)]
pub struct InitializeInstruction {
pub header: InstructionHeader,
}
#[repr(C, packed)]
pub struct InsertInstruction {
pub header: InstructionHeader,
pub key: u16,
pub value: u16,
}
#[repr(C, packed)]
pub struct RemoveInstruction {
pub header: InstructionHeader,
pub key: u16,
}
constant_group! {
/// Offsets for instruction processing.
instruction {
prefix = "INSN",
/// Offset to instruction discriminator byte.
offset!(DISCRIMINATOR, InstructionHeader.discriminator),
/// Initialize instruction discriminator.
DISCRIMINATOR_INITIALIZE: u8 = Instruction::Initialize as u8,
/// Insert instruction discriminator.
DISCRIMINATOR_INSERT: u8 = Instruction::Insert as u8,
/// Remove instruction discriminator.
DISCRIMINATOR_REMOVE: u8 = Instruction::Remove as u8,
/// Key field in insert instruction.
offset!(INSERT_KEY, InsertInstruction.key),
/// Value field in insert instruction.
offset!(INSERT_VALUE, InsertInstruction.value),
/// Key field in remove instruction.
offset!(REMOVE_KEY, RemoveInstruction.key),
}
}
// ANCHOR_END: instructions
#[repr(C, packed)]
pub struct InitInputBufferFooter {
pub instruction_data_len: u64,
pub instruction: InitializeInstruction,
pub program_id: Address,
}
#[repr(C)]
pub struct RuntimeAccount<const DATA_SIZE: usize> {
pub header: RuntimeAccountHeader,
pub data: [u8; DATA_SIZE],
pub rent_epoch: u64,
}
type EmptyRuntimeAccount = RuntimeAccount<{ runtime_data_size(data::DATA_LEN_ZERO as usize) }>;
type SystemProgramRuntimeAccount =
RuntimeAccount<{ runtime_data_size(input_buffer::SYSTEM_PROGRAM_DATA_LEN) }>;
type RentRuntimeAccount = RuntimeAccount<{ runtime_data_size(input_buffer::RENT_DATA_LEN) }>;
/// Compute the data buffer size for a runtime account with the given data length.
const fn runtime_data_size(data_len: usize) -> usize {
MAX_PERMITTED_DATA_INCREASE + data_len.next_multiple_of(data::BPF_ALIGN_OF_U128)
}rs
extern crate alloc;
use crate::bindings::{
SolAccountInfo, SolAccountMeta, SolInstruction, SolSignerSeed, SolSignerSeeds,
};
use crate::common::{
cpi, CreateAccountInstructionData, Direction, GeneralInputBufferHeader, InitInputBuffer,
InitializeInstruction, InputBufferHeader, InsertInstruction, Instruction, RemoveInstruction,
TreeHeader, TreeNode,
};
use macros::{asm_constant_group, extend_constant_group, pubkey_chunk_group, sizes, stack_frame};
use pinocchio::{
entrypoint::NON_DUP_MARKER,
sysvars::rent::{Rent, RENT_ID},
Address,
};
pubkey_chunk_group!();
sizes! {
u8,
u64,
Address,
u128,
TreeHeader,
InitializeInstruction,
InsertInstruction,
RemoveInstruction,
TreeNode,
}
extend_constant_group!(data {
/// No offset.
OFFSET_ZERO = 0,
/// Null pointer.
NULL = 0,
/// And mask for data length alignment.
DATA_LEN_AND_MASK = -8,
/// Maximum possible data length padding.
MAX_DATA_PAD = 7,
/// Boolean true value.
BOOL_TRUE = 1,
});
extend_constant_group!(input_buffer {
prefix = "IB",
/// User address field.
offset!(USER_ADDRESS, InputBufferHeader.user.header.address),
/// User data length field.
offset!(USER_DATA_LEN, InputBufferHeader.user.header.data_len),
/// Non-duplicate marker value.
NON_DUP_MARKER = NON_DUP_MARKER,
/// Tree non-duplicate marker field.
offset!(TREE_NON_DUP_MARKER, InputBufferHeader.tree_header.borrow_state),
/// Tree address field.
pubkey_offset!(TREE_ADDRESS, InputBufferHeader.tree_header.address),
/// Tree data length field.
offset!(TREE_DATA_LEN, InputBufferHeader.tree_header.data_len),
/// System Program non-duplicate marker field.
offset!(
SYSTEM_PROGRAM_NON_DUP_MARKER,
InitInputBuffer.header.system_program.header.borrow_state
),
/// System Program data length field.
offset!(SYSTEM_PROGRAM_DATA_LEN, InitInputBuffer.header.system_program.header.data_len),
/// Rent account non-duplicate marker field.
offset!(RENT_NON_DUP_MARKER, InitInputBuffer.header.rent.header.borrow_state),
/// Rent address field.
pubkey_offset!(RENT_ADDRESS, InitInputBuffer.header.rent.header.address),
/// Rent sysvar ID.
pubkey_value!(RENT_ID, RENT_ID),
/// Program ID field for initialize instruction.
offset_immediate!(INIT_PROGRAM_ID, InitInputBuffer.footer.program_id),
/// Tree top pointer field within tree data.
offset!(TREE_DATA_TOP, GeneralInputBufferHeader.tree_data.top),
/// Tree next pointer field within tree data.
offset!(TREE_DATA_NEXT, GeneralInputBufferHeader.tree_data.next),
/// Tree root pointer field within tree data.
offset!(TREE_DATA_ROOT, GeneralInputBufferHeader.tree_data.root),
/// Relative offset from user data field to tree pubkey field.
relative_offset_immediate!(
USER_DATA,
TREE_ADDRESS,
InputBufferHeader.user.data,
InputBufferHeader.tree_header.address
),
});
#[stack_frame]
struct InitStackFrame {
bump_seed: u8,
instruction_data: CreateAccountInstructionData,
instruction: SolInstruction,
account_metas: [SolAccountMeta; cpi::N_ACCOUNTS],
account_infos: [SolAccountInfo; cpi::N_ACCOUNTS],
signers_seeds: [SolSignerSeeds; cpi::N_PDA_SIGNERS],
signer_seeds: [SolSignerSeed; cpi::N_SEEDS_CREATE_ACCOUNT],
pda: Address,
rent: Rent,
/// Zero-initialized on stack.
system_program_address: Address,
}
asm_constant_group! {
/// Init stack frame layout.
init_stack_frame {
prefix = "SF_INIT",
/// Bump seed.
stack_frame_offset!(BUMP_SEED, InitStackFrame.bump_seed),
/// Bump signer seed address field.
stack_frame_offset!(SIGNER_SEED_ADDR, InitStackFrame.signer_seeds[0].addr),
/// Bump signer seed length field.
stack_frame_offset!(SIGNER_SEED_LEN, InitStackFrame.signer_seeds[0].len),
/// PDA address field.
stack_frame_offset!(PDA, InitStackFrame.pda),
/// Discriminator field in CPI instruction data.
stack_frame_offset_unaligned!(
CREATE_ACCOUNT_DISCRIMINATOR,
InitStackFrame.instruction_data.discriminator
),
/// Lamports field in CreateAccount instruction data.
stack_frame_offset_unaligned!(
CREATE_ACCOUNT_LAMPORTS,
InitStackFrame.instruction_data.lamports
),
/// Space address field in CreateAccount instruction data.
stack_frame_offset_unaligned!(CREATE_ACCOUNT_SPACE, InitStackFrame.instruction_data.space),
/// Owner field in CreateAccount instruction data.
stack_frame_pubkey_offset_unaligned!(
CREATE_ACCOUNT_OWNER,
InitStackFrame.instruction_data.owner
),
/// Signers seeds address field.
stack_frame_offset!(SIGNERS_SEEDS_ADDR, InitStackFrame.signers_seeds),
/// Signers seeds length field.
stack_frame_offset!(SIGNERS_SEEDS_LEN, InitStackFrame.signers_seeds[0].len),
/// System Program address.
stack_frame_offset!(SYSTEM_PROGRAM_ADDRESS, InitStackFrame.system_program_address),
/// SolInstruction program_id field.
stack_frame_offset!(INSN_PROGRAM_ID, InitStackFrame.instruction.program_id),
/// SolInstruction accounts field.
stack_frame_offset!(INSN_ACCOUNTS, InitStackFrame.instruction.accounts),
/// SolInstruction account_len field.
stack_frame_offset!(INSN_ACCOUNT_LEN, InitStackFrame.instruction.account_len),
/// SolInstruction data field.
stack_frame_offset!(INSN_DATA, InitStackFrame.instruction.data),
/// SolInstruction data_len field.
stack_frame_offset!(INSN_DATA_LEN, InitStackFrame.instruction.data_len),
/// SolAccountMeta is_writable field for user account.
stack_frame_offset!(
USER_META_IS_WRITABLE,
InitStackFrame.account_metas[cpi::USER_ACCOUNT_INDEX].is_writable
),
/// SolAccountMeta is_writable field for tree account.
stack_frame_offset!(
TREE_META_IS_WRITABLE,
InitStackFrame.account_metas[cpi::TREE_ACCOUNT_INDEX].is_writable
),
/// SolAccountInfo is_signer field for user account.
stack_frame_offset!(
USER_INFO_IS_SIGNER,
InitStackFrame.account_infos[cpi::USER_ACCOUNT_INDEX].is_signer
),
/// SolAccountMeta pubkey field for user account.
stack_frame_offset!(
USER_META_PUBKEY,
InitStackFrame.account_metas[cpi::USER_ACCOUNT_INDEX].pubkey
),
/// SolAccountInfo pubkey field for user account.
stack_frame_offset!(
USER_INFO_PUBKEY,
InitStackFrame.account_infos[cpi::USER_ACCOUNT_INDEX].key
),
/// SolAccountInfo owner field for user account.
stack_frame_offset!(
USER_INFO_OWNER,
InitStackFrame.account_infos[cpi::USER_ACCOUNT_INDEX].owner
),
/// SolAccountInfo lamports field for user account.
stack_frame_offset!(
USER_INFO_LAMPORTS,
InitStackFrame.account_infos[cpi::USER_ACCOUNT_INDEX].lamports
),
/// SolAccountInfo data_len field for user account.
stack_frame_offset!(
USER_INFO_DATA,
InitStackFrame.account_infos[cpi::USER_ACCOUNT_INDEX].data
),
/// SolAccountInfo data_len for tree account.
stack_frame_offset!(
TREE_INFO_DATA_LEN,
InitStackFrame.account_infos[cpi::TREE_ACCOUNT_INDEX].data_len
),
/// SolAccountInfo is_signer field for tree account.
stack_frame_offset!(
TREE_INFO_IS_SIGNER,
InitStackFrame.account_infos[cpi::TREE_ACCOUNT_INDEX].is_signer
),
/// SolAccountInfo is_writable field for tree account.
stack_frame_offset_unaligned!(
TREE_INFO_IS_WRITABLE,
InitStackFrame.account_infos[cpi::TREE_ACCOUNT_INDEX].is_writable
),
/// SolAccountMeta pubkey field for tree account.
stack_frame_offset!(
TREE_META_PUBKEY,
InitStackFrame.account_metas[cpi::TREE_ACCOUNT_INDEX].pubkey
),
/// SolAccountInfo pubkey field for tree account.
stack_frame_offset!(
TREE_INFO_PUBKEY,
InitStackFrame.account_infos[cpi::TREE_ACCOUNT_INDEX].key
),
/// SolAccountInfo owner field for tree account.
stack_frame_offset!(
TREE_INFO_OWNER,
InitStackFrame.account_infos[cpi::TREE_ACCOUNT_INDEX].owner
),
/// SolAccountInfo lamports field for tree account.
stack_frame_offset!(
TREE_INFO_LAMPORTS,
InitStackFrame.account_infos[cpi::TREE_ACCOUNT_INDEX].lamports
),
/// SolAccountInfo data_len field for tree account.
stack_frame_offset!(
TREE_INFO_DATA,
InitStackFrame.account_infos[cpi::TREE_ACCOUNT_INDEX].data
),
/// Relative offset from PDA on stack to System Program ID.
relative_offset_immediate!(
PDA,
SYSTEM_PROGRAM_ID,
InitStackFrame.pda,
InitStackFrame.system_program_address
),
/// Relative offset from System Program ID to first SolAccountMeta.
relative_offset_immediate!(
SYSTEM_PROGRAM_ID,
ACCT_METAS,
InitStackFrame.system_program_address,
InitStackFrame.account_metas
),
/// Relative offset from SolAccountMeta array to instruction data.
relative_offset_immediate!(
ACCT_METAS,
INSN_DATA,
InitStackFrame.account_metas,
InitStackFrame.instruction_data
),
/// Relative offset from instruction data to signer seeds.
relative_offset_immediate!(
INSN_DATA,
SIGNER_SEEDS,
InitStackFrame.instruction_data,
InitStackFrame.signer_seeds
),
/// Relative offset from signer seeds to signers seeds.
relative_offset_immediate!(
SIGNER_SEEDS,
SIGNERS_SEEDS,
InitStackFrame.signer_seeds,
InitStackFrame.signers_seeds
),
/// Account infos array.
stack_frame_offset!(ACCT_INFOS, InitStackFrame.account_infos),
}
}
extend_constant_group!(tree {
prefix = "TREE",
/// Discriminator for insert instruction.
DISCRIMINATOR_INSERT = Instruction::Insert as u8,
/// Node key field.
offset!(NODE_KEY, TreeNode.key),
/// Node value field.
offset!(NODE_VALUE, TreeNode.value),
/// Node left child field.
offset!(NODE_CHILD_L, TreeNode.child[Direction::Left as usize]),
/// Node right child field.
offset!(NODE_CHILD_R, TreeNode.child[Direction::Right as usize]),
/// Node parent field.
offset!(NODE_PARENT, TreeNode.parent),
/// Color field.
offset!(NODE_COLOR, TreeNode.color),
});rs
/// Generated from Agave using bindgen.
use pinocchio::Address;
/// SolInstruction from cpi.h.
#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct SolInstruction {
/// Pubkey of the instruction processor that executes this instruction.
pub program_id: *mut Address,
/// Metadata for what accounts should be passed to the instruction processor.
pub accounts: *mut SolAccountMeta,
/// Number of SolAccountMetas.
pub account_len: u64,
/// Opaque data passed to the instruction processor.
pub data: *mut u8,
/// Length of the data in bytes.
pub data_len: u64,
}
/// SolAccountMeta from cpi.h.
#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct SolAccountMeta {
/// An account's public key.
pub pubkey: *mut Address,
/// True if the `pubkey` can be loaded as a read-write account.
pub is_writable: bool,
/// True if an Instruction requires a Transaction signature matching `pubkey`.
pub is_signer: bool,
}
/// SolAccountInfo from entrypoint.h.
#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct SolAccountInfo {
/// Public key of the account.
pub key: *mut Address,
/// Number of lamports owned by this account.
pub lamports: *mut u64,
/// Length of data in bytes.
pub data_len: u64,
/// On-chain data within this account.
pub data: *mut u8,
/// Program that owns this account.
pub owner: *mut Address,
/// The epoch at which this account will next owe rent.
pub rent_epoch: u64,
/// Transaction was signed by this account's key?
pub is_signer: bool,
/// Is the account writable?
pub is_writable: bool,
/// This account's data contains a loaded program (and is now read-only).
pub executable: bool,
}
/// SolSignerSeed from pubkey.h.
#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct SolSignerSeed {
/// Seed bytes.
pub addr: *const u8,
/// Length of the seed bytes.
pub len: u64,
}
/// SolSignerSeeds from pubkey.h.
#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct SolSignerSeeds {
/// An array of a signer's seeds.
pub addr: *const SolSignerSeed,
/// Number of seeds.
pub len: u64,
}build.rs
rs
use std::{collections::HashSet, fs, path::Path};
use tree_interface::*;
const CONSTANTS_ANCHOR_START: &str = "# ANCHOR: constants";
const CONSTANTS_ANCHOR_END: &str = "# ANCHOR_END: constants";
macro_rules! asm_groups {
($($group:ident),* $(,)?) => {
[$($group::to_asm()),*]
};
}
fn main() {
// Collect all constant groups.
let groups = asm_groups![
error_codes,
sizes,
data,
pubkey_chunk,
input_buffer,
instruction,
init_stack_frame,
cpi,
tree
];
// Read in the assembly file.
let manifest_dir = env!("CARGO_MANIFEST_DIR");
let manifest_path = Path::new(manifest_dir);
let asm_path = manifest_path.join("src/tree/tree.s");
let content = fs::read_to_string(&asm_path).unwrap();
// Find the constants anchor region.
let anchor_start = content
.find(CONSTANTS_ANCHOR_START)
.expect("missing '# ANCHOR: constants' in assembly file");
let anchor_end = content
.find(CONSTANTS_ANCHOR_END)
.expect("missing '# ANCHOR_END: constants' in assembly file");
assert!(
anchor_start < anchor_end,
"ANCHOR: constants must come before ANCHOR_END: constants"
);
// Check for duplicate constant names.
let mut seen = HashSet::new();
for group in &groups {
for line in group.lines() {
if let Some(name) = line.strip_prefix(".equ ") {
if let Some(name) = name.split(',').next() {
if !seen.insert(name.to_string()) {
panic!("Duplicate constant name: {}", name);
}
}
}
}
}
// Generate the constants and insert inside the anchor region.
let constants = groups.join("\n");
let constants = constants.trim();
let before_anchor = &content[..anchor_start + CONSTANTS_ANCHOR_START.len()];
let after_anchor = &content[anchor_end..];
let new_content = format!("{}\n{}\n{}", before_anchor, constants, after_anchor);
if new_content != content {
fs::write(&asm_path, &new_content).unwrap();
}
// Extract ANCHOR snippets from source files (use new_content for asm since it's canonical).
extract_snippets(&new_content, manifest_path, "asm");
let rs_path = manifest_path.join("src/program.rs");
let rs_content = fs::read_to_string(&rs_path).unwrap();
extract_snippets(&rs_content, manifest_path, "rs");
let common_path = manifest_path.join("interface/src/common.rs");
let common_content = fs::read_to_string(&common_path).unwrap();
extract_snippets(&common_content, manifest_path, "interface");
}
/// Extract ANCHOR snippets from source content and write to artifacts/snippets/{kind}/.
fn extract_snippets(content: &str, manifest_dir: &Path, kind: &str) {
let snippets_dir = manifest_dir.join(format!("artifacts/snippets/{}", kind));
// Wipe existing snippets directory.
if snippets_dir.exists() {
fs::remove_dir_all(&snippets_dir).unwrap();
}
let mut current_anchor: Option<String> = None;
let mut current_lines: Vec<&str> = Vec::new();
let mut snippets: Vec<(String, String)> = Vec::new();
let mut seen_names: HashSet<String> = HashSet::new();
for line in content.lines() {
let trimmed = line.trim();
// Check for ANCHOR: start tag (supports # or // comments).
if let Some(rest) = trimmed
.strip_prefix("# ANCHOR:")
.or(trimmed.strip_prefix("// ANCHOR:"))
{
let name = rest.trim().to_string();
assert!(
current_anchor.is_none(),
"Nested ANCHOR not allowed: found '{}' while inside '{}'",
name,
current_anchor.unwrap()
);
assert!(
seen_names.insert(name.clone()),
"Duplicate ANCHOR name: '{}'",
name
);
current_anchor = Some(name);
current_lines.clear();
continue;
}
// Check for ANCHOR_END: end tag.
if let Some(rest) = trimmed
.strip_prefix("# ANCHOR_END:")
.or(trimmed.strip_prefix("// ANCHOR_END:"))
{
let name = rest.trim();
if let Some(ref anchor_name) = current_anchor {
assert_eq!(
anchor_name, name,
"ANCHOR_END mismatch: expected '{}', found '{}'",
anchor_name, name
);
let snippet_content = current_lines.join("\n");
snippets.push((anchor_name.clone(), snippet_content));
current_anchor = None;
current_lines.clear();
} else {
panic!("ANCHOR_END '{}' without matching ANCHOR", name);
}
continue;
}
// Collect lines inside an anchor.
if current_anchor.is_some() {
current_lines.push(line);
}
}
assert!(
current_anchor.is_none(),
"Unclosed ANCHOR: '{}'",
current_anchor.unwrap()
);
// Write snippets to files.
if !snippets.is_empty() {
fs::create_dir_all(&snippets_dir).unwrap();
for (name, content) in snippets {
let snippet_path = snippets_dir.join(format!("{}.txt", name));
fs::write(&snippet_path, content).unwrap();
}
}
}Macros
rs
// cspell:word strs
// cspell:word idents
use convert_case::{Case, Casing};
use proc_macro::TokenStream;
use quote::quote;
use syn::{
braced, bracketed,
parse::{discouraged::Speculative, Parse, ParseStream},
parse_macro_input, Ident, Lit, LitInt, Meta, Token,
};
/// Maximum line length for ASM output.
const MAX_LINE_LEN: usize = 75;
/// BPF alignment requirement (align_of::<u128>()).
const BPF_ALIGN: i64 = 8;
/// Maximum comment length (accounting for `# ` prefix).
const MAX_COMMENT_LEN: usize = MAX_LINE_LEN - "# ".len();
/// Error code entry: doc comment + snake_case name.
struct ErrorCodeEntry {
doc: String,
name: Ident,
}
/// Input for error_codes! macro.
struct ErrorCodesInput {
entries: Vec<ErrorCodeEntry>,
}
impl Parse for ErrorCodesInput {
fn parse(input: ParseStream) -> syn::Result<Self> {
let mut entries = Vec::new();
while !input.is_empty() {
let attrs = input.call(syn::Attribute::parse_outer)?;
let doc = extract_doc_comment(&attrs)
.ok_or_else(|| input.error("Error code must have a doc comment"))?;
if let Err(e) = validate_doc_comment(&doc) {
return Err(input.error(e));
}
let name: Ident = input.parse()?;
// Optional trailing comma.
let _ = input.parse::<Token![,]>();
entries.push(ErrorCodeEntry { doc, name });
}
Ok(ErrorCodesInput { entries })
}
}
/// Macro for defining error codes shared between Rust and ASM.
///
/// Creates an `Error` enum with `#[repr(u32)]` and auto-numbered variants starting at 1.
/// Variant names are SCREAMING_SNAKE_CASE. ASM names have `E_` prefix added.
///
/// # Example
/// ```ignore
/// error_codes! {
/// /// An invalid number of accounts were passed.
/// N_ACCOUNTS_INVALID,
/// /// The user account has nonzero data length.
/// USER_HAS_DATA,
/// }
/// ```
///
/// Generates:
/// - Rust: `enum Error { N_ACCOUNTS_INVALID, USER_HAS_DATA }` with `From<Error> for u32`
/// - ASM: `.equ E_N_ACCOUNTS_INVALID, 1` and `.equ E_USER_HAS_DATA, 2`
#[proc_macro]
pub fn error_codes(input: TokenStream) -> TokenStream {
let input = parse_macro_input!(input as ErrorCodesInput);
let mut variant_defs = Vec::new();
let mut asm_lines = Vec::new();
for (idx, entry) in input.entries.iter().enumerate() {
let doc = &entry.doc;
let variant_name = &entry.name;
// Just add E_ prefix for ASM.
let asm_name = format!("E_{}", entry.name);
let value = (idx + 1) as u32;
variant_defs.push(quote! {
#[doc = #doc]
#variant_name = #value
});
asm_lines.push(asm_equ_line(&asm_name, value, doc));
}
let header = asm_header("Error codes.");
let body = asm_lines.join("\n");
let n_codes = input.entries.len() as u32;
let expanded = quote! {
pub mod error_codes {
#[repr(u32)]
#[allow(non_camel_case_types)]
pub enum error {
#(#variant_defs),*
}
impl error {
/// Total number of error codes.
pub const N_CODES: u32 = #n_codes;
}
impl From<error> for u32 {
fn from(e: error) -> u32 {
e as u32
}
}
impl From<error> for u64 {
fn from(e: error) -> u64 {
e as u64
}
}
/// Generate ASM constants for error codes.
pub fn to_asm() -> alloc::string::String {
alloc::format!("{}\n{}\n", #header, #body)
}
}
};
TokenStream::from(expanded)
}
/// Generate an ASM section header with auto-width dashes.
fn asm_header(title: &str) -> String {
let dash_len = title.len().min(MAX_COMMENT_LEN);
format!("# {}\n# {}", title, "-".repeat(dash_len))
}
/// Validate a doc comment: must end with period and fit within max length.
fn validate_doc_comment(comment: &str) -> Result<(), String> {
if !comment.ends_with('.') {
return Err(format!("Doc comment must end with a period: {:?}", comment));
}
if comment.len() > MAX_COMMENT_LEN {
return Err(format!(
"Doc comment exceeds max length of {} chars (got {}): {:?}",
MAX_COMMENT_LEN,
comment.len(),
comment
));
}
Ok(())
}
/// Format an ASM .equ line. If inline comment would exceed max line length,
/// put the comment on its own line above.
fn asm_equ_line(name: &str, value: impl std::fmt::Display, comment: &str) -> String {
let inline = format!(".equ {}, {} # {}", name, value, comment);
if inline.len() <= MAX_LINE_LEN {
inline
} else {
format!("# {}\n.equ {}, {}", comment, name, value)
}
}
/// Extract the doc comment from attributes.
fn extract_doc_comment(attrs: &[syn::Attribute]) -> Option<String> {
let mut doc_parts = Vec::new();
for attr in attrs {
if attr.path().is_ident("doc") {
if let Meta::NameValue(meta) = &attr.meta {
if let syn::Expr::Lit(expr_lit) = &meta.value {
if let Lit::Str(lit_str) = &expr_lit.lit {
doc_parts.push(lit_str.value().trim().to_string());
}
}
}
}
}
if doc_parts.is_empty() {
None
} else {
Some(doc_parts.join(" "))
}
}
enum ConstantKind {
/// Regular constant with explicit type and value.
Value {
ty: syn::Type,
value: syn::Expr,
/// Original literal string for ASM output (preserves hex/binary).
literal_repr: Option<String>,
},
/// Offset derived from struct field path (i16 validated).
/// Name gets `_OFF` suffix appended.
/// Supports optional array indexing via `__<Struct>_fields` companion module.
Offset {
struct_name: Ident,
field_path_tokens: proc_macro2::TokenStream,
array_index: Option<ArrayIndexInfo>,
},
/// Offset derived from struct field path (i32 validated).
/// Name gets `_OFF_IMM` suffix appended. For use as BPF immediate operand.
/// Supports optional array indexing via `__<Struct>_fields` companion module.
OffsetImmediate {
struct_name: Ident,
field_path_tokens: proc_macro2::TokenStream,
array_index: Option<ArrayIndexInfo>,
},
/// Negative offset from end of a stack frame struct (i16 validated).
/// Computed as `offset_of!(Struct, field) - size_of::<Struct>()`.
/// Name gets `_OFF` suffix (aligned) or `_UOFF` suffix (unaligned).
/// Array element types are resolved via the `__<Struct>_fields` companion module
/// generated by `#[stack_frame]`.
StackFrameOffset {
struct_name: Ident,
field_path_tokens: proc_macro2::TokenStream,
array_index: Option<ArrayIndexInfo>,
aligned: bool,
},
/// Pubkey chunk offset (i16 validated).
/// Base offset + chunk_index * 8 for 8-byte register loads.
/// Name gets `_OFF_{chunk_index}` suffix appended.
PubkeyChunkOffset {
struct_name: Ident,
field_path: Vec<Ident>,
chunk_index: usize,
},
/// Negative stack frame offset for pubkey chunks (i16 validated).
/// Computed as `offset_of!(Struct, field) + chunk_index * 8 - size_of::<Struct>()`.
/// Always unaligned. Name gets `_UOFF_{chunk_index}` suffix appended.
StackFramePubkeyChunkOffset {
struct_name: Ident,
field_path_tokens: proc_macro2::TokenStream,
array_index: Option<ArrayIndexInfo>,
chunk_index: usize,
},
/// Relative offset between two fields (i32 validated).
/// Computed as `offset_of!(Struct, to_field) - offset_of!(Struct, from_field)`.
/// Name is `{FROM}_TO_{TO}_REL_OFF_IMM`.
RelativeOffsetImmediate {
from_struct_name: Ident,
from_field_path: Vec<Ident>,
to_struct_name: Ident,
to_field_path: Vec<Ident>,
},
}
/// Array index info for stack frame offset computation.
/// Supports expression indices (e.g., `CpiAccountIndex::User`) and
/// optional inner field access (e.g., `.is_writable`).
struct ArrayIndexInfo {
/// The array field name on the struct (for type alias lookup).
array_field_name: Ident,
/// The index expression (e.g., `0`, `CpiAccountIndex::User`).
index_expr: syn::Expr,
/// Optional inner field path after the array index (e.g., `is_writable`).
inner_field_path: Option<proc_macro2::TokenStream>,
}
struct ConstantDef {
doc: String,
name: Ident,
kind: ConstantKind,
}
struct ConstantGroup {
doc: String,
name: Ident,
prefix: Option<String>,
constants: Vec<ConstantDef>,
}
impl Parse for ConstantGroup {
fn parse(input: ParseStream) -> syn::Result<Self> {
// Parse doc comments for the module.
let attrs = input.call(syn::Attribute::parse_outer)?;
let doc = extract_doc_comment(&attrs)
.ok_or_else(|| input.error("Module must have a doc comment"))?;
// Validate group doc comment.
if let Err(e) = validate_doc_comment(&doc) {
return Err(input.error(format!("Group doc comment: {}", e)));
}
// Parse group name.
let name: Ident = input.parse()?;
// Parse module body.
let content;
braced!(content in input);
// Parse optional header parameter: prefix = "..."
let prefix = parse_group_params(&content)?;
// Parse constants.
let mut constants = Vec::new();
while !content.is_empty() {
let const_attrs = content.call(syn::Attribute::parse_outer)?;
let const_doc = extract_doc_comment(&const_attrs)
.ok_or_else(|| content.error("Constant must have a doc comment"))?;
// Validate constant doc comment.
if let Err(e) = validate_doc_comment(&const_doc) {
return Err(content.error(e));
}
let ident: Ident = content.parse()?;
// pubkey_offset!(NAME, Struct.field) expands to 4 chunk constants.
if ident == "pubkey_offset" {
content.parse::<Token![!]>()?;
let inner;
syn::parenthesized!(inner in content);
let base_name: Ident = inner.parse()?;
inner.parse::<Token![,]>()?;
let struct_name: Ident = inner.parse()?;
let mut field_path = Vec::new();
while inner.peek(Token![.]) {
inner.parse::<Token![.]>()?;
field_path.push(inner.parse::<Ident>()?);
}
if field_path.is_empty() {
return Err(inner.error("Expected at least one field after struct name"));
}
let _ = content.parse::<Token![,]>();
let base_doc = const_doc.trim_end_matches('.');
for chunk in 0..4usize {
let chunk_name =
Ident::new(&format!("{}_OFF_{}", base_name, chunk), base_name.span());
let chunk_doc = format!("{} (chunk index {}).", base_doc, chunk);
if let Err(e) = validate_doc_comment(&chunk_doc) {
return Err(content.error(e));
}
constants.push(ConstantDef {
doc: chunk_doc,
name: chunk_name,
kind: ConstantKind::PubkeyChunkOffset {
struct_name: struct_name.clone(),
field_path: field_path.clone(),
chunk_index: chunk,
},
});
}
continue;
}
// stack_frame_pubkey_offset_unaligned!(NAME, Struct.field) expands to 4 chunk constants.
if ident == "stack_frame_pubkey_offset_unaligned" {
content.parse::<Token![!]>()?;
let inner;
syn::parenthesized!(inner in content);
let base_name: Ident = inner.parse()?;
inner.parse::<Token![,]>()?;
let struct_name: Ident = inner.parse()?;
let parsed = parse_indexed_field_path(&inner)?;
let _ = content.parse::<Token![,]>();
let base_doc = const_doc.trim_end_matches('.');
for chunk in 0..4usize {
let chunk_name =
Ident::new(&format!("{}_UOFF_{}", base_name, chunk), base_name.span());
let chunk_doc = format!("{} (chunk index {}).", base_doc, chunk);
if let Err(e) = validate_doc_comment(&chunk_doc) {
return Err(content.error(e));
}
constants.push(ConstantDef {
doc: chunk_doc,
name: chunk_name,
kind: ConstantKind::StackFramePubkeyChunkOffset {
struct_name: struct_name.clone(),
field_path_tokens: parsed.field_path_tokens.clone(),
array_index: parsed.array_index.as_ref().map(|a| ArrayIndexInfo {
array_field_name: a.array_field_name.clone(),
index_expr: a.index_expr.clone(),
inner_field_path: a.inner_field_path.clone(),
}),
chunk_index: chunk,
},
});
}
continue;
}
// Support `offset!(NAME, Struct.field)`, `offset_immediate!(NAME, Struct.field)`,
// `relative_offset_immediate!(FROM, TO, Struct.from, Struct.to)`,
// `NAME: type = value`, and `NAME = expr as Type` forms.
let (const_name, kind) = if ident == "relative_offset_immediate" {
// relative_offset_immediate!(FROM_NAME, TO_NAME, Struct.from.path, Struct.to.path)
content.parse::<Token![!]>()?;
let inner;
syn::parenthesized!(inner in content);
let from_name: Ident = inner.parse()?;
inner.parse::<Token![,]>()?;
let to_name: Ident = inner.parse()?;
inner.parse::<Token![,]>()?;
let from_struct_name: Ident = inner.parse()?;
let mut from_field_path = Vec::new();
while inner.peek(Token![.]) {
inner.parse::<Token![.]>()?;
from_field_path.push(inner.parse::<Ident>()?);
}
if from_field_path.is_empty() {
return Err(inner.error("Expected at least one field after struct name"));
}
inner.parse::<Token![,]>()?;
let to_struct_name: Ident = inner.parse()?;
let mut to_field_path = Vec::new();
while inner.peek(Token![.]) {
inner.parse::<Token![.]>()?;
to_field_path.push(inner.parse::<Ident>()?);
}
if to_field_path.is_empty() {
return Err(inner.error("Expected at least one field after struct name"));
}
let full_name = Ident::new(
&format!("{}_TO_{}_REL_OFF_IMM", from_name, to_name),
from_name.span(),
);
(
full_name,
ConstantKind::RelativeOffsetImmediate {
from_struct_name,
from_field_path,
to_struct_name,
to_field_path,
},
)
} else if ident == "offset" {
// offset!(NAME, Struct.field[expr].nested)
content.parse::<Token![!]>()?;
let inner;
syn::parenthesized!(inner in content);
let base_name: Ident = inner.parse()?;
inner.parse::<Token![,]>()?;
let struct_name: Ident = inner.parse()?;
let parsed = parse_indexed_field_path(&inner)?;
let full_name = Ident::new(&format!("{}_OFF", base_name), base_name.span());
(
full_name,
ConstantKind::Offset {
struct_name,
field_path_tokens: parsed.field_path_tokens,
array_index: parsed.array_index,
},
)
} else if ident == "offset_immediate" {
// offset_immediate!(NAME, Struct.field[expr].nested)
content.parse::<Token![!]>()?;
let inner;
syn::parenthesized!(inner in content);
let base_name: Ident = inner.parse()?;
inner.parse::<Token![,]>()?;
let struct_name: Ident = inner.parse()?;
let parsed = parse_indexed_field_path(&inner)?;
let full_name = Ident::new(&format!("{}_OFF_IMM", base_name), base_name.span());
(
full_name,
ConstantKind::OffsetImmediate {
struct_name,
field_path_tokens: parsed.field_path_tokens,
array_index: parsed.array_index,
},
)
} else if ident == "stack_frame_offset" || ident == "stack_frame_offset_unaligned" {
let aligned = ident == "stack_frame_offset";
let suffix = if aligned { "_OFF" } else { "_UOFF" };
content.parse::<Token![!]>()?;
let inner;
syn::parenthesized!(inner in content);
let base_name: Ident = inner.parse()?;
inner.parse::<Token![,]>()?;
let struct_name: Ident = inner.parse()?;
let parsed = parse_indexed_field_path(&inner)?;
let full_name = Ident::new(&format!("{}{}", base_name, suffix), base_name.span());
(
full_name,
ConstantKind::StackFrameOffset {
struct_name,
field_path_tokens: parsed.field_path_tokens,
array_index: parsed.array_index,
aligned,
},
)
} else if content.peek(Token![:]) {
// NAME: type = value
content.parse::<Token![:]>()?;
let ty: syn::Type = content.parse()?;
content.parse::<Token![=]>()?;
// Try literal first to preserve hex/binary representation.
let fork = content.fork();
if let Ok(lit) = fork.parse::<LitInt>() {
content.advance_to(&fork);
let repr = lit.to_string();
let expr = syn::Expr::Lit(syn::ExprLit {
attrs: vec![],
lit: Lit::Int(lit),
});
(
ident,
ConstantKind::Value {
ty,
value: expr,
literal_repr: Some(repr),
},
)
} else {
let expr: syn::Expr = content.parse()?;
(
ident,
ConstantKind::Value {
ty,
value: expr,
literal_repr: None,
},
)
}
} else {
// NAME = expr (type inferred from `as Type` cast).
content.parse::<Token![=]>()?;
let expr: syn::Expr = content.parse()?;
let ty = if let syn::Expr::Cast(cast) = &expr {
(*cast.ty).clone()
} else {
return Err(content.error(
"Expression must include `as Type` when type annotation is omitted",
));
};
(
ident,
ConstantKind::Value {
ty,
value: expr,
literal_repr: None,
},
)
};
// Optional trailing comma.
let _ = content.parse::<Token![,]>();
constants.push(ConstantDef {
doc: const_doc,
name: const_name,
kind,
});
}
// Reject multiple groups in a single macro invocation.
if !input.is_empty() {
return Err(input.error(
"Only one constant group per macro invocation; use separate constant_group! calls",
));
}
Ok(ConstantGroup {
doc,
name,
prefix,
constants,
})
}
}
/// Macro for defining groups of constants shared between Rust and ASM.
///
/// Values are validated at compile time to fit within i32 range (sBPF immediate constraint).
/// The prefix is automatically joined with an underscore.
///
/// Two syntaxes are supported:
/// - `NAME: type = value` — explicit type, literal values preserve hex/binary in ASM.
/// - `NAME = expr as Type` — type inferred from `as` cast, value computed at build time.
///
/// # Example
/// ```ignore
/// constant_group! {
/// /// Input buffer layout.
/// input_buffer {
/// /// Number of accounts expected.
/// N_ACCOUNTS: u64 = 2,
/// /// Expected data length of system program account.
/// SYSTEM_PROGRAM_DATA_LEN = b"system_program".len() as u64,
/// }
/// }
/// ```
///
/// To extend a group with ASM-only constants, use `extend_constant_group!`.
#[proc_macro]
pub fn constant_group(input: TokenStream) -> TokenStream {
let group = parse_macro_input!(input as ConstantGroup);
let mod_name = &group.name;
let max_line_len = MAX_LINE_LEN;
let header = asm_header(&group.doc);
// Generate Rust constants with i32 bounds checking for ASM compatibility.
let mut const_defs = Vec::new();
let mut const_value_strs: Vec<Option<String>> = Vec::new();
for c in &group.constants {
let name = &c.name;
let doc = &c.doc;
match &c.kind {
ConstantKind::Value {
ty,
value,
literal_repr,
} => {
let assert_name = Ident::new(&format!("_ASSERT_{}_FITS_I32", name), name.span());
const_value_strs.push(literal_repr.clone());
if literal_repr.is_some() {
// Literal value - no scope wrapper needed.
const_defs.push(quote! {
#[doc = #doc]
pub const #name: #ty = #value;
const #assert_name: () = assert!(
(#value as i64) >= (i32::MIN as i64) && (#value as i64) <= (i32::MAX as i64),
"ASM immediate must fit in i32 range"
);
});
} else {
// Expression value - use super::* for scope access.
const_defs.push(quote! {
#[doc = #doc]
pub const #name: #ty = { use super::*; #value };
const #assert_name: () = assert!(
({ use super::*; #value } as i64) >= (i32::MIN as i64)
&& ({ use super::*; #value } as i64) <= (i32::MAX as i64),
"ASM immediate must fit in i32 range"
);
});
}
}
ConstantKind::Offset {
struct_name,
field_path_tokens,
array_index,
} => {
let assert_name = Ident::new(&format!("_ASSERT_{}_FITS", name), name.span());
const_value_strs.push(None);
let offset_expr = gen_offset_expr(struct_name, field_path_tokens, array_index);
const_defs.push(quote! {
#[doc = #doc]
pub const #name: i16 = { use super::*; (#offset_expr) as i16 };
const #assert_name: () = assert!(
{ use super::*; #offset_expr } >= (i16::MIN as i64)
&& { use super::*; #offset_expr } <= (i16::MAX as i64),
"Offset must fit in i16 range"
);
});
}
ConstantKind::OffsetImmediate {
struct_name,
field_path_tokens,
array_index,
} => {
let assert_name = Ident::new(&format!("_ASSERT_{}_FITS", name), name.span());
const_value_strs.push(None);
let offset_expr = gen_offset_expr(struct_name, field_path_tokens, array_index);
const_defs.push(quote! {
#[doc = #doc]
pub const #name: i32 = { use super::*; (#offset_expr) as i32 };
const #assert_name: () = assert!(
{ use super::*; #offset_expr } >= (i32::MIN as i64)
&& { use super::*; #offset_expr } <= (i32::MAX as i64),
"Offset immediate must fit in i32 range"
);
});
}
ConstantKind::StackFrameOffset {
struct_name,
field_path_tokens,
array_index,
aligned,
} => {
let (const_def, literal_repr) = gen_stack_frame_offset_code(
name,
doc,
struct_name,
field_path_tokens,
array_index,
*aligned,
);
const_value_strs.push(literal_repr);
const_defs.push(const_def);
}
ConstantKind::PubkeyChunkOffset {
struct_name,
field_path,
chunk_index,
} => {
const_value_strs.push(None);
const_defs.push(gen_pubkey_chunk_offset_code(
name,
doc,
struct_name,
field_path,
*chunk_index,
));
}
ConstantKind::StackFramePubkeyChunkOffset {
struct_name,
field_path_tokens,
array_index,
chunk_index,
} => {
let (const_def, literal_repr) = gen_stack_frame_pubkey_chunk_offset_code(
name,
doc,
struct_name,
field_path_tokens,
array_index,
*chunk_index,
);
const_value_strs.push(literal_repr);
const_defs.push(const_def);
}
ConstantKind::RelativeOffsetImmediate {
from_struct_name,
from_field_path,
to_struct_name,
to_field_path,
} => {
const_value_strs.push(None);
const_defs.push(gen_relative_offset_immediate_code(
name,
doc,
from_struct_name,
from_field_path,
to_struct_name,
to_field_path,
));
}
}
}
let const_names: Vec<String> = group.constants.iter().map(|c| c.name.to_string()).collect();
let const_idents: Vec<&Ident> = group.constants.iter().map(|c| &c.name).collect();
let const_docs: Vec<String> = group.constants.iter().map(|c| c.doc.clone()).collect();
let value_str_opts: Vec<_> = const_value_strs
.iter()
.map(|opt| match opt {
Some(s) => quote! { Some(#s) },
None => quote! { None },
})
.collect();
// Generate to_asm function signature based on whether prefix is baked in.
let to_asm_fn = if let Some(ref prefix) = group.prefix {
let name_format = quote! { format!("{}_{}", #prefix, names[i]) };
quote! {
/// Generate ASM constants for this module.
pub fn to_asm() -> alloc::string::String {
use alloc::string::String;
use alloc::format;
let mut result = String::from(#header);
result.push('\n');
let names = [#(#const_names),*];
let computed_values: &[i64] = &[#(#const_idents as i64),*];
let literal_values: &[Option<&str>] = &[#(#value_str_opts),*];
let docs = [#(#const_docs),*];
for i in 0..names.len() {
let full_name = #name_format;
let value_str = match literal_values[i] {
Some(lit) => String::from(lit),
None => format!("{}", computed_values[i]),
};
let inline = format!(".equ {}, {} # {}", full_name, value_str, docs[i]);
if inline.len() <= #max_line_len {
result.push_str(&inline);
} else {
result.push_str(&format!("# {}\n.equ {}, {}", docs[i], full_name, value_str));
}
result.push('\n');
}
result
}
}
} else {
quote! {
/// Generate ASM constants for this module with the given prefix.
/// Prefix is automatically joined with underscore (e.g., "IB" -> "IB_NAME").
pub fn to_asm(prefix: &str) -> alloc::string::String {
use alloc::string::String;
use alloc::format;
let mut result = String::from(#header);
result.push('\n');
let names = [#(#const_names),*];
let computed_values: &[i64] = &[#(#const_idents as i64),*];
let literal_values: &[Option<&str>] = &[#(#value_str_opts),*];
let docs = [#(#const_docs),*];
for i in 0..names.len() {
let full_name = if prefix.is_empty() {
String::from(names[i])
} else {
format!("{}_{}", prefix, names[i])
};
let value_str = match literal_values[i] {
Some(lit) => String::from(lit),
None => format!("{}", computed_values[i]),
};
let inline = format!(".equ {}, {} # {}", full_name, value_str, docs[i]);
if inline.len() <= #max_line_len {
result.push_str(&inline);
} else {
result.push_str(&format!("# {}\n.equ {}, {}", docs[i], full_name, value_str));
}
result.push('\n');
}
result
}
}
};
let expanded = quote! {
pub mod #mod_name {
#(#const_defs)*
#to_asm_fn
}
};
TokenStream::from(expanded)
}
/// ASM-only constant (no Rust type needed).
struct AsmConstantDef {
doc: String,
name: Ident,
kind: AsmConstantKind,
}
/// Kind of ASM constant.
enum AsmConstantKind {
/// Literal value (i32 validated) - preserves original representation (hex, etc.).
Literal(LitInt),
/// Expression value (i32 validated) - computed at runtime, shown as decimal.
Expr(syn::Expr),
/// Offset derived from struct field path (i16 validated).
/// Name gets `_OFF` suffix appended.
/// Supports optional array indexing via `__<Struct>_fields` companion module.
Offset {
struct_name: Ident,
field_path_tokens: proc_macro2::TokenStream,
array_index: Option<ArrayIndexInfo>,
},
/// Offset derived from struct field path (i32 validated).
/// Name gets `_OFF_IMM` suffix appended. For use as BPF immediate operand.
/// Supports optional array indexing via `__<Struct>_fields` companion module.
OffsetImmediate {
struct_name: Ident,
field_path_tokens: proc_macro2::TokenStream,
array_index: Option<ArrayIndexInfo>,
},
/// Negative offset from end of a stack frame struct (i16 validated).
/// Computed as `offset_of!(Struct, field) - size_of::<Struct>()`.
/// Name gets `_OFF` suffix (aligned) or `_UOFF` suffix (unaligned).
/// Array element types are resolved via the `__<Struct>_fields` companion module
/// generated by `#[stack_frame]`.
StackFrameOffset {
struct_name: Ident,
field_path_tokens: proc_macro2::TokenStream,
array_index: Option<ArrayIndexInfo>,
aligned: bool,
},
/// Pubkey chunk offset (i16 validated).
/// Base offset + chunk_index * 8 for 8-byte register loads.
/// Name gets `_OFF_{chunk_index}` suffix appended.
PubkeyChunkOffset {
struct_name: Ident,
field_path: Vec<Ident>,
chunk_index: usize,
},
/// Negative stack frame offset for pubkey chunks (i16 validated).
/// Computed as `offset_of!(Struct, field) + chunk_index * 8 - size_of::<Struct>()`.
/// Always unaligned. Name gets `_UOFF_{chunk_index}` suffix appended.
StackFramePubkeyChunkOffset {
struct_name: Ident,
field_path_tokens: proc_macro2::TokenStream,
array_index: Option<ArrayIndexInfo>,
chunk_index: usize,
},
/// Chunk of a known pubkey/address constant.
/// Extracts `width` bytes at `byte_offset` from the 32-byte address.
/// Width 4 produces i32 (for LO/HI halves), width 8 produces i64 (for lddw).
PubkeyValueChunk {
expr: syn::Expr,
byte_offset: usize,
width: usize,
},
/// Relative offset between two fields (i32 validated).
/// Computed as `offset_of!(Struct, to_field) - offset_of!(Struct, from_field)`.
/// Name is `{FROM}_TO_{TO}_REL_OFF_IMM`.
RelativeOffsetImmediate {
from_struct_name: Ident,
from_field_path: Vec<Ident>,
to_struct_name: Ident,
to_field_path: Vec<Ident>,
},
}
/// Input for asm_constant_group! macro.
struct AsmConstantGroupInput {
doc: String,
name: Ident,
prefix: Option<String>,
constants: Vec<AsmConstantDef>,
}
/// Parse a dot-separated field path for `offset_of!`.
///
/// Expects tokens like `.field` or `.field1.field2.field3`.
/// Returns a `TokenStream` suitable for use inside `offset_of!(Type, <path>)`.
fn parse_field_path(inner: ParseStream) -> syn::Result<proc_macro2::TokenStream> {
let mut tokens = proc_macro2::TokenStream::new();
let mut has_field = false;
while inner.peek(Token![.]) {
inner.parse::<Token![.]>()?;
let field: Ident = inner.parse()?;
if has_field {
tokens.extend(quote! { . #field });
} else {
tokens.extend(quote! { #field });
has_field = true;
}
}
if !has_field {
return Err(inner.error("Expected at least one field in path"));
}
Ok(tokens)
}
/// Result of parsing a field path with optional array indexing.
struct IndexedFieldPath {
/// Token stream for `offset_of!(Struct, <path>)` (fields before any bracket).
field_path_tokens: proc_macro2::TokenStream,
/// Optional array index info (bracket expression, inner field access).
array_index: Option<ArrayIndexInfo>,
}
/// Parse a field path with optional array indexing and inner field access.
///
/// Supports:
/// - `Struct.field` → simple field access
/// - `Struct.field[expr]` → array element access
/// - `Struct.field[expr].inner` → array element + inner field access
/// - `Struct.a.b.c` → nested field access (no array)
///
/// Array element types are resolved via `__<Struct>_fields::<array_field>` type aliases
/// generated by `#[stack_frame]` or `#[array_fields]`.
fn parse_indexed_field_path(inner: ParseStream) -> syn::Result<IndexedFieldPath> {
let mut fields: Vec<Ident> = Vec::new();
let mut tokens = proc_macro2::TokenStream::new();
// Parse dot-separated field segments until we hit a bracket or end.
while inner.peek(Token![.]) {
inner.parse::<Token![.]>()?;
let field: Ident = inner.parse()?;
if !fields.is_empty() {
tokens.extend(quote! { . });
}
tokens.extend(quote! { #field });
fields.push(field);
// Check for array index after this field.
if inner.peek(syn::token::Bracket) {
let bracket_content;
bracketed!(bracket_content in inner);
let index_expr: syn::Expr = bracket_content.parse()?;
// Parse optional inner field path.
let inner_field_path = if inner.peek(Token![.]) {
Some(parse_field_path(inner)?)
} else {
None
};
let array_field_name = fields.last().unwrap().clone();
return Ok(IndexedFieldPath {
field_path_tokens: tokens,
array_index: Some(ArrayIndexInfo {
array_field_name,
index_expr,
inner_field_path,
}),
});
}
}
if fields.is_empty() {
return Err(inner.error("Expected at least one field in path"));
}
Ok(IndexedFieldPath {
field_path_tokens: tokens,
array_index: None,
})
}
/// Extract the last path segment name from a type (for auto-generating constant names).
fn extract_type_name(ty: &syn::Type) -> Option<String> {
if let syn::Type::Path(type_path) = ty {
type_path
.path
.segments
.last()
.map(|seg| seg.ident.to_string())
} else {
None
}
}
/// Convert a type name to UPPER_SNAKE_CASE for ASM constants.
///
/// Uses `convert_case` but fixes primitive types like `u8`, `i16` etc.
/// where the default conversion inserts an unwanted underscore (`U_8` → `U8`).
fn type_name_to_upper_snake(name: &str) -> String {
let s = name.to_case(Case::UpperSnake);
// Primitive types: single letter + digits (e.g. U_8, I_16, F_32).
// Remove the underscore between the letter and digits.
if s.starts_with(|c: char| c.is_ascii_uppercase())
&& s.as_bytes().get(1) == Some(&b'_')
&& s[2..].chars().all(|c| c.is_ascii_digit())
{
format!("{}{}", &s[..1], &s[2..])
} else {
s
}
}
/// Parse optional `prefix = "..."` group header parameter.
fn parse_group_params(content: ParseStream) -> syn::Result<Option<String>> {
// Check for prefix = "..." before constants.
if content.peek(Ident) {
let fork = content.fork();
let ident: Ident = fork.parse()?;
if ident == "prefix" && fork.peek(Token![=]) {
content.parse::<Ident>()?;
content.parse::<Token![=]>()?;
let prefix_lit: syn::LitStr = content.parse()?;
content.parse::<Token![,]>()?;
return Ok(Some(prefix_lit.value()));
}
}
Ok(None)
}
/// Parse ASM constants (shared between asm_constant_group! and extend_constant_group!).
fn parse_asm_constants(content: ParseStream) -> syn::Result<Vec<AsmConstantDef>> {
let mut constants = Vec::new();
while !content.is_empty() {
let const_attrs = content.call(syn::Attribute::parse_outer)?;
let const_doc = extract_doc_comment(&const_attrs)
.ok_or_else(|| content.error("Constant must have a doc comment"))?;
if let Err(e) = validate_doc_comment(&const_doc) {
return Err(content.error(e));
}
// Check for identifier.
let lookahead = content.lookahead1();
if !lookahead.peek(Ident) {
return Err(lookahead.error());
}
let ident: Ident = content.parse()?;
// pubkey_offset!(NAME, Struct.field) expands to 4 chunk constants.
if ident == "pubkey_offset" {
content.parse::<Token![!]>()?;
let inner;
syn::parenthesized!(inner in content);
let base_name: Ident = inner.parse()?;
inner.parse::<Token![,]>()?;
let struct_name: Ident = inner.parse()?;
let mut field_path = Vec::new();
while inner.peek(Token![.]) {
inner.parse::<Token![.]>()?;
field_path.push(inner.parse::<Ident>()?);
}
if field_path.is_empty() {
return Err(inner.error("Expected at least one field after struct name"));
}
let _ = content.parse::<Token![,]>();
let base_doc = const_doc.trim_end_matches('.');
for chunk in 0..4usize {
let chunk_name =
Ident::new(&format!("{}_OFF_{}", base_name, chunk), base_name.span());
let chunk_doc = format!("{} (chunk index {}).", base_doc, chunk);
if let Err(e) = validate_doc_comment(&chunk_doc) {
return Err(content.error(e));
}
constants.push(AsmConstantDef {
doc: chunk_doc,
name: chunk_name,
kind: AsmConstantKind::PubkeyChunkOffset {
struct_name: struct_name.clone(),
field_path: field_path.clone(),
chunk_index: chunk,
},
});
}
continue;
}
// stack_frame_pubkey_offset_unaligned!(NAME, Struct.field) expands to 4 chunk constants.
if ident == "stack_frame_pubkey_offset_unaligned" {
content.parse::<Token![!]>()?;
let inner;
syn::parenthesized!(inner in content);
let base_name: Ident = inner.parse()?;
inner.parse::<Token![,]>()?;
let struct_name: Ident = inner.parse()?;
let parsed = parse_indexed_field_path(&inner)?;
let _ = content.parse::<Token![,]>();
let base_doc = const_doc.trim_end_matches('.');
for chunk in 0..4usize {
let chunk_name =
Ident::new(&format!("{}_UOFF_{}", base_name, chunk), base_name.span());
let chunk_doc = format!("{} (chunk index {}).", base_doc, chunk);
if let Err(e) = validate_doc_comment(&chunk_doc) {
return Err(content.error(e));
}
constants.push(AsmConstantDef {
doc: chunk_doc,
name: chunk_name,
kind: AsmConstantKind::StackFramePubkeyChunkOffset {
struct_name: struct_name.clone(),
field_path_tokens: parsed.field_path_tokens.clone(),
array_index: parsed.array_index.as_ref().map(|a| ArrayIndexInfo {
array_field_name: a.array_field_name.clone(),
index_expr: a.index_expr.clone(),
inner_field_path: a.inner_field_path.clone(),
}),
chunk_index: chunk,
},
});
}
continue;
}
// pubkey_value!(NAME, expr) expands to 8 chunk constants (4 LO/HI pairs).
if ident == "pubkey_value" {
content.parse::<Token![!]>()?;
let inner;
syn::parenthesized!(inner in content);
let base_name: Ident = inner.parse()?;
inner.parse::<Token![,]>()?;
let expr: syn::Expr = inner.parse()?;
let _ = content.parse::<Token![,]>();
let base_doc = const_doc.trim_end_matches('.');
for chunk in 0..4usize {
let byte_base = chunk * 8;
// Full i64 chunk (for lddw).
let full_name =
Ident::new(&format!("{}_CHUNK_{}", base_name, chunk), base_name.span());
let full_doc = format!("{} (chunk {}).", base_doc, chunk);
if let Err(e) = validate_doc_comment(&full_doc) {
return Err(content.error(e));
}
constants.push(AsmConstantDef {
doc: full_doc,
name: full_name,
kind: AsmConstantKind::PubkeyValueChunk {
expr: expr.clone(),
byte_offset: byte_base,
width: 8,
},
});
// Low i32 half (for jne reg, imm32 when sign-extension is safe).
let lo_name = Ident::new(
&format!("{}_CHUNK_{}_LO", base_name, chunk),
base_name.span(),
);
let lo_doc = format!("{} (chunk {} lo).", base_doc, chunk);
if let Err(e) = validate_doc_comment(&lo_doc) {
return Err(content.error(e));
}
constants.push(AsmConstantDef {
doc: lo_doc,
name: lo_name,
kind: AsmConstantKind::PubkeyValueChunk {
expr: expr.clone(),
byte_offset: byte_base,
width: 4,
},
});
// High i32 half (to check if lddw can be avoided).
let hi_name = Ident::new(
&format!("{}_CHUNK_{}_HI", base_name, chunk),
base_name.span(),
);
let hi_doc = format!("{} (chunk {} hi).", base_doc, chunk);
if let Err(e) = validate_doc_comment(&hi_doc) {
return Err(content.error(e));
}
constants.push(AsmConstantDef {
doc: hi_doc,
name: hi_name,
kind: AsmConstantKind::PubkeyValueChunk {
expr: expr.clone(),
byte_offset: byte_base + 4,
width: 4,
},
});
}
continue;
}
let (const_name, kind) = if ident == "relative_offset_immediate" {
// relative_offset_immediate!(FROM_NAME, TO_NAME, Struct.from.path, Struct.to.path)
content.parse::<Token![!]>()?;
let inner;
syn::parenthesized!(inner in content);
let from_name: Ident = inner.parse()?;
inner.parse::<Token![,]>()?;
let to_name: Ident = inner.parse()?;
inner.parse::<Token![,]>()?;
let from_struct_name: Ident = inner.parse()?;
let mut from_field_path = Vec::new();
while inner.peek(Token![.]) {
inner.parse::<Token![.]>()?;
from_field_path.push(inner.parse::<Ident>()?);
}
if from_field_path.is_empty() {
return Err(inner.error("Expected at least one field after struct name"));
}
inner.parse::<Token![,]>()?;
let to_struct_name: Ident = inner.parse()?;
let mut to_field_path = Vec::new();
while inner.peek(Token![.]) {
inner.parse::<Token![.]>()?;
to_field_path.push(inner.parse::<Ident>()?);
}
if to_field_path.is_empty() {
return Err(inner.error("Expected at least one field after struct name"));
}
let full_name = Ident::new(
&format!("{}_TO_{}_REL_OFF_IMM", from_name, to_name),
from_name.span(),
);
(
full_name,
AsmConstantKind::RelativeOffsetImmediate {
from_struct_name,
from_field_path,
to_struct_name,
to_field_path,
},
)
} else if ident == "offset" {
// Parse offset!(NAME, Struct.field[expr].nested)
content.parse::<Token![!]>()?;
let inner;
syn::parenthesized!(inner in content);
let base_name: Ident = inner.parse()?;
inner.parse::<Token![,]>()?;
let struct_name: Ident = inner.parse()?;
let parsed = parse_indexed_field_path(&inner)?;
let full_name = Ident::new(&format!("{}_OFF", base_name), base_name.span());
(
full_name,
AsmConstantKind::Offset {
struct_name,
field_path_tokens: parsed.field_path_tokens,
array_index: parsed.array_index,
},
)
} else if ident == "offset_immediate" {
// Parse offset_immediate!(NAME, Struct.field[expr].nested)
content.parse::<Token![!]>()?;
let inner;
syn::parenthesized!(inner in content);
let base_name: Ident = inner.parse()?;
inner.parse::<Token![,]>()?;
let struct_name: Ident = inner.parse()?;
let parsed = parse_indexed_field_path(&inner)?;
let full_name = Ident::new(&format!("{}_OFF_IMM", base_name), base_name.span());
(
full_name,
AsmConstantKind::OffsetImmediate {
struct_name,
field_path_tokens: parsed.field_path_tokens,
array_index: parsed.array_index,
},
)
} else if ident == "stack_frame_offset" || ident == "stack_frame_offset_unaligned" {
let aligned = ident == "stack_frame_offset";
let suffix = if aligned { "_OFF" } else { "_UOFF" };
content.parse::<Token![!]>()?;
let inner;
syn::parenthesized!(inner in content);
let base_name: Ident = inner.parse()?;
inner.parse::<Token![,]>()?;
let struct_name: Ident = inner.parse()?;
let parsed = parse_indexed_field_path(&inner)?;
let full_name = Ident::new(&format!("{}{}", base_name, suffix), base_name.span());
(
full_name,
AsmConstantKind::StackFrameOffset {
struct_name,
field_path_tokens: parsed.field_path_tokens,
array_index: parsed.array_index,
aligned,
},
)
} else {
// Regular NAME = value syntax.
content.parse::<Token![=]>()?;
// Try to parse as literal first (to preserve hex/binary representation),
// otherwise parse as expression (for constants like NON_DUP_MARKER).
let fork = content.fork();
if let Ok(lit) = fork.parse::<LitInt>() {
content.advance_to(&fork);
(ident, AsmConstantKind::Literal(lit))
} else {
let expr: syn::Expr = content.parse()?;
(ident, AsmConstantKind::Expr(expr))
}
};
// Optional trailing comma.
let _ = content.parse::<Token![,]>();
constants.push(AsmConstantDef {
doc: const_doc,
name: const_name,
kind,
});
}
Ok(constants)
}
impl Parse for AsmConstantGroupInput {
fn parse(input: ParseStream) -> syn::Result<Self> {
// Parse doc comment for the group.
let attrs = input.call(syn::Attribute::parse_outer)?;
let doc = extract_doc_comment(&attrs)
.ok_or_else(|| input.error("Group must have a doc comment"))?;
if let Err(e) = validate_doc_comment(&doc) {
return Err(input.error(format!("Group doc comment: {}", e)));
}
// Parse group name.
let name: Ident = input.parse()?;
// Parse body.
let content;
braced!(content in input);
// Parse optional header parameter: prefix = "..."
let prefix = parse_group_params(&content)?;
// Parse constants.
let constants = parse_asm_constants(&content)?;
// Reject multiple groups.
if !input.is_empty() {
return Err(input.error("Only one group per macro invocation"));
}
Ok(AsmConstantGroupInput {
doc,
name,
prefix,
constants,
})
}
}
/// Generate an offset expression for a struct field with optional array indexing.
///
/// Without array index: `offset_of!(super::Struct, field) as i64`.
/// With array index: adds `+ index * Struct::__FIELD_STRIDE` using the hidden
/// associated constant generated by `#[array_fields]`.
fn gen_offset_expr(
struct_name: &Ident,
field_path_tokens: &proc_macro2::TokenStream,
array_index: &Option<ArrayIndexInfo>,
) -> proc_macro2::TokenStream {
match array_index {
Some(info) => {
let array_field_name = &info.array_field_name;
let index_expr = &info.index_expr;
let stride_const = Ident::new(
&format!("__{}_STRIDE", array_field_name.to_string().to_uppercase()),
array_field_name.span(),
);
quote! {
core::mem::offset_of!(super::#struct_name, #field_path_tokens) as i64
+ (#index_expr) as i64 * (super::#struct_name::#stride_const) as i64
}
}
None => quote! {
core::mem::offset_of!(super::#struct_name, #field_path_tokens) as i64
},
}
}
/// Generate code for a `stack_frame_offset!` constant.
///
/// Returns `(const_def_tokens, literal_repr)` where literal_repr is always `None`
/// (offsets are always computed).
fn gen_stack_frame_offset_code(
name: &Ident,
doc: &str,
struct_name: &Ident,
field_path_tokens: &proc_macro2::TokenStream,
array_index: &Option<ArrayIndexInfo>,
aligned: bool,
) -> (proc_macro2::TokenStream, Option<String>) {
let assert_name = Ident::new(&format!("_ASSERT_{}_FITS", name), name.span());
let fields_mod = Ident::new(&format!("__{}_fields", struct_name), struct_name.span());
let offset_expr = match array_index {
Some(info) => {
let array_field_name = &info.array_field_name;
let index_expr = &info.index_expr;
let inner_offset = match &info.inner_field_path {
Some(path) => quote! {
+ core::mem::offset_of!(#fields_mod::#array_field_name, #path) as i64
},
None => quote! {},
};
quote! {
core::mem::offset_of!(#struct_name, #field_path_tokens) as i64
+ (#index_expr) as i64 * core::mem::size_of::<#fields_mod::#array_field_name>() as i64
#inner_offset
}
}
None => quote! {
core::mem::offset_of!(#struct_name, #field_path_tokens) as i64
},
};
let align_assertion = if aligned {
let bpf_align = BPF_ALIGN;
quote! {
assert!(
result % #bpf_align == 0,
"Stack frame offset must be aligned"
);
}
} else {
quote! {}
};
let const_def = quote! {
#[doc = #doc]
pub const #name: i16 = {
use super::*;
(#offset_expr - core::mem::size_of::<#struct_name>() as i64) as i16
};
const #assert_name: () = {
use super::*;
let result = #offset_expr - core::mem::size_of::<#struct_name>() as i64;
assert!(
result >= (i16::MIN as i64) && result <= (i16::MAX as i64),
"Stack frame offset must fit in i16"
);
assert!(result < 0, "Stack frame offset must be negative");
#align_assertion
};
};
(const_def, None)
}
/// Generate code for a `stack_frame_pubkey_offset_unaligned!` constant (single chunk).
///
/// Computes `offset_of!(Struct, field) + chunk_index * 8 - size_of::<Struct>()` as an i16 constant.
fn gen_stack_frame_pubkey_chunk_offset_code(
name: &Ident,
doc: &str,
struct_name: &Ident,
field_path_tokens: &proc_macro2::TokenStream,
array_index: &Option<ArrayIndexInfo>,
chunk_index: usize,
) -> (proc_macro2::TokenStream, Option<String>) {
let assert_name = Ident::new(&format!("_ASSERT_{}_FITS", name), name.span());
let fields_mod = Ident::new(&format!("__{}_fields", struct_name), struct_name.span());
let chunk_byte_offset = (chunk_index * BPF_ALIGN as usize) as i64;
let offset_expr = match array_index {
Some(info) => {
let array_field_name = &info.array_field_name;
let index_expr = &info.index_expr;
let inner_offset = match &info.inner_field_path {
Some(path) => quote! {
+ core::mem::offset_of!(#fields_mod::#array_field_name, #path) as i64
},
None => quote! {},
};
quote! {
core::mem::offset_of!(#struct_name, #field_path_tokens) as i64
+ (#index_expr) as i64 * core::mem::size_of::<#fields_mod::#array_field_name>() as i64
#inner_offset
}
}
None => quote! {
core::mem::offset_of!(#struct_name, #field_path_tokens) as i64
},
};
let const_def = quote! {
#[doc = #doc]
pub const #name: i16 = {
use super::*;
(#offset_expr + #chunk_byte_offset - core::mem::size_of::<#struct_name>() as i64) as i16
};
const #assert_name: () = {
use super::*;
let result = #offset_expr + #chunk_byte_offset - core::mem::size_of::<#struct_name>() as i64;
assert!(
result >= (i16::MIN as i64) && result <= (i16::MAX as i64),
"Stack frame offset must fit in i16"
);
assert!(result < 0, "Stack frame offset must be negative");
};
};
(const_def, None)
}
/// Generate code for a `pubkey_offset!` constant (single chunk).
///
/// Computes `offset_of!(Struct, field) + chunk_index * 8` as an i16 constant.
fn gen_pubkey_chunk_offset_code(
name: &Ident,
doc: &str,
struct_name: &Ident,
field_path: &[Ident],
chunk_index: usize,
) -> proc_macro2::TokenStream {
let assert_name = Ident::new(&format!("_ASSERT_{}_FITS", name), name.span());
let chunk_byte_offset = (chunk_index * BPF_ALIGN as usize) as i64;
quote! {
#[doc = #doc]
pub const #name: i16 = (core::mem::offset_of!(super::#struct_name, #(#field_path).*) as i64 + #chunk_byte_offset) as i16;
const #assert_name: () = assert!(
(core::mem::offset_of!(super::#struct_name, #(#field_path).*) as i64 + #chunk_byte_offset) >= (i16::MIN as i64)
&& (core::mem::offset_of!(super::#struct_name, #(#field_path).*) as i64 + #chunk_byte_offset) <= (i16::MAX as i64),
"Offset must fit in i16 range"
);
}
}
/// Generate code for a `relative_offset_immediate!` constant.
///
/// Computes `offset_of!(Struct, to_field) - offset_of!(Struct, from_field)` as an i32 constant.
fn gen_relative_offset_immediate_code(
name: &Ident,
doc: &str,
from_struct_name: &Ident,
from_field_path: &[Ident],
to_struct_name: &Ident,
to_field_path: &[Ident],
) -> proc_macro2::TokenStream {
let assert_name = Ident::new(&format!("_ASSERT_{}_FITS", name), name.span());
quote! {
#[doc = #doc]
pub const #name: i32 = (
core::mem::offset_of!(super::#to_struct_name, #(#to_field_path).*) as i64
- core::mem::offset_of!(super::#from_struct_name, #(#from_field_path).*) as i64
) as i32;
const #assert_name: () = assert!(
(core::mem::offset_of!(super::#to_struct_name, #(#to_field_path).*) as i64
- core::mem::offset_of!(super::#from_struct_name, #(#from_field_path).*) as i64) >= (i32::MIN as i64)
&& (core::mem::offset_of!(super::#to_struct_name, #(#to_field_path).*) as i64
- core::mem::offset_of!(super::#from_struct_name, #(#from_field_path).*) as i64) <= (i32::MAX as i64),
"Relative offset immediate must fit in i32 range"
);
}
}
/// Generate code for a `pubkey_value!` constant (single chunk).
///
/// Extracts `width` bytes at `byte_offset` from a 32-byte address constant.
/// Width 4 produces i32, width 8 produces i64.
fn gen_pubkey_value_chunk_code(
name: &Ident,
doc: &str,
expr: &syn::Expr,
byte_offset: usize,
width: usize,
) -> proc_macro2::TokenStream {
if width == 8 {
let b0 = byte_offset;
let b1 = byte_offset + 1;
let b2 = byte_offset + 2;
let b3 = byte_offset + 3;
let b4 = byte_offset + 4;
let b5 = byte_offset + 5;
let b6 = byte_offset + 6;
let b7 = byte_offset + 7;
quote! {
#[doc = #doc]
pub const #name: i64 = {
use super::*;
let bytes: [u8; 32] = unsafe { core::mem::transmute(#expr) };
i64::from_le_bytes([bytes[#b0], bytes[#b1], bytes[#b2], bytes[#b3], bytes[#b4], bytes[#b5], bytes[#b6], bytes[#b7]])
};
}
} else {
let b0 = byte_offset;
let b1 = byte_offset + 1;
let b2 = byte_offset + 2;
let b3 = byte_offset + 3;
quote! {
#[doc = #doc]
pub const #name: i32 = {
use super::*;
let bytes: [u8; 32] = unsafe { core::mem::transmute(#expr) };
i32::from_le_bytes([bytes[#b0], bytes[#b1], bytes[#b2], bytes[#b3]])
};
}
}
}
/// Macro for defining ASM-only constant groups.
///
/// Constants don't need types - values are `i64`, offsets are `i16`.
/// All values are validated at compile time to fit within i32 range (sBPF constraint).
///
/// Supports two constant syntaxes:
/// - `NAME = value` - direct value (i64, validated for i32 range)
/// - `offset!(NAME, Struct.field)` - offset (i16, validated for i16 range, `_OFF` suffix added)
///
/// # Example
/// ```ignore
/// asm_constant_group! {
/// /// Miscellaneous constants.
/// misc {
/// prefix = "M",
/// /// Data length of zero.
/// DATA_LENGTH_ZERO = 0,
/// }
/// }
/// // Creates `misc` module with:
/// // - pub const DATA_LENGTH_ZERO: i64 = 0;
/// // - to_asm() outputs ".equ M_DATA_LENGTH_ZERO, 0 # ..."
/// ```
#[proc_macro]
pub fn asm_constant_group(input: TokenStream) -> TokenStream {
let input = parse_macro_input!(input as AsmConstantGroupInput);
let mod_name = &input.name;
let max_line_len = MAX_LINE_LEN;
let header = asm_header(&input.doc);
// Generate constant definitions and collect info for ASM.
let mut const_defs = Vec::new();
let mut const_names = Vec::new();
let mut const_docs = Vec::new();
// Track value representations: Some(literal_str) for values, None for offsets.
let mut const_value_strs: Vec<Option<String>> = Vec::new();
for c in &input.constants {
let name = &c.name;
let doc = &c.doc;
let name_str = name.to_string();
let assert_name = Ident::new(&format!("_ASSERT_{}_FITS", name), name.span());
const_names.push(name_str);
const_docs.push(doc.clone());
match &c.kind {
AsmConstantKind::Literal(value) => {
// Preserve original literal representation (hex, binary, etc.).
const_value_strs.push(Some(value.to_string()));
const_defs.push(quote! {
#[doc = #doc]
pub const #name: i64 = #value;
const #assert_name: () = assert!(
(#value as i64) >= (i32::MIN as i64) && (#value as i64) <= (i32::MAX as i64),
"ASM immediate must fit in i32 range"
);
});
}
AsmConstantKind::Expr(expr) => {
// Expression (e.g., constant from another crate) - computed at runtime.
// Use super::* to access imports from parent scope.
const_value_strs.push(None);
const_defs.push(quote! {
#[doc = #doc]
pub const #name: i64 = { use super::*; #expr as i64 };
const #assert_name: () = assert!(
({ use super::*; #expr } as i64) >= (i32::MIN as i64)
&& ({ use super::*; #expr } as i64) <= (i32::MAX as i64),
"ASM immediate must fit in i32 range"
);
});
}
AsmConstantKind::Offset {
struct_name,
field_path_tokens,
array_index,
} => {
const_value_strs.push(None);
let offset_expr = gen_offset_expr(struct_name, field_path_tokens, array_index);
const_defs.push(quote! {
#[doc = #doc]
pub const #name: i16 = { use super::*; (#offset_expr) as i16 };
const #assert_name: () = assert!(
{ use super::*; #offset_expr } >= (i16::MIN as i64)
&& { use super::*; #offset_expr } <= (i16::MAX as i64),
"Offset must fit in i16 range"
);
});
}
AsmConstantKind::OffsetImmediate {
struct_name,
field_path_tokens,
array_index,
} => {
const_value_strs.push(None);
let offset_expr = gen_offset_expr(struct_name, field_path_tokens, array_index);
const_defs.push(quote! {
#[doc = #doc]
pub const #name: i32 = { use super::*; (#offset_expr) as i32 };
const #assert_name: () = assert!(
{ use super::*; #offset_expr } >= (i32::MIN as i64)
&& { use super::*; #offset_expr } <= (i32::MAX as i64),
"Offset immediate must fit in i32 range"
);
});
}
AsmConstantKind::StackFrameOffset {
struct_name,
field_path_tokens,
array_index,
aligned,
} => {
let (const_def, literal_repr) = gen_stack_frame_offset_code(
name,
doc,
struct_name,
field_path_tokens,
array_index,
*aligned,
);
const_value_strs.push(literal_repr);
const_defs.push(const_def);
}
AsmConstantKind::PubkeyChunkOffset {
struct_name,
field_path,
chunk_index,
} => {
const_value_strs.push(None);
const_defs.push(gen_pubkey_chunk_offset_code(
name,
doc,
struct_name,
field_path,
*chunk_index,
));
}
AsmConstantKind::StackFramePubkeyChunkOffset {
struct_name,
field_path_tokens,
array_index,
chunk_index,
} => {
let (const_def, literal_repr) = gen_stack_frame_pubkey_chunk_offset_code(
name,
doc,
struct_name,
field_path_tokens,
array_index,
*chunk_index,
);
const_value_strs.push(literal_repr);
const_defs.push(const_def);
}
AsmConstantKind::PubkeyValueChunk {
expr,
byte_offset,
width,
} => {
const_value_strs.push(None);
const_defs.push(gen_pubkey_value_chunk_code(
name,
doc,
expr,
*byte_offset,
*width,
));
}
AsmConstantKind::RelativeOffsetImmediate {
from_struct_name,
from_field_path,
to_struct_name,
to_field_path,
} => {
const_value_strs.push(None);
const_defs.push(gen_relative_offset_immediate_code(
name,
doc,
from_struct_name,
from_field_path,
to_struct_name,
to_field_path,
));
}
}
}
let const_idents: Vec<_> = input.constants.iter().map(|c| &c.name).collect();
// Generate name formatting logic based on whether prefix is present.
let name_format = match &input.prefix {
Some(prefix) => quote! { format!("{}_{}", #prefix, names[i]) },
None => quote! { String::from(names[i]) },
};
// Generate value string options for preserving hex/binary literals.
let value_str_opts: Vec<_> = const_value_strs
.iter()
.map(|opt| match opt {
Some(s) => quote! { Some(#s) },
None => quote! { None },
})
.collect();
let expanded = quote! {
pub mod #mod_name {
use alloc::string::String;
use alloc::format;
#(#const_defs)*
/// Generate ASM constants.
pub fn to_asm() -> String {
let mut result = String::from(#header);
result.push('\n');
let names: &[&str] = &[#(#const_names),*];
let computed_values: &[i64] = &[#(#const_idents as i64),*];
let literal_values: &[Option<&str>] = &[#(#value_str_opts),*];
let docs: &[&str] = &[#(#const_docs),*];
for i in 0..names.len() {
let full_name = #name_format;
// Use original literal if available, otherwise use computed value.
let value_str = match literal_values[i] {
Some(lit) => String::from(lit),
None => format!("{}", computed_values[i]),
};
let inline = format!(".equ {}, {} # {}", full_name, value_str, docs[i]);
if inline.len() <= #max_line_len {
result.push_str(&inline);
} else {
result.push_str(&format!("# {}\n.equ {}, {}", docs[i], full_name, value_str));
}
result.push('\n');
}
result
}
}
};
TokenStream::from(expanded)
}
/// Input for extend_constant_group! macro.
struct ExtendConstantGroupInput {
name: Ident,
prefix: Option<String>,
constants: Vec<AsmConstantDef>,
}
impl Parse for ExtendConstantGroupInput {
fn parse(input: ParseStream) -> syn::Result<Self> {
// Parse module name.
let name: Ident = input.parse()?;
// Parse body.
let content;
braced!(content in input);
// Parse optional header parameter: prefix = "..."
let prefix = parse_group_params(&content)?;
// Parse constants using shared parser.
let constants = parse_asm_constants(&content)?;
Ok(ExtendConstantGroupInput {
name,
prefix,
constants,
})
}
}
/// Macro for extending a constant group with ASM-only constants.
///
/// This creates a module that re-exports the base group's constants from
/// `crate::common::{name}` and adds ASM-only constants. The `to_asm()` function
/// combines both under one header. The prefix is automatically joined with an underscore.
///
/// Supports two constant syntaxes:
/// - `NAME = value` - direct value (validated for i32 range)
/// - `offset!(NAME, Struct.field)` - offset (validated for i16 range, `_OFF` suffix added)
///
/// # Example
/// ```ignore
/// extend_constant_group!(input_buffer {
/// prefix = "IB",
/// /// Offset to number of accounts field.
/// offset!(N_ACCOUNTS, InputBuffer.n_accounts),
/// });
/// // Creates `input_buffer` module that:
/// // - Re-exports all constants from crate::common::input_buffer
/// // - Adds N_ACCOUNTS_OFF constant derived from offset_of!(InputBuffer, n_accounts)
/// // - to_asm() outputs ".equ IB_N_ACCOUNTS_OFF, 0 # ..."
/// ```
#[proc_macro]
pub fn extend_constant_group(input: TokenStream) -> TokenStream {
let input = parse_macro_input!(input as ExtendConstantGroupInput);
let mod_name = &input.name;
let max_line_len = MAX_LINE_LEN;
// Generate constant definitions and collect info for ASM.
let mut const_defs = Vec::new();
let mut const_names = Vec::new();
let mut const_docs = Vec::new();
// Track value representations: Some(literal_str) for values, None for offsets.
let mut const_value_strs: Vec<Option<String>> = Vec::new();
for c in &input.constants {
let name = &c.name;
let doc = &c.doc;
let name_str = name.to_string();
let assert_name = Ident::new(&format!("_ASSERT_{}_FITS", name), name.span());
const_names.push(name_str);
const_docs.push(doc.clone());
match &c.kind {
AsmConstantKind::Literal(value) => {
// Preserve original literal representation (hex, binary, etc.).
const_value_strs.push(Some(value.to_string()));
const_defs.push(quote! {
#[doc = #doc]
pub const #name: i32 = #value;
const #assert_name: () = assert!(
(#value as i64) >= (i32::MIN as i64) && (#value as i64) <= (i32::MAX as i64),
"ASM immediate must fit in i32 range"
);
});
}
AsmConstantKind::Expr(expr) => {
// Expression (e.g., constant from another crate) - computed at runtime.
// Use super::* to access imports from parent scope.
const_value_strs.push(None);
const_defs.push(quote! {
#[doc = #doc]
pub const #name: i32 = { use super::*; #expr as i32 };
const #assert_name: () = assert!(
({ use super::*; #expr } as i64) >= (i32::MIN as i64)
&& ({ use super::*; #expr } as i64) <= (i32::MAX as i64),
"ASM immediate must fit in i32 range"
);
});
}
AsmConstantKind::Offset {
struct_name,
field_path_tokens,
array_index,
} => {
const_value_strs.push(None);
let offset_expr = gen_offset_expr(struct_name, field_path_tokens, array_index);
const_defs.push(quote! {
#[doc = #doc]
pub const #name: i16 = { use super::*; (#offset_expr) as i16 };
const #assert_name: () = assert!(
{ use super::*; #offset_expr } >= (i16::MIN as i64)
&& { use super::*; #offset_expr } <= (i16::MAX as i64),
"Offset must fit in i16 range"
);
});
}
AsmConstantKind::OffsetImmediate {
struct_name,
field_path_tokens,
array_index,
} => {
const_value_strs.push(None);
let offset_expr = gen_offset_expr(struct_name, field_path_tokens, array_index);
const_defs.push(quote! {
#[doc = #doc]
pub const #name: i32 = { use super::*; (#offset_expr) as i32 };
const #assert_name: () = assert!(
{ use super::*; #offset_expr } >= (i32::MIN as i64)
&& { use super::*; #offset_expr } <= (i32::MAX as i64),
"Offset immediate must fit in i32 range"
);
});
}
AsmConstantKind::StackFrameOffset {
struct_name,
field_path_tokens,
array_index,
aligned,
} => {
let (const_def, literal_repr) = gen_stack_frame_offset_code(
name,
doc,
struct_name,
field_path_tokens,
array_index,
*aligned,
);
const_value_strs.push(literal_repr);
const_defs.push(const_def);
}
AsmConstantKind::PubkeyChunkOffset {
struct_name,
field_path,
chunk_index,
} => {
const_value_strs.push(None);
const_defs.push(gen_pubkey_chunk_offset_code(
name,
doc,
struct_name,
field_path,
*chunk_index,
));
}
AsmConstantKind::StackFramePubkeyChunkOffset {
struct_name,
field_path_tokens,
array_index,
chunk_index,
} => {
let (const_def, literal_repr) = gen_stack_frame_pubkey_chunk_offset_code(
name,
doc,
struct_name,
field_path_tokens,
array_index,
*chunk_index,
);
const_value_strs.push(literal_repr);
const_defs.push(const_def);
}
AsmConstantKind::PubkeyValueChunk {
expr,
byte_offset,
width,
} => {
const_value_strs.push(None);
const_defs.push(gen_pubkey_value_chunk_code(
name,
doc,
expr,
*byte_offset,
*width,
));
}
AsmConstantKind::RelativeOffsetImmediate {
from_struct_name,
from_field_path,
to_struct_name,
to_field_path,
} => {
const_value_strs.push(None);
const_defs.push(gen_relative_offset_immediate_code(
name,
doc,
from_struct_name,
from_field_path,
to_struct_name,
to_field_path,
));
}
}
}
// Collect const idents for ASM output.
let const_idents: Vec<_> = input.constants.iter().map(|c| &c.name).collect();
// Generate name formatting logic based on whether prefix is present.
let base_prefix = match &input.prefix {
Some(prefix) => quote! { #prefix },
None => quote! { "" },
};
let name_format = match &input.prefix {
Some(prefix) => quote! { format!("{}_{}", #prefix, names[i]) },
None => quote! { String::from(names[i]) },
};
// Generate value string options for preserving hex/binary literals.
let value_str_opts: Vec<_> = const_value_strs
.iter()
.map(|opt| match opt {
Some(s) => quote! { Some(#s) },
None => quote! { None },
})
.collect();
let expanded = quote! {
pub mod #mod_name {
use alloc::string::String;
use alloc::format;
// Re-export base group's constants.
pub use crate::common::#mod_name::*;
#(#const_defs)*
/// Generate combined ASM (base + extension).
pub fn to_asm() -> String {
// Base group adds header and its constants.
let mut result = crate::common::#mod_name::to_asm(#base_prefix);
// Add extension constants (no separate header).
let names: &[&str] = &[#(#const_names),*];
let computed_values: &[i64] = &[#(#const_idents as i64),*];
let literal_values: &[Option<&str>] = &[#(#value_str_opts),*];
let docs: &[&str] = &[#(#const_docs),*];
for i in 0..names.len() {
let full_name = #name_format;
// Use original literal if available, otherwise use computed value.
let value_str = match literal_values[i] {
Some(lit) => String::from(lit),
None => format!("{}", computed_values[i]),
};
let inline = format!(".equ {}, {} # {}", full_name, value_str, docs[i]);
if inline.len() <= #max_line_len {
result.push_str(&inline);
} else {
result.push_str(&format!("# {}\n.equ {}, {}", docs[i], full_name, value_str));
}
result.push('\n');
}
result
}
}
};
TokenStream::from(expanded)
}
/// Input for sizes! macro.
struct SizesInput {
types: Vec<syn::Type>,
}
impl Parse for SizesInput {
fn parse(input: ParseStream) -> syn::Result<Self> {
let mut types = Vec::new();
while !input.is_empty() {
let ty: syn::Type = input.parse()?;
types.push(ty);
// Optional trailing comma.
let _ = input.parse::<Token![,]>();
}
Ok(SizesInput { types })
}
}
/// Macro for generating type size constants for ASM.
///
/// Takes a list of types and generates `SIZE_OF_<UPPER_SNAKE_NAME>` constants.
/// Creates a `sizes` module with a `to_asm()` function for build-time ASM generation.
///
/// # Example
/// ```ignore
/// sizes! {
/// u8,
/// SolAccountMeta,
/// Rent,
/// }
/// ```
///
/// Generates ASM:
/// ```text
/// # Type sizes.
/// # -----------
/// .equ SIZE_OF_U8, 1 # Size of u8.
/// .equ SIZE_OF_SOL_ACCOUNT_META, 24 # Size of SolAccountMeta.
/// .equ SIZE_OF_RENT, 17 # Size of Rent.
/// ```
#[proc_macro]
pub fn sizes(input: TokenStream) -> TokenStream {
let input = parse_macro_input!(input as SizesInput);
let max_line_len = MAX_LINE_LEN;
let header = asm_header("Type sizes.");
let mut const_defs = Vec::new();
let mut const_name_strs = Vec::new();
let mut const_doc_strs = Vec::new();
for ty in &input.types {
let type_name =
extract_type_name(ty).unwrap_or_else(|| panic!("Expected a named type path in sizes!"));
let screaming = type_name_to_upper_snake(&type_name);
let name_str = format!("SIZE_OF_{}", screaming);
let name = Ident::new(&name_str, proc_macro2::Span::call_site());
let assert_name = Ident::new(
&format!("_ASSERT_{}_FITS", name_str),
proc_macro2::Span::call_site(),
);
let doc = format!("Size of {}.", type_name);
const_defs.push(quote! {
#[doc = #doc]
pub const #name: i64 = { use super::*; core::mem::size_of::<#ty>() as i64 };
const #assert_name: () = assert!(
({ use super::*; core::mem::size_of::<#ty>() } as i64) >= (i32::MIN as i64)
&& ({ use super::*; core::mem::size_of::<#ty>() } as i64) <= (i32::MAX as i64),
"ASM immediate must fit in i32 range"
);
});
const_name_strs.push(name_str);
const_doc_strs.push(doc);
}
let const_idents: Vec<_> = const_name_strs
.iter()
.map(|n| Ident::new(n, proc_macro2::Span::call_site()))
.collect();
let expanded = quote! {
pub mod sizes {
use alloc::string::String;
use alloc::format;
#(#const_defs)*
/// Generate ASM constants for type sizes.
pub fn to_asm() -> String {
let mut result = String::from(#header);
result.push('\n');
let names: &[&str] = &[#(#const_name_strs),*];
let values: &[i64] = &[#(#const_idents as i64),*];
let docs: &[&str] = &[#(#const_doc_strs),*];
for i in 0..names.len() {
let inline = format!(".equ {}, {} # {}", names[i], values[i], docs[i]);
if inline.len() <= #max_line_len {
result.push_str(&inline);
} else {
result.push_str(&format!("# {}\n.equ {}, {}", docs[i], names[i], values[i]));
}
result.push('\n');
}
result
}
}
};
TokenStream::from(expanded)
}
/// Macro for generating pubkey chunking offset constants.
///
/// A pubkey (Address) is 32 bytes. For 8-byte register loads, it is
/// accessed in 4 chunks of 8 bytes each. This macro generates a
/// `pubkey_chunk` module with `OFF_0` through `OFF_3` constants
/// and a `to_asm()` function.
///
/// # Example
/// ```ignore
/// pubkey_chunk_group!();
/// ```
///
/// Generates ASM:
/// ```text
/// # Pubkey chunking offsets.
/// # ------------------------
/// .equ PUBKEY_CHUNK_OFF_0, 0 # Offset for the first 8 bytes.
/// .equ PUBKEY_CHUNK_OFF_1, 8 # Offset for the second 8 bytes.
/// .equ PUBKEY_CHUNK_OFF_2, 16 # Offset for the third 8 bytes.
/// .equ PUBKEY_CHUNK_OFF_3, 24 # Offset for the fourth 8 bytes.
/// ```
#[proc_macro]
pub fn pubkey_chunk_group(_input: TokenStream) -> TokenStream {
const PUBKEY_SIZE: usize = 32;
const CHUNK_SIZE: usize = BPF_ALIGN as usize;
const N_CHUNKS: usize = PUBKEY_SIZE / CHUNK_SIZE;
const ORDINALS: [&str; N_CHUNKS] = ["first", "second", "third", "fourth"];
let max_line_len = MAX_LINE_LEN;
let header = asm_header("Pubkey chunking offsets.");
let mut const_defs = Vec::new();
let mut const_name_strs = Vec::new();
let mut const_doc_strs = Vec::new();
for (i, ordinal) in ORDINALS.iter().enumerate() {
let offset = (i * CHUNK_SIZE) as i64;
let name_str = format!("PUBKEY_CHUNK_OFF_{}", i);
let name = Ident::new(&format!("OFF_{}", i), proc_macro2::Span::call_site());
let doc = format!("Offset for the {} 8 bytes.", ordinal);
const_defs.push(quote! {
#[doc = #doc]
pub const #name: i64 = #offset;
});
const_name_strs.push(name_str);
const_doc_strs.push(doc);
}
let const_idents: Vec<_> = (0..N_CHUNKS)
.map(|i| Ident::new(&format!("OFF_{}", i), proc_macro2::Span::call_site()))
.collect();
let expanded = quote! {
pub mod pubkey_chunk {
use alloc::string::String;
use alloc::format;
#(#const_defs)*
/// Generate ASM constants for pubkey chunking offsets.
pub fn to_asm() -> String {
let mut result = String::from(#header);
result.push('\n');
let names: &[&str] = &[#(#const_name_strs),*];
let values: &[i64] = &[#(#const_idents as i64),*];
let docs: &[&str] = &[#(#const_doc_strs),*];
for i in 0..names.len() {
let inline = format!(".equ {}, {} # {}", names[i], values[i], docs[i]);
if inline.len() <= #max_line_len {
result.push_str(&inline);
} else {
result.push_str(&format!("# {}\n.equ {}, {}", docs[i], names[i], values[i]));
}
result.push('\n');
}
result
}
}
};
TokenStream::from(expanded)
}
/// Attribute macro for stack frame structs.
///
/// Adds `#[repr(C, align(8))]` to ensure C layout and 8-byte alignment.
/// Any existing `#[repr(...)]` attributes are removed.
///
/// Also generates a companion module `__<StructName>_fields` containing type aliases
/// for each array field's element type. This allows `stack_frame_offset!` to resolve
/// element types without requiring them in the syntax.
///
/// # Example
/// ```ignore
/// #[stack_frame]
/// struct InitStackFrame {
/// data: [u8; 32],
/// metas: [SolAccountMeta; 2],
/// }
/// ```
///
/// Generates:
/// ```ignore
/// #[repr(C, align(8))]
/// struct InitStackFrame { ... }
///
/// mod __InitStackFrame_fields {
/// use super::*;
/// pub type metas = SolAccountMeta;
/// pub type data = u8;
/// }
/// ```
/// Generate the `__<StructName>_fields` companion module containing type
/// aliases for each array field's element type.
fn gen_fields_module(input: &syn::ItemStruct) -> proc_macro2::TokenStream {
let struct_name = &input.ident;
let fields_mod = Ident::new(&format!("__{}_fields", struct_name), struct_name.span());
let mut type_aliases = Vec::new();
if let syn::Fields::Named(ref fields) = input.fields {
for field in &fields.named {
if let Some(ref field_name) = field.ident {
if let syn::Type::Array(array_type) = &field.ty {
let elem_type = &*array_type.elem;
type_aliases.push(quote! {
#[allow(non_camel_case_types)]
pub type #field_name = #elem_type;
});
}
}
}
}
quote! {
#[doc(hidden)]
#[allow(non_snake_case)]
pub mod #fields_mod {
use super::*;
#(#type_aliases)*
}
}
}
#[proc_macro_attribute]
pub fn stack_frame(_attr: TokenStream, item: TokenStream) -> TokenStream {
let mut input = parse_macro_input!(item as syn::ItemStruct);
// Remove existing repr attributes to avoid conflicts.
input.attrs.retain(|attr| !attr.path().is_ident("repr"));
let fields_mod = gen_fields_module(&input);
let expanded = quote! {
#[repr(C, align(8))]
#input
#fields_mod
};
TokenStream::from(expanded)
}
/// Attribute macro that generates hidden associated constants for array field
/// element sizes.
///
/// Enables `offset!` and `offset_immediate!` to resolve element sizes when
/// using array index syntax (e.g., `Struct.field[expr]`).
///
/// Unlike `#[stack_frame]`, this does not modify the struct's `#[repr]`.
///
/// # Example
/// ```ignore
/// #[array_fields]
/// #[repr(C, packed)]
/// struct TreeNode {
/// child: [*mut TreeNode; 2],
/// key: u16,
/// }
/// ```
///
/// Generates:
/// ```ignore
/// impl TreeNode {
/// const __CHILD_STRIDE: usize = core::mem::size_of::<*mut TreeNode>();
/// }
/// ```
#[proc_macro_attribute]
pub fn array_fields(_attr: TokenStream, item: TokenStream) -> TokenStream {
let input = parse_macro_input!(item as syn::ItemStruct);
let struct_name = &input.ident;
let mut stride_consts = Vec::new();
if let syn::Fields::Named(ref fields) = input.fields {
for field in &fields.named {
if let Some(ref field_name) = field.ident {
if let syn::Type::Array(array_type) = &field.ty {
let elem_type = &*array_type.elem;
let const_name = Ident::new(
&format!("__{}_STRIDE", field_name.to_string().to_uppercase()),
field_name.span(),
);
stride_consts.push(quote! {
#[doc(hidden)]
pub const #const_name: usize = core::mem::size_of::<#elem_type>();
});
}
}
}
}
let expanded = quote! {
#input
#[doc(hidden)]
#[allow(dead_code)]
impl #struct_name {
#(#stride_consts)*
}
};
TokenStream::from(expanded)
}🔀 Entrypoint branching
The Rust implementation does not use pinocchio for the entrypoint. Instead, it uses C-style bindings with the SIMD-0321 r2 pointer. Note that the Rust implementation already introduces overhead at this point in program flow due to greedy tail call optimizations.
Implementations
asm
.globl entrypoint
entrypoint:
# Read instruction data length and discriminator.
# ---------------------------------------------------------------------
ldxdw r9, [r2 - SIZE_OF_U64] # Get instruction data length.
ldxdw r8, [r1 + IB_N_ACCOUNTS_OFF] # Get n input buffer accounts.
ldxb r7, [r2 + OFFSET_ZERO] # Get discriminator.
# Jump to branch for given discriminator.
# ---------------------------------------------------------------------
jeq r7, INSN_DISCRIMINATOR_INSERT, insert
jeq r7, INSN_DISCRIMINATOR_REMOVE, remove
jeq r7, INSN_DISCRIMINATOR_INITIALIZE, initialize
# Error if invalid discriminator provided.
mov64 r0, E_INSTRUCTION_DISCRIMINATOR
exitrs
no_allocator!();
nostd_panic_handler!();
#[no_mangle]
pub unsafe extern "C" fn entrypoint(input: *mut u8, instruction_data: *mut u8) -> u64 {
let instruction_data_len = ldxdw(instruction_data, -(size_of::<u64>() as i16));
let n_accounts = ldxdw(input, input_buffer::N_ACCOUNTS_OFF);
let instruction_discriminator = ldxb(instruction_data, instruction::DISCRIMINATOR_OFF);
if likely(instruction_discriminator == instruction::DISCRIMINATOR_INSERT) {
insert(input, instruction_data, instruction_data_len, n_accounts)
} else if likely(instruction_discriminator == instruction::DISCRIMINATOR_REMOVE) {
remove(input, instruction_data, instruction_data_len, n_accounts)
} else if likely(instruction_discriminator == instruction::DISCRIMINATOR_INITIALIZE) {
initialize(input, instruction_data, instruction_data_len, n_accounts)
} else {
error::INSTRUCTION_DISCRIMINATOR.into()
}
}Benchmarking
| Test case | ASM (CUs) | Rust (CUs) | Overhead | Overhead % |
|---|---|---|---|---|
| Invalid instruction discriminator | 8 | 11 | +3 | +37.5% |
🚀 Initialize
The initialize operation creates a tree PDA for the entire program, then invokes a CreateAccount CPI, with the same fixed costs as in the counter example.
🛡️ Input checks
Implementations
asm
initialize:
# Error if invalid instruction data length.
# ---------------------------------------------------------------------
jne r9, SIZE_OF_INITIALIZE_INSTRUCTION, e_instruction_data_len
# Error if invalid number of accounts.
# ---------------------------------------------------------------------
jne r8, IB_N_ACCOUNTS_INIT, e_n_accounts
# Error if user has data.
# ---------------------------------------------------------------------
ldxdw r9, [r1 + IB_USER_DATA_LEN_OFF]
jne r9, DATA_LEN_ZERO, e_user_data_len
# Error if tree is duplicate or has data.
# ---------------------------------------------------------------------
ldxb r9, [r1 + IB_TREE_NON_DUP_MARKER_OFF]
jne r9, IB_NON_DUP_MARKER, e_tree_duplicate
ldxdw r9, [r1 + IB_TREE_DATA_LEN_OFF]
jne r9, DATA_LEN_ZERO, e_tree_data_len
# Error if System Program is duplicate or has invalid data length.
# ---------------------------------------------------------------------
ldxb r9, [r1 + IB_SYSTEM_PROGRAM_NON_DUP_MARKER_OFF]
jne r9, IB_NON_DUP_MARKER, e_system_program_duplicate
ldxdw r9, [r1 + IB_SYSTEM_PROGRAM_DATA_LEN_OFF]
jne r9, IB_SYSTEM_PROGRAM_DATA_LEN, e_system_program_data_len
# Error if Rent account is duplicate or has incorrect address.
# ---------------------------------------------------------------------
ldxb r9, [r1 + IB_RENT_NON_DUP_MARKER_OFF]
jne r9, IB_NON_DUP_MARKER, e_rent_duplicate
ldxdw r9, [r1 + IB_RENT_ADDRESS_OFF_0]
lddw r8, IB_RENT_ID_CHUNK_0
jne r9, r8, e_rent_address
ldxdw r9, [r1 + IB_RENT_ADDRESS_OFF_1]
lddw r8, IB_RENT_ID_CHUNK_1
jne r9, r8, e_rent_address
ldxdw r9, [r1 + IB_RENT_ADDRESS_OFF_2]
lddw r8, IB_RENT_ID_CHUNK_2
jne r9, r8, e_rent_address
ldxdw r9, [r1 + IB_RENT_ADDRESS_OFF_3]
# Optimize out the following line, which costs two CUs due to two
# 32-bit immediate loads across two opcodes:
# ```
# lddw r8, IB_RENT_ID_CHUNK_3
# ```
# Instead, replace with mov32, which only loads one 32-bit immediate,
# since the rent sysvar address has all chunk 3 hi bits unset.
mov32 r8, IB_RENT_ID_CHUNK_3_LO
jne r9, r8, e_rent_addressrs
#[inline(always)]
unsafe fn initialize(
input: *mut u8,
_instruction_data: *mut u8,
instruction_data_len: u64,
n_accounts: u64,
) -> u64 {
check_instruction_data_len!(instruction_data_len, InitializeInstruction);
// Error if incorrect number of accounts.
if_err!(
n_accounts != input_buffer::N_ACCOUNTS_INIT,
error::N_ACCOUNTS
);
// Error if user has data.
let _user = user_account!(input);
// Error if tree is duplicate or has data.
let tree = account_non_dup!(input, input_buffer::TREE_ACCOUNT_OFF, error::TREE_DUPLICATE);
check_data_len!(tree, data::DATA_LEN_ZERO, error::TREE_DATA_LEN);
check_cpi_accounts!(input);Benchmarking
| Test case | ASM (CUs) | Rust (CUs) | Overhead | Overhead % |
|---|---|---|---|---|
| Invalid instruction data length | 9 | 12 | +3 | +33.3% |
| Too few accounts | 10 | 13 | +3 | +30.0% |
| Too many accounts | 10 | 13 | +3 | +30.0% |
| User has nonzero data length | 12 | 15 | +3 | +25.0% |
| Tree account is duplicate | 14 | 17 | +3 | +21.4% |
| Tree has nonzero data length | 16 | 19 | +3 | +18.8% |
| System program is duplicate | 18 | 21 | +3 | +16.7% |
| System program wrong data length | 20 | 23 | +3 | +15.0% |
| Rent sysvar is duplicate | 22 | 25 | +3 | +13.6% |
| Rent address mismatch word 0 | 25 | 28 | +3 | +12.0% |
| Rent address mismatch word 1 | 25 | 28 | +3 | +12.0% |
| Rent address mismatch word 2 | 28 | 32 | +4 | +14.3% |
| Rent address mismatch word 3 | 28 | 32 | +4 | +14.3% |
| Rent address mismatch word 4 | 31 | 36 | +5 | +16.1% |
| Rent address mismatch word 5 | 31 | 36 | +5 | +16.1% |
| Rent address mismatch word 6 | 34 | 40 | +6 | +17.6% |
| Rent address mismatch word 7 | 34 | 40 | +6 | +17.6% |
🔍 PDA checks
Implementations
asm
# Compute PDA.
# ---------------------------------------------------------------------
# Skip assignment for r1, since no seeds need to be parsed and this
# argument is effectively ignored.
# ---------------------------------------------------------------------
mov64 r2, CPI_N_SEEDS_TRY_FIND_PDA # Declare no seeds to parse.
mov64 r3, r1 # Get input buffer pointer.
add64 r3, IB_INIT_PROGRAM_ID_OFF_IMM # Point at program ID.
mov64 r4, r10 # Get stack frame pointer.
add64 r4, SF_INIT_PDA_OFF # Point to PDA region on stack.
mov64 r5, r10 # Get stack frame pointer.
add64 r5, SF_INIT_BUMP_SEED_OFF # Point to bump seed region on stack.
call sol_try_find_program_address # Find PDA.
# Compare computed PDA against passed account.
# ---------------------------------------------------------------------
ldxdw r9, [r1 + IB_TREE_ADDRESS_OFF_0]
ldxdw r8, [r4 + PUBKEY_CHUNK_OFF_0]
jne r9, r8, e_pda_mismatch
ldxdw r9, [r1 + IB_TREE_ADDRESS_OFF_1]
ldxdw r8, [r4 + PUBKEY_CHUNK_OFF_1]
jne r9, r8, e_pda_mismatch
ldxdw r9, [r1 + IB_TREE_ADDRESS_OFF_2]
ldxdw r8, [r4 + PUBKEY_CHUNK_OFF_2]
jne r9, r8, e_pda_mismatch
ldxdw r9, [r1 + IB_TREE_ADDRESS_OFF_3]
ldxdw r8, [r4 + PUBKEY_CHUNK_OFF_3]
jne r9, r8, e_pda_mismatchrs
#[cfg(target_os = "solana")]
// Invoke syscall.
let (pda, bump) = {
let mut pda = MaybeUninit::<Address>::uninit();
let mut bump = MaybeUninit::<u8>::uninit();
// Get input buffer footer pointer.
sol_try_find_program_address(
// Pass a declared pointer instead of null to prevent unnecessary register assignment.
input,
cpi::N_SEEDS_TRY_FIND_PDA,
input.add(input_buffer::INIT_PROGRAM_ID_OFF_IMM as usize),
pda.as_mut_ptr().cast(),
bump.as_mut_ptr(),
);
(pda.assume_init(), bump.assume_init())
};
#[cfg(not(target_os = "solana"))]
let (pda, bump) = (Address::default(), 0);
// Compare result with passed PDA.
if_err!(
!address_eq(
addr_of!(pda),
input.add(input_buffer::TREE_ADDRESS_OFF_0 as usize).cast()
),
error::PDA_MISMATCH
);Benchmarking
| Test case | Fixed CU costs | ASM (net CUs) | Rust (net CUs) | Overhead | Overhead % |
|---|---|---|---|---|---|
| PDA mismatch chunk 0 | 1500 | 45 | 55 | +10 | +22.2% |
| PDA mismatch chunk 1 | 1500 | 48 | 58 | +10 | +20.8% |
| PDA mismatch chunk 2 | 1500 | 51 | 61 | +10 | +19.6% |
| PDA mismatch chunk 3 | 1500 | 54 | 64 | +10 | +18.5% |
🛠️ Create account
The assembly implementation includes pointer walkthrough optimizations that are not available in Rust, since the compiler enforces instruction-level parallelism.
Implementations
asm
# Pack SolInstruction.
# ---------------------------------------------------------------------
# Packed later during bulk pointer load operation:
# - [x] System Program ID pointer.
# - [x] Account metas pointer.
# - [x] Instruction data pointer.
# ---------------------------------------------------------------------
stdw [r10 + SF_INIT_INSN_ACCOUNT_LEN_OFF], CPI_N_ACCOUNTS
stdw [r10 + SF_INIT_INSN_DATA_LEN_OFF], CPI_CREATE_ACCOUNT_INSN_DATA_LEN
# Pack CreateAccount instruction data.
# ---------------------------------------------------------------------
# - Discriminator is already set to 0 since stack is zero initialized.
# - Reuses r3 from PDA syscall.
# ---------------------------------------------------------------------
ldxdw r9, [r1 + IB_RENT_DATA_OFF] # Load lamports per byte
mul64 r9, CPI_ACCOUNT_DATA_SCALAR # Multiply to get rent-exempt cost.
# Store in instruction data.
stxdw [r10 + SF_INIT_CREATE_ACCOUNT_LAMPORTS_UOFF], r9
# Store new account data length.
stdw [r10 + SF_INIT_CREATE_ACCOUNT_SPACE_UOFF], CPI_TREE_DATA_LEN
# Copy in program ID to instruction data.
ldxdw r9, [r3 + PUBKEY_CHUNK_OFF_0]
stxdw [r10 + SF_INIT_CREATE_ACCOUNT_OWNER_UOFF_0], r9
ldxdw r9, [r3 + PUBKEY_CHUNK_OFF_1]
stxdw [r10 + SF_INIT_CREATE_ACCOUNT_OWNER_UOFF_1], r9
ldxdw r9, [r3 + PUBKEY_CHUNK_OFF_2]
stxdw [r10 + SF_INIT_CREATE_ACCOUNT_OWNER_UOFF_2], r9
ldxdw r9, [r3 + PUBKEY_CHUNK_OFF_3]
stxdw [r10 + SF_INIT_CREATE_ACCOUNT_OWNER_UOFF_3], r9
# Pack SolAccountMeta for user and tree.
# ---------------------------------------------------------------------
# Packed later during bulk pointer load operation:
# - [x] User pubkey pointer.
# - [x] Tree pubkey pointer.
# ---------------------------------------------------------------------
sth [r10 + SF_INIT_USER_META_IS_WRITABLE_OFF], CPI_WRITABLE_SIGNER
sth [r10 + SF_INIT_TREE_META_IS_WRITABLE_OFF], CPI_WRITABLE_SIGNER
# Pack SolAccountInfo for user and tree.
# ---------------------------------------------------------------------
# Packed later during bulk pointer load operation:
# - [x] User pubkey pointer.
# - [x] Tree pubkey pointer.
# - [x] User lamports pointer.
# - [x] Tree lamports pointer.
# - [x] User data pointer.
# - [x] Tree data pointer.
# - [x] User owner pointer.
# - [x] Tree owner pointer.
# Skipped due to zero-initialized stack memory:
# - User data length (already checked as zero).
# - Tree data length (already checked as zero).
# - User rent epoch.
# - Tree rent epoch.
# - User executable.
# - Tree executable.
# ---------------------------------------------------------------------
sth [r10 + SF_INIT_USER_INFO_IS_SIGNER_OFF], CPI_WRITABLE_SIGNER
sth [r10 + SF_INIT_TREE_INFO_IS_SIGNER_OFF], CPI_WRITABLE_SIGNER
# Initialize signer seed for PDA bump seed.
# ---------------------------------------------------------------------
# Reuses r5 from PDA derivation syscall.
# ---------------------------------------------------------------------
# Store pointer to bump seed.
stxdw [r10 + SF_INIT_SIGNER_SEED_ADDR_OFF], r5
stdw [r10 + SF_INIT_SIGNER_SEED_LEN_OFF], SIZE_OF_U8 # Store length.
# Initialize signers seeds for PDA.
# ---------------------------------------------------------------------
# Packed later during bulk pointer load operation:
# - [x] Signer seed pointer.
# ---------------------------------------------------------------------
stdw [r10 + SF_INIT_SIGNERS_SEEDS_LEN_OFF], CPI_N_SEEDS_CREATE_ACCOUNT
# Bulk assign/load pointers for account metas and infos.
# ---------------------------------------------------------------------
# Since pointers must be loaded from registers, this block steps
# through the input buffer in order to reduce intermediate loads.
# ---------------------------------------------------------------------
add64 r1, IB_USER_ADDRESS_OFF # Point to user address in input buffer.
stxdw [r10 + SF_INIT_USER_META_PUBKEY_OFF], r1 # Store in account meta.
stxdw [r10 + SF_INIT_USER_INFO_PUBKEY_OFF], r1 # Store in account info.
add64 r1, SIZE_OF_ADDRESS # Advance to user owner.
stxdw [r10 + SF_INIT_USER_INFO_OWNER_OFF], r1 # Store in account info.
add64 r1, SIZE_OF_ADDRESS # Advance to user lamports.
stxdw [r10 + SF_INIT_USER_INFO_LAMPORTS_OFF], r1 # Store in acct info.
add64 r1, SIZE_OF_U128 # Advance to user data.
stxdw [r10 + SF_INIT_USER_INFO_DATA_OFF], r1 # Store in account info.
# Advance to tree address field.
add64 r1, IB_USER_DATA_TO_TREE_ADDRESS_REL_OFF_IMM
stxdw [r10 + SF_INIT_TREE_META_PUBKEY_OFF], r1 # Store in account meta.
stxdw [r10 + SF_INIT_TREE_INFO_PUBKEY_OFF], r1 # Store in account info.
add64 r1, SIZE_OF_ADDRESS # Advance to tree owner.
stxdw [r10 + SF_INIT_TREE_INFO_OWNER_OFF], r1 # Store in account info.
add64 r1, SIZE_OF_ADDRESS # Advance to tree lamports.
stxdw [r10 + SF_INIT_TREE_INFO_LAMPORTS_OFF], r1 # Store in acct info.
add64 r1, SIZE_OF_U128 # Advance to tree data.
stxdw [r10 + SF_INIT_TREE_INFO_DATA_OFF], r1 # Store in account info.
mov64 r6, r1 # Store tree data pointer for later.
# Bulk assign/load pointers for CPI bindings.
# ---------------------------------------------------------------------
# This block steps through the stack frame, optimizing assignments in
# preparation for the impending CreateAccount CPI, which requires:
# - [x] r1 = pointer to instruction.
# - [x] r2 = pointer to account infos.
# - [x] r4 = pointer to signers seeds.
# Notably, it reuses r4 from the PDA derivation syscall to walk through
# pointers on the stack, before advancing it to its final value.
# ---------------------------------------------------------------------
# Advance to System Program ID pointer on zero-initialized stack.
add64 r4, SF_INIT_PDA_TO_SYSTEM_PROGRAM_ID_REL_OFF_IMM
# Store in SolInstruction.
stxdw [r10 + SF_INIT_INSN_PROGRAM_ID_OFF], r4
# Advance to SolAccountMeta array pointer.
add64 r4, SF_INIT_SYSTEM_PROGRAM_ID_TO_ACCT_METAS_REL_OFF_IMM
stxdw [r10 + SF_INIT_INSN_ACCOUNTS_OFF], r4 # Store in SolInstruction.
# Advance to instruction data pointer.
add64 r4, SF_INIT_ACCT_METAS_TO_INSN_DATA_REL_OFF_IMM
stxdw [r10 + SF_INIT_INSN_DATA_OFF], r4 # Store in SolInstruction.
# Advance to signer seeds pointer.
add64 r4, SF_INIT_INSN_DATA_TO_SIGNER_SEEDS_REL_OFF_IMM
stxdw [r10 + SF_INIT_SIGNERS_SEEDS_ADDR_OFF], r4
# Advance to signers seeds pointer.
add64 r4, SF_INIT_SIGNER_SEEDS_TO_SIGNERS_SEEDS_REL_OFF_IMM
# Assign remaining syscall pointers.
mov64 r1, r10
add64 r1, SF_INIT_INSN_PROGRAM_ID_OFF
mov64 r2, r10
add64 r2, SF_INIT_ACCT_INFOS_OFF
# Invoke CPI.
# ---------------------------------------------------------------------
mov64 r3, CPI_N_ACCOUNTS
mov64 r5, CPI_N_PDA_SIGNERS
call sol_invoke_signed_c
# Store next pointer in tree header.
# ---------------------------------------------------------------------
mov64 r7, r6 # Get copy of tree data pointer.
add64 r7, SIZE_OF_TREE_HEADER # Advance to next node.
stxdw [r6 + TREE_HEADER_NEXT_OFF], r7 # Store in next field.
exitrs
// Pack CreateAccount instruction data.
let create_account_instruction_data = CreateAccountInstructionData {
discriminator: cpi::CREATE_ACCOUNT_DISCRIMINATOR,
lamports: (cpi::ACCOUNT_DATA_SCALAR as u64) * ldxdw(input, input_buffer::RENT_DATA_OFF),
space: cpi::TREE_DATA_LEN as u64,
owner: read_unaligned(
input
.add(input_buffer::INIT_PROGRAM_ID_OFF_IMM as usize)
.cast(),
),
};
// Pack account metas and infos.
let user_key = input.add(input_buffer::USER_ADDRESS_OFF as usize).cast();
let tree_key = input.add(input_buffer::TREE_ADDRESS_OFF as usize).cast();
let sol_account_metas = [
SolAccountMeta {
pubkey: user_key,
is_writable: true,
is_signer: true,
},
SolAccountMeta {
pubkey: tree_key,
is_writable: true,
is_signer: true,
},
];
let sol_account_infos = [
SolAccountInfo {
key: user_key,
owner: input.add(input_buffer::USER_OWNER_OFF as usize).cast(),
lamports: input.add(input_buffer::USER_LAMPORTS_OFF as usize).cast(),
data: input.add(input_buffer::USER_DATA_OFF as usize),
data_len: data::DATA_LEN_ZERO,
rent_epoch: cpi::RENT_EPOCH_NULL,
is_signer: true,
is_writable: true,
executable: false,
},
SolAccountInfo {
key: tree_key,
owner: input.add(input_buffer::TREE_OWNER_OFF as usize).cast(),
lamports: input.add(input_buffer::TREE_LAMPORTS_OFF as usize).cast(),
data: input.add(input_buffer::TREE_DATA_OFF as usize),
data_len: data::DATA_LEN_ZERO,
rent_epoch: cpi::RENT_EPOCH_NULL,
is_signer: true,
is_writable: true,
executable: false,
},
];
// Pack instruction.
let system_program_address = Address::default();
let sol_instruction = SolInstruction {
program_id: addr_of!(system_program_address).cast_mut().cast(),
accounts: sol_account_metas.as_ptr().cast_mut().cast(),
account_len: sol_account_metas.len() as u64,
data: addr_of!(create_account_instruction_data).cast_mut().cast(),
data_len: cpi::CREATE_ACCOUNT_INSN_DATA_LEN as u64,
};
// Initialize signer seed for PDA bump.
let bump_seed = SolSignerSeed {
addr: addr_of!(bump).cast(),
len: size_of::<u8>() as u64,
};
// Initialize signer seeds for PDA.
let signers_seeds = SolSignerSeeds {
addr: addr_of!(bump_seed).cast(),
len: cpi::N_SEEDS_CREATE_ACCOUNT as u64,
};
#[cfg(target_os = "solana")]
sol_invoke_signed_c(
addr_of!(sol_instruction).cast(),
addr_of!(sol_account_infos).cast(),
cpi::N_ACCOUNTS as u64,
addr_of!(signers_seeds).cast(),
cpi::N_PDA_SIGNERS as u64,
);
#[cfg(not(target_os = "solana"))]
#[allow(path_statements)]
{
signers_seeds;
sol_account_infos;
sol_instruction;
}
// Store next pointer in tree header.
let tree_data: *mut TreeHeader = input.add(input_buffer::TREE_DATA_OFF as usize).cast();
(*tree_data).next = tree_data.add(1).cast();Benchmarking
| Test case | Fixed CU costs | ASM (net CUs) | Rust (net CUs) | Overhead | Overhead % |
|---|---|---|---|---|---|
| System Program is wrong address | 2446 | 108 | 142 | +34 | +31.5% |
| User has insufficient Lamports | 2596 | 108 | 142 | +34 | +31.5% |
| CreateAccount happy path | 2596 | 112 | 148 | +36 | +32.1% |
➕ Insert
🛡️ Input checks
Implementations
asm
insert:
# Error if invalid instruction data length.
# ---------------------------------------------------------------------
jne r9, SIZE_OF_INSERT_INSTRUCTION, e_instruction_data_len
# Error if too few accounts.
# ---------------------------------------------------------------------
jlt r8, IB_N_ACCOUNTS_GENERAL, e_n_accounts
# Error if user has data.
# ---------------------------------------------------------------------
ldxdw r9, [r1 + IB_USER_DATA_LEN_OFF]
jne r9, DATA_LEN_ZERO, e_user_data_len
# Error if tree is duplicate.
# ---------------------------------------------------------------------
ldxb r9, [r1 + IB_TREE_NON_DUP_MARKER_OFF]
jne r9, IB_NON_DUP_MARKER, e_tree_duplicaters
#[allow(unused_assignments)]
#[inline(always)]
unsafe fn insert(
input: *mut u8,
instruction_data: *mut u8,
instruction_data_len: u64,
n_accounts: u64,
) -> u64 {
check_instruction_data_len!(instruction_data_len, InsertInstruction);
// Error if too few accounts.
if_err!(
n_accounts < input_buffer::N_ACCOUNTS_GENERAL,
error::N_ACCOUNTS
);
// Error if user has data.
let _user = user_account!(input);
// Error if tree is duplicate.
let tree = account_non_dup!(input, input_buffer::TREE_ACCOUNT_OFF, error::TREE_DUPLICATE);Benchmarking
| Test case | ASM (CUs) | Rust (CUs) | Overhead | Overhead % |
|---|---|---|---|---|
| Instruction data too short | 7 | 9 | +2 | +28.6% |
| Instruction data too long | 7 | 9 | +2 | +28.6% |
| Too few accounts | 8 | 10 | +2 | +25.0% |
| User has nonzero data length | 10 | 12 | +2 | +20.0% |
| Tree account is duplicate | 12 | 14 | +2 | +16.7% |
📦 Allocate
See AccountView::resize_unchecked for reference implementation.
Implementations
asm
# Branch based on state of stack top.
# ---------------------------------------------------------------------
# Get stack top pointer.
ldxdw r9, [r1 + IB_TREE_DATA_TOP_OFF]
jne r9, NULL, insert_pop # Pop node from stack if non-null.
insert_allocate:
# Error if wrong number of accounts for allocation.
# ---------------------------------------------------------------------
jne r8, IB_N_ACCOUNTS_INIT, e_n_accounts_insert_allocation
# Compute shifted input buffer pointer based on tree data length.
# ---------------------------------------------------------------------
ldxdw r9, [r1 + IB_TREE_DATA_LEN_OFF] # Get tree data length.
# Store in account info for CPI.
stxdw [r10 + SF_INIT_TREE_INFO_DATA_LEN_OFF], r9
mov64 r7, r9 # Store copy for later.
add64 r9, MAX_DATA_PAD # Add max possible padding.
and64 r9, DATA_LEN_AND_MASK # Truncate to 8-byte alignment.
add64 r9, r1 # Increment by input buffer.
# Check system program is not duplicate and has correct data length.
# ---------------------------------------------------------------------
ldxb r8, [r9 + IB_SYSTEM_PROGRAM_NON_DUP_MARKER_OFF]
jne r8, IB_NON_DUP_MARKER, e_system_program_duplicate
ldxdw r8, [r9 + IB_SYSTEM_PROGRAM_DATA_LEN_OFF]
jne r8, IB_SYSTEM_PROGRAM_DATA_LEN, e_system_program_data_len
# Check rent sysvar is not duplicate and has correct address.
# ---------------------------------------------------------------------
ldxb r8, [r9 + IB_RENT_NON_DUP_MARKER_OFF]
jne r8, IB_NON_DUP_MARKER, e_rent_duplicate
ldxdw r8, [r9 + IB_RENT_ADDRESS_OFF_0]
lddw r4, IB_RENT_ID_CHUNK_0
jne r8, r4, e_rent_address
ldxdw r8, [r9 + IB_RENT_ADDRESS_OFF_1]
lddw r4, IB_RENT_ID_CHUNK_1
jne r8, r4, e_rent_address
ldxdw r8, [r9 + IB_RENT_ADDRESS_OFF_2]
lddw r4, IB_RENT_ID_CHUNK_2
jne r8, r4, e_rent_address
ldxdw r8, [r9 + IB_RENT_ADDRESS_OFF_3]
mov32 r4, IB_RENT_ID_CHUNK_3_LO
jne r8, r4, e_rent_address
# Calculate transfer lamports.
# ---------------------------------------------------------------------
ldxdw r8, [r9 + IB_RENT_DATA_OFF] # Load lamports per byte.
mul64 r8, SIZE_OF_TREE_NODE # Multiply to get transfer cost.
# Pack Transfer instruction data in CreateAccount slot on stack.
# ---------------------------------------------------------------------
stw [r10 + SF_INIT_CREATE_ACCOUNT_DISCRIMINATOR_UOFF], CPI_TRANSFER_DISCRIMINATOR
stxdw [r10 + SF_INIT_CREATE_ACCOUNT_LAMPORTS_UOFF], r8
# Pack SolInstruction.
# ---------------------------------------------------------------------
stdw [r10 + SF_INIT_INSN_ACCOUNT_LEN_OFF], CPI_N_ACCOUNTS
stdw [r10 + SF_INIT_INSN_DATA_LEN_OFF], CPI_TRANSFER_INSN_DATA_LEN
# Pack SolAccountMeta flags for user and tree.
# ---------------------------------------------------------------------
sth [r10 + SF_INIT_USER_META_IS_WRITABLE_OFF], CPI_WRITABLE_SIGNER
stb [r10 + SF_INIT_TREE_META_IS_WRITABLE_OFF], BOOL_TRUE
# Pack SolAccountInfo flags for user and tree.
# ---------------------------------------------------------------------
sth [r10 + SF_INIT_USER_INFO_IS_SIGNER_OFF], CPI_WRITABLE_SIGNER
stb [r10 + SF_INIT_TREE_INFO_IS_WRITABLE_UOFF], BOOL_TRUE
# Bulk assign/load pointers for account metas and infos.
# ---------------------------------------------------------------------
mov64 r6, r1 # Store input buffer pointer for later.
add64 r1, IB_USER_ADDRESS_OFF # Point to user address in input buffer.
stxdw [r10 + SF_INIT_USER_META_PUBKEY_OFF], r1 # Store in account meta.
stxdw [r10 + SF_INIT_USER_INFO_PUBKEY_OFF], r1 # Store in account info.
add64 r1, SIZE_OF_ADDRESS # Advance to user owner.
stxdw [r10 + SF_INIT_USER_INFO_OWNER_OFF], r1 # Store in account info.
add64 r1, SIZE_OF_ADDRESS # Advance to user lamports.
stxdw [r10 + SF_INIT_USER_INFO_LAMPORTS_OFF], r1 # Store in acct info.
add64 r1, SIZE_OF_U128 # Advance to user data.
stxdw [r10 + SF_INIT_USER_INFO_DATA_OFF], r1 # Store in account info.
# Advance to tree address field.
add64 r1, IB_USER_DATA_TO_TREE_ADDRESS_REL_OFF_IMM
stxdw [r10 + SF_INIT_TREE_META_PUBKEY_OFF], r1 # Store in account meta.
stxdw [r10 + SF_INIT_TREE_INFO_PUBKEY_OFF], r1 # Store in account info.
add64 r1, SIZE_OF_ADDRESS # Advance to tree owner.
stxdw [r10 + SF_INIT_TREE_INFO_OWNER_OFF], r1 # Store in account info.
add64 r1, SIZE_OF_ADDRESS # Advance to tree lamports.
stxdw [r10 + SF_INIT_TREE_INFO_LAMPORTS_OFF], r1 # Store in acct info.
add64 r1, SIZE_OF_U128 # Advance to tree data.
stxdw [r10 + SF_INIT_TREE_INFO_DATA_OFF], r1 # Store in account info.
# Bulk assign/load pointers for CPI bindings.
# ---------------------------------------------------------------------
# Point to System Program ID on zero-initialized stack.
mov64 r4, r10
add64 r4, SF_INIT_SYSTEM_PROGRAM_ADDRESS_OFF
# Store in SolInstruction.
stxdw [r10 + SF_INIT_INSN_PROGRAM_ID_OFF], r4
# Advance to SolAccountMeta array pointer.
add64 r4, SF_INIT_SYSTEM_PROGRAM_ID_TO_ACCT_METAS_REL_OFF_IMM
# Store in SolInstruction.
stxdw [r10 + SF_INIT_INSN_ACCOUNTS_OFF], r4
# Advance to instruction data pointer.
add64 r4, SF_INIT_ACCT_METAS_TO_INSN_DATA_REL_OFF_IMM
stxdw [r10 + SF_INIT_INSN_DATA_OFF], r4 # Store in SolInstruction.
# Invoke Transfer CPI.
# ---------------------------------------------------------------------
mov64 r1, r10
add64 r1, SF_INIT_INSN_PROGRAM_ID_OFF
mov64 r8, r2 # Save instruction data pointer for later.
mov64 r2, r10
add64 r2, SF_INIT_ACCT_INFOS_OFF
mov64 r3, CPI_N_ACCOUNTS
# Ignore PDA signer seeds pointer, since none required.
mov64 r5, CPI_N_PDA_SIGNERS_TRANSFER
call sol_invoke_signed_c
mov64 r2, r8 # Restore instruction data pointer.
mov64 r1, r6 # Restore input buffer pointer.
# Update tree data length.
# ---------------------------------------------------------------------
add64 r7, SIZE_OF_TREE_NODE # Increment data length.
stxdw [r1 + IB_TREE_DATA_LEN_OFF], r7 # Store in input buffer.
# Get node = next, then advance next by one TreeNode.
# ---------------------------------------------------------------------
ldxdw r7, [r1 + IB_TREE_DATA_NEXT_OFF] # Get pointer to next node.
mov64 r9, r7 # Store node pointer for later, the new node.
add64 r7, SIZE_OF_TREE_NODE # Increment to point to new next.
stxdw [r1 + IB_TREE_DATA_NEXT_OFF], r7 # Advance next.
# Continue insert.
# ---------------------------------------------------------------------
ja insert_store_key_value_pair
insert_pop:
# Pop node from free stack.
# ---------------------------------------------------------------------
ldxdw r8, [r9 + OFFSET_ZERO] # Load StackNode.next.
stxdw [r1 + IB_TREE_DATA_TOP_OFF], r8 # Update top in header.
insert_store_key_value_pair:
ldxw r4, [r2 + INSN_INSERT_KEY_OFF] # Load two fields together.
stxw [r9 + TREE_NODE_KEY_OFF], r4 # Store both fields together.rs
// Allocate or recycle a node.
let tree_header: *mut TreeHeader = input.add(input_buffer::TREE_DATA_OFF as usize).cast();
let mut node: *mut TreeNode = if (*tree_header).top.is_null() {
// Error if wrong number of accounts passed, since need extra accounts to allocate space.
if_err!(
n_accounts != input_buffer::N_ACCOUNTS_INIT,
error::N_ACCOUNTS_INSERT_ALLOCATION
);
// Get shifted input buffer pointer based on tree data length.
let tree_data_len: *mut u64 = addr_of_mut!((*tree).data_len);
let shifted_input =
input.add((*tree_data_len).next_multiple_of(data::BPF_ALIGN_OF_U128 as u64) as usize);
// Check system program and rent sysvar accounts using shifted input buffer pointer.
check_cpi_accounts!(shifted_input);
// Calculate additional lamports for rent exemption of one TreeNode.
let lamports_per_byte = ldxdw(shifted_input, input_buffer::RENT_DATA_OFF);
let transfer_lamports = size_of::<TreeNode>() as u64 * lamports_per_byte;
// Pack Transfer instruction data.
let transfer_instruction_data = TransferInstructionData {
discriminator: cpi::TRANSFER_DISCRIMINATOR,
lamports: transfer_lamports,
};
// Pack account metas and infos.
let user_key = input.add(input_buffer::USER_ADDRESS_OFF as usize).cast();
let tree_key = input.add(input_buffer::TREE_ADDRESS_OFF as usize).cast();
let sol_account_metas = [
SolAccountMeta {
pubkey: user_key,
is_writable: true,
is_signer: true,
},
SolAccountMeta {
pubkey: tree_key,
is_writable: true,
is_signer: false,
},
];
let sol_account_infos = [
SolAccountInfo {
key: user_key,
owner: input.add(input_buffer::USER_OWNER_OFF as usize).cast(),
lamports: input.add(input_buffer::USER_LAMPORTS_OFF as usize).cast(),
data: input.add(input_buffer::USER_DATA_OFF as usize),
data_len: data::DATA_LEN_ZERO,
rent_epoch: cpi::RENT_EPOCH_NULL,
is_signer: true,
is_writable: true,
executable: false,
},
SolAccountInfo {
key: tree_key,
owner: input.add(input_buffer::TREE_OWNER_OFF as usize).cast(),
lamports: input.add(input_buffer::TREE_LAMPORTS_OFF as usize).cast(),
data: input.add(input_buffer::TREE_DATA_OFF as usize),
data_len: *tree_data_len,
rent_epoch: cpi::RENT_EPOCH_NULL,
is_signer: false,
is_writable: true,
executable: false,
},
];
// Pack instruction.
let system_program_address = Address::default();
let sol_instruction = SolInstruction {
program_id: addr_of!(system_program_address).cast_mut().cast(),
accounts: sol_account_metas.as_ptr().cast_mut().cast(),
account_len: sol_account_metas.len() as u64,
data: addr_of!(transfer_instruction_data).cast_mut().cast(),
data_len: cpi::TRANSFER_INSN_DATA_LEN as u64,
};
// No signers needed, since user is already a signer on the transaction.
let empty_signers = SolSignerSeeds {
addr: core::ptr::null(),
len: 0,
};
#[cfg(target_os = "solana")]
sol_invoke_signed_c(
addr_of!(sol_instruction).cast(),
addr_of!(sol_account_infos).cast(),
cpi::N_ACCOUNTS as u64,
addr_of!(empty_signers).cast(),
cpi::N_PDA_SIGNERS_TRANSFER,
);
#[cfg(not(target_os = "solana"))]
#[allow(path_statements)]
{
empty_signers;
sol_account_infos;
sol_instruction;
}
// Increase tree data length by size of one TreeNode.
*tree_data_len += size_of::<TreeNode>() as u64;
// Advance next pointer by one TreeNode.
let node = (*tree_header).next;
(*tree_header).next = (*tree_header).next.add(1);
node
} else {
// Pop node from free stack.
let top = (*tree_header).top;
(*tree_header).top = (*top).next;
top.cast()
};
// Set key and value together as a single word.
*addr_of_mut!((*node).key).cast() = ldxw(instruction_data, instruction::INSERT_KEY_OFF);Benchmarking
| Test case | ASM (CUs) | Rust (CUs) | Overhead | Overhead % |
|---|---|---|---|---|
| Wrong N accounts for allocation | 15 | 19 | +4 | +26.7% |
| System program is duplicate | 23 | 27 | +4 | +17.4% |
| System program wrong data length | 25 | 29 | +4 | +16.0% |
| Rent sysvar is duplicate | 27 | 31 | +4 | +14.8% |
| Rent address mismatch chunk 0 | 30 | 34 | +4 | +13.3% |
| Rent address mismatch chunk 1 | 33 | 38 | +5 | +15.2% |
| Rent address mismatch chunk 2 | 36 | 42 | +6 | +16.7% |
| Rent address mismatch chunk 3 | 39 | 46 | +7 | +17.9% |
| Test case | Fixed CU costs | ASM (net CUs) | Rust (net CUs) | Overhead | Overhead % |
|---|---|---|---|---|---|
| Insert alloc happy path | 1096 | 100 | 129 | +29 | +29.0% |
| Alloc exceeds max data length | 1096 | 1398904 | 1398904 | +0 | +0.0% |
🔎 Search
Implementations
asm
insert_search: # r9 = node
ldxh r4, [r2 + INSN_INSERT_KEY_OFF] # r4 = insn.key;
ldxdw r3, [r1 + IB_TREE_DATA_ROOT_OFF] # r3 = cursor = root;
jeq r3, NULL, insert_root
insert_search_loop:
mov64 r2, r3 # r2 = parent = cursor;
ldxh r5, [r3 + TREE_NODE_KEY_OFF] # r5 = cursor.key;
jlt r4, r5, insert_search_branch_l
jgt r4, r5, insert_search_branch_r
mov64 r0, E_KEY_EXISTS # Error if key already exists.
exit
insert_root:
# Root is null: new node becomes root.
# ---------------------------------------------------------------------
stb [r9 + TREE_NODE_COLOR_OFF], TREE_COLOR_R # node.color = red;
mov64 r2, NULL # r2 = parent = null;
stxdw [r9 + TREE_NODE_PARENT_OFF], r2 # node.parent = null;
stxdw [r1 + IB_TREE_DATA_ROOT_OFF], r9 # root = node;
exit
insert_search_branch_l:
ldxdw r3, [r2 + TREE_NODE_CHILD_L_OFF] # r3 = parent.child[L];
jne r3, NULL, insert_search_loop
# Null child: insert node as left child of parent.
stb [r9 + TREE_NODE_COLOR_OFF], TREE_COLOR_R # node.color = red;
stxdw [r9 + TREE_NODE_PARENT_OFF], r2 # node.parent = parent;
stxdw [r2 + TREE_NODE_CHILD_L_OFF], r9 # parent.child[L] = node;
# Inline case 1: if parent is black, tree is valid.
ldxb r6, [r2 + TREE_NODE_COLOR_OFF] # r6 = parent.color;
jne r6, TREE_COLOR_B, insert_fixup_check_case_4
exit
insert_search_branch_r:
ldxdw r3, [r2 + TREE_NODE_CHILD_R_OFF] # r3 = parent.child[R];
jne r3, NULL, insert_search_loop
# Null child: insert node as right child of parent.
stb [r9 + TREE_NODE_COLOR_OFF], TREE_COLOR_R # node.color = red;
stxdw [r9 + TREE_NODE_PARENT_OFF], r2 # node.parent = parent;
stxdw [r2 + TREE_NODE_CHILD_R_OFF], r9 # parent.child[R] = node;rs
let key = ldxh(instruction_data, instruction::INSERT_KEY_OFF);
let mut cursor = (*tree_header).root;
// Root is null: new node becomes root.
if cursor.is_null() {
(*node).color = Color::Red;
(*node).parent = null_mut();
(*tree_header).root = node;
return SUCCESS;
}
let mut parent: *mut TreeNode;
loop {
parent = cursor;
let cursor_key = (*cursor).key;
if likely(key > cursor_key) {
cursor = (*parent).child[tree::DIR_R];
if cursor.is_null() {
(*node).color = Color::Red;
(*node).parent = parent;
(*parent).child[tree::DIR_R] = node;
if (*parent).color == Color::Black {
return SUCCESS;
}
break;
}
} else if likely(key < cursor_key) {
cursor = (*parent).child[tree::DIR_L];
if cursor.is_null() {
(*node).color = Color::Red;
(*node).parent = parent;
(*parent).child[tree::DIR_L] = node;
if (*parent).color == Color::Black {
return SUCCESS;
}
break;
}
} else {
return error::KEY_EXISTS.into();
}
}Benchmarking
| Test case | ASM (CUs) | Rust (CUs) | Overhead | Overhead % |
|---|---|---|---|---|
| Dup at root | 25 | 30 | +5 | +20.0% |
| Dup in left | 30 | 37 | +7 | +23.3% |
| Dup in right | 31 | 36 | +5 | +16.1% |
🔧 Insert fixup
Case 1
asm
insert_fixup_main:
# r2 := parent
# r5 := parent.key
# Case 1. # r9 := node
# ---------------------------------------------------------------------
ldxb r6, [r2 + TREE_NODE_COLOR_OFF] # r6 = parent.color;
jne r6, TREE_COLOR_B, insert_fixup_check_case_4
exit # If parent is black, tree is still valid, so exit.rs
// Main insert fixup.
loop {
// Case 1.
if (*parent).color == Color::Black {
return SUCCESS;
}Case 4
asm
insert_fixup_check_case_4:
# Check case 4.
# ---------------------------------------------------------------------
ldxdw r3, [r2 + TREE_NODE_PARENT_OFF] # r3 = grandparent;
jne r3, NULL, insert_fixup_check_case_5_6
stb [r2 + TREE_NODE_COLOR_OFF], TREE_COLOR_B # parent.color = black;
exitrs
let grandparent = (*parent).parent;
if grandparent.is_null() {
// Case 4.
(*parent).color = Color::Black;
return SUCCESS;
}Cases 5 and 6 (left)
asm
insert_fixup_check_case_5_6:
# Get uncle and check for case 5 or 6.
# ---------------------------------------------------------------------
ldxh r4, [r3 + TREE_NODE_KEY_OFF] # r4 = grandparent.key;
jgt r5, r4, insert_fixup_check_case_5_6_dir_r
insert_fixup_check_case_5_6_dir_l:
ldxdw r7, [r3 + TREE_NODE_CHILD_R_OFF] # r7 = uncle;
jeq r7, NULL, insert_fixup_case_5_6_dir_l
ldxb r8, [r7 + TREE_NODE_COLOR_OFF] # r8 = uncle.color;
jne r8, TREE_COLOR_B, insert_fixup_case_2
insert_fixup_case_5_6_dir_l:
ldxdw r6, [r2 + TREE_NODE_CHILD_R_OFF] # r6 = new_root = parent.child[R];
jne r9, r6, insert_fixup_case_6_dir_l
insert_fixup_case_5_dir_l:
ldxdw r8, [r6 + TREE_NODE_CHILD_L_OFF] # r8 = new_child = new_root.child[L];
stxdw [r2 + TREE_NODE_CHILD_R_OFF], r8 # parent.child[R] = new_child;
jeq r8, NULL, insert_fixup_case_5_dir_l_skip
stxdw [r8 + TREE_NODE_PARENT_OFF], r2 # new_child.parent = parent;
insert_fixup_case_5_dir_l_skip:
stxdw [r6 + TREE_NODE_CHILD_L_OFF], r2 # new_root.child[L] = parent;
stxdw [r6 + TREE_NODE_PARENT_OFF], r3 # new_root.parent = grandparent;
stxdw [r2 + TREE_NODE_PARENT_OFF], r6 # parent.parent = new_root;
stxdw [r3 + TREE_NODE_CHILD_L_OFF], r6 # grandparent.child[L] = new_root;
mov64 r9, r2 # node = old parent;
mov64 r2, r6 # parent = new_root;
insert_fixup_case_6_dir_l:
ldxdw r4, [r3 + TREE_NODE_PARENT_OFF] # r4 = great-grandparent;
ldxdw r8, [r2 + TREE_NODE_CHILD_R_OFF] # r8 = new_child = parent.child[R];
stxdw [r3 + TREE_NODE_CHILD_L_OFF], r8 # grandparent.child[L] = new_child;
jeq r8, NULL, insert_fixup_case_6_dir_l_skip
stxdw [r8 + TREE_NODE_PARENT_OFF], r3 # new_child.parent = grandparent;
insert_fixup_case_6_dir_l_skip:
stxdw [r2 + TREE_NODE_CHILD_R_OFF], r3 # parent.child[R] = grandparent;
stxdw [r2 + TREE_NODE_PARENT_OFF], r4 # parent.parent = great-grandparent;
stxdw [r3 + TREE_NODE_PARENT_OFF], r2 # grandparent.parent = parent;
jeq r4, NULL, insert_fixup_case_6_dir_l_root
ldxdw r8, [r4 + TREE_NODE_CHILD_R_OFF] # r8 = great-grandparent.child[R];
jne r3, r8, insert_fixup_case_6_dir_l_left
# Inline color + exit to eliminate `ja` (saves 1 CU).
stxdw [r4 + TREE_NODE_CHILD_R_OFF], r2 # great-grandparent.child[R] = parent;
stb [r2 + TREE_NODE_COLOR_OFF], TREE_COLOR_B # parent.color = black;
stb [r3 + TREE_NODE_COLOR_OFF], TREE_COLOR_R # grandparent.color = red;
exit
insert_fixup_case_6_dir_l_left:
# Inline color + exit to eliminate `ja` (saves 1 CU).
stxdw [r4 + TREE_NODE_CHILD_L_OFF], r2 # great-grandparent.child[L] = parent;
stb [r2 + TREE_NODE_COLOR_OFF], TREE_COLOR_B # parent.color = black;
stb [r3 + TREE_NODE_COLOR_OFF], TREE_COLOR_R # grandparent.color = red;
exit
insert_fixup_case_6_dir_l_root:
stxdw [r1 + IB_TREE_DATA_ROOT_OFF], r2 # tree.root = parent;
stb [r2 + TREE_NODE_COLOR_OFF], TREE_COLOR_B # parent.color = black;
stb [r3 + TREE_NODE_COLOR_OFF], TREE_COLOR_R # grandparent.color = red;
exitrs
// Determine direction and uncle with hardcoded child indices.
let uncle;
if parent == (*grandparent).child[tree::DIR_L] {
// dir_l: parent is left child of grandparent.
uncle = (*grandparent).child[tree::DIR_R];
if uncle.is_null() || (*uncle).color == Color::Black {
// Case 5 dir_l: rotate parent in DIR_L.
//
// Grandparent is guaranteed non-null by the case 4 check, so
// no root-replacement path is needed. Parent is known to be
// grandparent.child[DIR_L] from the dir_l branch, so the
// child pointer update is hardcoded without comparison.
if node == (*parent).child[tree::DIR_R] {
let new_root = (*parent).child[tree::DIR_R];
let new_child = (*new_root).child[tree::DIR_L];
(*parent).child[tree::DIR_R] = new_child;
if !new_child.is_null() {
(*new_child).parent = parent;
}
(*new_root).child[tree::DIR_L] = parent;
(*new_root).parent = grandparent;
(*parent).parent = new_root;
(*grandparent).child[tree::DIR_L] = new_root;
node = parent;
parent = new_root;
}
// Case 6 dir_l: rotate grandparent in DIR_R.
//
// The new root of this rotation is parent
// (= grandparent.child[DIR_L]), already in scope,
// eliminating the generic version's load of
// subtree.child[opposite(direction)].
//
// Great-grandparent may be null (grandparent could be root),
// so the null check and root-replacement path are retained.
// Grandparent's position under great-grandparent is unrelated
// to dir, so the pointer comparison is also retained.
{
let great_grandparent = (*grandparent).parent;
let new_child = (*parent).child[tree::DIR_R];
(*grandparent).child[tree::DIR_L] = new_child;
if !new_child.is_null() {
(*new_child).parent = grandparent;
}
(*parent).child[tree::DIR_R] = grandparent;
(*parent).parent = great_grandparent;
(*grandparent).parent = parent;
if !great_grandparent.is_null() {
if grandparent == (*great_grandparent).child[tree::DIR_R] {
(*great_grandparent).child[tree::DIR_R] = parent;
} else {
(*great_grandparent).child[tree::DIR_L] = parent;
}
} else {
(*tree_header).root = parent;
}
}
(*parent).color = Color::Black;
(*grandparent).color = Color::Red;
return SUCCESS;
}Cases 5 and 6 (right)
asm
insert_fixup_check_case_5_6_dir_r:
ldxdw r7, [r3 + TREE_NODE_CHILD_L_OFF] # r7 = uncle;
jeq r7, NULL, insert_fixup_case_5_6_dir_r
ldxb r8, [r7 + TREE_NODE_COLOR_OFF] # r8 = uncle.color;
jne r8, TREE_COLOR_B, insert_fixup_case_2
insert_fixup_case_5_6_dir_r:
ldxdw r6, [r2 + TREE_NODE_CHILD_L_OFF] # r6 = new_root = parent.child[L];
jne r9, r6, insert_fixup_case_6_dir_r
insert_fixup_case_5_dir_r:
ldxdw r8, [r6 + TREE_NODE_CHILD_R_OFF] # r8 = new_child = new_root.child[R];
stxdw [r2 + TREE_NODE_CHILD_L_OFF], r8 # parent.child[L] = new_child;
jeq r8, NULL, insert_fixup_case_5_dir_r_skip
stxdw [r8 + TREE_NODE_PARENT_OFF], r2 # new_child.parent = parent;
insert_fixup_case_5_dir_r_skip:
stxdw [r6 + TREE_NODE_CHILD_R_OFF], r2 # new_root.child[R] = parent;
stxdw [r6 + TREE_NODE_PARENT_OFF], r3 # new_root.parent = grandparent;
stxdw [r2 + TREE_NODE_PARENT_OFF], r6 # parent.parent = new_root;
stxdw [r3 + TREE_NODE_CHILD_R_OFF], r6 # grandparent.child[R] = new_root;
mov64 r9, r2 # node = old parent;
mov64 r2, r6 # parent = new_root;
insert_fixup_case_6_dir_r:
ldxdw r4, [r3 + TREE_NODE_PARENT_OFF] # r4 = great-grandparent;
ldxdw r8, [r2 + TREE_NODE_CHILD_L_OFF] # r8 = new_child = parent.child[L];
stxdw [r3 + TREE_NODE_CHILD_R_OFF], r8 # grandparent.child[R] = new_child;
jeq r8, NULL, insert_fixup_case_6_dir_r_skip
stxdw [r8 + TREE_NODE_PARENT_OFF], r3 # new_child.parent = grandparent;
insert_fixup_case_6_dir_r_skip:
stxdw [r2 + TREE_NODE_CHILD_L_OFF], r3 # parent.child[L] = grandparent;
stxdw [r2 + TREE_NODE_PARENT_OFF], r4 # parent.parent = great-grandparent;
stxdw [r3 + TREE_NODE_PARENT_OFF], r2 # grandparent.parent = parent;
jeq r4, NULL, insert_fixup_case_6_dir_r_root
ldxdw r8, [r4 + TREE_NODE_CHILD_R_OFF] # r8 = great-grandparent.child[R];
jne r3, r8, insert_fixup_case_6_dir_r_left
# Inline color + exit to eliminate `ja` (saves 1 CU).
stxdw [r4 + TREE_NODE_CHILD_R_OFF], r2 # great-grandparent.child[R] = parent;
stb [r2 + TREE_NODE_COLOR_OFF], TREE_COLOR_B # parent.color = black;
stb [r3 + TREE_NODE_COLOR_OFF], TREE_COLOR_R # grandparent.color = red;
exit
insert_fixup_case_6_dir_r_left:
# Inline color + exit to eliminate `ja` (saves 1 CU).
stxdw [r4 + TREE_NODE_CHILD_L_OFF], r2 # great-grandparent.child[L] = parent;
stb [r2 + TREE_NODE_COLOR_OFF], TREE_COLOR_B # parent.color = black;
stb [r3 + TREE_NODE_COLOR_OFF], TREE_COLOR_R # grandparent.color = red;
exit
insert_fixup_case_6_dir_r_root:
stxdw [r1 + IB_TREE_DATA_ROOT_OFF], r2 # tree.root = parent;
stb [r2 + TREE_NODE_COLOR_OFF], TREE_COLOR_B # parent.color = black;
stb [r3 + TREE_NODE_COLOR_OFF], TREE_COLOR_R # grandparent.color = red;
exitrs
} else {
// dir_r: parent is right child of grandparent.
uncle = (*grandparent).child[tree::DIR_L];
if uncle.is_null() || (*uncle).color == Color::Black {
// Case 5 dir_r: rotate parent in DIR_R.
//
// Grandparent is guaranteed non-null by the case 4 check, so
// no root-replacement path is needed. Parent is known to be
// grandparent.child[DIR_R] from the dir_r branch, so the
// child pointer update is hardcoded without comparison.
if node == (*parent).child[tree::DIR_L] {
let new_root = (*parent).child[tree::DIR_L];
let new_child = (*new_root).child[tree::DIR_R];
(*parent).child[tree::DIR_L] = new_child;
if !new_child.is_null() {
(*new_child).parent = parent;
}
(*new_root).child[tree::DIR_R] = parent;
(*new_root).parent = grandparent;
(*parent).parent = new_root;
(*grandparent).child[tree::DIR_R] = new_root;
node = parent;
parent = new_root;
}
// Case 6 dir_r: rotate grandparent in DIR_L.
//
// The new root of this rotation is parent
// (= grandparent.child[DIR_R]), already in scope,
// eliminating the generic version's load of
// subtree.child[opposite(direction)].
//
// Great-grandparent may be null (grandparent could be root),
// so the null check and root-replacement path are retained.
// Grandparent's position under great-grandparent is unrelated
// to dir, so the pointer comparison is also retained.
{
let great_grandparent = (*grandparent).parent;
let new_child = (*parent).child[tree::DIR_L];
(*grandparent).child[tree::DIR_R] = new_child;
if !new_child.is_null() {
(*new_child).parent = grandparent;
}
(*parent).child[tree::DIR_L] = grandparent;
(*parent).parent = great_grandparent;
(*grandparent).parent = parent;
if !great_grandparent.is_null() {
if grandparent == (*great_grandparent).child[tree::DIR_R] {
(*great_grandparent).child[tree::DIR_R] = parent;
} else {
(*great_grandparent).child[tree::DIR_L] = parent;
}
} else {
(*tree_header).root = parent;
}
}
(*parent).color = Color::Black;
(*grandparent).color = Color::Red;
return SUCCESS;
}
}Cases 2 and 3
asm
insert_fixup_case_2:
# r2 := parent
# r3 := grandparent
# r7 := uncle
# Case 2/3. # r9 := node
# ---------------------------------------------------------------------
stb [r2 + TREE_NODE_COLOR_OFF], TREE_COLOR_B # parent.color = black;
stb [r7 + TREE_NODE_COLOR_OFF], TREE_COLOR_B # uncle.color = black;
stb [r3 + TREE_NODE_COLOR_OFF], TREE_COLOR_R # grandparent.color = red;
mov64 r9, r3 # node = grandparent;
ldxdw r2, [r9 + TREE_NODE_PARENT_OFF] # parent = node.parent;
jne r2, NULL, insert_fixup_main
exit # Case 3.rs
// Case 2.
(*parent).color = Color::Black;
(*uncle).color = Color::Black;
(*grandparent).color = Color::Red;
node = grandparent;
parent = (*node).parent;
if parent.is_null() {
break;
}
}
// Case 3.
SUCCESS
}Benchmarking
| Test case | ASM (CUs) | Rust (CUs) | Overhead | Overhead % |
|---|---|---|---|---|
| Empty tree | 24 | 27 | +3 | +12.5% |
| Case 1: left child | 30 | 38 | +8 | +26.7% |
| Case 1: right child | 31 | 37 | +6 | +19.4% |
| Case 4: left child | 33 | 42 | +9 | +27.3% |
| Case 4: right child | 34 | 42 | +8 | +23.5% |
| Case 2+3: left-left | 49 | 60 | +11 | +22.4% |
| Case 2+3: left-right | 50 | 60 | +10 | +20.0% |
| Case 2+3: right-left | 50 | 58 | +8 | +16.0% |
| Case 2+3: right-right | 51 | 58 | +7 | +13.7% |
| Case 2+1: left | 56 | 69 | +13 | +23.2% |
| Case 2+1: right | 59 | 66 | +7 | +11.9% |
| Case 6: left-left null uncle | 54 | 67 | +13 | +24.1% |
| Case 6: right-right null uncle | 56 | 65 | +9 | +16.1% |
| Case 5+6: left-right null uncle | 64 | 74 | +10 | +15.6% |
| Case 5+6: right-left null uncle | 64 | 72 | +8 | +12.5% |
| Case 6: GGP non-null, LL GP-left | 61 | 78 | +17 | +27.9% |
| Case 6: GGP non-null, LL GP-right | 62 | 78 | +16 | +25.8% |
| Case 6: GGP non-null, RR GP-right | 64 | 76 | +12 | +18.8% |
| Case 6: GGP non-null, RR GP-left | 63 | 76 | +13 | +20.6% |
| Case 2+6: non-null new_child dir_l | 83 | 100 | +17 | +20.5% |
| Case 2+6: non-null new_child dir_r | 87 | 96 | +9 | +10.3% |
| Case 2+5+6: non-null new_child dir_l | 94 | 106 | +12 | +12.8% |
| Case 2+5+6: non-null new_child dir_r | 96 | 104 | +8 | +8.3% |
| Test case | ASM (CUs) | Rust (CUs) | Overhead | Overhead % |
|---|---|---|---|---|
| 3-node balanced (10,5,15) | 88 | 106 | +18 | +20.5% |
| Left-skew rotation (10,5,1) | 111 | 136 | +25 | +22.5% |
| Right-skew rotation (10,15,20) | 114 | 134 | +20 | +17.5% |
| Zigzag double rotation (10,5,7) | 121 | 143 | +22 | +18.2% |
| 7-node full tree (10,5,15,3,7,12,20) | 246 | 297 | +51 | +20.7% |
✂️ Remove
🛡️ Input checks
Implementations
asm
remove:
# Error if invalid instruction data length.
# ---------------------------------------------------------------------
jne r9, SIZE_OF_REMOVE_INSTRUCTION, e_instruction_data_len
# Error if too few accounts.
# ---------------------------------------------------------------------
jlt r8, IB_N_ACCOUNTS_GENERAL, e_n_accounts
# Error if user has data.
# ---------------------------------------------------------------------
ldxdw r9, [r1 + IB_USER_DATA_LEN_OFF]
jne r9, DATA_LEN_ZERO, e_user_data_len
# Error if tree is duplicate.
# ---------------------------------------------------------------------
ldxb r9, [r1 + IB_TREE_NON_DUP_MARKER_OFF]
jne r9, IB_NON_DUP_MARKER, e_tree_duplicaters
#[inline(always)]
unsafe fn remove(
input: *mut u8,
instruction_data: *mut u8,
instruction_data_len: u64,
n_accounts: u64,
) -> u64 {
check_instruction_data_len!(instruction_data_len, RemoveInstruction);
// Error if too few accounts.
if_err!(
n_accounts < input_buffer::N_ACCOUNTS_GENERAL,
error::N_ACCOUNTS
);
// Error if user has data.
let _user = user_account!(input);
// Error if tree is duplicate.
let _tree = account_non_dup!(input, input_buffer::TREE_ACCOUNT_OFF, error::TREE_DUPLICATE);Benchmarking
| Test case | ASM (CUs) | Rust (CUs) | Overhead | Overhead % |
|---|---|---|---|---|
| Data too short | 8 | 10 | +2 | +25.0% |
| Data too long | 8 | 10 | +2 | +25.0% |
| Too few accounts | 9 | 11 | +2 | +22.2% |
| User has data | 11 | 13 | +2 | +18.2% |
| Tree is duplicate | 13 | 15 | +2 | +15.4% |
🔎 Search
Implementations
asm
remove_search:
ldxdw r3, [r1 + IB_TREE_DATA_ROOT_OFF] # r3 = node = root;
jeq r3, NULL, e_key_does_not_exist
ldxh r4, [r2 + INSN_REMOVE_KEY_OFF] # r4 = key;
remove_search_loop:
ldxh r5, [r3 + TREE_NODE_KEY_OFF] # r5 = node.key;
jeq r4, r5, remove_found
jgt r4, r5, remove_search_r
remove_search_l:
ldxdw r3, [r3 + TREE_NODE_CHILD_L_OFF] # r3 = node.child[L];
jne r3, NULL, remove_search_loop
mov64 r0, E_KEY_DOES_NOT_EXIST
exit
remove_search_r:
ldxdw r3, [r3 + TREE_NODE_CHILD_R_OFF] # r3 = node.child[R];
jne r3, NULL, remove_search_loop
e_key_does_not_exist:
mov64 r0, E_KEY_DOES_NOT_EXIST
exitrs
let tree_header: *mut TreeHeader = input.add(input_buffer::TREE_DATA_OFF as usize).cast();
let mut node = (*tree_header).root;
if node.is_null() {
return error::KEY_DOES_NOT_EXIST.into();
}
let key = ldxh(instruction_data, instruction::REMOVE_KEY_OFF);
loop {
let node_key = (*node).key;
if key > node_key {
node = (*node).child[tree::DIR_R];
if node.is_null() {
return error::KEY_DOES_NOT_EXIST.into();
}
} else if key < node_key {
node = (*node).child[tree::DIR_L];
if node.is_null() {
return error::KEY_DOES_NOT_EXIST.into();
}
} else {
break;
}
}Benchmarking
| Test case | ASM (CUs) | Rust (CUs) | Overhead | Overhead % |
|---|---|---|---|---|
| Empty tree | 15 | 16 | +1 | +6.7% |
| Not found (left) | 21 | 28 | +7 | +33.3% |
| Not found (right) | 21 | 25 | +4 | +19.0% |
| Not found (deep) | 26 | 33 | +7 | +26.9% |
📑 Simple cases
Simple case 1: successor swap
asm
remove_found:
ldxdw r4, [r3 + TREE_NODE_CHILD_L_OFF] # r4 = node.child[L];
jeq r4, NULL, remove_check_child_r
ldxdw r5, [r3 + TREE_NODE_CHILD_R_OFF] # r5 = node.child[R];
jeq r5, NULL, remove_simple_2_child_l
# Simple case 1: successor swap.
# ---------------------------------------------------------------------
remove_successor_loop:
ldxdw r4, [r5 + TREE_NODE_CHILD_L_OFF] # r4 = successor.child[L];
jeq r4, NULL, remove_successor_copy
mov64 r5, r4 # successor = left;
ja remove_successor_loop
remove_successor_copy:
# Copy key/value pair as u32 from successor to found node.
# ---------------------------------------------------------------------
ldxw r4, [r5 + TREE_NODE_KEY_OFF] # r4 = successor.{key,value};
stxw [r3 + TREE_NODE_KEY_OFF], r4 # node.{key,value} = r4;
mov64 r3, r5 # node = successor;rs
if !(*node).child[tree::DIR_L].is_null() {
if !(*node).child[tree::DIR_R].is_null() {
// Simple case 1: successor swap.
let mut successor = (*node).child[tree::DIR_R];
loop {
let left = (*successor).child[tree::DIR_L];
if left.is_null() {
break;
}
successor = left;
}
// Copy successor's key/value to the found node as a u32
// pair. The successor's fields are left as-is (insert
// overwrites both when reusing the node from the stack).
let node_kv = addr_of_mut!((*node).key).cast::<u32>();
let successor_kv = addr_of!((*successor).key).cast::<u32>();
*node_kv = *successor_kv;
node = successor;Simple case 2: one-child replacement
asm
remove_check_child_r: # r3 = node
ldxdw r4, [r3 + TREE_NODE_CHILD_R_OFF] # r4 = node.child[R];
jeq r4, NULL, remove_check_simple_3_4
remove_simple_2_child_l: # r4 = child
# Simple case 2: replace node with child, recolor child black.
# ---------------------------------------------------------------------
ldxdw r5, [r3 + TREE_NODE_PARENT_OFF] # r5 = parent;
stxdw [r4 + TREE_NODE_PARENT_OFF], r5 # child.parent = parent;
stb [r4 + TREE_NODE_COLOR_OFF], TREE_COLOR_B # child.color = black;
jeq r5, NULL, remove_simple_2_root
ldxdw r6, [r5 + TREE_NODE_CHILD_R_OFF] # r6 = parent.child[R];
jne r3, r6, remove_simple_2_dir_l
stxdw [r5 + TREE_NODE_CHILD_R_OFF], r4 # parent.child[R] = child;
stdw [r3 + TREE_NODE_CHILD_L_OFF], NULL # node.child[L] = null;
stdw [r3 + TREE_NODE_CHILD_R_OFF], NULL # node.child[R] = null;
ldxdw r4, [r1 + IB_TREE_DATA_TOP_OFF] # r4 = header.top;
stxdw [r3 + TREE_NODE_PARENT_OFF], r4 # node.next = top;
stxdw [r1 + IB_TREE_DATA_TOP_OFF], r3 # header.top = node;
exit
remove_simple_2_dir_l:
stxdw [r5 + TREE_NODE_CHILD_L_OFF], r4 # parent.child[L] = child;
stdw [r3 + TREE_NODE_CHILD_L_OFF], NULL # node.child[L] = null;
stdw [r3 + TREE_NODE_CHILD_R_OFF], NULL # node.child[R] = null;
ldxdw r4, [r1 + IB_TREE_DATA_TOP_OFF] # r4 = header.top;
stxdw [r3 + TREE_NODE_PARENT_OFF], r4 # node.next = top;
stxdw [r1 + IB_TREE_DATA_TOP_OFF], r3 # header.top = node;
exit
remove_simple_2_root:
stxdw [r1 + IB_TREE_DATA_ROOT_OFF], r4 # tree.root = child;
stdw [r3 + TREE_NODE_CHILD_L_OFF], NULL # node.child[L] = null;
stdw [r3 + TREE_NODE_CHILD_R_OFF], NULL # node.child[R] = null;
ldxdw r4, [r1 + IB_TREE_DATA_TOP_OFF] # r4 = header.top;
stxdw [r3 + TREE_NODE_PARENT_OFF], r4 # node.next = top;
stxdw [r1 + IB_TREE_DATA_TOP_OFF], r3 # header.top = node;
exitrs
} else {
// Simple case 2: one child (L).
let child = (*node).child[tree::DIR_L];
remove_simple_2_child_replace!(node, child, tree_header);
}
};
if !(*node).child[tree::DIR_R].is_null() {
// Simple case 2: one child (R).
let child = (*node).child[tree::DIR_R];
remove_simple_2_child_replace!(node, child, tree_header);Simple case 3: root leaf
asm
remove_check_simple_3_4: # r3 = node
# Simple case 3: root leaf — clear root.
# ---------------------------------------------------------------------
ldxdw r5, [r3 + TREE_NODE_PARENT_OFF] # r5 = parent;
jne r5, NULL, remove_check_simple_4
stdw [r1 + IB_TREE_DATA_ROOT_OFF], NULL # tree.root = null;
stdw [r3 + TREE_NODE_CHILD_L_OFF], NULL # node.child[L] = null;
stdw [r3 + TREE_NODE_CHILD_R_OFF], NULL # node.child[R] = null;
ldxdw r4, [r1 + IB_TREE_DATA_TOP_OFF] # r4 = header.top;
stxdw [r3 + TREE_NODE_PARENT_OFF], r4 # node.next = top;
stxdw [r1 + IB_TREE_DATA_TOP_OFF], r3 # header.top = node;
exitrs
} else if unlikely((*node).parent.is_null()) {
// Simple case 3: root leaf.
(*tree_header).root = null_mut();
remove_recycle_node!(node, tree_header);Simple case 4: red leaf
asm
remove_check_simple_4: # r5 = parent
# Simple case 4: red leaf — detach from parent.
# ---------------------------------------------------------------------
ldxb r4, [r3 + TREE_NODE_COLOR_OFF] # r4 = node.color;
jne r4, TREE_COLOR_R, remove_complex
ldxdw r4, [r5 + TREE_NODE_CHILD_R_OFF] # r4 = parent.child[R];
jne r3, r4, remove_simple_4_dir_l
stdw [r5 + TREE_NODE_CHILD_R_OFF], NULL # parent.child[R] = null;
stdw [r3 + TREE_NODE_CHILD_L_OFF], NULL # node.child[L] = null;
stdw [r3 + TREE_NODE_CHILD_R_OFF], NULL # node.child[R] = null;
ldxdw r4, [r1 + IB_TREE_DATA_TOP_OFF] # r4 = header.top;
stxdw [r3 + TREE_NODE_PARENT_OFF], r4 # node.next = top;
stxdw [r1 + IB_TREE_DATA_TOP_OFF], r3 # header.top = node;
exit
remove_simple_4_dir_l:
stdw [r5 + TREE_NODE_CHILD_L_OFF], NULL # parent.child[L] = null;
stdw [r3 + TREE_NODE_CHILD_L_OFF], NULL # node.child[L] = null;
stdw [r3 + TREE_NODE_CHILD_R_OFF], NULL # node.child[R] = null;
ldxdw r4, [r1 + IB_TREE_DATA_TOP_OFF] # r4 = header.top;
stxdw [r3 + TREE_NODE_PARENT_OFF], r4 # node.next = top;
stxdw [r1 + IB_TREE_DATA_TOP_OFF], r3 # header.top = node;
exitrs
} else if (*node).color == Color::Red {
// Simple case 4: red leaf.
let parent = (*node).parent;
if node == (*parent).child[tree::DIR_R] {
(*parent).child[tree::DIR_R] = null_mut();
} else {
(*parent).child[tree::DIR_L] = null_mut();
}
remove_recycle_node!(node, tree_header);Benchmarking
| Test case | ASM (CUs) | Rust (CUs) | Overhead | Overhead % |
|---|---|---|---|---|
| Successor immediate R (SC 1) | 40 | 59 | +19 | +47.5% |
| Successor deep L descent (SC 1) | 44 | 63 | +19 | +43.2% |
| Successor with R child (SC 1) | 40 | 59 | +19 | +47.5% |
| One child root R (SC 2) | 31 | 48 | +17 | +54.8% |
| One child root L (SC 2) | 31 | 44 | +13 | +41.9% |
| One child non-root R,R (SC 2) | 38 | 57 | +19 | +50.0% |
| One child non-root L,L (SC 2) | 38 | 56 | +18 | +47.4% |
| One child non-root R,L (SC 2) | 38 | 60 | +22 | +57.9% |
| One child non-root L,R (SC 2) | 38 | 53 | +15 | +39.5% |
| Root leaf (SC 3) | 29 | 46 | +17 | +58.6% |
| Red leaf L (SC 4) | 38 | 60 | +22 | +57.9% |
| Red leaf R (SC 4) | 38 | 57 | +19 | +50.0% |