-
Notifications
You must be signed in to change notification settings - Fork 76
/
main.rs
178 lines (158 loc) · 5.54 KB
/
main.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
use std::{
collections::{
HashMap,
HashSet,
},
env,
ffi::OsStr,
fs::{
read_dir,
read_to_string,
remove_file,
},
io::Write as _,
path::{
Path,
PathBuf,
},
process::Command,
};
const OUT_DIR: &str = "../../crates/astria-core/src/generated";
const SRC_DIR: &str = "../../proto";
const INCLUDES: &[&str] = &[SRC_DIR];
fn main() {
let buf = get_buf_from_env();
let mut cmd = Command::new(buf.clone());
let buf_img = tempfile::NamedTempFile::new()
.expect("should be able to create a temp file to hold the buf image file descriptor set");
let crate_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
let src_dir = crate_dir.join(SRC_DIR);
let out_dir = crate_dir.join(OUT_DIR);
cmd.arg("build")
.arg("--output")
.arg(buf_img.path())
.arg("--as-file-descriptor-set")
.arg(&src_dir);
let buf_output = match cmd.output() {
Err(e) => {
panic!(
"failed creating file descriptor set from protobuf: failed to invoke buf (path: \
{buf:?}): {e:?}"
);
}
Ok(output) => output,
};
emit_buf_stdout(&buf_output.stdout).expect("able to write to stdout");
emit_buf_stderr(&buf_output.stderr).expect("able to write to stderr");
assert!(
buf_output.status.success(),
"failed creating file descriptor set from protobuf: `buf` returned non-zero exit code"
);
let files = find_protos(&src_dir);
tonic_build::configure()
.build_client(true)
.build_server(true)
.emit_rerun_if_changed(false)
.client_mod_attribute(".", "#[cfg(feature=\"client\")]")
.server_mod_attribute(".", "#[cfg(feature=\"server\")]")
.extern_path(".tendermint.abci", "::tendermint-proto::abci")
.extern_path(".tendermint.crypto", "::tendermint-proto::crypto")
.extern_path(".tendermint.version", "::tendermint-proto::version")
.extern_path(".tendermint.types", "::tendermint-proto::types")
.extern_path(".penumbra", "::penumbra-proto")
.type_attribute(".astria.primitive.v1.Uint128", "#[derive(Copy)]")
.out_dir(&out_dir)
.file_descriptor_set_path(buf_img.path())
.skip_protoc_run()
.compile(&files, INCLUDES)
.expect("should be able to compile protobuf using tonic");
let mut after_build = build_content_map(&out_dir);
clean_non_astria_code(&mut after_build);
}
fn emit_buf_stdout(buf: &[u8]) -> std::io::Result<()> {
if !buf.is_empty() {
std::io::stdout().lock().write_all(buf)?;
println!();
}
Ok(())
}
fn emit_buf_stderr(buf: &[u8]) -> std::io::Result<()> {
if !buf.is_empty() {
std::io::stderr().lock().write_all(buf)?;
eprintln!();
}
Ok(())
}
fn clean_non_astria_code(generated: &mut ContentMap) {
let mut foreign_file_names: HashSet<_> = generated
.files
.keys()
.filter(|name| !name.starts_with("astria."))
.cloned()
.collect();
// also mask mod.rs because we need are defining it
foreign_file_names.remove("mod.rs");
for name in foreign_file_names {
let _ = generated.codes.remove(&name);
let file = generated
.files
.remove(&name)
.expect("file should exist under the name");
let _ = remove_file(file);
}
}
fn find_protos<P: AsRef<Path>>(dir: P) -> Vec<PathBuf> {
use walkdir::{
DirEntry,
WalkDir,
};
WalkDir::new(dir)
.into_iter()
.filter_map(Result::ok)
.filter(|e| e.file_type().is_file() && e.path().extension() == Some(OsStr::new("proto")))
.map(DirEntry::into_path)
.collect()
}
struct ContentMap {
files: HashMap<String, PathBuf>,
codes: HashMap<String, String>,
}
fn build_content_map(path: impl AsRef<Path>) -> ContentMap {
let mut files = HashMap::new();
let mut codes = HashMap::new();
for entry in read_dir(path)
.expect("should be able to read target folder for generated files")
.flatten()
{
let path = entry.path();
let name = path
.file_name()
.expect("generated file should have a file name")
.to_string_lossy()
.to_string();
let contents = read_to_string(&path)
.expect("should be able to read the contents of an existing generated file");
files.insert(name.clone(), path);
codes.insert(name.clone(), contents);
}
ContentMap {
files,
codes,
}
}
fn get_buf_from_env() -> PathBuf {
let os_specific_hint = match env::consts::OS {
"macos" => "You could try running `brew install buf` or downloading a recent release from https://github.com/bufbuild/buf/releases",
"linux" => "You can download it from https://github.com/bufbuild/buf/releases; if you are on Arch Linux, install it from the AUR with `rua install buf` or another helper",
_other => "Check if there is a precompiled version for your OS at https://github.com/bufbuild/buf/releases"
};
let error_msg = "Could not find `buf` installation and this build crate cannot proceed without
this knowledge. If `buf` is installed and this crate had trouble finding
it, you can set the `BUF` environment variable with the specific path to your
installed `buf` binary.";
let msg = format!("{error_msg} {os_specific_hint}");
env::var_os("BUF")
.map(PathBuf::from)
.or_else(|| which::which("buf").ok())
.expect(&msg)
}