Fix UseDeclaredVarsMoreThanAssignment rule

Rules/UseDeclaredVarsMoreThanAssignments.cs

@@ -27,6 +27,12 @@ namespace Microsoft.Windows.PowerShell.ScriptAnalyzer.BuiltinRules
 #endif
     public class UseDeclaredVarsMoreThanAssignments : IScriptRule
     {
+        private Dictionary<ScriptBlockAst, Dictionary<string, AssignmentStatementAst>> scriptBlockAssignmentMap;
+        private Dictionary<ScriptBlockAst, Dictionary<string, bool>> scriptblockVariableUsageMap;
+        private Dictionary<ScriptBlockAst, ScriptBlockAst> scriptBlockAstParentMap;
+        private Ast ast;
+        private string fileName;
+
         /// <summary>
         /// AnalyzeScript: Analyzes the ast to check that variables are used in more than just there assignment.
         /// </summary>
@@ -46,13 +52,20 @@ public IEnumerable<DiagnosticRecord> AnalyzeScript(Ast ast, string fileName)
                 yield break;
             }
 
+            scriptBlockAssignmentMap = new Dictionary<ScriptBlockAst, Dictionary<string, AssignmentStatementAst>>();
+            scriptblockVariableUsageMap = new Dictionary<ScriptBlockAst, Dictionary<string, bool>>();
+            scriptBlockAstParentMap = new Dictionary<ScriptBlockAst, ScriptBlockAst>();
+            this.ast = ast;
+            this.fileName = fileName;
             foreach (var scriptBlockAst in scriptBlockAsts)
             {
                 var sbAst = scriptBlockAst as ScriptBlockAst;
-                foreach (var diagnosticRecord in AnalyzeScriptBlockAst(sbAst, fileName))
-                {
-                    yield return diagnosticRecord;
-                }
+                AnalyzeScriptBlockAst(sbAst, fileName);
+            }
+
+            foreach (var diagnosticRecord in GetViolations())
+            {
+                yield return diagnosticRecord;
             }
         }
 
@@ -114,20 +127,18 @@ public string GetSourceName()
         /// <param name="scriptBlockAst">Ast of type ScriptBlock</param>
         /// <param name="fileName">Name of file containing the ast</param>
         /// <returns>An enumerable containing diagnostic records</returns>
-        private IEnumerable<DiagnosticRecord> AnalyzeScriptBlockAst(ScriptBlockAst scriptBlockAst, string fileName)
+        private void AnalyzeScriptBlockAst(ScriptBlockAst scriptBlockAst, string fileName)
         {
             IEnumerable<Ast> assignmentAsts = scriptBlockAst.FindAll(testAst => testAst is AssignmentStatementAst, false);
             IEnumerable<Ast> varAsts = scriptBlockAst.FindAll(testAst => testAst is VariableExpressionAst, true);
-            IEnumerable<Ast> varsInAssignment;
-
             Dictionary<string, AssignmentStatementAst> assignments = new Dictionary<string, AssignmentStatementAst>(StringComparer.OrdinalIgnoreCase);
-
+            Dictionary<string, bool> isVariableUsed = new Dictionary<string, bool>(StringComparer.OrdinalIgnoreCase);
             string varKey;
             bool inAssignment;
 
             if (assignmentAsts == null)
             {
-                yield break;
+                return;
             }
 
             foreach (AssignmentStatementAst assignmentAst in assignmentAsts)
@@ -147,6 +158,7 @@ private IEnumerable<DiagnosticRecord> AnalyzeScriptBlockAst(ScriptBlockAst scrip
                         if (!assignments.ContainsKey(variableName))
                         {
                             assignments.Add(variableName, assignmentAst);
+                            isVariableUsed.Add(variableName, false);
                         }
                     }
                 }
@@ -161,38 +173,93 @@ private IEnumerable<DiagnosticRecord> AnalyzeScriptBlockAst(ScriptBlockAst scrip
 
                     if (assignments.ContainsKey(varKey))
                     {
-                        varsInAssignment = assignments[varKey].Left.FindAll(testAst => testAst is VariableExpressionAst, true);
+                        // Check if varAst is part of an AssignmentStatementAst
+                        inAssignment = varAst.Parent is AssignmentStatementAst;
 
+                        // Check if variable belongs to PowerShell built-in variables
                         // Checks if this variableAst is part of the logged assignment
-                        foreach (VariableExpressionAst varInAssignment in varsInAssignment)
+                        if (!inAssignment || Helper.Instance.HasSpecialVars(varKey))
                         {
-                            inAssignment |= varInAssignment.Equals(varAst);
+                            isVariableUsed[varKey] = true;
                         }
+                    }
+                }
+            }
 
-                        if (!inAssignment)
-                        {
-                            assignments.Remove(varKey);
-                        }
+            scriptBlockAssignmentMap[scriptBlockAst] = assignments;
+            scriptblockVariableUsageMap[scriptBlockAst] = isVariableUsed;
+            scriptBlockAstParentMap[scriptBlockAst] = GetScriptBlockAstParent(scriptBlockAst);
+        }
 
-                        // Check if variable belongs to PowerShell built-in variables
-                        if (Helper.Instance.HasSpecialVars(varKey))
-                        {
-                            assignments.Remove(varKey);
-                        }
+        /// <summary>
+        /// Gets the first upstream node away from the input argument that is of type ScriptBlockAst
+        /// </summary>
+        /// <param name="scriptBlockAst">Ast</param>
+        /// <returns>Null if the input argument's Parent is null
+        /// or if Parent is ast, the input to AnalyzeAst</returns>
+        private ScriptBlockAst GetScriptBlockAstParent(Ast scriptBlockAst)
+        {
+            if (scriptBlockAst == this.ast
+                || scriptBlockAst.Parent == null)
+            {
+                return null;
+            }
+
+            var parent = scriptBlockAst.Parent as ScriptBlockAst;
+            if (parent == null)
+            {
+                return GetScriptBlockAstParent(scriptBlockAst.Parent);
+            }
+
+            return parent;
+        }
+
+        /// <summary>
+        /// Returns the violations in the given ast
+        /// </summary>
+        private IEnumerable<DiagnosticRecord> GetViolations()
+        {
+            foreach (var sbAst in scriptblockVariableUsageMap.Keys)
+            {
+                foreach (var variable in scriptblockVariableUsageMap[sbAst].Keys)
+                {
+                    if (!DoesScriptBlockUseVariable(sbAst, variable))
+                    {
+                        yield return new DiagnosticRecord(
+                            string.Format(CultureInfo.CurrentCulture, Strings.UseDeclaredVarsMoreThanAssignmentsError, variable),
+                            scriptBlockAssignmentMap[sbAst][variable].Left.Extent,
+                            GetName(),
+                            DiagnosticSeverity.Warning,
+                            fileName,
+                            variable);
                     }
                 }
             }
+        }
 
-            foreach (string key in assignments.Keys)
+        /// <summary>
+        /// Returns true if the input variable argument is used more than just declaration in
+        /// the input scriptblock or in its parent scriptblock, otherwise returns false
+        /// </summary>
+        private bool DoesScriptBlockUseVariable(ScriptBlockAst scriptBlockAst, string variable)
+        {
+            if (scriptblockVariableUsageMap[scriptBlockAst].ContainsKey(variable))
             {
-                yield return new DiagnosticRecord(
-                    string.Format(CultureInfo.CurrentCulture, Strings.UseDeclaredVarsMoreThanAssignmentsError, key),
-                    assignments[key].Left.Extent,
-                    GetName(),
-                    DiagnosticSeverity.Warning,
-                    fileName,
-                    key);
+                if (!scriptblockVariableUsageMap[scriptBlockAst][variable])
+                {
+                    if (scriptBlockAstParentMap[scriptBlockAst] == null)
+                    {
+                        return false;
+                    }
+
+                    return DoesScriptBlockUseVariable(scriptBlockAstParentMap[scriptBlockAst], variable);
+                }
+
+                return true;
             }
+
+            return false;
         }
+
     }
 }

Tests/Rules/UseDeclaredVarsMoreThanAssignments.tests.ps1

@@ -9,6 +9,17 @@ $violationName = "PSUseDeclaredVarsMoreThanAssignments"
 $violations = Invoke-ScriptAnalyzer $directory\UseDeclaredVarsMoreThanAssignments.ps1 | Where-Object {$_.RuleName -eq $violationName}
 $noViolations = Invoke-ScriptAnalyzer $directory\UseDeclaredVarsMoreThanAssignmentsNoViolations.ps1 | Where-Object {$_.RuleName -eq $violationName}
 
+Function Test-UseDeclaredVarsMoreThanAssignments
+{
+    param(
+        [string] $targetScript,
+        [int] $expectedNumViolations
+    )
+            Invoke-ScriptAnalyzer -ScriptDefinition $targetScript -IncludeRule $violationName | `
+            Get-Count | `
+            Should Be $expectedNumViolations
+}
+
 Describe "UseDeclaredVarsMoreThanAssignments" {
     Context "When there are violations" {
         It "has 2 use declared vars more than assignments violations" {
@@ -33,27 +44,52 @@ function MyFunc2() {
     $a + $a
 }
 '@
-            Invoke-ScriptAnalyzer -ScriptDefinition $target -IncludeRule $violationName | `
-            Get-Count | `
-            Should Be 1
-        }
-
-        It "does not flag `$InformationPreference variable" {
-            Invoke-ScriptAnalyzer -ScriptDefinition '$InformationPreference=Stop' -IncludeRule $violationName  | `
-            Get-Count | `
-            Should Be 0
+            Test-UseDeclaredVarsMoreThanAssignments $target 1
         }
 
-        It "does not flag `$PSModuleAutoLoadingPreference variable" {
-            Invoke-ScriptAnalyzer -ScriptDefinition '$PSModuleAutoLoadingPreference=None' -IncludeRule $violationName | `
-            Get-Count | `
-            Should Be 0
+        It "flags variables assigned in downstream scopes" {
+$target = @'
+function Get-Directory() {
+    $a = 1
+    1..10 | ForEach-Object { $a = $_ }
+}
+'@
+            Test-UseDeclaredVarsMoreThanAssignments $target 2
         }
     }
 
     Context "When there are no violations" {
         It "returns no violations" {
             $noViolations.Count | Should Be 0
         }
+
+        It "does not flag `$InformationPreference variable" {
+            Test-UseDeclaredVarsMoreThanAssignments '$InformationPreference=Stop' 0
+        }
+
+        It "does not flag `$PSModuleAutoLoadingPreference variable" {
+            Test-UseDeclaredVarsMoreThanAssignments '$PSModuleAutoLoadingPreference=None' 0
+        }
+
+        It "does not flags variables used in downstream scopes" {
+$target = @'
+function Get-Directory() {
+    $a = 1
+    1..10 | ForEach-Object { Write-Output $a }
+}
+'@
+            Test-UseDeclaredVarsMoreThanAssignments $target 0
+        }
+
+        It "does not flag variables assigned in downstream scope but used in parent scope" {
+$target = @'
+function Get-Directory() {
+    $a = 1
+    1..10 | ForEach-Object { $a = $_ }
+    $a
+}
+'@
+            Test-UseDeclaredVarsMoreThanAssignments $target 0
+        }
     }
 }