Simplify GetTopVariableAssignment using new extension methods

Author: JustinGrote

README.md

@@ -151,6 +151,7 @@ The focus of the rename support is on quick updates to variables or functions wi
 ❌ Dynamically constructed splat parameters will not be renamed/updated (e.g. `$splat = @{};$splat.a = 5;Do-Thing @a`)
 ❌ Scoped variables (e.g. $SCRIPT:test) are not currently supported
 ❌ Renaming a variable inside of a scriptblock that is used in unscoped operations like `Foreach-Parallel` or `Start-Job` and the variable is not defined within the scriptblock is not supported and can have unexpected results.
+❌ Scriptblocks part of an assignment are considered isolated scopes. For example `$a = 5; $x = {$a}; & $x` does not consider the two $a to be related, even though in execution this reference matches.
 
 📄📄 Filing a Rename Issue
 

src/PowerShellEditorServices/Services/TextDocument/RenameService.cs

@@ -460,7 +460,9 @@ private bool ShouldRename(Ast candidate)
         }
 
         if (candidate == VariableDefinition) { return true; }
+        // Performance optimization
         if (VariableDefinition.IsAfter(candidate)) { return false; }
+
         if (candidate.GetTopVariableAssignment() == VariableDefinition) { return true; }
 
         return false;

src/PowerShellEditorServices/Utility/AstExtensions.cs

@@ -121,14 +121,17 @@ internal static bool StartsAfter(this Ast ast, IScriptPosition other)
     /// </summary>
     /// <param name="target">The target Ast to search from</param>
     /// <param name="predicate">The predicate to match the Ast against</param>
-    /// <param name="crossScopeBoundaries">If true, the search will continue until the topmost scope boundary is reached</param>
-    internal static Ast? FindStartsBefore(this Ast target, Func<Ast, bool> predicate, bool crossScopeBoundaries = false)
+    /// <param name="crossScopeBoundaries">If true, the search will continue until the topmost scope boundary is
+    /// <param name="searchNestedScriptBlocks">Searches scriptblocks within the parent at each level. This can be helpful to find "side" scopes but affects performance</param>
+    internal static Ast? FindStartsBefore(this Ast target, Func<Ast, bool> predicate, bool crossScopeBoundaries = false, bool searchNestedScriptBlocks = false)
     {
         Ast? scope = target.GetScopeBoundary();
         do
         {
-            Ast? result = scope?.Find(ast => ast.StartsBefore(target) && predicate(ast)
-            , searchNestedScriptBlocks: false);
+            Ast? result = scope?.Find(ast =>
+                ast.StartsBefore(target)
+                && predicate(ast)
+            , searchNestedScriptBlocks);
 
             if (result is not null)
             {
@@ -141,6 +144,12 @@ internal static bool StartsAfter(this Ast ast, IScriptPosition other)
         return null;
     }
 
+    internal static T? FindStartsBefore<T>(this Ast target, Func<T, bool> predicate, bool crossScopeBoundaries = false, bool searchNestedScriptBlocks = false) where T : Ast
+        => target.FindStartsBefore
+        (
+            ast => ast is T type && predicate(type), crossScopeBoundaries, searchNestedScriptBlocks
+        ) as T;
+
     /// <summary>
     /// Finds all AST items that start before the target and match the predicate within the scope. Items are returned in order from closest to furthest. Returns an empty list if none found. Useful for finding definitions of variable/function references
     /// </summary>
@@ -252,21 +261,19 @@ public static Ast GetHighestParent(this Ast ast, params Type[] type)
         => FindParents(ast, types).FirstOrDefault();
 
     /// <summary>
-    /// Returns an array of parents in order from closest to furthest
+    /// Returns an enumerable of parents, in order of closest to furthest, that match the specified types.
     /// </summary>
-    public static Ast[] FindParents(this Ast ast, params Type[] types)
+    public static IEnumerable<Ast> FindParents(this Ast ast, params Type[] types)
     {
-        List<Ast> parents = new();
         Ast parent = ast.Parent;
         while (parent is not null)
         {
             if (types.Contains(parent.GetType()))
             {
-                parents.Add(parent);
+                yield return parent;
             }
             parent = parent.Parent;
         }
-        return parents.ToArray();
     }
 
     /// <summary>
@@ -442,10 +449,8 @@ public static bool TryGetFunction(this ParameterAst ast, out FunctionDefinitionA
             StringConstantExpressionAst stringConstantExpressionAst => stringConstantExpressionAst.Value,
             _ => throw new NotSupportedException("The provided reference is not a variable reference type.")
         };
-
-        Ast? scope = reference.GetScopeBoundary();
-
         VariableExpressionAst? varAssignment = null;
+        Ast? scope = reference;
 
         while (scope is not null)
         {
@@ -481,34 +486,29 @@ public static bool TryGetFunction(this ParameterAst ast, out FunctionDefinitionA
                     }
                 }
             }
-            // Will find the outermost assignment that matches the reference.
+
+            // Will find the outermost assignment within the scope that matches the reference.
             varAssignment = reference switch
             {
-                VariableExpressionAst var => scope.Find
-                (
-                    ast => ast is VariableExpressionAst var
-                            && ast.IsBefore(reference)
-                            &&
-                            (
-                                (var.IsVariableAssignment() && !var.IsOperatorAssignment())
-                                || var.IsScopedVariableAssignment()
-                            )
-                            && var.GetUnqualifiedName().ToLower() == name.ToLower()
-                    , searchNestedScriptBlocks: false
-                ) as VariableExpressionAst,
-
-                CommandParameterAst param => scope.Find
-                (
-                    ast => ast is VariableExpressionAst var
-                            && ast.IsBefore(reference)
-                            && var.GetUnqualifiedName().ToLower() == name.ToLower()
-                            && var.Parent is ParameterAst paramAst
-                            && paramAst.TryGetFunction(out FunctionDefinitionAst? foundFunction)
-                            && foundFunction?.Name.ToLower()
-                                == (param.Parent as CommandAst)?.GetCommandName()?.ToLower()
-                            && foundFunction?.Parent?.Parent == scope
-                    , searchNestedScriptBlocks: true //This might hit side scopes...
-                ) as VariableExpressionAst,
+                VariableExpressionAst => scope.FindStartsBefore<VariableExpressionAst>(var =>
+                    var.GetUnqualifiedName().ToLower() == name.ToLower()
+                    && (
+                        (var.IsVariableAssignment() && !var.IsOperatorAssignment())
+                        || var.IsScopedVariableAssignment()
+                    )
+                    , crossScopeBoundaries: false, searchNestedScriptBlocks: false
+                ),
+
+                CommandParameterAst param => scope.FindStartsBefore<VariableExpressionAst>(var =>
+                    var.GetUnqualifiedName().ToLower() == name.ToLower()
+                    && var.Parent is ParameterAst paramAst
+                    && paramAst.TryGetFunction(out FunctionDefinitionAst? foundFunction)
+                    && foundFunction?.Name.ToLower()
+                        == (param.Parent as CommandAst)?.GetCommandName()?.ToLower()
+                    && foundFunction?.Parent?.Parent == scope
+                ),
+
+
                 _ => null
             };