realloc/steel - zero init, assertion

Author: thewuhxyz

basics/realloc/steel/api/src/instruction/extend.rs

@@ -5,7 +5,7 @@ use sysvar::rent::Rent;
 use super::SteelInstruction;
 
 instruction!(SteelInstruction, ExtendAddressInfo);
-/// create address info
+/// extend address info
 #[derive(Debug, Clone, Copy, PartialEq, Pod, Zeroable)]
 #[repr(C, packed)]
 pub struct ExtendAddressInfo {
@@ -25,50 +25,33 @@ impl ExtendAddressInfo {
         //
         let address_info_extended_data = Self::try_from_bytes(data)?.address_info;
 
-        let address_info = address_info_account.as_account::<AddressInfo>(&crate::ID)?;
-
-        // We collect our current address data and combime with the
-        //  data extension
-        //
-        let extended_address_info_data = ExtendedAddressInfo {
-            name: address_info.name,
-            house_number: address_info.house_number,
-            street: address_info.street,
-            city: address_info.city,
-            zip: address_info_extended_data.zip,
-            state: address_info_extended_data.state,
-        };
-
         // We get pay for the extra space we need before allocating the space
         //
-        let account_span = 8 + (extended_address_info_data.to_bytes()).len();
+        let account_span = 8 + std::mem::size_of::<ExtendedAddressInfo>(); // 8 byte Account Discriminator + sizeof ExtendedAddressInfo
         let lamports_required = (Rent::get()?).minimum_balance(account_span);
-        let diff = lamports_required - address_info_account.lamports();
-
-        // transfer the difference
+        let diff = lamports_required.saturating_sub(address_info_account.lamports());
         address_info_account.collect(diff, payer)?;
 
         // we reallocate new space to accomodate the new data
         //
-        address_info_account.realloc(account_span, false)?;
+        address_info_account.realloc(account_span, false)?; // no zero init
 
-        // we set the discriminator to the `ExtendedAccountInfo`, so Steel can deserialize.
-        // the account as such.
+        // we set the discriminator to the `ExtendedAccountInfo`, so Steel can deserialize the account as such.
         //
         {
             let mut data = address_info_account.data.borrow_mut();
-            // reset the account discriminator
-            data[0] = ExtendAddressInfo::discriminator();
+            data[0] = ExtendedAddressInfo::discriminator();
         }
 
-        // now we reset the account discriminator, we can deserialise as `ExtendedAddressInfo
+        // now we reset the account discriminator, we can deserialise as `ExtendedAddressInfo`
         //
         let extended_address_info =
             address_info_account.as_account_mut::<ExtendedAddressInfo>(&crate::ID)?;
 
         // set the extended address info
         //
-        *extended_address_info = extended_address_info_data;
+        extended_address_info.state = address_info_extended_data.state;
+        extended_address_info.zip = address_info_extended_data.zip;
 
         Ok(())
     }

basics/realloc/steel/api/src/instruction/mod.rs

@@ -4,6 +4,9 @@ pub use create::*;
 pub mod extend;
 pub use extend::*;
 
+pub mod zero_init;
+pub use zero_init::*;
+
 use steel::*;
 
 #[repr(u8)]
@@ -12,4 +15,5 @@ use steel::*;
 pub enum SteelInstruction {
     CreateAddressInfo = 0,
     ExtendAddressInfo = 1,
+    ZeroInit = 2,
 }

basics/realloc/steel/api/src/instruction/zero_init.rs

@@ -0,0 +1,56 @@
+use crate::state::*;
+use steel::*;
+use sysvar::rent::Rent;
+
+use super::SteelInstruction;
+
+instruction!(SteelInstruction, ZeroInit);
+/// work info
+#[derive(Debug, Clone, Copy, PartialEq, Pod, Zeroable)]
+#[repr(C)]
+pub struct ZeroInit {
+    pub work_info: WorkInfo,
+}
+
+impl ZeroInit {
+    pub fn process(accounts: &[AccountInfo<'_>], data: &[u8]) -> ProgramResult {
+        // Steel uses zero_copy types, so the sizes are fixed, however we can move from one account
+        // type to another, e.g ExtendedAddressInfo -> WorkInfo
+
+        let [payer, address_info_account, _system_program] = accounts else {
+            return Err(ProgramError::NotEnoughAccountKeys);
+        };
+
+        // we collect the extended data
+        //
+        let work_info_data = Self::try_from_bytes(data)?.work_info;
+
+        // We get pay for the extra space we need before allocating the space
+        //
+        let account_span = 8 + std::mem::size_of::<WorkInfo>(); // 8 byte Account Discriminator + sizeof WorkInfo
+        let lamports_required = (Rent::get()?).minimum_balance(account_span);
+        let diff = lamports_required.saturating_sub(address_info_account.lamports());
+        address_info_account.collect(diff, payer)?;
+
+        // we reallocate new space to accomodate the new data
+        //
+        address_info_account.realloc(account_span, true)?; // zero init
+
+        // we set the discriminator to the `WorkInfo`, so Steel can deserialize the account as such.
+        //
+        {
+            let mut data = address_info_account.data.borrow_mut();
+            data[0] = WorkInfo::discriminator();
+        }
+
+        // now we reset the account discriminator, we can deserialise as `WorkInfo``
+        //
+        let work_info = address_info_account.as_account_mut::<WorkInfo>(&crate::ID)?;
+
+        // set the extended address info
+        //
+        *work_info = work_info_data;
+
+        Ok(())
+    }
+}

basics/realloc/steel/api/src/state.rs

@@ -5,10 +5,11 @@ use steel::*;
 pub enum SteelAccount {
     AddressInfo,
     ExtendedAddressInfo,
+    WorkInfo,
 }
 
 account!(SteelAccount, AddressInfo);
-#[repr(C, packed)]
+#[repr(C)]
 #[derive(Clone, Copy, Debug, PartialEq, Pod, Zeroable)]
 pub struct AddressInfo {
     pub name: [u8; 48],
@@ -29,7 +30,17 @@ pub struct ExtendedAddressInfo {
     pub zip: u32,
 }
 
-#[repr(C, packed)]
+account!(SteelAccount, WorkInfo);
+#[repr(C)]
+#[derive(Clone, Copy, Debug, PartialEq, Pod, Zeroable)]
+pub struct WorkInfo {
+    pub name: [u8; 48],
+    pub position: [u8; 48],
+    pub company: [u8; 48],
+    pub years_employed: u8,
+}
+
+#[repr(C)]
 #[derive(Clone, Copy, Debug, PartialEq, Pod, Zeroable)]
 pub struct EnhancedAddressInfoExtender {
     pub state: [u8; 48],

basics/realloc/steel/program/src/lib.rs

@@ -13,5 +13,6 @@ pub fn process_instruction(
     match ix {
         SteelInstruction::CreateAddressInfo => CreateAddressInfo::process(accounts, data),
         SteelInstruction::ExtendAddressInfo => ExtendAddressInfo::process(accounts, data),
+        SteelInstruction::ZeroInit => ZeroInit::process(accounts, data),
     }
 }

basics/realloc/steel/tests/schema.ts

@@ -4,6 +4,7 @@ import * as borsh from 'borsh';
 export enum ReallocInstruction {
   Create = 0,
   Extend = 1,
+  ZeroInit = 2,
 }
 
 export class AddressInfo {
@@ -165,3 +166,70 @@ const AddressInfoExtenderSchema = new Map([
     },
   ],
 ]);
+
+export class WorkInfo {
+  name: Uint8Array;
+  position: Uint8Array;
+  company: Uint8Array;
+  years_employed: number;
+
+  constructor(info: {
+    name: Uint8Array;
+    position: Uint8Array;
+    company: Uint8Array;
+    years_employed: number;
+  }) {
+    this.name = info.name;
+    this.position = info.position;
+    this.company = info.company;
+    this.years_employed = info.years_employed;
+  }
+
+  toBuffer() {
+    return Buffer.from(borsh.serialize(WorkInfoSchema, this));
+  }
+
+  static fromAccountData(buffer: Buffer) {
+    const _accountData = Uint8Array.from(buffer).slice(8); // remove 8 byte discriminator
+
+    return borsh.deserialize(WorkInfoSchema, WorkInfo, Buffer.from(_accountData));
+  }
+
+  static fromData(info: {
+    name: string;
+    position: string;
+    company: string;
+    years_employed: number;
+  }) {
+    return new WorkInfo({
+      name: Uint8Array.from(Buffer.from(info.name.padEnd(48, '\0'))),
+      position: Uint8Array.from(Buffer.from(info.position.padEnd(48, '\0'))),
+      company: Uint8Array.from(Buffer.from(info.company.padEnd(48, '\0'))),
+      years_employed: info.years_employed,
+    });
+  }
+
+  toData() {
+    return {
+      name: Buffer.from(this.name).toString(),
+      position: Buffer.from(this.position).toString(),
+      company: Buffer.from(this.company).toString(),
+      years_employed: this.years_employed,
+    };
+  }
+}
+
+const WorkInfoSchema = new Map([
+  [
+    WorkInfo,
+    {
+      kind: 'struct',
+      fields: [
+        ['name', [48]], // Fixed array of 48 bytes
+        ['position', [48]], // Fixed array of 48 bytes
+        ['company', [48]], // Fixed array of 48 bytes
+        ['years_employed', 'u8'],
+      ],
+    },
+  ],
+]);

basics/realloc/steel/tests/test.ts

@@ -1,8 +1,9 @@
 import { Buffer } from 'node:buffer';
 import { describe, test } from 'node:test';
 import { Keypair, PublicKey, SystemProgram, Transaction, TransactionInstruction } from '@solana/web3.js';
+import { assert } from 'chai';
 import { start } from 'solana-bankrun';
-import { AddressInfo, AddressInfoExtender, ExtendedAddressInfo, ReallocInstruction } from './schema';
+import { AddressInfo, AddressInfoExtender, ExtendedAddressInfo, ReallocInstruction, WorkInfo } from './schema';
 
 describe('Realloc!', async () => {
   const PROGRAM_ID = new PublicKey('z7msBPQHDJjTvdQRoEcKyENgXDhSRYeHieN1ZMTqo35');
@@ -18,15 +19,14 @@ describe('Realloc!', async () => {
     console.log(`Payer Address      : ${payer.publicKey}`);
     console.log(`Address Info Acct  : ${addressInfoAccount.publicKey}`);
 
-    const ixData = Buffer.concat([
-      Buffer.from([ReallocInstruction.Create]),
-      AddressInfo.fromData({
-        name: 'Joe C',
-        house_number: 136,
-        street: 'Mile High Dr.',
-        city: 'Solana Beach',
-      }).toBuffer(),
-    ]);
+    const addressInfo = {
+      name: 'Joe C',
+      house_number: 136,
+      street: 'Mile High Dr.',
+      city: 'Solana Beach',
+    };
+
+    const ixData = Buffer.concat([Buffer.from([ReallocInstruction.Create]), AddressInfo.fromData(addressInfo).toBuffer()]);
 
     const ix = new TransactionInstruction({
       keys: [
@@ -48,9 +48,7 @@ describe('Realloc!', async () => {
     tx.recentBlockhash = blockhash;
     tx.add(ix).sign(payer, addressInfoAccount);
     await client.processTransaction(tx);
-  });
 
-  test("Read the new account's data", async () => {
     const accountInfo = await client.getAccount(addressInfoAccount.publicKey);
 
     const readAddressInfo = AddressInfo.fromAccountData(Buffer.from(accountInfo.data)).toData();
@@ -59,16 +57,20 @@ describe('Realloc!', async () => {
     console.log(`House Num: ${readAddressInfo.house_number}`);
     console.log(`Street   : ${readAddressInfo.street}`);
     console.log(`City     : ${readAddressInfo.city}`);
+
+    assert(readAddressInfo.name.slice(0, addressInfo.name.length) === addressInfo.name, 'name does not match');
+    assert(readAddressInfo.house_number === addressInfo.house_number, 'house number does not match');
+    assert(readAddressInfo.street.slice(0, addressInfo.street.length) === addressInfo.street, 'street does not match');
+    assert(readAddressInfo.city.slice(0, addressInfo.city.length) === addressInfo.city, 'city does not match');
   });
 
   test('Extend the address info account', async () => {
-    const ixData = Buffer.concat([
-      Buffer.from([ReallocInstruction.Extend]),
-      AddressInfoExtender.fromData({
-        state: 'Illinois',
-        zip: 12345,
-      }).toBuffer(),
-    ]);
+    const addressInfoExtender = {
+      state: 'Illinois',
+      zip: 12345,
+    };
+
+    const ixData = Buffer.concat([Buffer.from([ReallocInstruction.Extend]), AddressInfoExtender.fromData(addressInfoExtender).toBuffer()]);
 
     const ix = new TransactionInstruction({
       keys: [
@@ -90,9 +92,7 @@ describe('Realloc!', async () => {
     tx.recentBlockhash = blockhash;
     tx.add(ix).sign(payer, addressInfoAccount);
     await client.processTransaction(tx);
-  });
 
-  test("Read the new account's data", async () => {
     const accountInfo = await client.getAccount(addressInfoAccount.publicKey);
 
     const readAddressInfo = ExtendedAddressInfo.fromAccountData(Buffer.from(accountInfo.data)).toData();
@@ -103,5 +103,54 @@ describe('Realloc!', async () => {
     console.log(`City     : ${readAddressInfo.city}`);
     console.log(`State    : ${readAddressInfo.state}`);
     console.log(`Zip      : ${readAddressInfo.zip}`);
+
+    assert(readAddressInfo.state.slice(0, addressInfoExtender.state.length) === addressInfoExtender.state, 'state does not match');
+    assert(readAddressInfo.zip === addressInfoExtender.zip, 'zip does not match');
+  });
+
+  test('zero init work info account', async () => {
+    const workInfo = {
+      name: 'Pete',
+      company: 'Solana Labs',
+      position: 'Engineer',
+      years_employed: 2,
+    };
+
+    const ixData = Buffer.concat([Buffer.from([ReallocInstruction.ZeroInit]), WorkInfo.fromData(workInfo).toBuffer()]);
+
+    const ix = new TransactionInstruction({
+      keys: [
+        { pubkey: payer.publicKey, isSigner: true, isWritable: true },
+        {
+          pubkey: addressInfoAccount.publicKey,
+          isSigner: true,
+          isWritable: true,
+        },
+        { pubkey: SystemProgram.programId, isSigner: false, isWritable: false },
+      ],
+      programId: PROGRAM_ID,
+      data: ixData,
+    });
+
+    const blockhash = context.lastBlockhash;
+
+    const tx = new Transaction();
+    tx.recentBlockhash = blockhash;
+    tx.add(ix).sign(payer, addressInfoAccount);
+    await client.processTransaction(tx);
+
+    const accountInfo = await client.getAccount(addressInfoAccount.publicKey);
+
+    const readWorkInfo = WorkInfo.fromAccountData(Buffer.from(accountInfo.data)).toData();
+
+    console.log(`Name          : ${readWorkInfo.name}`);
+    console.log(`Position      : ${readWorkInfo.position}`);
+    console.log(`Company       : ${readWorkInfo.company}`);
+    console.log(`Years Employed: ${readWorkInfo.years_employed}`);
+
+    assert(readWorkInfo.name.slice(0, workInfo.name.length) === workInfo.name, 'name does not match');
+    assert(readWorkInfo.position.slice(0, workInfo.position.length) === workInfo.position, 'position does not match');
+    assert(readWorkInfo.company.slice(0, workInfo.company.length) === workInfo.company, 'company does not match');
+    assert(readWorkInfo.years_employed === workInfo.years_employed, 'years employed does not match');
   });
 });