1 files changed, 57 insertions, 20 deletions

diff --git a/src/stages/compiler.rs b/src/stages/compiler.rs
index 7dfebd0..1476ec8 100644
--- a/src/stages/compiler.rs
+++ b/src/stages/compiler.rs
@@ -1,8 +1,6 @@
 use crate::*;
 
 use assembler::{Symbol, SymbolRole, DefinitionType};
-use SymbolRole::*;
-use DefinitionType::*;
 
 use std::path::Path;
 
@@ -22,7 +20,27 @@ pub fn parse_symbols(source_code: &str, path: Option<&Path>) -> Option<Vec<Symbo
         Ok(semantic) => semantic,
         Err(_errors) => return None,
     };
-    Some(SymbolParser::new().parse(&semantic))
+    // Convert symbols to the format required by the assembler library.
+    let parsed = SymbolParser::new().parse(&semantic);
+    let mut symbols = Vec::new();
+    for symbol in parsed {
+        let name = format!("{}::{}", symbol.name, symbol.arg_count);
+        let namespace = match symbol.macro_name {
+            Some(macro_name) => vec![macro_name],
+            None => vec![],
+        };
+        let source = symbol.source;
+        let role = match symbol.role {
+            SymbolRoleDetailed::MacroDefinition =>
+                SymbolRole::Definition(DefinitionType::CanFollowReference),
+            SymbolRoleDetailed::LabelDefinition =>
+                SymbolRole::Definition(DefinitionType::CanFollowReference),
+            SymbolRoleDetailed::Invocation =>
+                SymbolRole::Reference,
+        };
+        symbols.push(Symbol { name, namespace, source, role });
+    }
+    Some(symbols)
 }
 
 /// Push source code to a source compilation string.
@@ -43,10 +61,27 @@ pub fn push_code(compilation: &mut String, source_file: &SourceFile) {
 }
 
 
-// Extract symbol definitions from a list of semantic tokens.
+// Track additional information for each symbol.
+pub struct SymbolDetailed {
+    pub name: String,
+    pub macro_name: Option<String>,
+    pub arg_count: usize,
+    pub role: SymbolRoleDetailed,
+    pub source: SourceSpan,
+}
+
+pub enum SymbolRoleDetailed {
+    MacroDefinition,
+    LabelDefinition,
+    Invocation,
+}
+
+
+// Extract symbol definitions and invocations from a list of semantic tokens.
 pub struct SymbolParser {
+    /// Current macro definition name.
     pub macro_name: Option<String>,
-    pub symbols: Vec<Symbol>,
+    pub symbols: Vec<SymbolDetailed>,
 }
 
 impl SymbolParser {
@@ -57,18 +92,19 @@ impl SymbolParser {
         }
     }
 
-    fn record_symbol(&mut self, name: &str, arg_count: usize, source: &SourceSpan, role: SymbolRole) {
-        let name = format!("{name}:{arg_count}");
-        let namespace = match &self.macro_name {
-            Some(macro_name) => vec![macro_name.to_owned()],
-            None => vec![],
-        };
-        let source = source.to_owned();
-        self.symbols.push(Symbol { name, namespace, source, role });
-
+    fn record_symbol(&mut self, name: &str, arg_count: usize, source: &SourceSpan, role: SymbolRoleDetailed) {
+        self.symbols.push(
+            SymbolDetailed {
+                name: name.to_string(),
+                macro_name: self.macro_name.clone(),
+                arg_count,
+                role,
+                source: source.clone(),
+            }
+        );
     }
 
-    pub fn parse(mut self, semantic: &[Tracked<SemanticToken>]) -> Vec<Symbol> {
+    pub fn parse(mut self, semantic: &[Tracked<SemanticToken>]) -> Vec<SymbolDetailed> {
         for token in semantic {
             let source = &token.source;
             match &token.value {
@@ -78,8 +114,9 @@ impl SymbolParser {
                         &definition.name,
                         definition.arguments.len(),
                         &definition.name.source,
-                        Definition(CanFollowReference),
+                        SymbolRoleDetailed::MacroDefinition,
                     );
+                    // Track that we're currently inside a macro definition.
                     self.macro_name = Some(definition.name.to_string());
 
                     for argument in &definition.arguments {
@@ -87,7 +124,7 @@ impl SymbolParser {
                             &argument.name,
                             0,
                             &argument.source,
-                            Definition(MustPrecedeReference),
+                            SymbolRoleDetailed::MacroDefinition,
                         );
                     }
                     match &definition.body {
@@ -139,7 +176,7 @@ impl SymbolParser {
             &invocation.name,
             invocation.arguments.len(),
             &source,
-            Reference,
+            SymbolRoleDetailed::Invocation,
         );
 
         for argument in &invocation.arguments {
@@ -168,7 +205,7 @@ impl SymbolParser {
                     &name,
                     0,
                     &source,
-                    Definition(CanFollowReference),
+                    SymbolRoleDetailed::LabelDefinition,
                 );
             }
             BlockToken::PinnedAddress(integer) => {
@@ -184,7 +221,7 @@ impl SymbolParser {
                         &field.name.to_string(),
                         0,
                         &field.source,
-                        Reference,
+                        SymbolRoleDetailed::Invocation,
                     );
                 }
             }