Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
186 changes: 142 additions & 44 deletions sv2/binary-sv2/derive_codec/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ use alloc::{
vec::Vec,
};
use core::iter::FromIterator;
use proc_macro::{Group, TokenStream, TokenTree};
use proc_macro::{Delimiter, Group, Ident, Literal, Punct, Spacing, Span, TokenStream, TokenTree};

// Reserved field names to avoid conflicts
const RESERVED_FIELDS: [&str; 2] = ["__decodable_internal_data", "__decodable_internal_offset"];
Expand Down Expand Up @@ -245,16 +245,24 @@ struct ParsedStruct {
struct ParsedField {
// Name of the field.
name: String,
// Span of the field name.
span: Span,
// Type of the field.
type_: String,
// Generics associated with the field, if any.
generics: String,
}

struct MacroError {
span: Span,
message: String,
}

impl ParsedField {
pub fn new() -> Self {
ParsedField {
name: "".to_string(),
span: Span::call_site(),
type_: "".to_string(),
generics: "".to_string(),
}
Expand All @@ -274,6 +282,46 @@ impl ParsedField {
".into_static()".to_string()
}
}

pub fn is_empty(&self) -> bool {
self.name.is_empty() && self.type_.is_empty() && self.generics.is_empty()
}
}

// Return `compile_error!(...)` tokens so the downstream crate gets a normal
// compiler error at the derive site instead of a proc-macro panic.
fn emit_compile_error(message: &str) -> TokenStream {
emit_compile_error_at(Span::call_site(), message)
}

fn emit_compile_error_at(span: Span, message: &str) -> TokenStream {
let mut inner = TokenStream::new();
let mut literal = Literal::string(message);
literal.set_span(span);
inner.extend([TokenTree::Literal(literal)]);

let mut group = Group::new(Delimiter::Parenthesis, inner);
group.set_span(span);
let mut bang = Punct::new('!', Spacing::Alone);
bang.set_span(span);
let mut semicolon = Punct::new(';', Spacing::Alone);
semicolon.set_span(span);

TokenStream::from_iter([
TokenTree::Ident(Ident::new("compile_error", span)),
TokenTree::Punct(bang),
TokenTree::Group(group),
TokenTree::Punct(semicolon),
])
}

fn parse_generated_tokens(result: &str) -> TokenStream {
result.parse().unwrap_or_else(|err| {
emit_compile_error(&format!(
"derive_codec_sv2 generated invalid Rust tokens: {}",
err
))
})
}

// Extracts properties of a struct, including its name, generics, and fields.
Expand Down Expand Up @@ -303,61 +351,82 @@ impl ParsedField {
//
// This example demonstrates how `get_struct_properties` identifies the struct's
// name, any generic parameters, and its fields, including types and generic parameters.
fn get_struct_properties(item: TokenStream) -> ParsedStruct {
fn get_struct_properties(item: TokenStream) -> Result<ParsedStruct, MacroError> {
let item = remove_attributes(item);
let mut stream = item.into_iter();

// Check if the stream is a struct
loop {
match stream.next().expect("Stream not a struct") {
TokenTree::Ident(i) => {
if i.to_string() == "struct" {
break;
}
match stream.next() {
Some(TokenTree::Ident(i)) if i.to_string() == "struct" => break,
Some(_) => continue,
None => {
return Err(MacroError {
span: Span::call_site(),
message: "Expected a struct definition".to_string(),
})
}
_ => continue,
}
}

// Get the struct name
let struct_name = match stream.next().expect("Struct has no name") {
TokenTree::Ident(i) => i.to_string(),
// Never executed at runtime it ok to panic
_ => panic!("Strcut has no name"),
let struct_name = match stream.next() {
Some(TokenTree::Ident(i)) => i.to_string(),
Some(token) => {
return Err(MacroError {
span: token.span(),
message: format!("Expected a struct name, found '{}'", token),
})
}
None => {
return Err(MacroError {
span: Span::call_site(),
message: "Struct has no name".to_string(),
})
}
};

let mut struct_generics = "".to_string();
let group: Vec<TokenTree>;

// Get the struct generics if any
loop {
match stream
.next()
// Never executed at runtime it ok to panic
.unwrap_or_else(|| panic!("Struct {} has no fields", struct_name))
{
TokenTree::Group(g) => {
match stream.next() {
Some(TokenTree::Group(g)) => {
group = g.stream().into_iter().collect();
break;
}
TokenTree::Punct(p) => {
Some(TokenTree::Punct(p)) => {
struct_generics = format!("{struct_generics}{p}");
}
TokenTree::Ident(i) => {
Some(TokenTree::Ident(i)) => {
struct_generics = format!("{struct_generics}{i}");
}
// Never executed at runtime it ok to panic
_ => panic!("Struct {} has no fields", struct_name),
Some(token) => {
return Err(MacroError {
span: token.span(),
message: format!(
"Struct '{}' has invalid tokens before its fields: '{}'",
struct_name, token
),
});
}
None => {
return Err(MacroError {
span: Span::call_site(),
message: format!("Struct '{}' has no fields", struct_name),
})
}
};
}

let fields = parse_struct_fields(group);
let fields = parse_struct_fields(group)?;

ParsedStruct {
Ok(ParsedStruct {
name: struct_name,
generics: struct_generics,
fields,
}
})
}

// Parses the fields of a struct, scanning tokens to identify field names, types, and generics.
Expand All @@ -384,17 +453,20 @@ fn get_struct_properties(item: TokenStream) -> ParsedStruct {
//
// This example shows how `parse_struct_fields` handles both a primitive field (`id`) and a
// generic field (`data`) within a struct, including parsing of generic parameters.
fn parse_struct_fields(group: Vec<TokenTree>) -> Vec<ParsedField> {
fn parse_struct_fields(group: Vec<TokenTree>) -> Result<Vec<ParsedField>, MacroError> {
let mut fields = Vec::new();
let mut field_ = ParsedField::new();
let mut field_parser_state = ParserState::Name;
let mut last_span = Span::call_site();
for token in group {
last_span = token.span();
match (token, &field_parser_state) {
(TokenTree::Ident(i), ParserState::Name) => {
if i.to_string() == "pub" {
continue;
} else {
field_.name = i.to_string();
field_.span = i.span();
}
}
(TokenTree::Ident(i), ParserState::Type) => {
Expand All @@ -407,8 +479,10 @@ fn parse_struct_fields(group: Vec<TokenTree>) -> Vec<ParsedField> {
if p.to_string() == ":" {
field_parser_state = ParserState::Type
} else {
// Never executed at runtime it ok to panic
panic!("Unexpected token '{}' in parsing {:#?}", p, field_);
return Err(MacroError {
span: p.span(),
message: format!("Unexpected token '{}' in parsing {:#?}", p, field_),
});
}
}
(TokenTree::Punct(p), ParserState::Type) => match p.to_string().as_ref() {
Expand All @@ -421,8 +495,12 @@ fn parse_struct_fields(group: Vec<TokenTree>) -> Vec<ParsedField> {
field_.generics = "<".to_string();
field_parser_state = ParserState::Generics(0);
}
// Never executed at runtime it ok to panic
_ => panic!("Unexpected token '{}' in parsing {:#?}", p, field_),
_ => {
return Err(MacroError {
span: p.span(),
message: format!("Unexpected token '{}' in parsing {:#?}", p, field_),
})
}
},
(TokenTree::Punct(p), ParserState::Generics(open_brackets)) => {
match p.to_string().as_ref() {
Expand All @@ -446,11 +524,29 @@ fn parse_struct_fields(group: Vec<TokenTree>) -> Vec<ParsedField> {
}
}
}
// Never executed at runtime it ok to panic
_ => panic!("Unexpected token"),
(token, _) => {
return Err(MacroError {
span: token.span(),
message: format!("Unexpected token '{}' while parsing struct fields", token),
});
}
}
}
fields

if !field_.is_empty() {
if field_.name.is_empty() || field_.type_.is_empty() {
return Err(MacroError {
span: last_span,
message: format!(
"Incomplete field definition near '{}'; expected a named field like `name: Type`",
field_.name
),
});
}
fields.push(field_);
}

Ok(fields)
}

/// Derives the `Decodable` trait, generating implementations for deserializing a struct from a
Expand Down Expand Up @@ -555,18 +651,19 @@ fn parse_struct_fields(group: Vec<TokenTree>) -> Vec<ParsedField> {
/// ownership and lifetime management of decoded fields in the struct.
#[proc_macro_derive(Decodable)]
pub fn decodable(item: TokenStream) -> TokenStream {
let parsed_struct = get_struct_properties(item);
let parsed_struct = match get_struct_properties(item) {
Ok(parsed_struct) => parsed_struct,
Err(err) => return emit_compile_error_at(err.span, &err.message),
};
let data_ident = RESERVED_FIELDS[0];
let offset_ident = RESERVED_FIELDS[1];

for field in &parsed_struct.fields {
if RESERVED_FIELDS.contains(&field.name.as_str()) {
return format!(
"compile_error!(\"Field name '{}' is reserved and cannot be used in struct '{}'. Rename it to avoid conflicts.\");",
return emit_compile_error_at(field.span, &format!(
"Field name '{}' is reserved and cannot be used in struct '{}'. Rename it to avoid conflicts.",
field.name, parsed_struct.name
)
.parse()
.unwrap();
));
}
}

Expand Down Expand Up @@ -703,8 +800,7 @@ pub fn decodable(item: TokenStream) -> TokenStream {

);

// Never executed at runtime it ok to panic
result.parse().unwrap()
parse_generated_tokens(&result)
}

fn get_static_generics(gen: &str) -> &str {
Expand Down Expand Up @@ -787,7 +883,10 @@ fn get_static_generics(gen: &str) -> &str {
#[proc_macro_derive(Encodable, attributes(already_sized))]
pub fn encodable(item: TokenStream) -> TokenStream {
let is_already_sized = is_already_sized(item.clone());
let parsed_struct = get_struct_properties(item);
let parsed_struct = match get_struct_properties(item) {
Ok(parsed_struct) => parsed_struct,
Err(err) => return emit_compile_error_at(err.span, &err.message),
};
let fields = parsed_struct.fields.clone();

let mut field_into_decoded_field = String::new();
Expand Down Expand Up @@ -871,6 +970,5 @@ pub fn encodable(item: TokenStream) -> TokenStream {
get_size,
);

// Never executed at runtime it ok to panic
result.parse().unwrap()
parse_generated_tokens(&result)
}
Loading