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