xtask: rework code generator for bpf helpers

xtask/Cargo.toml

@@ -8,7 +8,7 @@ edition = "2018"
 aya-gen = { path = "../aya-gen" }
 structopt = {version = "0.3", default-features = false }
 anyhow = "1"
-syn = {version = "1", features = ["visit-mut", "extra-traits"] }
+syn = "1"
 quote = "1"
 proc-macro2 = "1"
 indexmap = "1.6"

xtask/src/codegen/aya_bpf_bindings.rs

@@ -1,23 +1,15 @@
-use std::path::PathBuf;
-
 use anyhow::anyhow;
 use proc_macro2::TokenStream;
-use quote::{quote, ToTokens};
+use quote::ToTokens;
+use std::path::PathBuf;
 use structopt::StructOpt;
-use syn::{
-    self, parse_str,
-    punctuated::Punctuated,
-    token::Comma,
-    visit_mut::{self, VisitMut},
-    AngleBracketedGenericArguments, ForeignItemStatic, GenericArgument, Ident, Item, Path,
-    PathArguments::AngleBracketed,
-    Type,
-};
 
 use aya_gen::getters::{generate_getters_for_items, probe_read_getter};
+use syn::{parse_str, Item};
 
 use crate::codegen::{
     bindings::{self, bindgen},
+    helpers::{expand_helpers, extract_helpers},
     Architecture,
 };
 
@@ -49,13 +41,12 @@ pub fn codegen(opts: CodegenOptions) -> Result<(), anyhow::Error> {
         return Err(anyhow!("bindgen failed: {}", output.status));
     }
 
-    // delete the helpers, then rewrite them in helpers.rs
     let mut tree = parse_str::<syn::File>(bindings).unwrap();
-
-    let mut tx = RewriteBpfHelpers {
-        helpers: Vec::new(),
-    };
-    tx.visit_file_mut(&mut tree);
+    let (indexes, helpers) = extract_helpers(&tree.items);
+    let helpers = expand_helpers(&helpers);
+    for index in indexes {
+        tree.items[index] = Item::Verbatim(TokenStream::new())
+    }
 
     bindings::write(
         &tree.to_token_stream().to_string(),
@@ -64,12 +55,12 @@ pub fn codegen(opts: CodegenOptions) -> Result<(), anyhow::Error> {
     )?;
 
     bindings::write(
-        &tx.helpers.join(""),
+        &helpers.to_string(),
         "use super::bindings::*;",
         &generated.join("helpers.rs"),
     )?;
 
-    let bpf_probe_read = syn::parse_str::<Path>("crate::bpf_probe_read").unwrap();
+    let bpf_probe_read = syn::parse_str("crate::bpf_probe_read").unwrap();
     bindings::write(
         &generate_getters_for_items(&tree.items, |getter| {
             probe_read_getter(getter, &bpf_probe_read)
@@ -81,60 +72,3 @@ pub fn codegen(opts: CodegenOptions) -> Result<(), anyhow::Error> {
 
     Ok(())
 }
-
-struct RewriteBpfHelpers {
-    helpers: Vec<String>,
-}
-
-impl VisitMut for RewriteBpfHelpers {
-    fn visit_item_mut(&mut self, item: &mut Item) {
-        visit_mut::visit_item_mut(self, item);
-        if let Item::ForeignMod(_) = item {
-            *item = Item::Verbatim(TokenStream::new())
-        }
-    }
-    fn visit_foreign_item_static_mut(&mut self, static_item: &mut ForeignItemStatic) {
-        if let Type::Path(path) = &*static_item.ty {
-            let ident = &static_item.ident;
-            let ident_str = ident.to_string();
-            let last = path.path.segments.last().unwrap();
-            let ty_ident = last.ident.to_string();
-            if ident_str.starts_with("bpf_") && ty_ident == "Option" {
-                let fn_ty = match &last.arguments {
-                    AngleBracketed(AngleBracketedGenericArguments { args, .. }) => {
-                        args.first().unwrap()
-                    }
-                    _ => panic!(),
-                };
-                let mut ty_s = quote! {
-                    #[inline(always)]
-                    pub #fn_ty
-                }
-                .to_string();
-                ty_s = ty_s.replace("fn (", &format!("fn {} (", ident_str));
-                let call_idx = self.helpers.len() + 1;
-                let args: Punctuated<Ident, Comma> = match fn_ty {
-                    GenericArgument::Type(Type::BareFn(f)) => f
-                        .inputs
-                        .iter()
-                        .map(|arg| arg.name.clone().unwrap().0)
-                        .collect(),
-                    _ => unreachable!(),
-                };
-                let body = quote! {
-                    {
-                        let f: #fn_ty = ::core::mem::transmute(#call_idx);
-                        f(#args)
-                    }
-                }
-                .to_string();
-                ty_s.push_str(&body);
-                let mut helper = ty_s;
-                if helper.contains("printk") {
-                    helper = format!("/* {} */", helper);
-                }
-                self.helpers.push(helper);
-            }
-        }
-    }
-}

xtask/src/codegen/helpers.rs

@@ -0,0 +1,100 @@
+use proc_macro2::TokenStream;
+use quote::{quote, TokenStreamExt};
+use syn::{
+    punctuated::Punctuated, AngleBracketedGenericArguments, BareFnArg, ForeignItem,
+    ForeignItemStatic, GenericArgument, Ident, Item, Path, PathArguments, ReturnType, Token, Type,
+    TypeBareFn, TypePath,
+};
+
+pub fn extract_helpers(items: &[Item]) -> (Vec<usize>, Vec<Helper<'_>>) {
+    let mut helpers = Vec::new();
+    let mut indexes = Vec::new();
+    for (item_index, item) in items.iter().enumerate() {
+        if let Item::ForeignMod(module) = item {
+            for i in &module.items {
+                if let ForeignItem::Static(s_item) = i {
+                    let ident_s = s_item.ident.to_string();
+                    if ident_s.starts_with("bpf_") {
+                        helpers.push(
+                            helper_from_item(s_item, helpers.len() + 1)
+                                .expect("unexpected bindgen helper signature"),
+                        );
+                        indexes.push(item_index);
+                    }
+                }
+            }
+        }
+    }
+
+    (indexes, helpers)
+}
+
+pub fn helper_from_item(item: &ForeignItemStatic, call_index: usize) -> Option<Helper<'_>> {
+    if let Type::Path(TypePath {
+        path: Path { segments, .. },
+        ..
+    }) = &*item.ty
+    {
+        let generics = &segments.last().unwrap().arguments;
+        if let PathArguments::AngleBracketed(AngleBracketedGenericArguments { args, .. }) = generics
+        {
+            if let Some(GenericArgument::Type(ty)) = args.first() {
+                if let Type::BareFn(TypeBareFn { inputs, output, .. }) = ty {
+                    return Some(Helper {
+                        ident: &item.ident,
+                        ty,
+                        inputs,
+                        output,
+                        call_index,
+                    });
+                }
+            }
+        };
+    }
+
+    None
+}
+
+pub fn expand_helpers(helpers: &[Helper<'_>]) -> TokenStream {
+    let mut tokens = TokenStream::new();
+    tokens.append_all(
+        helpers
+            .iter()
+            .filter(|h| h.ident.to_string() != "bpf_trace_printk")
+            .map(expand_helper),
+    );
+
+    return tokens;
+}
+
+pub fn expand_helper(helper: &Helper<'_>) -> TokenStream {
+    let Helper {
+        ident,
+        ty,
+        inputs,
+        output,
+        call_index,
+    } = helper;
+
+    let args = inputs
+        .iter()
+        .map(|arg| &arg.name.as_ref().unwrap().0)
+        .collect::<Vec<_>>();
+
+    let helper = quote! {
+        pub unsafe fn #ident(#inputs) #output {
+            let fun: #ty = ::core::mem::transmute(#call_index);
+            fun(#(#args),*)
+        }
+    };
+
+    helper
+}
+
+pub struct Helper<'a> {
+    ident: &'a Ident,
+    ty: &'a Type,
+    inputs: &'a Punctuated<BareFnArg, Token![,]>,
+    output: &'a ReturnType,
+    call_index: usize,
+}

xtask/src/codegen/mod.rs

@@ -1,5 +1,6 @@
 mod aya_bpf_bindings;
 mod bindings;
+mod helpers;
 
 use structopt::StructOpt;