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 {} function comments and {} struct/enum comments to vmlinux.h",
151        matches_found, struct_matches_found
152    );
153    Ok(output)
154}
155
156fn build_kernel_struct_map(kernel_dir: &str) -> HashMap<String, String> {
157    // Get all .c and .h files
158    let files: Vec<_> = walkdir::WalkDir::new(kernel_dir)
159        .into_iter()
160        .filter_map(|e| e.ok())
161        .filter(|e| {
162            let path = e.path();
163            path.is_file() && path.extension().is_some_and(|ext| ext == "c" || ext == "h")
164        })
165        .map(|e| e.path().to_path_buf())
166        .collect();
167
168    println!("Found {} files to process", files.len());
169
170    // Process files sequentially
171    let mut struct_map = HashMap::new();
172
173    for file in files {
174        let content = match fs::read_to_string(&file) {
175            Ok(content) => content,
176            Err(e) => {
177                eprintln!("Error reading {}: {}", file.display(), e);
178                continue;
179            }
180        };
181
182        let mut current_comment = String::new();
183        let mut in_comment = false;
184        let mut last_comment_line = 0;
185        let mut line_number = 0;
186
187        for line in content.lines() {
188            line_number += 1;
189
190            if line.contains("/**") {
191                in_comment = true;
192                current_comment.clear();
193                current_comment.push_str(line);
194                current_comment.push('\n');
195            } else if in_comment {
196                if line.contains("*/") {
197                    current_comment.push_str(line);
198                    in_comment = false;
199                    last_comment_line = line_number;
200                } else {
201                    current_comment.push_str(line);
202                    current_comment.push('\n');
203                }
204            } else if !current_comment.is_empty()
205                && (line.trim().starts_with("struct") || line.trim().starts_with("enum"))
206                && line.trim().ends_with(" {")
207                && line
208                    .trim()
209                    .split_terminator(' ')
210                    .collect::<Vec<&str>>()
211                    .len()
212                    == 3
213                && (line_number - last_comment_line <= 1)
214            {
215                let mut key =
216                    line.trim().split_terminator(' ').collect::<Vec<&str>>()[1].to_string();
217                if line.contains("struct") {
218                    key = "struct ".to_string() + key.as_str() + " {";
219                } else if line.contains("enum") {
220                    key = "enum ".to_string() + key.as_str() + " {";
221                }
222                struct_map.insert(key.clone(), current_comment.clone());
223                current_comment.clear();
224            } else if !line.is_empty()
225                && !current_comment.is_empty()
226                && (line_number - last_comment_line > 1)
227            {
228                // Only clear the comment if we're far from the last comment
229                current_comment.clear();
230            }
231        }
232    }
233
234    println!("Processing complete");
235    struct_map
236}
237
238fn main() {
239    let args = Args::parse();
240
241    // Build the map of __bpf_kfunc declarations to their comments
242    let comments_map = build_bpf_kfunc_map(&args.kernel_dir);
243    println!(
244        "Found {} __bpf_kfunc declarations with comments",
245        comments_map.len()
246    );
247
248    // Build the map of structs and enums to their comments and definitions
249    let struct_map = build_kernel_struct_map(&args.kernel_dir);
250    println!("Found {} structs and enums with comments", struct_map.len());
251
252    // Annotate vmlinux.h with comments
253    match annotate_vmlinux(&args.vmlinux_h, &comments_map, &struct_map) {
254        Ok(annotated_content) => {
255            // Write the annotated output
256            if let Err(e) = fs::write(&args.output, annotated_content) {
257                eprintln!("Error writing annotated file: {}", e);
258                process::exit(1);
259            }
260            println!("Successfully wrote annotated vmlinux.h to {}", args.output);
261        }
262        Err(e) => {
263            eprintln!("Error processing vmlinux.h: {}", e);
264            process::exit(1);
265        }
266    }
267}
vmlinux_docify/main.rs

vmlinux_docify/
main.rs
