spec-gen: add support for type params

ethereum-consensus/src/capella/spec/mod.rs

@@ -3195,6 +3195,13 @@ pub fn state_transition_block_in_slot<
     const MAX_TRANSACTIONS_PER_PAYLOAD: usize,
     const MAX_WITHDRAWALS_PER_PAYLOAD: usize,
     const MAX_BLS_TO_EXECUTION_CHANGES: usize,
+    E: ExecutionEngine<
+        BYTES_PER_LOGS_BLOOM,
+        MAX_EXTRA_DATA_BYTES,
+        MAX_BYTES_PER_TRANSACTION,
+        MAX_TRANSACTIONS_PER_PAYLOAD,
+        MAX_WITHDRAWALS_PER_PAYLOAD,
+    >,
 >(
     state: &mut BeaconState<
         SLOTS_PER_HISTORICAL_ROOT,
@@ -3262,6 +3269,13 @@ pub fn state_transition<
     const MAX_TRANSACTIONS_PER_PAYLOAD: usize,
     const MAX_WITHDRAWALS_PER_PAYLOAD: usize,
     const MAX_BLS_TO_EXECUTION_CHANGES: usize,
+    E: ExecutionEngine<
+        BYTES_PER_LOGS_BLOOM,
+        MAX_EXTRA_DATA_BYTES,
+        MAX_BYTES_PER_TRANSACTION,
+        MAX_TRANSACTIONS_PER_PAYLOAD,
+        MAX_WITHDRAWALS_PER_PAYLOAD,
+    >,
 >(
     state: &mut BeaconState<
         SLOTS_PER_HISTORICAL_ROOT,

spec-gen/src/generator.rs

@@ -1,14 +1,14 @@
 use crate::visitors::{
-    collate_generics_from, collect_lifetimes, generics_to_arguments, ArgumentsEditor,
-    TypeNameVisitor,
+    collate_generics_from, collect_lifetimes, collect_type_params, generics_to_arguments,
+    ArgumentsEditor, TypeNameVisitor,
 };
 use std::{
     collections::{BTreeMap, HashMap},
     fs,
     path::PathBuf,
     rc::Rc,
 };
-use syn::{parse_quote, Ident, Item};
+use syn::{parse_quote, visit_mut::VisitMut, Ident, Item};
 
 const SOURCE_ROOT: &str = "ethereum-consensus/src";
 
@@ -405,7 +405,24 @@ impl Spec {
 
                 let lifetimes = collect_lifetimes(&fragment);
 
-                let generics = collate_generics_from(&all_arguments, &lifetimes);
+                let (mut type_params, bounds) = collect_type_params(&fragment.sig.generics);
+                for (name, type_param) in bounds.iter().zip(type_params.iter_mut()) {
+                    if let Some(target_module) = index.get(name) {
+                        let target_module = self.diff.modules.get(target_module).unwrap();
+                        let trait_def = target_module
+                            .trait_defs
+                            .iter()
+                            .find(|&c| &c.name == name)
+                            .expect("internal state integrity");
+
+                        let arguments = generics_to_arguments(&trait_def.item.generics);
+                        let mut editor = ArgumentsEditor::new(&trait_def.name, &arguments);
+                        editor.visit_type_param_mut(type_param);
+
+                        f.fork = self.fork;
+                    }
+                }
+                let generics = collate_generics_from(&all_arguments, &lifetimes, &type_params);
                 fragment.sig.generics = generics;
 
                 f.item = fragment;

spec-gen/src/visitors.rs

@@ -23,6 +23,45 @@ pub fn collect_lifetimes(f: &ItemFn) -> Vec<syn::Lifetime> {
     visitor.items
 }
 
+#[derive(Default)]
+struct TypeParamVisitor {
+    in_context: bool,
+    items: Vec<syn::TypeParam>,
+    bounds: Vec<String>,
+}
+
+impl<'ast> Visit<'ast> for TypeParamVisitor {
+    fn visit_path_segment(&mut self, i: &'ast syn::PathSegment) {
+        if self.in_context {
+            self.bounds.push(i.ident.to_string());
+        }
+    }
+
+    fn visit_type_param_bound(&mut self, i: &'ast syn::TypeParamBound) {
+        self.in_context = true;
+        visit::visit_type_param_bound(self, i);
+        self.in_context = false;
+    }
+
+    fn visit_type_param(&mut self, i: &'ast syn::TypeParam) {
+        if !self.items.contains(i) {
+            // NOTE: only supports a single bound for now
+            i.bounds.iter().for_each(|i| {
+                visit::visit_type_param_bound(self, i);
+            });
+            self.items.push(i.clone());
+        }
+        visit::visit_type_param(self, i);
+    }
+}
+
+// returns type params and the name of their trait bound
+pub fn collect_type_params(g: &Generics) -> (Vec<syn::TypeParam>, Vec<String>) {
+    let mut visitor = TypeParamVisitor::default();
+    visitor.visit_generics(g);
+    (visitor.items, visitor.bounds)
+}
+
 #[derive(Default)]
 struct ToGenericsVisitor {
     bounds: Vec<Ident>,
@@ -40,6 +79,7 @@ impl<'ast> Visit<'ast> for ToGenericsVisitor {
 pub fn collate_generics_from(
     arguments: &[AngleBracketedGenericArguments],
     lifetimes: &[syn::Lifetime],
+    type_params: &[syn::TypeParam],
 ) -> Generics {
     let mut visitor = ToGenericsVisitor::default();
 
@@ -49,22 +89,27 @@ pub fn collate_generics_from(
 
     let bounds = visitor.bounds;
 
-    // NOTE: macro failed when trying to combine both arms into one invocation
-    // it also seems like this will only suppport one explicit lifetime
-    if lifetimes.is_empty() {
-        parse_quote! {
-            <
-            #(const #bounds: usize),*
-            >
-        }
-    } else {
-        parse_quote! {
-            <
+    let mut generics: syn::Generics = parse_quote! {
+        <
+        #(const #bounds: usize),*
+        >
+    };
+
+    if !lifetimes.is_empty() {
+        let lifetimes: syn::punctuated::Punctuated<syn::GenericParam, syn::Token![,]> = parse_quote! {
             #(#lifetimes)*,
-            #(const #bounds: usize),*
-            >
-        }
+        };
+        generics.params.extend(lifetimes);
     }
+
+    if !type_params.is_empty() {
+        let type_params: syn::punctuated::Punctuated<syn::GenericParam, syn::Token![,]> = parse_quote! {
+            #(#type_params)*,
+        };
+        generics.params.extend(type_params);
+    }
+
+    generics
 }
 
 #[derive(Default)]