Experimental defmacro changes (#126)

* first pass for better defmacro support

crates/steel-core/src/compiler/code_gen.rs

@@ -542,7 +542,7 @@ impl<'a> VisitorMut for CodeGenerator<'a> {
                     // directly from the stack
                     self.push(
                         LabeledInstruction::builder(OpCode::COPYCAPTURESTACK)
-                            .payload(var.stack_offset.unwrap())
+                            .payload(var.stack_offset.ok_or_else(crate::throw!(Generic => format!("Error compiling this function - are you missing an expression after a local define?"); lambda_function.location.span))?)
                             .contents(SyntaxObject::default(TokenType::Identifier(*key))),
                     );
                 }

crates/steel-core/src/compiler/compiler.rs

@@ -6,13 +6,13 @@ use crate::{
         map::SymbolMap,
         passes::{
             analysis::SemanticAnalysis, begin::flatten_begins_and_expand_defines,
-            reader::MultipleArityFunctions, shadow::RenameShadowedVariables,
+            reader::MultipleArityFunctions, shadow::RenameShadowedVariables, VisitorMutRefUnit,
         },
     },
     core::labels::Expr,
     parser::{
         ast::AstTools,
-        expand_visitor::{expand_kernel, expand_kernel_in_env},
+        expand_visitor::{expand_kernel, expand_kernel_in_env, expand_kernel_in_env_with_change},
         interner::InternedString,
         kernel::Kernel,
         parser::{lower_entire_ast, lower_macro_and_require_definitions},
@@ -57,7 +57,7 @@ use crate::steel_vm::const_evaluation::ConstantEvaluatorManager;
 use super::{
     constants::SerializableConstantMap,
     modules::{CompiledModule, ModuleManager},
-    passes::analysis::Analysis,
+    passes::{analysis::Analysis, mangle::NameMangler},
     program::RawProgramWithSymbols,
 };
 
@@ -278,6 +278,13 @@ pub enum OptLevel {
     Three,
 }
 
+#[derive(Clone)]
+pub(crate) struct KernelDefMacroSpec {
+    pub(crate) env: String,
+    pub(crate) exported: Option<HashSet<InternedString>>,
+    pub(crate) name_mangler: NameMangler,
+}
+
 #[derive(Clone)]
 pub struct Compiler {
     pub(crate) symbol_map: SymbolMap,
@@ -288,6 +295,12 @@ pub struct Compiler {
     pub(crate) kernel: Option<Kernel>,
     memoization_table: MemoizationTable,
     mangled_identifiers: HashSet<InternedString>,
+    // Try this out?
+    lifted_kernel_environments: HashMap<String, KernelDefMacroSpec>,
+    // Macros that... we need to compile against directly at the top level
+    // This is really just a hack, but it solves cases for interactively
+    // running at the top level using private macros.
+    lifted_macro_environments: HashSet<PathBuf>,
 }
 
 #[derive(Serialize, Deserialize)]
@@ -352,6 +365,8 @@ impl Compiler {
             kernel: None,
             memoization_table: MemoizationTable::new(),
             mangled_identifiers: HashSet::new(),
+            lifted_kernel_environments: HashMap::new(),
+            lifted_macro_environments: HashSet::new(),
         }
     }
 
@@ -371,6 +386,8 @@ impl Compiler {
             kernel: Some(kernel),
             memoization_table: MemoizationTable::new(),
             mangled_identifiers: HashSet::new(),
+            lifted_kernel_environments: HashMap::new(),
+            lifted_macro_environments: HashSet::new(),
         }
     }
 
@@ -539,6 +556,8 @@ impl Compiler {
             exprs,
             path,
             builtin_modules,
+            &mut self.lifted_kernel_environments,
+            &mut self.lifted_macro_environments,
         );
 
         #[cfg(not(feature = "modules"))]
@@ -626,16 +645,46 @@ impl Compiler {
         expanded_statements = expanded_statements
             .into_iter()
             .map(|x| {
+                let mut x = x;
+
+                for module in &self.lifted_macro_environments {
+                    if let Some(macro_env) = self.modules().get(module).map(|x| &x.macro_map) {
+                        x = crate::parser::expand_visitor::expand(x, macro_env)?
+                    }
+                }
+
+                // Lift all of the kernel macros as well?
+                for (module, lifted_env) in &mut self.lifted_kernel_environments {
+                    let mut changed = false;
+
+                    (x, changed) = expand_kernel_in_env_with_change(
+                        x,
+                        self.kernel.as_mut(),
+                        builtin_modules.clone(),
+                        module.to_string(),
+                    )?;
+
+                    if changed {
+                        lifted_env.name_mangler.visit(&mut x);
+                    }
+                }
+
                 expand_kernel_in_env(
                     x,
                     self.kernel.as_mut(),
                     builtin_modules.clone(),
                     "top-level".to_string(),
                 )
+                // TODO: If we have this, then we have to lower all of the expressions again
                 .and_then(|x| crate::parser::expand_visitor::expand(x, &self.macro_env))
             })
             .collect::<Result<Vec<_>>>()?;
 
+        expanded_statements = expanded_statements
+            .into_iter()
+            .map(lower_entire_ast)
+            .collect::<std::result::Result<Vec<_>, ParseError>>()?;
+
         log::debug!(target: "expansion-phase", "Beginning constant folding");
 
         let mut expanded_statements =
@@ -750,7 +799,22 @@ impl Compiler {
 
         expanded_statements = expanded_statements
             .into_iter()
-            .map(|x| {
+            .map(|mut x| {
+                // // Expanded any / all lifted environments:
+                // let mut x = x;
+
+                // // Lift all of the kernel macros as well?
+                // for (module, lifted_env) in &self.lifted_kernel_environments {
+                //     dbg!(&module);
+
+                //     x = expand_kernel_in_env(
+                //         x,
+                //         self.kernel.as_mut(),
+                //         builtin_modules.clone(),
+                //         module.to_string(),
+                //     )?;
+                // }
+
                 expand_kernel_in_env(
                     x,
                     self.kernel.as_mut(),
@@ -769,6 +833,30 @@ impl Compiler {
         expanded_statements = expanded_statements
             .into_iter()
             .map(|x| {
+                let mut x = x;
+
+                for module in &self.lifted_macro_environments {
+                    if let Some(macro_env) = self.modules().get(module).map(|x| &x.macro_map) {
+                        x = crate::parser::expand_visitor::expand(x, macro_env)?
+                    }
+                }
+
+                // Lift all of the kernel macros as well?
+                for (module, lifted_env) in &mut self.lifted_kernel_environments {
+                    let mut changed = false;
+
+                    (x, changed) = expand_kernel_in_env_with_change(
+                        x,
+                        self.kernel.as_mut(),
+                        builtin_modules.clone(),
+                        module.to_string(),
+                    )?;
+
+                    if changed {
+                        lifted_env.name_mangler.visit(&mut x);
+                    }
+                }
+
                 expand_kernel_in_env(
                     x,
                     self.kernel.as_mut(),
@@ -780,7 +868,10 @@ impl Compiler {
             })
             .collect::<Result<Vec<_>>>()?;
 
-        // println!("{:#?}", expanded_statements);
+        expanded_statements = expanded_statements
+            .into_iter()
+            .map(lower_entire_ast)
+            .collect::<std::result::Result<Vec<_>, ParseError>>()?;
 
         // Let's do the @doc macro here?
 

crates/steel-core/src/compiler/modules.rs

@@ -29,6 +29,7 @@ use crate::{
 };
 use crate::{parser::expand_visitor::Expander, rvals::Result};
 
+use once_cell::sync::Lazy;
 use steel_parser::expr_list;
 
 use std::{
@@ -50,6 +51,7 @@ use log::{debug, info, log_enabled};
 use crate::parser::ast::IteratorExtensions;
 
 use super::{
+    compiler::KernelDefMacroSpec,
     passes::{
         begin::FlattenBegin,
         mangle::{collect_globals, NameMangler, NameUnMangler},
@@ -107,6 +109,8 @@ create_prelude!(
     for_syntax "#%private/steel/contract"
 );
 
+pub static STEEL_HOME: Lazy<Option<String>> = Lazy::new(|| std::env::var("STEEL_HOME").ok());
+
 /// Manages the modules
 /// keeps some visited state on the manager for traversal
 /// Also keeps track of the metadata for each file in order to determine
@@ -195,6 +199,8 @@ impl ModuleManager {
         exprs: Vec<ExprKind>,
         path: Option<PathBuf>,
         builtin_modules: ModuleContainer,
+        lifted_kernel_environments: &mut HashMap<String, KernelDefMacroSpec>,
+        lifted_macro_environments: &mut HashSet<PathBuf>,
     ) -> Result<Vec<ExprKind>> {
         // Wipe the visited set on entry
         self.visited.clear();
@@ -588,14 +594,47 @@ impl ModuleManager {
                     // }
 
                     if expander.changed || changed {
-                        expand(first_round_expanded, &module.macro_map)
+                        let fully_expanded = expand(first_round_expanded, &module.macro_map)?;
+
+                        let module_name = module.name.to_str().unwrap().to_string();
+
+                        // Expanding the kernel with only these macros...
+                        let (mut first_round_expanded, changed) = expand_kernel_in_env_with_change(
+                            fully_expanded,
+                            kernel.as_mut(),
+                            // We don't need to expand those here
+                            ModuleContainer::default(),
+                            module_name.clone(),
+                            // &kernel_macros_in_scope,
+                        )?;
+
+                        if changed {
+                            name_mangler.visit(&mut first_round_expanded);
+                        }
+
+                        lifted_kernel_environments.insert(
+                            module_name.clone(),
+                            KernelDefMacroSpec {
+                                env: module_name,
+                                exported: None,
+                                name_mangler: name_mangler.clone(),
+                            },
+                        );
+
+                        Ok(first_round_expanded)
                     } else {
                         Ok(first_round_expanded)
                     }
                 })
                 .collect::<Result<_>>()?;
 
+            // Global macro map - also need to expand with ALL macros
+            // post expansion in the target environment, which means we can't _just_
+            // extend the global macro map with the target in scope macros, we need to
+            // do something like the two pass expansion
             global_macro_map.extend(in_scope_macros);
+
+            lifted_macro_environments.insert(module.name.clone());
         }
 
         // Include the defines from the modules now imported
@@ -1762,7 +1801,36 @@ impl<'a> ModuleBuilder<'a> {
                     // }
 
                     if expander.changed || changed {
-                        expand(first_round_expanded, &module.macro_map)
+                        // expand(first_round_expanded, &module.macro_map)
+
+                        let fully_expanded = expand(first_round_expanded, &module.macro_map)?;
+
+                        let module_name = module.name.to_str().unwrap().to_string();
+
+                        // Expanding the kernel with only these macros...
+                        let (mut first_round_expanded, changed) = expand_kernel_in_env_with_change(
+                            fully_expanded,
+                            self.kernel.as_mut(),
+                            // We don't need to expand those here
+                            ModuleContainer::default(),
+                            module_name.clone(),
+                            // &kernel_macros_in_scope,
+                        )?;
+
+                        if changed {
+                            name_mangler.visit(&mut first_round_expanded);
+                        }
+
+                        // lifted_kernel_environments.insert(
+                        //     module_name.clone(),
+                        //     KernelDefMacroSpec {
+                        //         env: module_name,
+                        //         exported: None,
+                        //         name_mangler: name_mangler.clone(),
+                        //     },
+                        // );
+
+                        Ok(first_round_expanded)
                     } else {
                         Ok(first_round_expanded)
                     }
@@ -2192,13 +2260,10 @@ impl<'a> ModuleBuilder<'a> {
         let mut exprs_without_requires = Vec::new();
         let exprs = std::mem::take(&mut self.source_ast);
 
-        let home = std::env::var("STEEL_HOME")
-            .map(PathBuf::from)
-            .map(|mut x| {
-                x.push("cogs");
-                x
-            })
-            .ok();
+        let home = STEEL_HOME.clone().map(PathBuf::from).map(|mut x| {
+            x.push("cogs");
+            x
+        });
 
         fn walk(
             module_builder: &mut ModuleBuilder,

crates/steel-core/src/compiler/passes/mangle.rs

@@ -128,6 +128,7 @@ impl<'a> VisitorMutRefUnit for NameUnMangler<'a> {
     }
 }
 
+#[derive(Clone)]
 pub struct NameMangler {
     pub(crate) globals: HashSet<InternedString>,
     prefix: String,