From 935255d06b1e1ac1a1818ae9b278bc47dc277f48 Mon Sep 17 00:00:00 2001
From: Ayaz Hafiz <ayaz@rwx.com>
Date: Wed, 7 Dec 2022 10:48:41 -0600
Subject: [PATCH] Support compiling programs with test mode in test_mono

---
 crates/compiler/test_mono/src/tests.rs      | 35 ++++++++++++---------
 crates/compiler/test_mono_macros/src/lib.rs | 26 ++++++++++++---
 2 files changed, 41 insertions(+), 20 deletions(-)

diff --git a/crates/compiler/test_mono/src/tests.rs b/crates/compiler/test_mono/src/tests.rs
index 9793a85052..fc6952ebee 100644
--- a/crates/compiler/test_mono/src/tests.rs
+++ b/crates/compiler/test_mono/src/tests.rs
@@ -76,10 +76,16 @@ fn promote_expr_to_module(src: &str) -> String {
     buffer
 }
 
-fn compiles_to_ir(test_name: &str, src: &str, no_check: bool) {
+fn compiles_to_ir(test_name: &str, src: &str, mode: &str, no_check: bool) {
     use roc_packaging::cache::RocCacheDir;
     use std::path::PathBuf;
 
+    let exec_mode = match mode {
+        "exec" => ExecutionMode::Executable,
+        "test" => ExecutionMode::Test,
+        _ => panic!("Invalid test_mono exec mode {mode}"),
+    };
+
     let arena = &Bump::new();
 
     let filename = PathBuf::from("Test.roc");
@@ -87,7 +93,7 @@ fn compiles_to_ir(test_name: &str, src: &str, no_check: bool) {
 
     let module_src;
     let temp;
-    if src.starts_with("app") {
+    if src.starts_with("app") || src.starts_with("interface") {
         // this is already a module
         module_src = src;
     } else {
@@ -101,7 +107,7 @@ fn compiles_to_ir(test_name: &str, src: &str, no_check: bool) {
         threading: Threading::Single,
         render: roc_reporting::report::RenderTarget::Generic,
         palette: roc_reporting::report::DEFAULT_PALETTE,
-        exec_mode: ExecutionMode::Executable,
+        exec_mode,
     };
     let loaded = roc_load::load_and_monomorphize_from_str(
         arena,
@@ -143,9 +149,7 @@ fn compiles_to_ir(test_name: &str, src: &str, no_check: bool) {
 
     assert!(type_problems.is_empty());
 
-    debug_assert_eq!(exposed_to_host.values.len(), 1);
-
-    let main_fn_symbol = exposed_to_host.values.keys().copied().next().unwrap();
+    let main_fn_symbol = exposed_to_host.values.keys().copied().next();
 
     if !no_check {
         check_procedures(arena, &interns, &layout_interner, &procedures);
@@ -173,22 +177,23 @@ fn verify_procedures<'a>(
     test_name: &str,
     interner: STLayoutInterner<'a>,
     procedures: MutMap<(Symbol, ProcLayout<'a>), Proc<'a>>,
-    main_fn_symbol: Symbol,
+    opt_main_fn_symbol: Option<Symbol>,
 ) {
-    let index = procedures
-        .keys()
-        .position(|(s, _)| *s == main_fn_symbol)
-        .unwrap();
-
     let mut procs_string = procedures
         .values()
         .map(|proc| proc.to_pretty(&interner, 200, false))
         .collect::<Vec<_>>();
 
-    let main_fn = procs_string.swap_remove(index);
-
     procs_string.sort();
-    procs_string.push(main_fn);
+
+    if let Some(main_fn_symbol) = opt_main_fn_symbol {
+        let index = procedures
+            .keys()
+            .position(|(s, _)| *s == main_fn_symbol)
+            .unwrap();
+        let main_fn = procs_string.swap_remove(index);
+        procs_string.push(main_fn);
+    }
 
     let result = procs_string.join("\n");
 
diff --git a/crates/compiler/test_mono_macros/src/lib.rs b/crates/compiler/test_mono_macros/src/lib.rs
index 0f373bd06e..bb7d8845df 100644
--- a/crates/compiler/test_mono_macros/src/lib.rs
+++ b/crates/compiler/test_mono_macros/src/lib.rs
@@ -1,14 +1,30 @@
 //! Macros for use in `test_mono`.
 extern crate proc_macro;
 
-use proc_macro::{TokenStream, TokenTree};
+use proc_macro::TokenStream;
 use quote::quote;
 
 #[proc_macro_attribute]
 pub fn mono_test(args: TokenStream, item: TokenStream) -> TokenStream {
-    let no_check = args
-        .into_iter()
-        .any(|tok| matches!(tok, TokenTree::Ident(id) if id.to_string() == "no_check"));
+    let mut no_check = false;
+    let mut mode = "exec".to_owned();
+    for arg in syn::parse_macro_input!(args as syn::AttributeArgs) {
+        use syn::{Lit, Meta, MetaNameValue, NestedMeta};
+        if matches!(&arg, NestedMeta::Meta(Meta::Path(p)) if p.is_ident("no_check")) {
+            no_check = true;
+        }
+        if let NestedMeta::Meta(Meta::NameValue(MetaNameValue {
+            path,
+            eq_token: _,
+            lit: Lit::Str(s),
+        })) = arg
+        {
+            if path.is_ident("mode") {
+                mode = s.value();
+            }
+        }
+    }
+
     let task_fn = syn::parse_macro_input!(item as syn::ItemFn);
 
     let args = task_fn.sig.inputs.clone();
@@ -24,7 +40,7 @@ pub fn mono_test(args: TokenStream, item: TokenStream) -> TokenStream {
         #[test]
         #(#attributes)*
         #visibility fn #name(#args) {
-            compiles_to_ir(#name_str, #body, #no_check);
+            compiles_to_ir(#name_str, #body, &#mode, #no_check);
 
         }
     };