Implement an intelligent source merging strategy

The previous source merging strategy was to concatenate source units
in the reverse order that they were added to the resolver, which
generally only worked when each source unit had at most one
macro-resolving parent.

An issue arose when some macros in a source unit were resolved by a
source unit which had been added earlier in the order, as the required
macro definitions would then be merged after they were referenced,
preventing the program from assembling.

The new source merging strategy finds an optimal merge order by first
recording for a given source unit the ID of each unit which resolves a
macro referenced by the given unit, and then only merging those source
units whose macro-defining dependencies have already been merged. In the
case that a cycle is detected, where two or more source units depend on
one another, a message is printed and the assembly is aborted.

1 files changed, 76 insertions, 12 deletions

diff --git a/src/symbol_resolver.rs b/src/symbol_resolver.rs
index cced994..e19a7bf 100644
--- a/src/symbol_resolver.rs
+++ b/src/symbol_resolver.rs
@@ -6,6 +6,9 @@ use std::mem::take;
 /// Resolve symbol references across source units.
 pub struct SymbolResolver {
     pub definitions: Vec<TrackedSymbol>,
+    /// All resolved references.
+    pub resolved: Vec<TrackedSymbol>,
+    /// All unresolved references.
     pub unresolved: Vec<TrackedSymbol>,
     /// Contains the ID of the owner of the original definition.
     pub redefinitions: Vec<(TrackedSymbol, usize)>,
@@ -20,6 +23,7 @@ impl SymbolResolver {
     pub fn from_source_unit(source_unit: SourceUnit) -> Self {
         let mut new = Self {
             definitions: Vec::new(),
+            resolved: Vec::new(),
             unresolved: Vec::new(),
             redefinitions: Vec::new(),
             source_units: Vec::new(),
@@ -47,6 +51,20 @@ impl SymbolResolver {
             }
             break;
         }
+
+        // For every macro reference in every unit, find the ID of the unit which
+        // resolves that reference and add it to the .parent_ids field of the
+        // referencing unit.
+        for reference in &self.resolved {
+            let predicate = |d: &&TrackedSymbol| d.symbol.name == reference.symbol.name;
+            if let Some(definition) = self.definitions.iter().find(predicate) {
+                let is_self = reference.source_id == definition.source_id;
+                let is_label = definition.symbol.variant == SymbolVariant::LabelDefinition;
+                if  is_self || is_label { continue; }
+                let referencing_unit = &mut self.source_units[reference.source_id];
+                referencing_unit.parent_ids.push(definition.source_id);
+            };
+        }
     }
 
     /// Add a source unit to the resolver and link it to a parent unit.
@@ -83,15 +101,21 @@ impl SymbolResolver {
             self.root_unit_ids.push(source_id);
         }
 
-        let source_unit = HeirarchicalSourceUnit { source_unit, child_ids: Vec::new() };
-        self.source_units.push(source_unit);
+        self.source_units.push(
+            HeirarchicalSourceUnit {
+                source_unit,
+                child_ids: Vec::new(),
+                parent_ids: Vec::new(),
+            }
+        );
     }
 
     fn add_references(&mut self, references: Vec<Symbol>, source_id: usize, source_role: SourceRole) {
         for symbol in references {
             let reference = TrackedSymbol { symbol, source_id, source_role };
-            if !self.definitions.contains(&reference) {
-                self.unresolved.push(reference);
+            match self.definitions.contains(&reference) {
+                true => self.resolved.push(reference),
+                false => self.unresolved.push(reference),
             }
         }
     }
@@ -104,6 +128,10 @@ impl SymbolResolver {
                 let redefinition = (definition, def.source_id);
                 self.redefinitions.push(redefinition);
             } else {
+                let predicate = |s: &mut TrackedSymbol| s.symbol.name == symbol.name;
+                for symbol in self.unresolved.extract_if(predicate) {
+                    self.resolved.push(symbol);
+                }
                 self.unresolved.retain(|s| s.symbol.name != symbol.name);
                 let definition = TrackedSymbol { symbol, source_id, source_role };
                 self.definitions.push(definition);
@@ -162,28 +190,61 @@ impl SymbolResolver {
     }
 
     /// Create a source file by concatenating all source units.
-    pub fn get_merged_source_code(&self) -> String {
-        // The first source unit is guaranteed to be the root unit, so we can
-        // just push source files in their current order.
+    /// If the source unit dependency graph contains a cycle, the IDs of the
+    /// source units involved in the cycle will be returned.
+    pub fn get_merged_source_code(&self) -> Result<String, Vec<usize>> {
+        // The ID of a given source unit will come after the IDs of all
+        // source units which define at least one symbol referenced in the
+        // given source unit.
+        let source_order = {
+            let mut included_source_ids: Vec<usize> = Vec::new();
+            let mut remaining_source_ids: Vec<usize> = Vec::new();
+            for i in 0..self.source_units.len() {
+                remaining_source_ids.push(i);
+            }
+
+            'restart: while !remaining_source_ids.is_empty() {
+                'next: for (i, id) in remaining_source_ids.iter().enumerate() {
+                    let unit = &self.source_units[*id];
+                    for parent_id in &unit.parent_ids {
+                        if !included_source_ids.contains(&parent_id) {
+                            continue 'next;
+                        }
+                    }
+                    included_source_ids.push(*id);
+                    remaining_source_ids.remove(i);
+                    continue 'restart;
+                }
+                // All remaining source units depend on at least one remaining
+                // source unit, indicating a dependency cycle.
+                return Err(remaining_source_ids);
+            }
+            included_source_ids
+        };
+
         let mut source_code = String::new();
 
         // Push head source code.
-        for source_unit in self.source_units.iter().rev() {
+        for id in &source_order {
+            let source_unit = &self.source_units[*id];
             if let Some(head) = &source_unit.source_unit.head {
                 push_source_code_to_string(&mut source_code, head);
             }
         }
         // Push main source code.
-        for source_unit in self.source_units.iter() {
-            push_source_code_to_string(&mut source_code, &source_unit.source_unit.main);
+        for id in source_order.iter().rev() {
+            let source_unit = &self.source_units[*id];
+            let main = &source_unit.source_unit.main;
+            push_source_code_to_string(&mut source_code, &main);
         }
         // Push tail source code.
-        for source_unit in self.source_units.iter().rev() {
+        for id in &source_order {
+            let source_unit = &self.source_units[*id];
             if let Some(tail) = &source_unit.source_unit.tail {
                 push_source_code_to_string(&mut source_code, tail);
             }
         }
-        return source_code;
+        return Ok(source_code);
     }
 }
 
@@ -204,7 +265,10 @@ fn push_source_code_to_string(string: &mut String, source_file: &SourceFile) {
 
 pub struct HeirarchicalSourceUnit {
     pub source_unit: SourceUnit,
+    /// IDs of units which were added to resolve symbol references this unit.
     pub child_ids: Vec<usize>,
+    /// IDs of units which resolve macro references in this unit.
+    pub parent_ids: Vec<usize>,
 }