1use std::fs::{self, File};
2use std::io::{Read, Seek, SeekFrom, Write};
3use std::path::{Path, PathBuf};
4
5#[cfg(any(feature = "deploy", feature = "maelstrom"))]
6use dfir_lang::diagnostic::Diagnostics;
7#[cfg(any(feature = "deploy", feature = "maelstrom"))]
8use dfir_lang::graph::DfirGraph;
9use sha2::{Digest, Sha256};
10#[cfg(any(feature = "deploy", feature = "maelstrom"))]
11use stageleft::internal::quote;
12#[cfg(any(feature = "deploy", feature = "maelstrom"))]
13use syn::visit_mut::VisitMut;
14use trybuild_internals_api::cargo::{self, Metadata};
15use trybuild_internals_api::env::Update;
16use trybuild_internals_api::run::{PathDependency, Project};
17use trybuild_internals_api::{Runner, dependencies, features, path};
18
19#[cfg(any(feature = "deploy", feature = "maelstrom"))]
20use super::rewriters::UseTestModeStaged;
21
22pub const HYDRO_RUNTIME_FEATURES: &[&str] = &[
23 "deploy_integration",
24 "runtime_measure",
25 "docker_runtime",
26 "ecs_runtime",
27 "maelstrom_runtime",
28];
29
30#[cfg(any(feature = "deploy", feature = "maelstrom"))]
31#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35pub enum LinkingMode {
36 Static,
37 #[cfg(feature = "deploy")]
38 Dynamic,
39}
40
41#[cfg(any(feature = "deploy", feature = "maelstrom"))]
42#[derive(Debug, Clone, Copy, PartialEq, Eq)]
44pub enum DeployMode {
45 #[cfg(feature = "deploy")]
46 HydroDeploy,
48 #[cfg(any(feature = "docker_deploy", feature = "ecs_deploy"))]
49 Containerized,
51 #[cfg(feature = "maelstrom")]
52 Maelstrom,
54}
55
56pub(crate) static IS_TEST: std::sync::atomic::AtomicBool =
57 std::sync::atomic::AtomicBool::new(false);
58
59pub(crate) static CONCURRENT_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
60
61pub fn init_test() {
75 IS_TEST.store(true, std::sync::atomic::Ordering::Relaxed);
76}
77
78#[cfg(any(feature = "deploy", feature = "maelstrom"))]
79fn clean_bin_name_prefix(bin_name_prefix: &str) -> String {
80 bin_name_prefix
81 .replace("::", "__")
82 .replace(" ", "_")
83 .replace(",", "_")
84 .replace("<", "_")
85 .replace(">", "")
86 .replace("(", "")
87 .replace(")", "")
88}
89
90#[derive(Debug, Clone)]
91pub struct TrybuildConfig {
92 pub project_dir: PathBuf,
93 pub target_dir: PathBuf,
94 pub features: Option<Vec<String>>,
95 #[cfg(feature = "deploy")]
96 pub linking_mode: LinkingMode,
100}
101
102#[cfg(any(feature = "deploy", feature = "maelstrom"))]
103pub fn create_graph_trybuild(
104 graph: DfirGraph,
105 extra_stmts: &[syn::Stmt],
106 sidecars: &[syn::Expr],
107 bin_name_prefix: Option<&str>,
108 deploy_mode: DeployMode,
109 linking_mode: LinkingMode,
110) -> (String, TrybuildConfig) {
111 let source_dir = cargo::manifest_dir().unwrap();
112 let source_manifest = dependencies::get_manifest(&source_dir).unwrap();
113 let crate_name = source_manifest.package.name.replace("-", "_");
114
115 let is_test = IS_TEST.load(std::sync::atomic::Ordering::Relaxed);
116
117 let generated_code = compile_graph_trybuild(
118 graph,
119 extra_stmts,
120 sidecars,
121 &crate_name,
122 is_test,
123 deploy_mode,
124 );
125
126 let inlined_staged = if is_test {
127 let raw_toml_manifest = toml::from_str::<toml::Value>(
128 &fs::read_to_string(path!(source_dir / "Cargo.toml")).unwrap(),
129 )
130 .unwrap();
131
132 let maybe_custom_lib_path = raw_toml_manifest
133 .get("lib")
134 .and_then(|lib| lib.get("path"))
135 .and_then(|path| path.as_str());
136
137 let mut gen_staged = stageleft_tool::gen_staged_trybuild(
138 &maybe_custom_lib_path
139 .map(|s| path!(source_dir / s))
140 .unwrap_or_else(|| path!(source_dir / "src" / "lib.rs")),
141 &path!(source_dir / "Cargo.toml"),
142 &crate_name,
143 Some("hydro___test".to_owned()),
144 );
145
146 gen_staged.attrs.insert(
147 0,
148 syn::parse_quote! {
149 #![allow(
150 unused,
151 ambiguous_glob_reexports,
152 clippy::suspicious_else_formatting,
153 unexpected_cfgs,
154 reason = "generated code"
155 )]
156 },
157 );
158
159 Some(prettyplease::unparse(&gen_staged))
160 } else {
161 None
162 };
163
164 let source = prettyplease::unparse(&generated_code);
165
166 let hash = format!("{:X}", Sha256::digest(&source))
167 .chars()
168 .take(8)
169 .collect::<String>();
170
171 let bin_name = if let Some(bin_name_prefix) = &bin_name_prefix {
172 format!("{}_{}", clean_bin_name_prefix(bin_name_prefix), &hash)
173 } else {
174 hash
175 };
176
177 let (project_dir, target_dir, mut cur_bin_enabled_features) = create_trybuild().unwrap();
178
179 let examples_dir = match linking_mode {
181 LinkingMode::Static => path!(project_dir / "examples"),
182 #[cfg(feature = "deploy")]
183 LinkingMode::Dynamic => path!(project_dir / "dylib-examples" / "examples"),
184 };
185
186 fs::create_dir_all(&examples_dir).unwrap();
188
189 let out_path = path!(examples_dir / format!("{bin_name}.rs"));
190 {
191 let _concurrent_test_lock = CONCURRENT_TEST_LOCK.lock().unwrap();
192 write_atomic(source.as_ref(), &out_path).unwrap();
193 }
194
195 if let Some(inlined_staged) = inlined_staged {
196 let staged_path = path!(project_dir / "src" / "__staged.rs");
197 {
198 let _concurrent_test_lock = CONCURRENT_TEST_LOCK.lock().unwrap();
199 write_atomic(inlined_staged.as_bytes(), &staged_path).unwrap();
200 }
201 }
202
203 if is_test {
204 if cur_bin_enabled_features.is_none() {
205 cur_bin_enabled_features = Some(vec![]);
206 }
207
208 cur_bin_enabled_features
209 .as_mut()
210 .unwrap()
211 .push("hydro___test".to_owned());
212 }
213
214 (
215 bin_name,
216 TrybuildConfig {
217 project_dir,
218 target_dir,
219 features: cur_bin_enabled_features,
220 #[cfg(feature = "deploy")]
221 linking_mode,
222 },
223 )
224}
225
226#[cfg(any(feature = "deploy", feature = "maelstrom"))]
227pub fn compile_graph_trybuild(
228 partitioned_graph: DfirGraph,
229 extra_stmts: &[syn::Stmt],
230 sidecars: &[syn::Expr],
231 crate_name: &str,
232 is_test: bool,
233 deploy_mode: DeployMode,
234) -> syn::File {
235 use crate::staging_util::get_this_crate;
236
237 let mut diagnostics = Diagnostics::new();
238 let mut dfir_expr: syn::Expr = syn::parse2(
239 partitioned_graph
240 .as_code("e! { __root_dfir_rs }, true, quote!(), &mut diagnostics)
241 .expect("DFIR code generation failed with diagnostics."),
242 )
243 .unwrap();
244
245 if is_test {
246 UseTestModeStaged { crate_name }.visit_expr_mut(&mut dfir_expr);
247 }
248
249 let orig_crate_name = quote::format_ident!("{}", crate_name);
250 let trybuild_crate_name_ident = quote::format_ident!("{}_hydro_trybuild", crate_name);
251 let root = get_this_crate();
252 let tokio_main_ident = format!("{}::runtime_support::tokio", root);
253 let dfir_ident = quote::format_ident!("{}", crate::compile::DFIR_IDENT);
254
255 let source_ast: syn::File = match deploy_mode {
256 #[cfg(any(feature = "docker_deploy", feature = "ecs_deploy"))]
257 DeployMode::Containerized => {
258 syn::parse_quote! {
259 #![allow(unused_imports, unused_crate_dependencies, missing_docs, non_snake_case)]
260 use #trybuild_crate_name_ident::__root as #orig_crate_name;
261 use #trybuild_crate_name_ident::__staged::__deps::*;
262 use #root::prelude::*;
263 use #root::runtime_support::dfir_rs as __root_dfir_rs;
264 pub use #trybuild_crate_name_ident::__staged;
265
266 #[allow(unused)]
267 async fn __hydro_runtime<'a>() -> #root::runtime_support::dfir_rs::scheduled::graph::Dfir<'a> {
268 #( #extra_stmts )*
270
271 #dfir_expr
273 }
274
275 #[#root::runtime_support::tokio::main(crate = #tokio_main_ident, flavor = "current_thread")]
276 async fn main() {
277 #root::telemetry::initialize_tracing();
278
279 let mut #dfir_ident = __hydro_runtime().await;
280
281 let local_set = #root::runtime_support::tokio::task::LocalSet::new();
282 #(
283 let _ = local_set.spawn_local( #sidecars ); )*
285
286 let _ = local_set.run_until(#dfir_ident.run()).await;
287 }
288 }
289 }
290 #[cfg(feature = "deploy")]
291 DeployMode::HydroDeploy => {
292 syn::parse_quote! {
293 #![allow(unused_imports, unused_crate_dependencies, missing_docs, non_snake_case)]
294 use #trybuild_crate_name_ident::__root as #orig_crate_name;
295 use #trybuild_crate_name_ident::__staged::__deps::*;
296 use #root::prelude::*;
297 use #root::runtime_support::dfir_rs as __root_dfir_rs;
298 pub use #trybuild_crate_name_ident::__staged;
299
300 #[allow(unused)]
301 fn __hydro_runtime<'a>(
302 __hydro_lang_trybuild_cli: &'a #root::runtime_support::hydro_deploy_integration::DeployPorts<#root::__staged::deploy::deploy_runtime::HydroMeta>
303 )
304 -> #root::runtime_support::dfir_rs::scheduled::graph::Dfir<'a>
305 {
306 #( #extra_stmts )*
307
308 #dfir_expr
309 }
310
311 #[#root::runtime_support::tokio::main(crate = #tokio_main_ident, flavor = "current_thread")]
312 async fn main() {
313 let ports = #root::runtime_support::launch::init_no_ack_start().await;
314 let #dfir_ident = __hydro_runtime(&ports);
315 println!("ack start");
316
317 let local_set = #root::runtime_support::tokio::task::LocalSet::new();
321 #(
322 let _ = local_set.spawn_local( #sidecars ); )*
324
325 let _ = local_set.run_until(#root::runtime_support::launch::run_stdin_commands(#dfir_ident)).await;
326 }
327 }
328 }
329 #[cfg(feature = "maelstrom")]
330 DeployMode::Maelstrom => {
331 syn::parse_quote! {
332 #![allow(unused_imports, unused_crate_dependencies, missing_docs, non_snake_case)]
333 use #trybuild_crate_name_ident::__root as #orig_crate_name;
334 use #trybuild_crate_name_ident::__staged::__deps::*;
335 use #root::prelude::*;
336 use #root::runtime_support::dfir_rs as __root_dfir_rs;
337 pub use #trybuild_crate_name_ident::__staged;
338
339 #[allow(unused)]
340 fn __hydro_runtime<'a>(
341 __hydro_lang_maelstrom_meta: &'a #root::__staged::deploy::maelstrom::deploy_runtime_maelstrom::MaelstromMeta
342 )
343 -> #root::runtime_support::dfir_rs::scheduled::graph::Dfir<'a>
344 {
345 #( #extra_stmts )*
346
347 #dfir_expr
348 }
349
350 #[#root::runtime_support::tokio::main(crate = #tokio_main_ident, flavor = "current_thread")]
351 async fn main() {
352 #root::telemetry::initialize_tracing();
353
354 let __hydro_lang_maelstrom_meta = #root::__staged::deploy::maelstrom::deploy_runtime_maelstrom::maelstrom_init();
356
357 let mut #dfir_ident = __hydro_runtime(&__hydro_lang_maelstrom_meta);
358
359 __hydro_lang_maelstrom_meta.start_receiving(); let local_set = #root::runtime_support::tokio::task::LocalSet::new();
362 #(
363 let _ = local_set.spawn_local( #sidecars ); )*
365
366 let _ = local_set.run_until(#dfir_ident.run()).await;
367 }
368 }
369 }
370 };
371 source_ast
372}
373
374pub fn create_trybuild()
375-> Result<(PathBuf, PathBuf, Option<Vec<String>>), trybuild_internals_api::error::Error> {
376 let Metadata {
377 target_directory: target_dir,
378 workspace_root: workspace,
379 packages,
380 } = cargo::metadata()?;
381
382 let source_dir = cargo::manifest_dir()?;
383 let mut source_manifest = dependencies::get_manifest(&source_dir)?;
384
385 let mut dev_dependency_features = vec![];
386 source_manifest.dev_dependencies.retain(|k, v| {
387 if source_manifest.dependencies.contains_key(k) {
388 for feat in &v.features {
390 dev_dependency_features.push(format!("{}/{}", k, feat));
391 }
392
393 false
394 } else {
395 dev_dependency_features.push(format!("dep:{k}"));
397
398 v.optional = true;
399 true
400 }
401 });
402
403 let mut features = features::find();
404
405 let path_dependencies = source_manifest
406 .dependencies
407 .iter()
408 .filter_map(|(name, dep)| {
409 let path = dep.path.as_ref()?;
410 if packages.iter().any(|p| &p.name == name) {
411 None
413 } else {
414 Some(PathDependency {
415 name: name.clone(),
416 normalized_path: path.canonicalize().ok()?,
417 })
418 }
419 })
420 .collect();
421
422 let crate_name = source_manifest.package.name.clone();
423 let project_dir = path!(target_dir / "hydro_trybuild" / crate_name /);
424 fs::create_dir_all(&project_dir)?;
425
426 let project_name = format!("{}-hydro-trybuild", crate_name);
427 let mut manifest = Runner::make_manifest(
428 &workspace,
429 &project_name,
430 &source_dir,
431 &packages,
432 &[],
433 source_manifest,
434 )?;
435
436 if let Some(enabled_features) = &mut features {
437 enabled_features
438 .retain(|feature| manifest.features.contains_key(feature) || feature == "default");
439 }
440
441 for runtime_feature in HYDRO_RUNTIME_FEATURES {
442 manifest.features.insert(
443 format!("hydro___feature_{runtime_feature}"),
444 vec![format!("hydro_lang/{runtime_feature}")],
445 );
446 }
447
448 manifest
449 .dependencies
450 .get_mut("hydro_lang")
451 .unwrap()
452 .features
453 .push("runtime_support".to_owned());
454
455 manifest
456 .features
457 .insert("hydro___test".to_owned(), dev_dependency_features);
458
459 if manifest
460 .workspace
461 .as_ref()
462 .is_some_and(|w| w.dependencies.is_empty())
463 {
464 manifest.workspace = None;
465 }
466
467 let project = Project {
468 dir: project_dir,
469 source_dir,
470 target_dir,
471 name: project_name.clone(),
472 update: Update::env()?,
473 has_pass: false,
474 has_compile_fail: false,
475 features,
476 workspace,
477 path_dependencies,
478 manifest,
479 keep_going: false,
480 };
481
482 {
483 let _concurrent_test_lock = CONCURRENT_TEST_LOCK.lock().unwrap();
484
485 let project_lock = File::create(path!(project.dir / ".hydro-trybuild-lock"))?;
486 project_lock.lock()?;
487
488 fs::create_dir_all(path!(project.dir / "src"))?;
489 fs::create_dir_all(path!(project.dir / "examples"))?;
490
491 let crate_name_ident = syn::Ident::new(
492 &crate_name.replace("-", "_"),
493 proc_macro2::Span::call_site(),
494 );
495
496 write_atomic(
497 prettyplease::unparse(&syn::parse_quote! {
498 #![allow(unused_imports, unused_crate_dependencies, missing_docs, non_snake_case)]
499
500 pub use #crate_name_ident as __root;
501
502 #[cfg(feature = "hydro___test")]
503 pub mod __staged;
504
505 #[cfg(not(feature = "hydro___test"))]
506 pub use #crate_name_ident::__staged;
507 })
508 .as_bytes(),
509 &path!(project.dir / "src" / "lib.rs"),
510 )
511 .unwrap();
512
513 let base_manifest = toml::to_string(&project.manifest)?;
514
515 let feature_names: Vec<_> = project.manifest.features.keys().cloned().collect();
517
518 let dylib_dir = path!(project.dir / "dylib");
520 fs::create_dir_all(path!(dylib_dir / "src"))?;
521
522 let trybuild_crate_name_ident = syn::Ident::new(
523 &project_name.replace("-", "_"),
524 proc_macro2::Span::call_site(),
525 );
526 write_atomic(
527 prettyplease::unparse(&syn::parse_quote! {
528 #![allow(unused_imports, unused_crate_dependencies, missing_docs, non_snake_case)]
529 pub use #trybuild_crate_name_ident::*;
530 })
531 .as_bytes(),
532 &path!(dylib_dir / "src" / "lib.rs"),
533 )?;
534
535 let serialized_edition = toml::to_string(
536 &vec![("edition", &project.manifest.package.edition)]
537 .into_iter()
538 .collect::<std::collections::HashMap<_, _>>(),
539 )
540 .unwrap();
541
542 let dylib_manifest = format!(
546 r#"[package]
547name = "{project_name}-dylib"
548version = "0.0.0"
549{}
550
551[lib]
552crate-type = ["{}"]
553
554[dependencies]
555{project_name} = {{ path = "..", default-features = false }}
556"#,
557 serialized_edition,
558 if cfg!(target_os = "windows") {
559 "rlib"
560 } else {
561 "dylib"
562 }
563 );
564 write_atomic(dylib_manifest.as_ref(), &path!(dylib_dir / "Cargo.toml"))?;
565
566 let dylib_examples_dir = path!(project.dir / "dylib-examples");
567 fs::create_dir_all(path!(dylib_examples_dir / "src"))?;
568 fs::create_dir_all(path!(dylib_examples_dir / "examples"))?;
569
570 write_atomic(
571 b"#![allow(unused_crate_dependencies)]\n",
572 &path!(dylib_examples_dir / "src" / "lib.rs"),
573 )?;
574
575 let features_section = feature_names
577 .iter()
578 .map(|f| format!("{f} = [\"{project_name}/{f}\"]"))
579 .collect::<Vec<_>>()
580 .join("\n");
581
582 let dylib_examples_manifest = format!(
584 r#"[package]
585name = "{project_name}-dylib-examples"
586version = "0.0.0"
587{}
588
589[dev-dependencies]
590{project_name} = {{ path = "..", default-features = false }}
591{project_name}-dylib = {{ path = "../dylib", default-features = false }}
592
593[features]
594{features_section}
595
596[[example]]
597name = "sim-dylib"
598crate-type = ["cdylib"]
599"#,
600 serialized_edition
601 );
602 write_atomic(
603 dylib_examples_manifest.as_ref(),
604 &path!(dylib_examples_dir / "Cargo.toml"),
605 )?;
606
607 let sim_dylib_contents = prettyplease::unparse(&syn::parse_quote! {
609 #![allow(unused_imports, unused_crate_dependencies, missing_docs, non_snake_case)]
610 include!(std::concat!(env!("TRYBUILD_LIB_NAME"), ".rs"));
611 });
612 write_atomic(
613 sim_dylib_contents.as_bytes(),
614 &path!(project.dir / "examples" / "sim-dylib.rs"),
615 )?;
616 write_atomic(
617 sim_dylib_contents.as_bytes(),
618 &path!(dylib_examples_dir / "examples" / "sim-dylib.rs"),
619 )?;
620
621 let workspace_manifest = format!(
622 r#"{}
623[[example]]
624name = "sim-dylib"
625crate-type = ["cdylib"]
626
627[workspace]
628members = ["dylib", "dylib-examples"]
629"#,
630 base_manifest,
631 );
632
633 write_atomic(
634 workspace_manifest.as_ref(),
635 &path!(project.dir / "Cargo.toml"),
636 )?;
637
638 let manifest_hash = format!("{:X}", Sha256::digest(&workspace_manifest))
640 .chars()
641 .take(8)
642 .collect::<String>();
643
644 let workspace_cargo_lock = path!(project.workspace / "Cargo.lock");
645 let workspace_cargo_lock_contents_and_hash = if workspace_cargo_lock.exists() {
646 let cargo_lock_contents = fs::read_to_string(&workspace_cargo_lock)?;
647
648 let hash = format!("{:X}", Sha256::digest(&cargo_lock_contents))
649 .chars()
650 .take(8)
651 .collect::<String>();
652
653 Some((cargo_lock_contents, hash))
654 } else {
655 None
656 };
657
658 let trybuild_hash = format!(
659 "{}-{}",
660 manifest_hash,
661 workspace_cargo_lock_contents_and_hash
662 .as_ref()
663 .map(|(_contents, hash)| &**hash)
664 .unwrap_or_default()
665 );
666
667 if !check_contents(
668 trybuild_hash.as_bytes(),
669 &path!(project.dir / ".hydro-trybuild-manifest"),
670 )
671 .is_ok_and(|b| b)
672 {
673 if let Some((cargo_lock_contents, _)) = workspace_cargo_lock_contents_and_hash {
675 write_atomic(
678 cargo_lock_contents.as_ref(),
679 &path!(project.dir / "Cargo.lock"),
680 )?;
681 } else {
682 let _ = cargo::cargo(&project).arg("generate-lockfile").status();
683 }
684
685 std::process::Command::new("cargo")
687 .current_dir(&project.dir)
688 .args(["update", "-w"]) .stdout(std::process::Stdio::null())
690 .stderr(std::process::Stdio::null())
691 .status()
692 .unwrap();
693
694 write_atomic(
695 trybuild_hash.as_bytes(),
696 &path!(project.dir / ".hydro-trybuild-manifest"),
697 )?;
698 }
699
700 let examples_folder = path!(project.dir / "examples");
702 fs::create_dir_all(&examples_folder)?;
703
704 let workspace_dot_cargo_config_toml = path!(project.workspace / ".cargo" / "config.toml");
705 if workspace_dot_cargo_config_toml.exists() {
706 let dot_cargo_folder = path!(project.dir / ".cargo");
707 fs::create_dir_all(&dot_cargo_folder)?;
708
709 write_atomic(
710 fs::read_to_string(&workspace_dot_cargo_config_toml)?.as_ref(),
711 &path!(dot_cargo_folder / "config.toml"),
712 )?;
713 }
714
715 let vscode_folder = path!(project.dir / ".vscode");
716 fs::create_dir_all(&vscode_folder)?;
717 write_atomic(
718 include_bytes!("./vscode-trybuild.json"),
719 &path!(vscode_folder / "settings.json"),
720 )?;
721 }
722
723 Ok((
724 project.dir.as_ref().into(),
725 project.target_dir.as_ref().into(),
726 project.features,
727 ))
728}
729
730fn check_contents(contents: &[u8], path: &Path) -> Result<bool, std::io::Error> {
731 let mut file = File::options()
732 .read(true)
733 .write(false)
734 .create(false)
735 .truncate(false)
736 .open(path)?;
737 file.lock()?;
738
739 let mut existing_contents = Vec::new();
740 file.read_to_end(&mut existing_contents)?;
741 Ok(existing_contents == contents)
742}
743
744pub(crate) fn write_atomic(contents: &[u8], path: &Path) -> Result<(), std::io::Error> {
745 let mut file = File::options()
746 .read(true)
747 .write(true)
748 .create(true)
749 .truncate(false)
750 .open(path)?;
751
752 let mut existing_contents = Vec::new();
753 file.read_to_end(&mut existing_contents)?;
754 if existing_contents != contents {
755 file.lock()?;
756 file.seek(SeekFrom::Start(0))?;
757 file.set_len(0)?;
758 file.write_all(contents)?;
759 }
760
761 Ok(())
762}