Skip to main content

huff_neo_utils/
types.rs

1use crate::bytes_util::*;
2use alloy_dyn_abi::{DynSolValue, Word};
3use alloy_primitives::{Address, I256, U256};
4use lazy_static::lazy_static;
5use regex::Regex;
6use std::{fmt, str::FromStr};
7
8/// Primitive EVM types
9#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy)]
10pub enum PrimitiveEVMType {
11    /// String type
12    String,
13    /// Bytes array type
14    DynBytes,
15    /// Boolean type
16    Bool,
17    /// Uint ; should start from uint8 to uint256 by step of 8
18    Uint(usize),
19    /// Int ; should start from int8 to int256 by step of 8
20    Int(usize),
21    /// Address type
22    Address,
23    /// Bytes type ; bytes1, bytes2, ..., bytes32
24    Bytes(usize),
25}
26
27/// Automatically converts an input string to a PrimitiveEVMType.
28/// Example : PrimitiveEVMType::from("uint256") => PrimitiveEVMType::Uint(256)
29impl TryFrom<&String> for PrimitiveEVMType {
30    type Error = String;
31
32    fn try_from(input: &String) -> Result<Self, Self::Error> {
33        if input.starts_with("uint") {
34            // Default to 256 if no size
35            let size = match input.get(4..input.len()) {
36                Some(s) => match s.is_empty() {
37                    false => match s.parse::<usize>() {
38                        Ok(s) => s,
39                        Err(_) => return Err(format!("Invalid uint size : {s}")),
40                    },
41                    true => 256,
42                },
43                None => 256,
44            };
45            return Ok(PrimitiveEVMType::Uint(size));
46        }
47        if input.starts_with("int") {
48            // Default to 256 if no size
49            let size = match input.get(3..input.len()) {
50                Some(s) => match s.is_empty() {
51                    false => match s.parse::<usize>() {
52                        Ok(s) => s,
53                        Err(_) => return Err(format!("Invalid int size : {s}")),
54                    },
55                    true => 256,
56                },
57                None => 256,
58            };
59            return Ok(PrimitiveEVMType::Int(size));
60        }
61        if input.starts_with("bytes") && input.len() != 5 {
62            let remaining = input.get(5..input.len()).unwrap();
63            let size = match remaining.parse::<usize>() {
64                Ok(s) => s,
65                Err(_) => return Err(format!("Invalid bytes size : {remaining}")),
66            };
67            return Ok(PrimitiveEVMType::Bytes(size));
68        }
69        if input.eq("bool") {
70            return Ok(PrimitiveEVMType::Bool);
71        }
72        if input.eq("address") {
73            return Ok(PrimitiveEVMType::Address);
74        }
75        if input.eq("string") {
76            return Ok(PrimitiveEVMType::String);
77        }
78        if input == "bytes" { Ok(PrimitiveEVMType::DynBytes) } else { Err(format!("Invalid PrimitiveEVMType type: {input}")) }
79    }
80}
81
82impl fmt::Display for PrimitiveEVMType {
83    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
84        let x = match *self {
85            PrimitiveEVMType::Address => "address",
86            PrimitiveEVMType::String => "string",
87            PrimitiveEVMType::Bool => "bool",
88            PrimitiveEVMType::DynBytes => "bytes",
89            PrimitiveEVMType::Uint(s) => return write!(f, "uint{s}"),
90            PrimitiveEVMType::Int(s) => return write!(f, "int{s}"),
91            PrimitiveEVMType::Bytes(s) => return write!(f, "bytes{s}"),
92        };
93
94        write!(f, "{x}")
95    }
96}
97
98// Array of regex to matching fancier EVM types
99lazy_static! {
100    /// Array of regex to matching fancier EVM types
101    pub static ref EVM_TYPE_ARRAY_REGEX: Regex = Regex::new(r"((u|)int[0-9]*|address|bool|bytes|string|bytes[0-9]*)\[[0-9]*\]").unwrap();
102}
103
104/// Wrap ether-rs Token to allow to derive the TryFrom trait
105#[derive(Clone)]
106pub struct EToken(pub DynSolValue);
107
108impl TryFrom<String> for EToken {
109    type Error = String;
110
111    fn try_from(input: String) -> Result<Self, Self::Error> {
112        if input.starts_with("0x") {
113            // remove 0x prefix
114            let cleaned_input = input.get(2..input.len()).unwrap();
115            // either address or fixed bytes
116            if cleaned_input.len() <= 64 {
117                // could be either address or fixed bytes
118                // if length is 42, assume it's an address
119                match input.len() {
120                    42 => return Ok(EToken(DynSolValue::Address(Address::from_str(cleaned_input).map_err(|e| e.to_string())?))),
121                    _ => return Ok(EToken(DynSolValue::FixedBytes(Word::from(&str_to_bytes32(cleaned_input)), 32))),
122                }
123            } else {
124                // dyn bytes array
125                return Ok(EToken(DynSolValue::Bytes(str_to_vec(cleaned_input).map_err(|e| e.to_string())?)));
126            }
127        }
128        // array
129        if input.starts_with('[') {
130            let trimmed_input = input.trim_start_matches('[').trim_end_matches(']');
131            let v: Vec<String> = trimmed_input.split(',').map(|x| x.replace([' ', '"', '\''], "")).collect();
132            let etokens: Result<Vec<EToken>, _> = v.iter().map(|x| EToken::try_from(x.to_owned())).collect();
133            let tokens: Vec<DynSolValue> = etokens?.iter().map(move |x| x.clone().0).collect();
134            return Ok(EToken(DynSolValue::Array(tokens)));
135        }
136        if input.starts_with('-') || input.starts_with('+') {
137            return Ok(EToken(DynSolValue::Int(I256::from_str(&input).map_err(|e| e.to_string())?, 256)));
138        }
139        if input == "true" || input == "false" {
140            return Ok(EToken(DynSolValue::Bool(input == "true")));
141        }
142        if input.chars().all(|x| x.is_ascii_digit()) {
143            return Ok(EToken(DynSolValue::Uint(U256::from_str_radix(input.as_str(), 10).map_err(|e| e.to_string())?, 256)));
144        }
145        if input.chars().all(|x| x.is_alphanumeric()) {
146            Ok(EToken(DynSolValue::String(input)))
147        } else if input.contains(',') {
148            // Try to unwrap something like "100,0x123,20" without brackets
149            let e_tokens: Result<Vec<EToken>, _> =
150                input.split(',').map(|x| x.replace([' ', '"', '\''], "")).map(EToken::try_from).collect();
151            let tokens: Vec<DynSolValue> = e_tokens?.into_iter().map(|x| x.0).collect();
152            Ok(EToken(DynSolValue::Array(tokens)))
153        } else {
154            Err(format!("Invalid input: {input}"))
155        }
156    }
157}