cherry-svm-decode 0.2.0

use anyhow::{anyhow, Context, Result};
use arrow::array::{Array, BinaryArray};
use arrow::{array::RecordBatch, datatypes::*};
use std::sync::Arc;
mod deserialize;
pub use deserialize::{deserialize_data, DynType, DynValue, ParamInput};
mod arrow_converter;
use arrow_converter::{to_arrow, to_arrow_dtype};

#[derive(Debug, Clone)]
pub struct InstructionSignature {
    pub discriminator: Vec<u8>,
    pub params: Vec<ParamInput>,
    pub accounts_names: Vec<String>,
}

#[cfg(feature = "pyo3")]
impl<'py> pyo3::FromPyObject<'py> for InstructionSignature {
    fn extract_bound(ob: &pyo3::Bound<'py, pyo3::PyAny>) -> pyo3::PyResult<Self> {
        use pyo3::types::PyAnyMethods;
        use pyo3::types::PyTypeMethods;

        let discriminator_ob = ob.getattr("discriminator")?;

        let discriminator_ob_type: String = discriminator_ob.get_type().name()?.to_string();
        let discriminator = match discriminator_ob_type.as_str() {
            "str" => {
                let s: &str = discriminator_ob.extract()?;
                hex_to_bytes(s).context("failed to decode hex")?
            }
            "bytes" => discriminator_ob.extract()?,
            _ => return Err(anyhow!("unknown type: {}", discriminator_ob_type).into()),
        };

        let params = ob.getattr("params")?.extract::<Vec<ParamInput>>()?;
        let accounts_names = ob.getattr("accounts_names")?.extract::<Vec<String>>()?;

        Ok(InstructionSignature {
            discriminator,
            params,
            accounts_names,
        })
    }
}

fn hex_to_bytes(hex_string: &str) -> Result<Vec<u8>> {
    let hex_string = hex_string.strip_prefix("0x").unwrap_or(hex_string);
    let hex_string = if hex_string.len() % 2 == 1 {
        format!("0{}", hex_string)
    } else {
        hex_string.to_string()
    };
    let out = (0..hex_string.len())
        .step_by(2)
        .map(|i| {
            u8::from_str_radix(&hex_string[i..i + 2], 16)
                .context("failed to parse hexstring to bytes")
        })
        .collect::<Result<Vec<_>, _>>()?;

    Ok(out)
}

pub fn svm_decode_instructions(
    signature: InstructionSignature,
    batch: &RecordBatch,
    allow_decode_fail: bool,
) -> Result<RecordBatch> {
    let data_col = batch.column_by_name("data").unwrap();
    let data_array = data_col.as_any().downcast_ref::<BinaryArray>().unwrap();

    let account_arrays: Vec<&BinaryArray> = (0..10)
        .map(|i| {
            let col_name = format!("a{}", i);
            let col = batch.column_by_name(&col_name).unwrap();
            col.as_any().downcast_ref::<BinaryArray>().unwrap()
        })
        .collect();

    decode_instructions(signature, &account_arrays, data_array, allow_decode_fail)
}

pub fn decode_instructions(
    signature: InstructionSignature,
    accounts: &[&BinaryArray],
    data: &BinaryArray,
    allow_decode_fail: bool,
) -> Result<RecordBatch> {
    let num_params = signature.params.len();

    let mut decoded_params_vec: Vec<Vec<Option<DynValue>>> =
        (0..num_params).map(|_| Vec::new()).collect();

    for row_idx in 0..data.len() {
        if data.is_null(row_idx) {
            if allow_decode_fail {
                log::debug!("Instruction data is null");
                decoded_params_vec.iter_mut().for_each(|v| v.push(None));
                continue;
            } else {
                return Err(anyhow::anyhow!("Instruction data is null"));
            }
        }

        let instruction_data = data.value(row_idx).to_vec();
        let data_result = match_discriminators(&instruction_data, &signature.discriminator);
        let data = match data_result {
            Ok(data) => data,
            Err(e) if allow_decode_fail => {
                log::debug!("Error matching discriminators: {:?}", e);
                decoded_params_vec.iter_mut().for_each(|v| v.push(None));
                continue;
            }
            Err(e) => {
                return Err(anyhow::anyhow!("Error matching discriminators: {:?}", e));
            }
        };

        let decoded_ix_result = deserialize_data(&data, &signature.params);
        let decoded_ix = match decoded_ix_result {
            Ok(ix) => ix,
            Err(e) if allow_decode_fail => {
                log::debug!("Error deserializing instruction: {:?}", e);
                decoded_params_vec.iter_mut().for_each(|v| v.push(None));
                continue;
            }
            Err(e) => {
                return Err(anyhow::anyhow!("Error deserializing instruction: {:?}", e));
            }
        };

        for (i, value) in decoded_ix.into_iter().enumerate() {
            decoded_params_vec[i].push(Some(value));
        }
    }

    let data_arrays: Vec<Arc<dyn Array>> = decoded_params_vec
        .iter()
        .enumerate()
        .map(|(i, v)| to_arrow(&signature.params[i].param_type, v.clone()).unwrap())
        .collect::<Vec<_>>();

    let data_fields = signature
        .params
        .iter()
        .map(|p| Field::new(p.name.clone(), to_arrow_dtype(&p.param_type).unwrap(), true))
        .collect::<Vec<_>>();

    let acc_names_len = signature.accounts_names.len();

    let mut accounts: Vec<Arc<dyn Array>> = accounts
        .iter()
        .map(|arr| {
            let owned_array = arr.slice(0, arr.len());
            owned_array as Arc<dyn Array>
        })
        .collect();

    let mut acc_fields = Vec::new();
    if acc_names_len < 10 {
        let _ = accounts.split_off(acc_names_len);
        for i in 0..acc_names_len {
            let field = Field::new(signature.accounts_names[i].clone(), DataType::Binary, true);
            acc_fields.push(field);
        }
    } else {
        for i in 0..10 {
            let field = Field::new(signature.accounts_names[i].clone(), DataType::Binary, true);
            acc_fields.push(field);
        }
    }

    let decoded_instructions_array = data_arrays.into_iter().chain(accounts).collect::<Vec<_>>();
    let decoded_instructions_fields = data_fields
        .into_iter()
        .chain(acc_fields.clone())
        .collect::<Vec<_>>();

    let schema = Arc::new(Schema::new(decoded_instructions_fields));
    let batch = RecordBatch::try_new(schema, decoded_instructions_array)
        .context("Failed to create record batch from data arrays")
        .unwrap();

    Ok(batch)
}

pub fn match_discriminators(instr_data: &[u8], discriminator: &[u8]) -> Result<Vec<u8>> {
    let discriminator_len = discriminator.len();
    if instr_data.len() < discriminator_len {
        return Err(anyhow::anyhow!(
            "Instruction data is too short to contain discriminator. Expected at least {} bytes, got {} bytes",
            discriminator_len,
            instr_data.len()
        ));
    }
    let disc = &instr_data[..discriminator_len].to_vec();
    let ix_data = &instr_data[discriminator_len..];
    if !disc.eq(discriminator) {
        return Err(anyhow::anyhow!(
            "Instruction data discriminator doesn't match signature discriminator"
        ));
    }
    Ok(ix_data.to_vec())
}

pub fn instruction_signature_to_arrow_schema(signature: &InstructionSignature) -> Result<Schema> {
    let mut fields = Vec::new();

    for param in &signature.params {
        let field = Field::new(
            param.name.clone(),
            to_arrow_dtype(&param.param_type).unwrap(),
            true,
        );
        fields.push(field);
    }

    for account in &signature.accounts_names {
        let field = Field::new(account.clone(), DataType::Binary, true);
        fields.push(field);
    }

    Ok(Schema::new(fields))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::deserialize::{DynType, ParamInput};
    use std::fs::File;

    #[test]
    #[ignore]
    fn read_parquet_with_real_data() {
        use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;

        let builder = ParquetRecordBatchReaderBuilder::try_new(
            File::open("instruction_exemple.parquet").unwrap(),
        )
        .unwrap();
        let mut reader = builder.build().unwrap();
        let instructions = reader.next().unwrap().unwrap();
        let ix_signature = InstructionSignature {
            // // SPL Token Transfer
            // discriminator: &[3],
            // params: vec![ParamInput {
            //     name: "Amount".to_string(),
            //     param_type: DynType::U64,
            // }],
            // accounts: vec![
            //     "Source".to_string(),
            //     "Destination".to_string(),
            //     "Authority".to_string(),
            // ],

            // // JUP SwapEvent
            // discriminator: &[
            //     228, 69, 165, 46, 81, 203, 154, 29, 64, 198, 205, 232, 38, 8, 113, 226,
            // ],
            // params: vec![
            //     ParamInput {
            //         name: "Amm".to_string(),
            //         param_type: DynType::Pubkey,
            //     },
            //     ParamInput {
            //         name: "InputMint".to_string(),
            //         param_type: DynType::Pubkey,
            //     },
            //     ParamInput {
            //         name: "InputAmount".to_string(),
            //         param_type: DynType::U64,
            //     },
            //     ParamInput {
            //         name: "OutputMint".to_string(),
            //         param_type: DynType::Pubkey,
            //     },
            //     ParamInput {
            //         name: "OutputAmount".to_string(),
            //         param_type: DynType::U64,
            //     },
            // ],
            // accounts: vec![],

            // JUP Route
            discriminator: vec![229, 23, 203, 151, 122, 227, 173, 42],
            params: vec![
                ParamInput {
                    name: "RoutePlan".to_string(),
                    param_type: DynType::Array(Box::new(DynType::Struct(vec![
                        (
                            "Swap".to_string(),
                            DynType::Enum(vec![
                                ("Saber".to_string(), None),
                                ("SaberAddDecimalsDeposit".to_string(), None),
                                ("SaberAddDecimalsWithdraw".to_string(), None),
                                ("TokenSwap".to_string(), None),
                                ("Sencha".to_string(), None),
                                ("Step".to_string(), None),
                                ("Cropper".to_string(), None),
                                ("Raydium".to_string(), None),
                                (
                                    "Crema".to_string(),
                                    Some(DynType::Struct(vec![(
                                        "a_to_b".to_string(),
                                        DynType::Bool,
                                    )])),
                                ),
                                ("Lifinity".to_string(), None),
                                ("Mercurial".to_string(), None),
                                ("Cykura".to_string(), None),
                                (
                                    "Serum".to_string(),
                                    Some(DynType::Struct(vec![(
                                        "side".to_string(),
                                        DynType::Enum(vec![
                                            ("Bid".to_string(), None),
                                            ("Ask".to_string(), None),
                                        ]),
                                    )])),
                                ),
                                ("MarinadeDeposit".to_string(), None),
                                ("MarinadeUnstake".to_string(), None),
                                (
                                    "Aldrin".to_string(),
                                    Some(DynType::Struct(vec![(
                                        "side".to_string(),
                                        DynType::Enum(vec![
                                            ("Bid".to_string(), None),
                                            ("Ask".to_string(), None),
                                        ]),
                                    )])),
                                ),
                                (
                                    "AldrinV2".to_string(),
                                    Some(DynType::Struct(vec![(
                                        "side".to_string(),
                                        DynType::Enum(vec![
                                            ("Bid".to_string(), None),
                                            ("Ask".to_string(), None),
                                        ]),
                                    )])),
                                ),
                                (
                                    "Whirlpool".to_string(),
                                    Some(DynType::Struct(vec![(
                                        "a_to_b".to_string(),
                                        DynType::Bool,
                                    )])),
                                ),
                                (
                                    "Invariant".to_string(),
                                    Some(DynType::Struct(vec![(
                                        "x_to_y".to_string(),
                                        DynType::Bool,
                                    )])),
                                ),
                                ("Meteora".to_string(), None),
                                ("GooseFX".to_string(), None),
                                (
                                    "DeltaFi".to_string(),
                                    Some(DynType::Struct(vec![(
                                        "stable".to_string(),
                                        DynType::Bool,
                                    )])),
                                ),
                                ("Balansol".to_string(), None),
                                (
                                    "MarcoPolo".to_string(),
                                    Some(DynType::Struct(vec![(
                                        "x_to_y".to_string(),
                                        DynType::Bool,
                                    )])),
                                ),
                                (
                                    "Dradex".to_string(),
                                    Some(DynType::Struct(vec![(
                                        "side".to_string(),
                                        DynType::Enum(vec![
                                            ("Bid".to_string(), None),
                                            ("Ask".to_string(), None),
                                        ]),
                                    )])),
                                ),
                                ("LifinityV2".to_string(), None),
                                ("RaydiumClmm".to_string(), None),
                                (
                                    "Openbook".to_string(),
                                    Some(DynType::Struct(vec![(
                                        "side".to_string(),
                                        DynType::Enum(vec![
                                            ("Bid".to_string(), None),
                                            ("Ask".to_string(), None),
                                        ]),
                                    )])),
                                ),
                                (
                                    "Phoenix".to_string(),
                                    Some(DynType::Struct(vec![(
                                        "side".to_string(),
                                        DynType::Enum(vec![
                                            ("Bid".to_string(), None),
                                            ("Ask".to_string(), None),
                                        ]),
                                    )])),
                                ),
                                (
                                    "Symmetry".to_string(),
                                    Some(DynType::Struct(vec![
                                        ("from_token_id".to_string(), DynType::U64),
                                        ("to_token_id".to_string(), DynType::U64),
                                    ])),
                                ),
                                ("TokenSwapV2".to_string(), None),
                                ("HeliumTreasuryManagementRedeemV0".to_string(), None),
                                ("StakeDexStakeWrappedSol".to_string(), None),
                                (
                                    "StakeDexSwapViaStake".to_string(),
                                    Some(DynType::Struct(vec![(
                                        "bridge_stake_seed".to_string(),
                                        DynType::U32,
                                    )])),
                                ),
                                ("GooseFXV2".to_string(), None),
                                ("Perps".to_string(), None),
                                ("PerpsAddLiquidity".to_string(), None),
                                ("PerpsRemoveLiquidity".to_string(), None),
                                ("MeteoraDlmm".to_string(), None),
                                (
                                    "OpenBookV2".to_string(),
                                    Some(DynType::Struct(vec![(
                                        "side".to_string(),
                                        DynType::Enum(vec![
                                            ("Bid".to_string(), None),
                                            ("Ask".to_string(), None),
                                        ]),
                                    )])),
                                ),
                                ("RaydiumClmmV2".to_string(), None),
                                (
                                    "StakeDexPrefundWithdrawStakeAndDepositStake".to_string(),
                                    Some(DynType::Struct(vec![(
                                        "bridge_stake_seed".to_string(),
                                        DynType::U32,
                                    )])),
                                ),
                                (
                                    "Clone".to_string(),
                                    Some(DynType::Struct(vec![
                                        ("pool_index".to_string(), DynType::U8),
                                        ("quantity_is_input".to_string(), DynType::Bool),
                                        ("quantity_is_collateral".to_string(), DynType::Bool),
                                    ])),
                                ),
                                (
                                    "SanctumS".to_string(),
                                    Some(DynType::Struct(vec![
                                        ("src_lst_value_calc_accs".to_string(), DynType::U8),
                                        ("dst_lst_value_calc_accs".to_string(), DynType::U8),
                                        ("src_lst_index".to_string(), DynType::U32),
                                        ("dst_lst_index".to_string(), DynType::U32),
                                    ])),
                                ),
                                (
                                    "SanctumSAddLiquidity".to_string(),
                                    Some(DynType::Struct(vec![
                                        ("lst_value_calc_accs".to_string(), DynType::U8),
                                        ("lst_index".to_string(), DynType::U32),
                                    ])),
                                ),
                                (
                                    "SanctumSRemoveLiquidity".to_string(),
                                    Some(DynType::Struct(vec![
                                        ("lst_value_calc_accs".to_string(), DynType::U8),
                                        ("lst_index".to_string(), DynType::U32),
                                    ])),
                                ),
                                ("RaydiumCP".to_string(), None),
                                (
                                    "WhirlpoolSwapV2".to_string(),
                                    Some(DynType::Struct(vec![
                                        ("a_to_b".to_string(), DynType::Bool),
                                        (
                                            "remaining_accounts_info".to_string(),
                                            DynType::Struct(vec![(
                                                "slices".to_string(),
                                                DynType::Array(Box::new(DynType::Struct(vec![(
                                                    "remaining_accounts_slice".to_string(),
                                                    DynType::Struct(vec![
                                                        ("accounts_type".to_string(), DynType::U8),
                                                        ("length".to_string(), DynType::U8),
                                                    ]),
                                                )]))),
                                            )]),
                                        ),
                                    ])),
                                ),
                                ("OneIntro".to_string(), None),
                                ("PumpdotfunWrappedBuy".to_string(), None),
                                ("PumpdotfunWrappedSell".to_string(), None),
                                ("PerpsV2".to_string(), None),
                                ("PerpsV2AddLiquidity".to_string(), None),
                                ("PerpsV2RemoveLiquidity".to_string(), None),
                                ("MoonshotWrappedBuy".to_string(), None),
                                ("MoonshotWrappedSell".to_string(), None),
                                ("StabbleStableSwap".to_string(), None),
                                ("StabbleWeightedSwap".to_string(), None),
                                (
                                    "Obric".to_string(),
                                    Some(DynType::Struct(vec![(
                                        "x_to_y".to_string(),
                                        DynType::Bool,
                                    )])),
                                ),
                                ("FoxBuyFromEstimatedCost".to_string(), None),
                                (
                                    "FoxClaimPartial".to_string(),
                                    Some(DynType::Struct(vec![(
                                        "is_y".to_string(),
                                        DynType::Bool,
                                    )])),
                                ),
                                (
                                    "SolFi".to_string(),
                                    Some(DynType::Struct(vec![(
                                        "is_quote_to_base".to_string(),
                                        DynType::Bool,
                                    )])),
                                ),
                                ("SolayerDelegateNoInit".to_string(), None),
                                ("SolayerUndelegateNoInit".to_string(), None),
                                (
                                    "TokenMill".to_string(),
                                    Some(DynType::Struct(vec![(
                                        "side".to_string(),
                                        DynType::Enum(vec![
                                            ("Bid".to_string(), None),
                                            ("Ask".to_string(), None),
                                        ]),
                                    )])),
                                ),
                                ("DaosFunBuy".to_string(), None),
                                ("DaosFunSell".to_string(), None),
                                ("ZeroFi".to_string(), None),
                                ("StakeDexWithdrawWrappedSol".to_string(), None),
                                ("VirtualsBuy".to_string(), None),
                                ("VirtualsSell".to_string(), None),
                                (
                                    "Peren".to_string(),
                                    Some(DynType::Struct(vec![
                                        ("in_index".to_string(), DynType::U8),
                                        ("out_index".to_string(), DynType::U8),
                                    ])),
                                ),
                                ("PumpdotfunAmmBuy".to_string(), None),
                                ("PumpdotfunAmmSell".to_string(), None),
                                ("Gamma".to_string(), None),
                            ]),
                        ),
                        ("Percent".to_string(), DynType::U8),
                        ("InputIndex".to_string(), DynType::U8),
                        ("OutputIndex".to_string(), DynType::U8),
                    ]))),
                },
                ParamInput {
                    name: "InAmount".to_string(),
                    param_type: DynType::U64,
                },
                ParamInput {
                    name: "QuotedOutAmount".to_string(),
                    param_type: DynType::U64,
                },
                ParamInput {
                    name: "SlippageBps".to_string(),
                    param_type: DynType::U16,
                },
                ParamInput {
                    name: "PlatformFeeBps".to_string(),
                    param_type: DynType::U8,
                },
            ],
            accounts_names: vec![
                "TokenProgram".to_string(),
                "UserTransferAuthority".to_string(),
                "UserSourceTokenAccount".to_string(),
                "UserDestinationTokenAccount".to_string(),
                "DestinationTokenAccount".to_string(),
                "PlatformFeeAccount".to_string(),
                "EventAuthority".to_string(),
                "Program".to_string(),
            ],
        };

        let result = svm_decode_instructions(ix_signature, &instructions, true)
            .context("decode failed")
            .unwrap();

        // Save the filtered instructions to a new parquet file
        let mut file = File::create("decoded_instructions.parquet").unwrap();
        let mut writer =
            parquet::arrow::ArrowWriter::try_new(&mut file, result.schema(), None).unwrap();
        writer.write(&result).unwrap();
        writer.close().unwrap();
    }

    #[test]
    // #[ignore]
    fn test_instruction_signature_to_arrow_schema() {
        // Create a test instruction signature
        let signature = InstructionSignature {
            discriminator: vec![],
            params: vec![
                ParamInput {
                    name: "amount".to_string(),
                    param_type: DynType::U64,
                },
                ParamInput {
                    name: "is_valid".to_string(),
                    param_type: DynType::Bool,
                },
                ParamInput {
                    name: "amm".to_string(),
                    param_type: DynType::FixedArray(Box::new(DynType::U8), 32),
                },
            ],
            accounts_names: vec!["source".to_string(), "destination".to_string()],
        };

        // Convert to schema
        let schema = instruction_signature_to_arrow_schema(&signature).unwrap();

        // Verify the schema has the correct number of fields
        assert_eq!(schema.fields().len(), 5); // 2 params + 2 accounts

        // Verify param fields
        let amount_field = schema.field_with_name("amount").unwrap();
        assert_eq!(amount_field.name(), "amount");
        assert!(amount_field.is_nullable());

        let is_valid_field = schema.field_with_name("is_valid").unwrap();
        assert_eq!(is_valid_field.name(), "is_valid");
        assert!(is_valid_field.is_nullable());

        let amm_field = schema.field_with_name("amm").unwrap();
        assert_eq!(amm_field.name(), "amm");
        assert!(amm_field.is_nullable());

        // Verify account fields
        let source_field = schema.field_with_name("source").unwrap();
        assert_eq!(source_field.name(), "source");
        assert_eq!(source_field.data_type(), &DataType::Binary);
        assert!(source_field.is_nullable());

        let dest_field = schema.field_with_name("destination").unwrap();
        assert_eq!(dest_field.name(), "destination");
        assert_eq!(dest_field.data_type(), &DataType::Binary);
        assert!(dest_field.is_nullable());
    }
}