vmlinux_docify/
main.rs

1use clap::{ColorChoice, Parser};
2use std::collections::HashMap;
3use std::fs;
4use std::process;
5
6/// A tool to annotate vmlinux.h with documentation from kernel sources
7#[derive(Parser, Debug)]
8#[command(author, version, about, long_about = "A tool to annotate vmlinux.h with documentation from kernel sources", color = ColorChoice::Always)]
9struct Args {
10    /// Path to the kernel source directory
11    #[arg(short, long)]
12    kernel_dir: String,
13
14    /// Path to the vmlinux.h file to annotate
15    #[arg(short, long)]
16    vmlinux_h: String,
17
18    /// Path to the output file (default: vmlinux_annotated.h)
19    #[arg(short, long, default_value = "vmlinux_annotated.h")]
20    output: String,
21}
22
23/// Builds a map of BPF kfunc signatures to their comments from kernel source files
24fn build_bpf_kfunc_map(kernel_dir: &str) -> HashMap<String, String> {
25    // Get all .c and .h files
26    let files: Vec<_> = walkdir::WalkDir::new(kernel_dir)
27        .into_iter()
28        .filter_map(|e| e.ok())
29        .filter(|e| {
30            let path = e.path();
31            path.is_file() && path.extension().is_some_and(|ext| ext == "c" || ext == "h")
32        })
33        .map(|e| e.path().to_path_buf())
34        .collect();
35
36    println!("Found {} files to process", files.len());
37
38    // Process files sequentially
39    let mut comments_map = HashMap::new();
40
41    for file in files {
42        let content = match fs::read_to_string(&file) {
43            Ok(content) => content,
44            Err(e) => {
45                eprintln!("Error reading {}: {}", file.display(), e);
46                continue;
47            }
48        };
49
50        let mut current_comment = String::new();
51        let mut in_comment = false;
52        let mut prev_was_close = false;
53        let mut last_comment_line = 0;
54        let mut line_number = 0;
55
56        for line in content.lines() {
57            line_number += 1;
58
59            if line.contains("/**") {
60                in_comment = true;
61                current_comment.clear();
62                current_comment.push_str(line);
63                current_comment.push('\n');
64            } else if in_comment {
65                if line.contains("*/") {
66                    current_comment.push_str(line);
67                    in_comment = false;
68                    prev_was_close = true;
69                    last_comment_line = line_number;
70                } else {
71                    current_comment.push_str(line);
72                    current_comment.push('\n');
73                }
74            } else if !current_comment.is_empty() && line.contains("__bpf_kfunc") {
75                if line.contains("(") && line.contains(" ") {
76                    let sig = (" ".to_owned()
77                        + line
78                            .trim()
79                            .to_string()
80                            .split_terminator('(')
81                            .collect::<Vec<&str>>()[0]
82                            .split_terminator(' ')
83                            .collect::<Vec<&str>>()
84                            .last()
85                            .unwrap()
86                            .to_string()
87                            .to_owned()
88                            .as_str()
89                        + "(")
90                        .to_string();
91                    comments_map.insert(sig, current_comment.trim().to_string());
92                }
93                current_comment.clear();
94                prev_was_close = false;
95            } else if !line.is_empty()
96                && prev_was_close
97                && (line_number - last_comment_line > 1 || line.contains("__bpf_kptr"))
98            {
99                // Only clear the comment if we're far from the last comment or if we encounter __bpf_kptr
100                current_comment.clear();
101                prev_was_close = false;
102            }
103        }
104    }
105
106    println!("Processing complete");
107    comments_map
108}
109
110/// Adds comments to vmlinux.h based on the provided comments map
111fn annotate_vmlinux(
112    vmlinux_h: &str,
113    comments_map: &HashMap<String, String>,
114    struct_map: &HashMap<String, String>,
115) -> Result<String, std::io::Error> {
116    // Read vmlinux.h
117    let content = fs::read_to_string(vmlinux_h)?;
118
119    // Process the file and add comments
120    let mut output = String::new();
121    let mut matches_found = 0;
122    let mut struct_matches_found = 0;
123
124    for line in content.lines() {
125        // Check if this line contains a function that matches one of our keys
126        for (key, comment) in comments_map {
127            if line.contains("extern") && line.contains(key) {
128                output.push_str(comment);
129                output.push('\n');
130                matches_found += 1;
131                break;
132            }
133        }
134
135        // Check if this line contains a struct or enum that matches one of our keys
136        for (key, comment) in struct_map {
137            if line.contains(key) && (line.contains("struct") || line.contains("enum")) {
138                output.push_str(comment);
139                output.push('\n');
140                struct_matches_found += 1;
141                break;
142            }
143        }
144
145        output.push_str(line);
146        output.push('\n');
147    }
148
149    println!(
150        "Added {matches_found} function comments and {struct_matches_found} struct/enum comments to vmlinux.h"
151    );
152    Ok(output)
153}
154
155fn build_kernel_struct_map(kernel_dir: &str) -> HashMap<String, String> {
156    // Get all .c and .h files
157    let files: Vec<_> = walkdir::WalkDir::new(kernel_dir)
158        .into_iter()
159        .filter_map(|e| e.ok())
160        .filter(|e| {
161            let path = e.path();
162            path.is_file() && path.extension().is_some_and(|ext| ext == "c" || ext == "h")
163        })
164        .map(|e| e.path().to_path_buf())
165        .collect();
166
167    println!("Found {} files to process", files.len());
168
169    // Process files sequentially
170    let mut struct_map = HashMap::new();
171
172    for file in files {
173        let content = match fs::read_to_string(&file) {
174            Ok(content) => content,
175            Err(e) => {
176                eprintln!("Error reading {}: {}", file.display(), e);
177                continue;
178            }
179        };
180
181        let mut current_comment = String::new();
182        let mut in_comment = false;
183        let mut last_comment_line = 0;
184        let mut line_number = 0;
185
186        for line in content.lines() {
187            line_number += 1;
188
189            if line.contains("/**") {
190                in_comment = true;
191                current_comment.clear();
192                current_comment.push_str(line);
193                current_comment.push('\n');
194            } else if in_comment {
195                if line.contains("*/") {
196                    current_comment.push_str(line);
197                    in_comment = false;
198                    last_comment_line = line_number;
199                } else {
200                    current_comment.push_str(line);
201                    current_comment.push('\n');
202                }
203            } else if !current_comment.is_empty()
204                && (line.trim().starts_with("struct") || line.trim().starts_with("enum"))
205                && line.trim().ends_with(" {")
206                && line
207                    .trim()
208                    .split_terminator(' ')
209                    .collect::<Vec<&str>>()
210                    .len()
211                    == 3
212                && (line_number - last_comment_line <= 1)
213            {
214                let mut key =
215                    line.trim().split_terminator(' ').collect::<Vec<&str>>()[1].to_string();
216                if line.contains("struct") {
217                    key = "struct ".to_string() + key.as_str() + " {";
218                } else if line.contains("enum") {
219                    key = "enum ".to_string() + key.as_str() + " {";
220                }
221                struct_map.insert(key.clone(), current_comment.clone());
222                current_comment.clear();
223            } else if !line.is_empty()
224                && !current_comment.is_empty()
225                && (line_number - last_comment_line > 1)
226            {
227                // Only clear the comment if we're far from the last comment
228                current_comment.clear();
229            }
230        }
231    }
232
233    println!("Processing complete");
234    struct_map
235}
236
237fn main() {
238    let args = Args::parse();
239
240    // Build the map of __bpf_kfunc declarations to their comments
241    let comments_map = build_bpf_kfunc_map(&args.kernel_dir);
242    println!(
243        "Found {} __bpf_kfunc declarations with comments",
244        comments_map.len()
245    );
246
247    // Build the map of structs and enums to their comments and definitions
248    let struct_map = build_kernel_struct_map(&args.kernel_dir);
249    println!("Found {} structs and enums with comments", struct_map.len());
250
251    // Annotate vmlinux.h with comments
252    match annotate_vmlinux(&args.vmlinux_h, &comments_map, &struct_map) {
253        Ok(annotated_content) => {
254            // Write the annotated output
255            if let Err(e) = fs::write(&args.output, annotated_content) {
256                eprintln!("Error writing annotated file: {e}");
257                process::exit(1);
258            }
259            println!("Successfully wrote annotated vmlinux.h to {}", args.output);
260        }
261        Err(e) => {
262            eprintln!("Error processing vmlinux.h: {e}");
263            process::exit(1);
264        }
265    }
266}
vmlinux_docify/main.rs

vmlinux_docify/
main.rs
