Defer evaluation of block arguments

When a block value is evaluated, the address tracker is immediately
incremented for each word in the evaluated block. This was causing an
issue when a block value was passed as an argument to a macro
invocation, because argument values would all be evaluated before the
macro body was evaluated. This meant that the value of every label in
the macro body would be pushed forward by the size of the block
argument, even if the block argument is inserted later in the macro.
Furthermore, the address tracker would be incremented exactly once
for each word in the block, even if the block is never inserted into
the macro or is inserted multiple times.

To fix this, we now defer evaluation of a block argument until exactly
when that block argument is inserted into a macro. Because a block
argument could be an invocation with a set of arguments passed to it,
we also store a copy of the current environment alongside the block
value so that it can be evaluated correctly at any time in the future.

2 files changed, 109 insertions, 62 deletions

diff --git a/src/stages/intermediate.rs b/src/stages/intermediate.rs
index b890cd4..3144b77 100644
--- a/src/stages/intermediate.rs
+++ b/src/stages/intermediate.rs
@@ -14,29 +14,6 @@ macro_rules! next_id {
     };
 }
 
-pub struct LabelAddress {
-    pub previous: usize,
-    pub current: usize,
-    pub touched: bool,
-}
-
-#[derive(Eq, Hash, PartialEq, Clone)]
-pub struct SymbolSignature {
-    pub name: String,
-    pub arg_count: usize,
-}
-
-impl std::fmt::Display for SymbolSignature {
-    fn fmt(&self, f: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
-        write!(f, "{}::{}", self.name, self.arg_count)
-    }
-}
-
-struct Environment {
-    arguments: IndexMap<String, Tracked<IntermediateValue>>,
-    id: usize,
-}
-
 
 pub fn parse_intermediate(semantic: Vec<Tracked<SemanticToken>>) -> Result<Vec<Tracked<IntermediateToken>>, Vec<Tracked<IntermediateError>>> {
     IntermediateParser::new(semantic).parse()
@@ -294,14 +271,14 @@ impl IntermediateParser {
     fn parse_integer_invocation(&mut self, invocation: &Invocation, source: &SourceSpan) -> Option<Tracked<IntermediateInteger>> {
         let result = self.parse_invocation(invocation, source)?;
         match result.value {
-            IntermediateValue::Integer(integer) => {
+            DeferredValue::Integer(integer) => {
                 Some(Tracked::from(integer, source.clone()))
             }
-            IntermediateValue::List(mut integers) if integers.len() == 1 => {
+            DeferredValue::List(mut integers) if integers.len() == 1 => {
                 let integer = integers.pop().unwrap();
                 Some(Tracked::from(integer, source.clone()))
             }
-            IntermediateValue::Block(_) | IntermediateValue::List(_) => {
+            DeferredValue::Block(_) | DeferredValue::List(_) => {
                 let error = IntermediateError::ExpectedInteger;
                 self.errors.push(Tracked::from(error, source.clone()));
                 None
@@ -313,8 +290,18 @@ impl IntermediateParser {
         let result = self.parse_invocation(invocation, source)?;
         let source = result.source;
         match result.value {
-            IntermediateValue::Block(tokens) => Some(tokens),
-            IntermediateValue::Integer(_) | IntermediateValue::List(_) => {
+            DeferredValue::Block(blocks) => {
+                // Convert each deferred block value to a concrete value, using
+                // the environment that was stored alongside each value.
+                let mut values = Vec::new();
+                for block in blocks {
+                    self.environment_stack.push(block.env);
+                    values.append(&mut self.parse_block_token(&block.block, &source));
+                    self.environment_stack.pop().unwrap();
+                }
+                Some(values)
+            }
+            DeferredValue::Integer(_) | DeferredValue::List(_) => {
                 let error = IntermediateError::ExpectedBlock;
                 self.errors.push(Tracked::from(error, source));
                 None
@@ -322,7 +309,7 @@ impl IntermediateParser {
         }
     }
 
-    fn parse_invocation(&mut self, invocation: &Invocation, source: &SourceSpan) -> Option<Tracked<IntermediateValue>> {
+    fn parse_invocation(&mut self, invocation: &Invocation, source: &SourceSpan) -> Option<Tracked<DeferredValue>> {
         let received_count = invocation.arguments.len();
         let signature = SymbolSignature {
             name: invocation.name.clone(),
@@ -350,7 +337,7 @@ impl IntermediateParser {
                 let address = self.label_addresses.get(&tagged_name)
                     .and_then(|a| Some(a.current)).or(Some(0)).unwrap();
                 let tracked = Tracked::from(address as isize, label.source.clone());
-                let integer = IntermediateValue::Integer(tracked);
+                let integer = DeferredValue::Integer(tracked);
                 Some(Tracked::from(integer, source.clone()))
             }
         } else if let Some(definition) = self.macro_definitions.get(&signature) {
@@ -416,22 +403,23 @@ impl IntermediateParser {
                     // If the macro was invoked once, return the value.
                     values.pop()
                 } else {
-                    // If the macro was invoked multiple times, create a list of
-                    // block tokens from the returned values.
-                    let mut block = Vec::new();
+                    // If the macro was invoked multiple times, create a single
+                    // deferred block value from the returned values.
+                    let mut deferred_blocks = Vec::new();
                     for value in values {
                         match value.value {
-                            IntermediateValue::Integer(_) | IntermediateValue::List(_) => {
+                            DeferredValue::Integer(_) | DeferredValue::List(_) => {
                                 let error = IntermediateError::ExpectedBlock;
                                 self.errors.push(Tracked::from(error, value.source));
                                 return None;
                             }
-                            IntermediateValue::Block(mut tokens) => {
-                                block.append(&mut tokens);
+                            DeferredValue::Block(mut blocks) => {
+                                deferred_blocks.append(&mut blocks);
                             }
                         }
                     }
-                    Some(Tracked::from(IntermediateValue::Block(block), source.clone()))
+                    let deferred_value = DeferredValue::Block(deferred_blocks);
+                    Some(Tracked::from(deferred_value, source.clone()))
                 }
             }
         } else {
@@ -448,22 +436,25 @@ impl IntermediateParser {
             }
             InvocationArgument::IntegerToken(integer) => {
                 let tracked = self.parse_integer_token(&integer, &argument.source)?;
-                let value = IntermediateValue::Integer(tracked.value);
+                let value = DeferredValue::Integer(tracked.value);
                 (RepeatedArgument::Loop(value), ArgumentType::Integer)
             }
             InvocationArgument::BlockToken(block) => {
-                let tokens = self.parse_block_token(&block, &argument.source);
-                let value = IntermediateValue::Block(tokens);
+                // Store each block argument as a deferred value alongside a copy of the
+                // current environment so that they can be converted to concrete values later.
+                let tracked = Tracked::from(block.clone(), source.clone());
+                let deferred_block = DeferredBlock { block: tracked, env: self.env().clone() };
+                let value = DeferredValue::Block(vec![deferred_block]);
                 (RepeatedArgument::Loop(value), ArgumentType::Block)
             }
             InvocationArgument::Invocation(invocation) => {
-                let value = self.parse_invocation(&invocation, &argument.source)?;
+                let value = self.parse_invocation(&invocation, &source)?;
                 match value.value {
-                    IntermediateValue::Integer(_) =>
+                    DeferredValue::Integer(_) =>
                         (RepeatedArgument::Loop(value.value), ArgumentType::Integer),
-                    IntermediateValue::Block(_) =>
+                    DeferredValue::Block(_) =>
                         (RepeatedArgument::Loop(value.value), ArgumentType::Block),
-                    IntermediateValue::List(list) =>
+                    DeferredValue::List(list) =>
                         self.parse_invocation_list_argument(Tracked::from(list, value.source), expected_type)?
                 }
             }
@@ -480,36 +471,37 @@ impl IntermediateParser {
         if let ArgumentType::Integer = expected_type {
             let mut values = Vec::new();
             for value in &list.value {
-                values.push(IntermediateValue::Integer(value.clone()));
+                values.push(DeferredValue::Integer(value.clone()));
             }
             Some((RepeatedArgument::List(values), ArgumentType::Integer))
         } else {
-            let value = IntermediateValue::List(list.value);
+            let value = DeferredValue::List(list.value);
             Some((RepeatedArgument::Loop(value), ArgumentType::List))
         }
     }
 
-    fn parse_macro_definition_body(&mut self, body: &MacroDefinitionBody, source: &SourceSpan) -> Option<Tracked<IntermediateValue>> {
+    fn parse_macro_definition_body(&mut self, body: &MacroDefinitionBody, source: &SourceSpan) -> Option<Tracked<DeferredValue>> {
         match &body {
             MacroDefinitionBody::Integer(integer) => {
                 let token = self.parse_integer_token(&integer, &source)?;
-                let integer = IntermediateValue::Integer(token.value);
+                let integer = DeferredValue::Integer(token.value);
                 Some(Tracked::from(integer, source.clone()))
             }
             MacroDefinitionBody::Invocation(invocation) => {
                 self.parse_invocation(&invocation, &invocation.source)
             }
             MacroDefinitionBody::Block(blocks) => {
-                let mut tokens = Vec::new();
+                // Store each block value alongside a copy of the current environment
+                // so that they can be converted to concrete values later.
+                let mut values = Vec::new();
                 for block in blocks {
-                    tokens.append(&mut self.parse_block_token(block, &block.source));
+                    values.push(DeferredBlock { block: block.clone(), env: self.env().clone() });
                 }
-                let value = IntermediateValue::Block(tokens);
-                Some(Tracked::from(value, source.clone()))
+                Some(Tracked::from(DeferredValue::Block(values), source.clone()))
             }
             MacroDefinitionBody::List(list) => {
                 let list = self.parse_list_token(list, &source)?;
-                let integer = IntermediateValue::List(list.value);
+                let integer = DeferredValue::List(list.value);
                 Some(Tracked::from(integer, source.clone()))
             }
         }
@@ -530,13 +522,17 @@ impl IntermediateParser {
                 }
                 ExpressionToken::Invocation(invocation) => {
                     if let Some(value) = self.parse_invocation(invocation, source) {
-                        if let IntermediateValue::Block(_) = &value.value {
-                            let error = IntermediateError::InvalidBlockInExpression;
-                            let new_source = source.clone().wrap(value.source.clone());
-                            self.errors.push(Tracked::from(error, new_source));
-                            return None;
-                        }
-                        stack.push(value.value);
+                        let value = match value.value {
+                            DeferredValue::Block(_) => {
+                                let error = IntermediateError::InvalidBlockInExpression;
+                                let new_source = source.clone().wrap(value.source.clone());
+                                self.errors.push(Tracked::from(error, new_source));
+                                return None;
+                            }
+                            DeferredValue::List(list) => IntermediateValue::List(list),
+                            DeferredValue::Integer(integer) => IntermediateValue::Integer(integer),
+                        };
+                        stack.push(value);
                     }
                 }
                 ExpressionToken::Operator(operator) => {
diff --git a/src/stages/intermediate_tokens.rs b/src/stages/intermediate_tokens.rs
index d796299..116cec5 100644
--- a/src/stages/intermediate_tokens.rs
+++ b/src/stages/intermediate_tokens.rs
@@ -1,5 +1,33 @@
 use crate::*;
 
+use indexmap::IndexMap;
+
+
+pub struct LabelAddress {
+    pub previous: usize,
+    pub current: usize,
+    pub touched: bool,
+}
+
+#[derive(Eq, Hash, PartialEq, Clone)]
+pub struct SymbolSignature {
+    pub name: String,
+    pub arg_count: usize,
+}
+
+impl std::fmt::Display for SymbolSignature {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
+        write!(f, "{}::{}", self.name, self.arg_count)
+    }
+}
+
+// Arguments stored in the environment are lazily evaluated. See the
+// comment on DeferredValue for more information.
+#[derive(Clone)]
+pub struct Environment {
+    pub arguments: IndexMap<String, Tracked<DeferredValue>>,
+    pub id: usize,
+}
 
 #[derive(Clone)]
 pub enum IntermediateToken {
@@ -41,13 +69,36 @@ pub enum IntermediateValue {
     Block(IntermediateBlock),
 }
 
+/// Arguments are passed around as deferred values. This is to prevent
+/// block arguments values from being evaluated before they are used,
+/// which would cause the address tracker to have an incorrect value
+/// (because the address is incremented every time a word is assembled,
+/// and the block argument isn't necessarily used at the start of the
+/// macro).
+///
+/// This will also fix the issue of a block being 'assembled'
+/// exactly once when it's passed to a macro invocation, but being used
+/// twice or never in that macro, throwing off the address tracker.
+#[derive(Clone)]
+pub enum DeferredValue {
+    Integer(IntermediateInteger),
+    List(IntermediateList),
+    Block(Vec<DeferredBlock>),
+}
+
 pub type IntermediateInteger = Tracked<isize>;
 pub type IntermediateList = Vec<Tracked<isize>>;
 pub type IntermediateBlock = Vec<Tracked<IntermediateToken>>;
 
+#[derive(Clone)]
+pub struct DeferredBlock {
+    pub block: Tracked<BlockToken>,
+    pub env: Environment,
+}
+
 pub enum RepeatedArgument {
-    Loop(IntermediateValue),
-    List(Vec<IntermediateValue>),
+    Loop(DeferredValue),
+    List(Vec<DeferredValue>),
 }
 
 impl RepeatedArgument {