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