From c12c6791a1f7c471817e08874311a03a1181dff8 Mon Sep 17 00:00:00 2001
From: MartinGC94 <mgc94@outlook.dk>
Date: Fri, 9 Dec 2022 14:29:04 +0100
Subject: [PATCH] Fix type inference for all scope vars

---
 .../engine/SpecialVariables.cs                | 33 ++++++++++---------
 .../engine/parser/TypeInferenceVisitor.cs     | 19 ++++++++++-
 .../engine/Api/TypeInference.Tests.ps1        | 17 ++++++++++
 3 files changed, 52 insertions(+), 17 deletions(-)

diff --git a/src/System.Management.Automation/engine/SpecialVariables.cs b/src/System.Management.Automation/engine/SpecialVariables.cs
index d093cdbbb50..d1d90e3507c 100644
--- a/src/System.Management.Automation/engine/SpecialVariables.cs
+++ b/src/System.Management.Automation/engine/SpecialVariables.cs
@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 using System.Collections.Generic;
+using System.Management.Automation.Internal;
 
 namespace System.Management.Automation
 {
@@ -357,22 +358,22 @@ internal static class SpecialVariables
         // see an assignment to any of these variables so that they get handled properly (either throwing an exception
         // because they are constant/readonly, or having the value persist in parent scopes where the allscope variable
         // also exists.
-        internal static readonly string[] AllScopeVariables = {
-                                                                  SpecialVariables.Question,
-                                                                  SpecialVariables.ExecutionContext,
-                                                                  SpecialVariables.False,
-                                                                  SpecialVariables.Home,
-                                                                  SpecialVariables.Host,
-                                                                  SpecialVariables.PID,
-                                                                  SpecialVariables.PSCulture,
-                                                                  SpecialVariables.PSHome,
-                                                                  SpecialVariables.PSUICulture,
-                                                                  SpecialVariables.PSVersionTable,
-                                                                  SpecialVariables.PSEdition,
-                                                                  SpecialVariables.ShellId,
-                                                                  SpecialVariables.True,
-                                                                  SpecialVariables.EnabledExperimentalFeatures,
-                                                              };
+        internal static readonly Dictionary<string, Type> AllScopeVariables = new(StringComparer.OrdinalIgnoreCase) {
+            { Question, typeof(bool) },
+            { ExecutionContext, typeof(EngineIntrinsics) },
+            { False, typeof(bool) },
+            { Home, typeof(string) },
+            { Host, typeof(object) },
+            { PID, typeof(int) },
+            { PSCulture, typeof(string) },
+            { PSHome, typeof(string) },
+            { PSUICulture, typeof(string) },
+            { PSVersionTable, typeof(PSVersionHashTable) },
+            { PSEdition, typeof(string) },
+            { ShellId, typeof(string) },
+            { True, typeof(bool) },
+            { EnabledExperimentalFeatures, typeof(ReadOnlyBag<string>) }
+        };
 
         private static readonly HashSet<string> s_classMethodsAccessibleVariables = new HashSet<string>
             (
diff --git a/src/System.Management.Automation/engine/parser/TypeInferenceVisitor.cs b/src/System.Management.Automation/engine/parser/TypeInferenceVisitor.cs
index f21fb0f4f6d..206c0737c10 100644
--- a/src/System.Management.Automation/engine/parser/TypeInferenceVisitor.cs
+++ b/src/System.Management.Automation/engine/parser/TypeInferenceVisitor.cs
@@ -1846,7 +1846,7 @@ private void InferTypeFrom(VariableExpressionAst variableExpressionAst, List<PST
             // We don't need to handle drive qualified variables, we can usually get those values
             // without needing to "guess" at the type.
             var astVariablePath = variableExpressionAst.VariablePath;
-            if (!astVariablePath.IsVariable)
+            if (!astVariablePath.IsVariable || astVariablePath.UserPath.EqualsOrdinalIgnoreCase(SpecialVariables.Null))
             {
                 // Not a variable - the caller should have already tried going to session state
                 // to get the item and hence it's type, but that must have failed.  Don't try again.
@@ -1961,6 +1961,23 @@ private void InferTypeFrom(VariableExpressionAst variableExpressionAst, List<PST
                 var isThis = astVariablePath.UserPath.EqualsOrdinalIgnoreCase(SpecialVariables.This);
                 if (!isThis || (_context.CurrentTypeDefinitionAst == null && _context.CurrentThisType == null))
                 {
+                    if (SpecialVariables.AllScopeVariables.TryGetValue(astVariablePath.UserPath, out Type knownType))
+                    {
+                        if (knownType == typeof(object))
+                        {
+                            if (_context.TryGetRepresentativeTypeNameFromExpressionSafeEval(variableExpressionAst, out var psType))
+                            {
+                                inferredTypes.Add(psType);
+                            }
+                        }
+                        else
+                        {
+                            inferredTypes.Add(new PSTypeName(knownType));
+                        }
+
+                        return;
+                    }
+
                     for (int i = 0; i < SpecialVariables.AutomaticVariables.Length; i++)
                     {
                         if (!astVariablePath.UserPath.EqualsOrdinalIgnoreCase(SpecialVariables.AutomaticVariables[i]))
diff --git a/test/powershell/engine/Api/TypeInference.Tests.ps1 b/test/powershell/engine/Api/TypeInference.Tests.ps1
index 8a77945d03b..0376abdadab 100644
--- a/test/powershell/engine/Api/TypeInference.Tests.ps1
+++ b/test/powershell/engine/Api/TypeInference.Tests.ps1
@@ -1366,6 +1366,23 @@ Describe "Type inference Tests" -tags "CI" {
         $res.Count | Should -Be 2
         $res.Name -join ' ' | Should -Be "System.IO.FileInfo System.IO.DirectoryInfo"
     }
+
+    It 'Infers type of $null after variable assignment' {
+        $res = [AstTypeInference]::InferTypeOf( { $null = "Hello";$null }.Ast)
+        $res.Count | Should -Be 0
+    }
+
+    It 'Infers type of all scope variable after variable assignment' {
+        $res = [AstTypeInference]::InferTypeOf( { $true = "Hello";$true }.Ast)
+        $res.Count | Should -Be 1
+        $res.Name | Should -Be 'System.Boolean'
+    }
+
+    It 'Infers type of all scope variable host after variable assignment' {
+        $res = [AstTypeInference]::InferTypeOf( { $Host = "Hello";$Host }.Ast, [TypeInferenceRuntimePermissions]::AllowSafeEval)
+        $res.Count | Should -Be 1
+        $res.Name | Should -Be 'System.Management.Automation.Internal.Host.InternalHost'
+    }
 }
 
 Describe "AstTypeInference tests" -Tags CI {