simulator-api 0.23.0

//! Caller-supplied direct-fill templates: one instruction the server rewrites per probe.
//!
//! Split from the crate root because the wire shape, its parse errors and its round-trip
//! tests are one subject, and `lib.rs` is long enough that adding a ninth unrelated block
//! to it makes every one of them harder to find.

use std::{collections::HashSet, str::FromStr};

use serde::{Deserialize, Serialize};
use solana_address::Address;

use crate::{Base64DecodeError, BinaryEncoding, EncodedBinary, SwapMode};

/// The width of the little-endian amount a template's `amountOffset` addresses.
pub const AMOUNT_WIDTH: usize = 8;

/// One instruction the caller serialized, with the slots the server rewrites per probe and the
/// offset it writes the amount at. Matches hops of `mode` originals trading `inputMint` →
/// `outputMint` and is priced against each such hop's realized fill.
///
/// Every floor, fee, direction and venue payload is the caller's: the server writes the amount,
/// rewrites the taker, funds, simulates and reads the fill.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(try_from = "DirectFillTemplateWire", into = "DirectFillTemplateWire")]
pub struct DirectFillTemplate {
    /// The template's identity on rows, events, notifications, report groups and stats keys.
    /// Unique per session, non-blank after trim.
    pub label: String,
    /// Which originals match, and whether the hop's input (`ExactIn`) or output (`ExactOut`)
    /// amount is written.
    pub mode: SwapMode,
    pub input_mint: Address,
    pub output_mint: Address,
    pub instruction: TemplateInstruction,
    /// Offset of the [`AMOUNT_WIDTH`]-byte little-endian amount in `instruction.data`.
    pub amount_offset: u32,
    pub taker: TakerSlots,
}

/// The template's single instruction, sent verbatim except the role slots' addresses and the
/// amount bytes.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TemplateInstruction {
    pub program_id: Address,
    /// Metas in program order. Flags are never touched by the server; duplicate addresses are
    /// legal; a role slot's address is a placeholder (by convention the zero address).
    pub accounts: Vec<TemplateAccount>,
    pub data: Vec<u8>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct TemplateAccount {
    pub address: Address,
    pub writable: bool,
    pub signer: bool,
}

/// The account slots a probe rewrites, as indices into `instruction.accounts`.
///
/// Each role is a list because one instruction can name the same account more than once: a
/// router's own swap instruction carries the taker's token accounts for itself and again inside
/// the venue's account run, and rewriting only the first leaves the venue pointed at the
/// exemplar's accounts. A bare index decodes as a one-element list.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct TakerSlots {
    /// Every slot rewritten to the original's fee payer. Non-empty; each must be a signer.
    pub user: Vec<u8>,
    /// Rewritten to the fee payer's associated account for `inputMint` and seeded with the
    /// trade. Non-empty; each must be writable.
    #[serde(with = "slot_list")]
    pub input_token_account: Vec<u8>,
    /// Rewritten to the fee payer's associated account for `outputMint` and seeded to zero, so
    /// what lands there is the fill. Non-empty; each must be writable.
    #[serde(with = "slot_list")]
    pub output_token_account: Vec<u8>,
}

/// One slot or a list of them on the wire, always a list in memory.
mod slot_list {
    use serde::{Deserialize, Deserializer, Serialize, Serializer};

    #[derive(Deserialize)]
    #[serde(untagged)]
    enum OneOrMany {
        One(u8),
        Many(Vec<u8>),
    }

    pub(super) fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<Vec<u8>, D::Error> {
        Ok(match OneOrMany::deserialize(d)? {
            OneOrMany::One(slot) => vec![slot],
            OneOrMany::Many(slots) => slots,
        })
    }

    pub(super) fn serialize<S: Serializer>(slots: &[u8], s: S) -> Result<S::Ok, S::Error> {
        match slots {
            [only] => only.serialize(s),
            many => many.serialize(s),
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
struct DirectFillTemplateWire {
    label: String,
    mode: SwapMode,
    input_mint: String,
    output_mint: String,
    instruction: TemplateInstructionWire,
    amount_offset: u32,
    taker: TakerSlots,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
struct TemplateInstructionWire {
    program_id: String,
    accounts: Vec<TemplateAccountWire>,
    data: EncodedBinary,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
struct TemplateAccountWire {
    address: String,
    #[serde(default)]
    writable: bool,
    #[serde(default)]
    signer: bool,
}

/// Why a supplied template could not be used.
#[derive(Debug, thiserror::Error, PartialEq, Eq)]
pub enum ParseDirectFillTemplateError {
    #[error("invalid address {value} in template field {field}")]
    Address { field: String, value: String },
    #[error("a template's label must not be blank")]
    BlankLabel,
    #[error("a template must trade two different mints")]
    SameMint,
    /// The per-instruction runtime ceiling; an empty list names no slot to rewrite.
    #[error("a template instruction must hold 1..=255 accounts, got {count}")]
    AccountCount { count: usize },
    #[error("template instruction data is not valid base64: {0}")]
    Data(#[from] Base64DecodeError),
    #[error("amount offset {offset} leaves no room for an amount in {len} bytes of data")]
    AmountOffset { offset: u32, len: usize },
    #[error("a template must name at least one user slot")]
    NoUserSlot,
    #[error("a template must name at least one slot for each token account")]
    NoTokenAccountSlot,
    #[error("{role} slot {slot} is past the instruction's accounts")]
    SlotOutOfRange { role: &'static str, slot: u8 },
    #[error("slot {slot} is named by more than one taker role")]
    DuplicateRoleSlot { slot: u8 },
    #[error("user slot {slot} must be a signer")]
    UserNotSigner { slot: u8 },
    #[error("token account slot {slot} must be writable")]
    TokenAccountReadOnly { slot: u8 },
}

impl TryFrom<DirectFillTemplateWire> for DirectFillTemplate {
    type Error = ParseDirectFillTemplateError;

    fn try_from(wire: DirectFillTemplateWire) -> Result<Self, Self::Error> {
        let address = |field: String, value: &str| {
            Address::from_str(value).map_err(|_| ParseDirectFillTemplateError::Address {
                field,
                value: value.to_owned(),
            })
        };
        let accounts = wire
            .instruction
            .accounts
            .iter()
            .enumerate()
            .map(|(index, account)| {
                Ok(TemplateAccount {
                    address: address(
                        format!("instruction.accounts[{index}].address"),
                        &account.address,
                    )?,
                    writable: account.writable,
                    signer: account.signer,
                })
            })
            .collect::<Result<Vec<_>, ParseDirectFillTemplateError>>()?;
        let template = Self {
            label: wire.label,
            mode: wire.mode,
            input_mint: address("inputMint".to_owned(), &wire.input_mint)?,
            output_mint: address("outputMint".to_owned(), &wire.output_mint)?,
            instruction: TemplateInstruction {
                program_id: address(
                    "instruction.programId".to_owned(),
                    &wire.instruction.program_id,
                )?,
                accounts,
                data: wire.instruction.data.decode()?,
            },
            amount_offset: wire.amount_offset,
            taker: wire.taker,
        };
        template.validate()?;
        Ok(template)
    }
}

impl DirectFillTemplate {
    /// The account slots a probe rewrites: every user slot, then the input and output token
    /// accounts.
    pub fn role_slots(&self) -> impl Iterator<Item = u8> {
        self.taker
            .user
            .iter()
            .chain(&self.taker.input_token_account)
            .chain(&self.taker.output_token_account)
            .copied()
    }

    /// The rules `TryFrom<DirectFillTemplateWire>` enforces, for a template built in code:
    /// [`simulator_trading::template::instantiate`] indexes on them without checking.
    pub fn validate(&self) -> Result<(), ParseDirectFillTemplateError> {
        use ParseDirectFillTemplateError as Error;

        if self.label.trim().is_empty() {
            return Err(Error::BlankLabel);
        }
        if self.input_mint == self.output_mint {
            return Err(Error::SameMint);
        }
        let accounts = &self.instruction.accounts;
        if !(1..=usize::from(u8::MAX)).contains(&accounts.len()) {
            return Err(Error::AccountCount {
                count: accounts.len(),
            });
        }
        let len = self.instruction.data.len();
        let end = usize::try_from(self.amount_offset)
            .ok()
            .and_then(|offset| offset.checked_add(AMOUNT_WIDTH));
        if end.is_none_or(|end| end > len) {
            return Err(Error::AmountOffset {
                offset: self.amount_offset,
                len,
            });
        }
        if self.taker.user.is_empty() {
            return Err(Error::NoUserSlot);
        }
        if self.taker.input_token_account.is_empty() || self.taker.output_token_account.is_empty() {
            return Err(Error::NoTokenAccountSlot);
        }
        let mut role_slots = HashSet::new();
        fn labelled(role: &'static str, slots: &[u8]) -> impl Iterator<Item = (&'static str, u8)> {
            slots.iter().map(move |&slot| (role, slot))
        }
        self.taker
            .user
            .iter()
            .map(|&slot| ("user", slot))
            .chain(labelled(
                "inputTokenAccount",
                &self.taker.input_token_account,
            ))
            .chain(labelled(
                "outputTokenAccount",
                &self.taker.output_token_account,
            ))
            .try_for_each(|(role, slot)| {
                if usize::from(slot) >= accounts.len() {
                    return Err(Error::SlotOutOfRange { role, slot });
                }
                role_slots
                    .insert(usize::from(slot))
                    .then_some(())
                    .ok_or(Error::DuplicateRoleSlot { slot })
            })?;
        let flags_of = |slot: u8| accounts[usize::from(slot)];
        if let Some(&slot) = self.taker.user.iter().find(|&&slot| !flags_of(slot).signer) {
            return Err(Error::UserNotSigner { slot });
        }
        if let Some(&slot) = self
            .taker
            .input_token_account
            .iter()
            .chain(&self.taker.output_token_account)
            .find(|&&slot| !flags_of(slot).writable)
        {
            return Err(Error::TokenAccountReadOnly { slot });
        }
        Ok(())
    }
}

impl From<DirectFillTemplate> for DirectFillTemplateWire {
    fn from(template: DirectFillTemplate) -> Self {
        Self {
            label: template.label,
            mode: template.mode,
            input_mint: template.input_mint.to_string(),
            output_mint: template.output_mint.to_string(),
            instruction: TemplateInstructionWire {
                program_id: template.instruction.program_id.to_string(),
                accounts: template
                    .instruction
                    .accounts
                    .into_iter()
                    .map(|account| TemplateAccountWire {
                        address: account.address.to_string(),
                        writable: account.writable,
                        signer: account.signer,
                    })
                    .collect(),
                data: EncodedBinary::from_bytes(&template.instruction.data, BinaryEncoding::Base64),
            },
            amount_offset: template.amount_offset,
            taker: template.taker,
        }
    }
}

#[cfg(test)]
mod tests {
    use rstest::rstest;

    use super::*;

    fn address(byte: u8) -> Address {
        Address::from([byte; 32])
    }

    /// A captured Tempest SOL → USDC template: 16 Titan `swap_route_v3` prefix slots, the
    /// 13-account Tempest run, 37 data bytes.
    fn tempest_template_json() -> serde_json::Value {
        serde_json::json!({
            "label": "tempest/sol-usdc/sol-to-usdc",
            "mode": "ExactIn",
            "inputMint": "So11111111111111111111111111111111111111112",
            "outputMint": "EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v",
            "instruction": {
                "programId": "T1TANpTeScyeqVzzgNViGDNrkQ6qHz9KrSBS4aNXvGT",
                "accounts": [
                    {"address": "11111111111111111111111111111111", "writable": true, "signer": true},
                    {"address": "11111111111111111111111111111111", "writable": true, "signer": true},
                    {"address": "D5YqVMoSxnqeZAKAUUE1Dm3bmjtdxQ5DCF356ozqN9cM"},
                    {"address": "11111111111111111111111111111111", "writable": true},
                    {"address": "11111111111111111111111111111111", "writable": true},
                    {"address": "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA"},
                    {"address": "TokenzQdBNbLqP5VEhdkAS6EPFLC1PHnBqCXEpPxuEb"},
                    {"address": "11111111111111111111111111111111"},
                    {"address": "ATokenGPvbdGVxr1b2hvZbsiqW5xWH25efTNsLJA8knL"},
                    {"address": "T1TANpTeScyeqVzzgNViGDNrkQ6qHz9KrSBS4aNXvGT"},
                    {"address": "T1TANpTeScyeqVzzgNViGDNrkQ6qHz9KrSBS4aNXvGT"},
                    {"address": "T1TANpTeScyeqVzzgNViGDNrkQ6qHz9KrSBS4aNXvGT"},
                    {"address": "EW9diL91VgHY5i9qYScz53W3PihQPSnPoMtAKCo1Bs7J", "writable": true},
                    {"address": "3nnVbsCfN1mwUk2XSLCnjzY3bDDTdzpnKjvRmd8nESS2", "writable": true},
                    {"address": "So11111111111111111111111111111111111111112"},
                    {"address": "EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v"},
                    {"address": "D5YqVMoSxnqeZAKAUUE1Dm3bmjtdxQ5DCF356ozqN9cM", "writable": true},
                    {"address": "FQmFVQ7i8GCLqwE7EriUA4WQCvyCKxAeY4XUr55ZD6D7", "writable": true},
                    {"address": "EW9diL91VgHY5i9qYScz53W3PihQPSnPoMtAKCo1Bs7J", "writable": true},
                    {"address": "3nnVbsCfN1mwUk2XSLCnjzY3bDDTdzpnKjvRmd8nESS2", "writable": true},
                    {"address": "4kHHmecPiRNJtAQug2cHFChovM4wqdBesWQckrXBWinD", "writable": true},
                    {"address": "6vNWbfi3HwDtoANb13CszPwEdAiyaUyZQ4A5A3hw4pqt", "writable": true},
                    {"address": "So11111111111111111111111111111111111111112"},
                    {"address": "EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v"},
                    {"address": "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA"},
                    {"address": "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA"},
                    {"address": "FUQcTWJjXtRxk6XHQBS6gqZ6VCgYeQkGs1H7inC2YKVe"},
                    {"address": "Sysvar1nstructions1111111111111111111111111"},
                    {"address": "2G71FYychcDbuNvMmVcFzU8J9N3giFvgp3a5PS2LmgZm"}
                ],
                "data": {"data": "KgAAAAAAAAAAAAAAAAAAAAAAMgACAAAAAQAAAFsAAAEAypo7DQ==", "encoding": "base64"}
            },
            "amountOffset": 2,
            "taker": {"user": [0, 1], "inputTokenAccount": 3, "outputTokenAccount": 4}
        })
    }

    #[test]
    fn a_direct_fill_template_round_trips_through_its_wire_shape() {
        let json = tempest_template_json();

        let template: DirectFillTemplate = serde_json::from_value(json.clone()).expect("decode");

        assert_eq!(template.label, "tempest/sol-usdc/sol-to-usdc");
        assert_eq!(template.mode, SwapMode::ExactIn);
        assert_eq!(
            template.input_mint.to_string(),
            "So11111111111111111111111111111111111111112"
        );
        assert_eq!(
            template.output_mint.to_string(),
            "EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v"
        );
        assert_eq!(
            template.instruction.program_id.to_string(),
            "T1TANpTeScyeqVzzgNViGDNrkQ6qHz9KrSBS4aNXvGT"
        );
        assert_eq!(template.instruction.accounts.len(), 29);
        assert_eq!(
            template.instruction.accounts[0],
            TemplateAccount {
                address: Address::default(),
                writable: true,
                signer: true,
            }
        );
        assert_eq!(
            template.instruction.accounts[2],
            TemplateAccount {
                address: "D5YqVMoSxnqeZAKAUUE1Dm3bmjtdxQ5DCF356ozqN9cM"
                    .parse()
                    .unwrap(),
                writable: false,
                signer: false,
            }
        );
        assert_eq!(template.instruction.data.len(), 37);
        assert_eq!(template.instruction.data[0], 42);
        assert_eq!(template.instruction.data[2..10], [0; 8]);
        assert_eq!(template.instruction.data[28], 91);
        assert_eq!(template.amount_offset, 2);
        assert_eq!(
            template.taker,
            TakerSlots {
                user: vec![0, 1],
                input_token_account: vec![3],
                output_token_account: vec![4],
            }
        );

        let encoded = serde_json::to_value(&template).expect("encode");
        assert_eq!(encoded["mode"], "ExactIn");
        assert_eq!(encoded["amountOffset"], 2);
        assert_eq!(encoded["instruction"]["data"], json["instruction"]["data"]);
        assert_eq!(encoded["instruction"]["accounts"][2]["signer"], false);
        assert_eq!(encoded["taker"]["inputTokenAccount"], 3);
        assert_eq!(
            serde_json::from_value::<DirectFillTemplate>(encoded).expect("decode again"),
            template
        );
    }

    fn tempest_wire() -> DirectFillTemplateWire {
        serde_json::from_value(tempest_template_json()).expect("wire shape")
    }

    /// Each case corrupts exactly one thing about the Tempest template, and the parse names it.
    #[rstest]
    #[case::address(
        |wire: &mut DirectFillTemplateWire| wire.instruction.program_id = "not-base58".to_owned(),
        ParseDirectFillTemplateError::Address {
            field: "instruction.programId".to_owned(),
            value: "not-base58".to_owned(),
        }
    )]
    #[case::blank_label(
        |wire: &mut DirectFillTemplateWire| wire.label = "  ".to_owned(),
        ParseDirectFillTemplateError::BlankLabel
    )]
    #[case::same_mint(
        |wire: &mut DirectFillTemplateWire| wire.output_mint = wire.input_mint.clone(),
        ParseDirectFillTemplateError::SameMint
    )]
    #[case::no_accounts(
        |wire: &mut DirectFillTemplateWire| wire.instruction.accounts.clear(),
        ParseDirectFillTemplateError::AccountCount { count: 0 }
    )]
    #[case::too_many_accounts(
        |wire: &mut DirectFillTemplateWire| {
            let filler = wire.instruction.accounts[2].clone();
            wire.instruction.accounts.resize(256, filler);
        },
        ParseDirectFillTemplateError::AccountCount { count: 256 }
    )]
    #[case::data(
        |wire: &mut DirectFillTemplateWire| wire.instruction.data.data = "!!!".to_owned(),
        ParseDirectFillTemplateError::Data(Base64DecodeError::InvalidByte(0, b'!'))
    )]
    #[case::amount_offset(
        |wire: &mut DirectFillTemplateWire| wire.amount_offset = 30,
        ParseDirectFillTemplateError::AmountOffset { offset: 30, len: 37 }
    )]
    #[case::no_user_slot(
        |wire: &mut DirectFillTemplateWire| wire.taker.user.clear(),
        ParseDirectFillTemplateError::NoUserSlot
    )]
    #[case::slot_out_of_range(
        |wire: &mut DirectFillTemplateWire| wire.taker.output_token_account = vec![29],
        ParseDirectFillTemplateError::SlotOutOfRange { role: "outputTokenAccount", slot: 29 }
    )]
    #[case::duplicate_role_slot(
        |wire: &mut DirectFillTemplateWire| wire.taker.input_token_account = vec![1],
        ParseDirectFillTemplateError::DuplicateRoleSlot { slot: 1 }
    )]
    #[case::user_not_signer(
        |wire: &mut DirectFillTemplateWire| wire.instruction.accounts[1].signer = false,
        ParseDirectFillTemplateError::UserNotSigner { slot: 1 }
    )]
    #[case::token_account_read_only(
        |wire: &mut DirectFillTemplateWire| wire.instruction.accounts[4].writable = false,
        ParseDirectFillTemplateError::TokenAccountReadOnly { slot: 4 }
    )]
    fn a_malformed_template_is_refused_at_parse(
        #[case] corrupt: fn(&mut DirectFillTemplateWire),
        #[case] expected: ParseDirectFillTemplateError,
    ) {
        let mut wire = tempest_wire();
        corrupt(&mut wire);

        assert_eq!(DirectFillTemplate::try_from(wire), Err(expected));
    }

    #[test]
    fn a_255_account_template_is_admitted() {
        let mut wire = tempest_wire();
        wire.instruction.accounts = (0..u8::MAX)
            .map(|byte| TemplateAccountWire {
                address: address(byte).to_string(),
                writable: true,
                signer: true,
            })
            .collect();
        wire.taker.user = vec![0];

        let template = DirectFillTemplate::try_from(wire).expect("at the bound");

        assert_eq!(template.role_slots().collect::<Vec<_>>(), vec![0, 3, 4]);
    }
}