Update AST read/written variable finders.

Washi1337 · Jan 4, 2024 · 09e113f · 09e113f
1 parent 9f85b59
commit 09e113f
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 5 deletions.
diff --git a/src/Core/Echo.Ast/Analysis/ReadVariableFinder.cs b/src/Core/Echo.Ast/Analysis/ReadVariableFinder.cs
@@ -1,16 +1,41 @@
-using System.Collections.Generic;
+using System.Buffers;
+using System.Collections.Generic;
 using Echo.Code;
 
 namespace Echo.Ast.Analysis
 {
     internal sealed class ReadVariableFinder<TInstruction> : AstNodeListener<TInstruction>
     {
+        private readonly IArchitecture<TInstruction> _architecture;
+
+        public ReadVariableFinder(IArchitecture<TInstruction> architecture)
+        {
+            _architecture = architecture;
+        }
+
         internal HashSet<IVariable> Variables { get; } = new();
 
         public override void ExitVariableExpression(VariableExpression<TInstruction> expression)
         {
             base.ExitVariableExpression(expression);
             Variables.Add(expression.Variable);
         }
+
+        public override void ExitInstructionExpression(InstructionExpression<TInstruction> expression)
+        {
+            int count = _architecture.GetReadVariablesCount(expression.Instruction);
+            if (count == 0)
+                return;
+
+            var variables = ArrayPool<IVariable>.Shared.Rent(count);
+
+            int actualCount = _architecture.GetReadVariables(expression.Instruction, variables);
+            for (int i = 0; i < actualCount; i++)
+                Variables.Add(variables[i]);
+
+            ArrayPool<IVariable>.Shared.Return(variables);
+
+            base.ExitInstructionExpression(expression);
+        }
     }
 }
diff --git a/src/Core/Echo.Ast/Analysis/WrittenVariableFinder.cs b/src/Core/Echo.Ast/Analysis/WrittenVariableFinder.cs
@@ -1,10 +1,18 @@
+using System.Buffers;
 using System.Collections.Generic;
 using Echo.Code;
 
 namespace Echo.Ast.Analysis
 {
     internal sealed class WrittenVariableFinder<TInstruction> : AstNodeListener<TInstruction>
     {
+        private readonly IArchitecture<TInstruction> _architecture;
+
+        public WrittenVariableFinder(IArchitecture<TInstruction> architecture)
+        {
+            _architecture = architecture;
+        }
+
         internal HashSet<IVariable> Variables { get; } = new();
 
         public override void ExitAssignmentStatement(AssignmentStatement<TInstruction> statement)
@@ -19,5 +27,22 @@ public override void ExitPhiStatement(PhiStatement<TInstruction> phiStatement)
             base.ExitPhiStatement(phiStatement);
             Variables.Add(phiStatement.Representative);
         }
+
+        public override void ExitInstructionExpression(InstructionExpression<TInstruction> expression)
+        {
+            int count = _architecture.GetWrittenVariablesCount(expression.Instruction);
+            if (count == 0)
+                return;
+
+            var variables = ArrayPool<IVariable>.Shared.Rent(count);
+
+            int actualCount = _architecture.GetWrittenVariables(expression.Instruction, variables);
+            for (int i = 0; i < actualCount; i++)
+                Variables.Add(variables[i]);
+
+            ArrayPool<IVariable>.Shared.Return(variables);
+
+            base.ExitInstructionExpression(expression);
+        }
     }
 }
diff --git a/src/Core/Echo.Ast/AstArchitecture.cs b/src/Core/Echo.Ast/AstArchitecture.cs
@@ -11,6 +11,7 @@ namespace Echo.Ast
     public class AstArchitecture<TInstruction>
         : IArchitecture<Statement<TInstruction>>
     {
+        private readonly IArchitecture<TInstruction> _baseArchitecture;
         private readonly FlowControlDeterminer<TInstruction> _flowControlDeterminer;
 
         /// <summary>
@@ -19,6 +20,7 @@ public class AstArchitecture<TInstruction>
         /// <param name="baseArchitecture">The <see cref="IArchitecture{TInstruction}"/> to decorate</param>
         public AstArchitecture(IArchitecture<TInstruction> baseArchitecture)
         {
+            _baseArchitecture = baseArchitecture;
             _flowControlDeterminer = new FlowControlDeterminer<TInstruction>(baseArchitecture);
         }
 
@@ -43,15 +45,15 @@ public InstructionFlowControl GetFlowControl(in Statement<TInstruction> instruct
         /// <inheritdoc />
         public int GetReadVariablesCount(in Statement<TInstruction> instruction)
         {
-            var finder = new ReadVariableFinder<TInstruction>();
+            var finder = new ReadVariableFinder<TInstruction>(_baseArchitecture);
             AstNodeWalker<TInstruction>.Walk(finder, instruction);
             return finder.Variables.Count;
         }
 
         /// <inheritdoc />
         public int GetReadVariables(in Statement<TInstruction> instruction, Span<IVariable> variablesBuffer)
         {
-            var finder = new ReadVariableFinder<TInstruction>();
+            var finder = new ReadVariableFinder<TInstruction>(_baseArchitecture);
             AstNodeWalker<TInstruction>.Walk(finder, instruction);
 
             int i = 0;
@@ -64,15 +66,15 @@ public int GetReadVariables(in Statement<TInstruction> instruction, Span<IVariab
         /// <inheritdoc />
         public int GetWrittenVariablesCount(in Statement<TInstruction> instruction)
         {
-            var finder = new WrittenVariableFinder<TInstruction>();
+            var finder = new WrittenVariableFinder<TInstruction>(_baseArchitecture);
             AstNodeWalker<TInstruction>.Walk(finder, instruction);
             return finder.Variables.Count;
         }
 
         /// <inheritdoc />
         public int GetWrittenVariables(in Statement<TInstruction> instruction, Span<IVariable> variablesBuffer)
         {
-            var finder = new WrittenVariableFinder<TInstruction>();
+            var finder = new WrittenVariableFinder<TInstruction>(_baseArchitecture);
             AstNodeWalker<TInstruction>.Walk(finder, instruction);
 
             int i = 0;