Add UserData to AstNode, add more var tracking in CompilationUnit

Washi1337 · Jan 17, 2025 · 5d121df · 5d121df
1 parent 313ee9c
commit 5d121df
Show file tree

Hide file tree

Showing 5 changed files with 42 additions and 8 deletions.
diff --git a/src/Core/Echo.Ast/AstNode.cs b/src/Core/Echo.Ast/AstNode.cs
@@ -23,6 +23,15 @@ public AddressRange? OriginalRange
             get;
             set;
         }
+
+        /// <summary>
+        /// Gets or sets additional user data that is associated to the AST node. 
+        /// </summary>
+        public object? UserData
+        {
+            get;
+            set;
+        }
 
         /// <summary>
         /// Obtains the parent compilation unit the AST is added to (if available).

diff --git a/src/Core/Echo.Ast/AstVariableExtensions.cs b/src/Core/Echo.Ast/AstVariableExtensions.cs
@@ -14,7 +14,7 @@ public static class AstVariableExtensions
     /// <param name="self">The variable to cross-reference.</param>
     /// <param name="unit">The compilation unit to cross-reference in.</param>
     /// <returns>The expressions referencing the variable.</returns>
-    public static IEnumerable<VariableExpression<TInstruction>> GetIsUsedBy<TInstruction>(
+    public static IReadOnlyList<VariableExpression<TInstruction>> GetIsUsedBy<TInstruction>(
         this IVariable self,
         CompilationUnit<TInstruction> unit)
         where TInstruction : notnull
@@ -28,7 +28,7 @@ public static IEnumerable<VariableExpression<TInstruction>> GetIsUsedBy<TInstruc
     /// <param name="self">The variable to cross-reference.</param>
     /// <param name="unit">The compilation unit to cross-reference in.</param>
     /// <returns>The statements writing to the variable.</returns>
-    public static IEnumerable<Statement<TInstruction>> GetIsWrittenBy<TInstruction>(
+    public static IReadOnlyList<Statement<TInstruction>> GetIsWrittenBy<TInstruction>(
         this IVariable self,
         CompilationUnit<TInstruction> unit)
         where TInstruction : notnull

diff --git a/src/Core/Echo.Ast/CompilationUnit.cs b/src/Core/Echo.Ast/CompilationUnit.cs
@@ -15,6 +15,7 @@ public class CompilationUnit<TInstruction> : AstNode<TInstruction>
 {
     private readonly Dictionary<IVariable, List<VariableExpression<TInstruction>>> _variableUses = new();
     private readonly Dictionary<IVariable, List<Statement<TInstruction>>> _variableWrites = new();
+    private readonly HashSet<IVariable> _variables = new();
 
     private BlockStatement<TInstruction> _root = null!;
 
@@ -35,15 +36,20 @@ public BlockStatement<TInstruction> Root
         internal set => UpdateChildNotNull(ref _root, value);
     }
 
+    /// <summary>
+    /// Gets a collection of all referenced variables in the compilation unit.
+    /// </summary>
+    public IReadOnlyCollection<IVariable> Variables => _variables;
+
     /// <summary>
     /// Gets all expressions in the compilation unit that reference the provided variable.
     /// </summary>
     /// <param name="variable">The variable to cross-reference.</param>
     /// <returns>The expressions referencing the variable.</returns>
-    public IEnumerable<VariableExpression<TInstruction>> GetVariableUses(IVariable variable)
+    public IReadOnlyList<VariableExpression<TInstruction>> GetVariableUses(IVariable variable)
     {
         return !_variableUses.TryGetValue(variable, out var uses)
-            ? Enumerable.Empty<VariableExpression<TInstruction>>()
+            ? Array.Empty<VariableExpression<TInstruction>>()
             : uses.ToArray(); // Clone to prevent list being modified after returning.
     }
 
@@ -52,10 +58,10 @@ public IEnumerable<VariableExpression<TInstruction>> GetVariableUses(IVariable v
     /// </summary>
     /// <param name="variable">The variable to cross-reference.</param>
     /// <returns>The statements writing to the variable.</returns>
-    public IEnumerable<Statement<TInstruction>> GetVariableWrites(IVariable variable)
+    public IReadOnlyList<Statement<TInstruction>> GetVariableWrites(IVariable variable)
     {
         return !_variableWrites.TryGetValue(variable, out var writes) 
-            ? Enumerable.Empty<Statement<TInstruction>>()
+            ? Array.Empty<Statement<TInstruction>>()
             : writes.ToArray(); // Clone to prevent list being modified after returning.
     }
 
@@ -68,6 +74,7 @@ internal void RegisterVariableUse(VariableExpression<TInstruction> expression)
         }
 
         uses.Add(expression);
+        _variables.Add(expression.Variable);
     }
 
     internal void UnregisterVariableUse(VariableExpression<TInstruction> expression)
@@ -76,7 +83,11 @@ internal void UnregisterVariableUse(VariableExpression<TInstruction> expression)
         {
             uses.Remove(expression);
             if (uses.Count == 0)
+            {
                 _variableUses.Remove(expression.Variable);
+                if (!_variableWrites.ContainsKey(expression.Variable))
+                    _variables.Remove(expression.Variable);
+            }
         }
     }
 
@@ -89,6 +100,7 @@ internal void RegisterVariableWrite(IVariable variable, Statement<TInstruction>
         }
 
         writes.Add(statement);
+        _variables.Add(variable);
     }
 
     internal void UnregisterVariableWrite(IVariable variable, Statement<TInstruction> statement)
@@ -97,7 +109,11 @@ internal void UnregisterVariableWrite(IVariable variable, Statement<TInstruction
         {
             writes.Remove(statement);
             if (writes.Count == 0)
-                _variableUses.Remove(variable);
+            {
+                _variableWrites.Remove(variable);
+                if (!_variableUses.ContainsKey(variable))
+                    _variables.Remove(variable);
+            }
         }
     }
 

diff --git a/src/Core/Echo.Ast/SyntheticVariable.cs b/src/Core/Echo.Ast/SyntheticVariable.cs
@@ -55,6 +55,15 @@ public string Name
     {
         get;
     }
+
+    /// <summary>
+    /// Gets or sets additional user data associated to the synthetic variable.
+    /// </summary>
+    public object? UserData
+    {
+        get;
+        set;
+    }
 
     private string GenerateName()
     {

diff --git a/src/Core/Echo.Ast/VariableExpression.cs b/src/Core/Echo.Ast/VariableExpression.cs
@@ -28,7 +28,7 @@ public VariableExpression(IVariable variable)
         public IVariable Variable
         {
             get => _variable;
-            internal set
+            set
             {
                 if (_variable != value)
                 {