use std::collections::HashMap;
pub fn canonical_param_type(ty: &str) -> String {
let mut parts = ty.split_whitespace();
match parts.next() {
Some("struct" | "enum") => parts.next().unwrap_or_default().to_string(),
Some(first) => first.to_string(),
None => String::new(),
}
}
pub fn canonical_param_type_simple(ty: &str) -> String {
ty.split_whitespace().next().unwrap_or_default().to_string()
}
pub fn canonical_param_type_with_structs(
ty: &str,
struct_fields: &HashMap<String, Vec<(String, String)>>,
) -> String {
let trimmed = ty.trim();
let without_location = trimmed
.replace(" memory", "")
.replace(" calldata", "")
.replace(" storage", "");
let t = without_location.trim();
if t.ends_with(']') {
if let Some(open) = t.rfind('[') {
let (element, suffix) = t.split_at(open);
let inner = canonical_param_type_with_structs(element.trim(), struct_fields);
return format!("{inner}{suffix}");
}
}
let mut parts = t.split_whitespace();
match parts.next() {
Some("struct") => {
let name = parts.next().unwrap_or_default();
if name.is_empty() {
return String::new();
}
let mut visited = std::collections::HashSet::new();
expand_struct_canonical(name, struct_fields, &mut visited)
}
Some("enum") => "uint8".to_string(),
Some(first) => {
if struct_fields.contains_key(first) {
let mut visited = std::collections::HashSet::new();
expand_struct_canonical(first, struct_fields, &mut visited)
} else {
first.to_string()
}
}
None => String::new(),
}
}
fn expand_struct_canonical(
name: &str,
struct_fields: &HashMap<String, Vec<(String, String)>>,
visited: &mut std::collections::HashSet<String>,
) -> String {
if !visited.insert(name.to_string()) {
return name.to_string();
}
let Some(fields) = struct_fields.get(name) else {
visited.remove(name);
return name.to_string();
};
let field_sigs: Vec<String> = fields
.iter()
.map(|(_, field_ty)| {
let tt = field_ty
.replace(" memory", "")
.replace(" calldata", "")
.replace(" storage", "");
let tt = tt.trim();
let mut parts = tt.split_whitespace();
match parts.next() {
Some("struct") => {
let sub_name = parts.next().unwrap_or_default();
if sub_name.is_empty() {
String::new()
} else {
expand_struct_canonical(sub_name, struct_fields, visited)
}
}
Some("enum") => "uint8".to_string(),
Some(first) => {
if struct_fields.contains_key(first) {
expand_struct_canonical(first, struct_fields, visited)
} else {
first.to_string()
}
}
None => String::new(),
}
})
.collect();
visited.remove(name);
format!("({})", field_sigs.join(","))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_canonical_param_type() {
assert_eq!(canonical_param_type("uint256"), "uint256");
assert_eq!(canonical_param_type("struct MyStruct"), "MyStruct");
assert_eq!(canonical_param_type("enum MyEnum"), "MyEnum");
assert_eq!(canonical_param_type("address payable"), "address");
assert_eq!(canonical_param_type(""), "");
}
#[test]
fn test_canonical_param_type_simple() {
assert_eq!(canonical_param_type_simple("uint256"), "uint256");
assert_eq!(canonical_param_type_simple("struct MyStruct"), "struct");
assert_eq!(canonical_param_type_simple("address payable"), "address");
}
#[test]
fn test_canonical_param_type_with_structs_flat() {
let mut map: HashMap<String, Vec<(String, String)>> = HashMap::new();
map.insert(
"P".to_string(),
vec![
("a".to_string(), "uint256".to_string()),
("b".to_string(), "bool".to_string()),
],
);
assert_eq!(
canonical_param_type_with_structs("struct P", &map),
"(uint256,bool)"
);
assert_eq!(
canonical_param_type_with_structs("struct P memory", &map),
"(uint256,bool)"
);
assert_eq!(
canonical_param_type_with_structs("uint256", &map),
"uint256"
);
}
#[test]
fn test_canonical_param_type_with_structs_nested() {
let mut map: HashMap<String, Vec<(String, String)>> = HashMap::new();
map.insert(
"Inner".to_string(),
vec![("x".to_string(), "uint256".to_string())],
);
map.insert(
"Outer".to_string(),
vec![
("inner".to_string(), "struct Inner".to_string()),
("who".to_string(), "address".to_string()),
],
);
assert_eq!(
canonical_param_type_with_structs("struct Outer", &map),
"((uint256),address)"
);
}
#[test]
fn test_canonical_param_type_with_structs_unknown_fallback() {
let map: HashMap<String, Vec<(String, String)>> = HashMap::new();
assert_eq!(
canonical_param_type_with_structs("struct Unknown", &map),
"Unknown"
);
}
#[test]
fn test_canonical_param_type_with_structs_enum() {
let map: HashMap<String, Vec<(String, String)>> = HashMap::new();
assert_eq!(
canonical_param_type_with_structs("enum MyEnum", &map),
"uint8"
);
}
}