From b5cb1f479f7731bc71ab67005ca4e6a19473b05f Mon Sep 17 00:00:00 2001
From: Alessandro Decina <alessandro.d@gmail.com>
Date: Sun, 21 Feb 2021 20:59:44 +0000
Subject: [PATCH] xask: add bindings::write

---
 xtask/src/codegen/aya_bpf.rs  | 47 ++++++++++++++---------------------
 xtask/src/codegen/bindings.rs | 16 +++++++++++-
 2 files changed, 34 insertions(+), 29 deletions(-)

diff --git a/xtask/src/codegen/aya_bpf.rs b/xtask/src/codegen/aya_bpf.rs
index 1808ddb0..a79a9ec1 100644
--- a/xtask/src/codegen/aya_bpf.rs
+++ b/xtask/src/codegen/aya_bpf.rs
@@ -1,4 +1,4 @@
-use std::{fs::File, io::Write, path::PathBuf, process::Command};
+use std::path::PathBuf;
 
 use anyhow::anyhow;
 use proc_macro2::TokenStream;
@@ -15,7 +15,7 @@ use syn::{
 };
 
 use crate::codegen::{
-    bindings::bindgen,
+    bindings::{self, bindgen},
     getters::{generate_getters_for_items, Getter},
 };
 
@@ -52,32 +52,23 @@ pub fn codegen(opts: CodegenOptions) -> Result<(), anyhow::Error> {
     };
     tx.visit_file_mut(&mut tree);
 
-    let bindings = tree.to_token_stream().to_string();
-    let filename = generated.join("bindings.rs");
-    {
-        let mut file = File::create(&filename)?;
-        write!(file, "{}", bindings)?;
-    }
-    Command::new("rustfmt").arg(filename).status()?;
-
-    let filename = generated.join("helpers.rs");
-    {
-        let mut file = File::create(&filename)?;
-        write!(file, "use crate::bpf::generated::bindings::*;")?;
-        for helper in &tx.helpers {
-            file.write(helper.as_bytes())?;
-        }
-    }
-    Command::new("rustfmt").arg(filename).status()?;
-
-    let getters = generate_getters_for_items(&tree.items, gen_probe_read_getter);
-    let filename = generated.join("getters.rs");
-    {
-        let mut file = File::create(&filename)?;
-        write!(file, "use crate::bpf::generated::bindings::*;")?;
-        write!(file, "{}", getters)?;
-    }
-    Command::new("rustfmt").arg(filename).status()?;
+    bindings::write(
+        &tree.to_token_stream().to_string(),
+        "",
+        &generated.join("bindings.rs"),
+    )?;
+
+    bindings::write(
+        &tx.helpers.join(""),
+        "use crate::bpf::generated::bindings::*;",
+        &generated.join("helpers.rs"),
+    )?;
+
+    bindings::write(
+        &generate_getters_for_items(&tree.items, gen_probe_read_getter).to_string(),
+        "use crate::bpf::generated::bindings::*;",
+        &generated.join("getters.rs"),
+    )?;
 
     Ok(())
 }
diff --git a/xtask/src/codegen/bindings.rs b/xtask/src/codegen/bindings.rs
index 30be31ea..4ac9125b 100644
--- a/xtask/src/codegen/bindings.rs
+++ b/xtask/src/codegen/bindings.rs
@@ -1,4 +1,9 @@
-use std::process::Command;
+use std::{
+    fs::File,
+    io::{self, Write},
+    path::Path,
+    process::Command,
+};
 
 pub fn bindgen(types: &[&str], vars: &[&str]) -> Command {
     let mut cmd = Command::new("bindgen");
@@ -20,3 +25,12 @@ pub fn bindgen(types: &[&str], vars: &[&str]) -> Command {
 
     cmd
 }
+
+pub fn write(bindings: &str, header: &str, filename: &Path) -> io::Result<()> {
+    let mut file = File::create(&filename)?;
+    file.write(header.as_bytes())?;
+    file.write(bindings.as_bytes())?;
+
+    Command::new("rustfmt").arg(filename).status()?;
+    Ok(())
+}