Allow label references in pinned address calculations

This is a relaxation of the rule where a label reference could not
be used in any context that could change the length of an assembled
program.

We implement this in the bytecode stage by naively calculating an
initial address for each label as before. If a pinned address is
calculated from a label reference, some of the calculated addresses
could be incorrect. We then attempt to run the bytecode stage, which
will calculate a more accurate address for each label based on how
pinned addresses are calculated. If the address of any label was
changed by running this stage, we re-run the stage up to three more
times until all labels stabilise. If the labels fail to stabilise, we
return an error.

5 files changed, 137 insertions, 46 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 420bb61..f0de207 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -7,3 +7,5 @@ pub use types::*;
 pub use formats::*;
 
 pub use assembler::*;
+
+use log::*;
diff --git a/src/stages/bytecode.rs b/src/stages/bytecode.rs
index 1d27b28..71f1ff0 100644
--- a/src/stages/bytecode.rs
+++ b/src/stages/bytecode.rs
@@ -34,26 +34,82 @@ impl BytecodeParser {
     }
 
     pub fn parse(mut self, tokens: Vec<Tracked<IntermediateToken>>) -> Result<Vec<Segment>, Vec<Tracked<BytecodeError>>> {
-        // Calculate all label addresses ahead of time.
+        // Register all labels with address 0.
+        for token in &tokens {
+            if let IntermediateToken::LabelDefinition(name) = &token.value {
+                let tracked = Tracked::from(0, token.source.clone());
+                if let Some(_) = self.addresses.insert(name.clone(), tracked) {
+                    unreachable!("Uncaught duplicate label definition '{name}'");
+                }
+            }
+        }
+        // Attempt to calculate all label addresses naively ahead of time.
+        // This will give false results if we pin an address calculated from a label address.
         let mut address = 0;
         for token in &tokens {
             let source = &token.source;
             match &token.value {
                 IntermediateToken::LabelDefinition(name) => {
                     let tracked = Tracked::from(address, source.clone());
-                    if let Some(_) = self.addresses.insert(name.clone(), tracked) {
-                        unreachable!("Uncaught duplicate label definition '{name}'");
-                    }
+                    self.addresses.insert(name.clone(), tracked);
                 }
                 IntermediateToken::Word(_) => {
                     address += 1;
                 }
                 IntermediateToken::PinnedAddress(pinned) => {
-                    address = pinned.value;
+                    // Attempt to calculate a sane initial value for a pinned address.
+                    match &pinned.value {
+                        IntermediateInteger::Integer(value) => {
+                            address = (*value).try_into().unwrap_or(0);
+                        }
+                        IntermediateInteger::Expression(expression) => {
+                            let result = self.evaluate_expression(&expression, &pinned.source);
+                            address = result.try_into().unwrap_or(0);
+                        }
+                        IntermediateInteger::LabelReference(_) => {
+                            let error = BytecodeError::PinnedLabel;
+                            self.errors.push(Tracked::from(error, source.clone()));
+                        }
+                    }
                 }
             }
         }
-        for token in &tokens {
+        // Return unrecoverable errors.
+        if !self.errors.is_empty() {
+            return Err(self.errors);
+        }
+
+       for i in 0..4 {
+            info!("Attempting iteration {} of bytecode assembly stage", i+1);
+            // Erase the previous parse attempt.
+            self.segments.clear();
+            self.errors.clear();
+            // Attempt to parse the program.
+            let previous_addresses = self.addresses.clone();
+            self.parse_iteration(&tokens);
+            // Return unrecoverable errors.
+            if !self.errors.is_empty() {
+                return Err(self.errors);
+            }
+            // Check label stability
+            if self.check_for_instability(&previous_addresses) {
+                continue;
+            }
+            // Check for backtrack
+            if self.check_for_backtrack() {
+                continue;
+            };
+            // Program is stable, return.
+            info!("Stabilised in {} iteration of bytecode assembly stage", i+1);
+            return Ok(self.segments);
+        }
+
+        return Err(self.errors);
+    }
+
+    /// Attempt to parse the full program using the current label values.
+    fn parse_iteration(&mut self, tokens: &[Tracked<IntermediateToken>]) {
+        for token in tokens {
             let source = &token.source;
             match &token.value {
                 IntermediateToken::Word(word) => {
@@ -68,26 +124,37 @@ impl BytecodeParser {
                     self.words.push(word);
                     self.address += 1;
                 }
-                IntermediateToken::PinnedAddress(address) => {
-                    let current = self.address;
-                    let pinned = address.value;
-                    if current > pinned {
-                        let error = BytecodeError::PinnedAddressBacktrack(pinned, current);
-                        self.errors.push(Tracked::from(error, address.source.clone()));
-                    } else {
-                        let words = std::mem::take(&mut self.words);
-                        if !words.is_empty() {
-                            let address = self.segment_address;
-                            let source = std::mem::take(&mut self.segment_source);
-                            let segment = Segment { address, source, words };
-                            self.segments.push(segment);
+                IntermediateToken::PinnedAddress(integer) => {
+                    // Calculate the address of the new segment.
+                    let pinned = match &integer.value {
+                        IntermediateInteger::Integer(value) => {
+                            (*value).try_into().unwrap_or(0)
                         }
-                        self.segment_source = Some(address.source.clone());
-                        self.address = pinned;
-                        self.segment_address = pinned;
+                        IntermediateInteger::Expression(expression) => {
+                            let result = self.evaluate_expression(&expression, &integer.source);
+                            result.try_into().unwrap_or(0)
+                        }
+                        IntermediateInteger::LabelReference(_) =>
+                            // Already handled when registering initial label values.
+                            unreachable!(),
+                    };
+                    // Start a new segment.
+                    let words = std::mem::take(&mut self.words);
+                    if !words.is_empty() {
+                        let address = self.segment_address;
+                        let source = std::mem::take(&mut self.segment_source);
+                        let segment = Segment { address, source, words };
+                        self.segments.push(segment);
                     }
+                    self.segment_source = Some(integer.source.clone());
+                    self.address = pinned;
+                    self.segment_address = pinned;
+                }
+                IntermediateToken::LabelDefinition(name) => {
+                    // Record the latest known address of this label.
+                    let address = self.addresses.get_mut(name).unwrap();
+                    address.value = self.address;
                 }
-                IntermediateToken::LabelDefinition(_) => (),
             }
         }
         // Finish final segment.
@@ -98,11 +165,6 @@ impl BytecodeParser {
             let segment = Segment { address, source, words };
             self.segments.push(segment);
         }
-
-        match self.errors.is_empty() {
-            true => Ok(self.segments),
-            false => Err(self.errors),
-        }
     }
 
     fn evaluate_expression(&mut self, expression: &IntermediateExpression, source: &SourceSpan) -> isize {
@@ -175,4 +237,36 @@ impl BytecodeParser {
         let word = Word { width: word.width, value: word_value };
         return Tracked::from(word, source.clone());
     }
+
+    fn check_for_instability(&mut self, previous_addresses: &HashMap<String, Tracked<usize>>) -> bool {
+        let mut instability_occurred = false;
+        for (name, previous_address) in previous_addresses.iter() {
+            let current_address = &self.addresses[name];
+            if current_address != previous_address {
+                info!("Label '{name}' was unstable, moving from address 0x{:04x} to 0x{:04x}",
+                    previous_address.value, current_address.value);
+                let error = BytecodeError::UnstableLabel(name.to_string());
+                self.errors.push(Tracked::from(error, previous_address.source.clone()));
+                instability_occurred = true;
+            }
+        }
+        return instability_occurred;
+    }
+
+    fn check_for_backtrack(&mut self) -> bool {
+        let mut backtrack_occurred = false;
+        let mut current_address = 0;
+        for segment in &self.segments {
+            if segment.address < current_address {
+                let error = BytecodeError::PinnedAddressBacktrack(segment.address, current_address);
+                if let Some(source) = &segment.source {
+                    self.errors.push(Tracked::from(error, source.clone()));
+                }
+                info!("Backtrack occurred with segment at address 0x{:04x}", segment.address);
+                backtrack_occurred = true;
+            }
+            current_address = segment.address + segment.words.len();
+        }
+        return backtrack_occurred;
+    }
 }
diff --git a/src/stages/bytecode_tokens.rs b/src/stages/bytecode_tokens.rs
index b54cb0e..17b13b8 100644
--- a/src/stages/bytecode_tokens.rs
+++ b/src/stages/bytecode_tokens.rs
@@ -40,6 +40,8 @@ pub enum BytecodeError {
     PinnedAddressBacktrack(usize, usize),
     /// expected, received
     ValueTooWide(u32, u32),
+    PinnedLabel,
+    UnstableLabel(String),
     StackError(Tracked<StackError>),
 }
 
@@ -57,6 +59,10 @@ fn report_bytecode_error(error: &Tracked<BytecodeError>, source_code: &str) {
             &format!("Word is {received} bits wide, but was expected to have a fixed width of {expected} bits"),
         BytecodeError::PinnedAddressBacktrack(pinned, real) =>
             &format!("Cannot pin to address {pinned} when address is already {real}"),
+        BytecodeError::PinnedLabel =>
+            &format!("Cannot pin directly to a label"),
+        BytecodeError::UnstableLabel(name) =>
+            &format!("Label '{name}' never stabilised"),
         BytecodeError::StackError(stack_error) => {
             report_stack_error(stack_error, source_code); return; },
         BytecodeError::ValueTooWide(expected, received) =>
diff --git a/src/stages/intermediate.rs b/src/stages/intermediate.rs
index c4bb74d..cfe33b7 100644
--- a/src/stages/intermediate.rs
+++ b/src/stages/intermediate.rs
@@ -197,21 +197,8 @@ impl<'a> Environment<'a> {
             }
             BlockToken::PinnedAddress(address) => {
                 if let Some(integer) = self.parse_integer_token(address, &address.source) {
-                    if let Some(source) = integer_contains_label_reference(&integer) {
-                        let error = IntermediateError::LabelReferenceInPinnedAddress;
-                        let new_source = address.source.clone().wrap(source);
-                        self.errors.push(Tracked::from(error, new_source));
-                    } else {
-                        match evaluate_integer(&integer, source) {
-                            Ok(value) => {
-                                let value = usize::try_from(value).unwrap_or(0);
-                                let tracked = Tracked::from(value, address.source.clone());
-                                let token = IntermediateToken::PinnedAddress(tracked);
-                                intermediate.push(Tracked::from(token, source.clone()));
-                            }
-                            Err(error) => self.errors.push(error),
-                        }
-                    }
+                    let token = IntermediateToken::PinnedAddress(integer);
+                    intermediate.push(Tracked::from(token, source.clone()));
                 }
             }
             BlockToken::ConditionalBlock(cond) => {
diff --git a/src/stages/intermediate_tokens.rs b/src/stages/intermediate_tokens.rs
index 4c4a345..8b7cfaf 100644
--- a/src/stages/intermediate_tokens.rs
+++ b/src/stages/intermediate_tokens.rs
@@ -4,7 +4,7 @@ use crate::*;
 #[derive(Clone)]
 pub enum IntermediateToken {
     Word(IntermediateWord),
-    PinnedAddress(Tracked<usize>),
+    PinnedAddress(Tracked<IntermediateInteger>),
     LabelDefinition(String),
 }
 
@@ -125,8 +125,10 @@ pub fn print_intermediate_token(i: usize, token: &IntermediateToken) {
                 print_intermediate_integer(i+1, &field.value.value);
             }
         }
-        IntermediateToken::PinnedAddress(address) =>
-            indent!(i, "PinnedAddress({address})"),
+        IntermediateToken::PinnedAddress(address) => {
+            indent!(i, "PinnedAddress");
+            print_intermediate_integer(i+1, address);
+        }
         IntermediateToken::LabelDefinition(name) =>
             indent!(i, "LabelDefinition({name})"),
     }