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