From 0185fcc0b0c97b0de7de5249dbbadb972add6fcb Mon Sep 17 00:00:00 2001
From: MartinGC94 <mgc94@outlook.dk>
Date: Thu, 23 Mar 2023 20:42:36 +0100
Subject: [PATCH] Fix type inference crash

---
 .../engine/parser/TypeInferenceVisitor.cs     | 24 +++++++++++++------
 .../engine/Api/TypeInference.Tests.ps1        |  8 +++++++
 2 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/src/System.Management.Automation/engine/parser/TypeInferenceVisitor.cs b/src/System.Management.Automation/engine/parser/TypeInferenceVisitor.cs
index 5f2288e0106..afa33d0e918 100644
--- a/src/System.Management.Automation/engine/parser/TypeInferenceVisitor.cs
+++ b/src/System.Management.Automation/engine/parser/TypeInferenceVisitor.cs
@@ -739,19 +739,28 @@ object ICustomAstVisitor.VisitSubExpression(SubExpressionAst subExpressionAst)
         object ICustomAstVisitor.VisitErrorStatement(ErrorStatementAst errorStatementAst)
         {
             var inferredTypes = new List<PSTypeName>();
-            foreach (var ast in errorStatementAst.Conditions)
+            if (errorStatementAst.Conditions is not null)
             {
-                inferredTypes.AddRange(InferTypes(ast));
+                foreach (var ast in errorStatementAst.Conditions)
+                {
+                    inferredTypes.AddRange(InferTypes(ast));
+                }
             }
 
-            foreach (var ast in errorStatementAst.Bodies)
+            if (errorStatementAst.Bodies is not null)
             {
-                inferredTypes.AddRange(InferTypes(ast));
+                foreach (var ast in errorStatementAst.Bodies)
+                {
+                    inferredTypes.AddRange(InferTypes(ast));
+                }
             }
 
-            foreach (var ast in errorStatementAst.NestedAst)
+            if (errorStatementAst.NestedAst is not null)
             {
-                inferredTypes.AddRange(InferTypes(ast));
+                foreach (var ast in errorStatementAst.NestedAst)
+                {
+                    inferredTypes.AddRange(InferTypes(ast));
+                }
             }
 
             return inferredTypes;
@@ -1944,7 +1953,8 @@ private void InferTypeFrom(VariableExpressionAst variableExpressionAst, List<PST
                     {
                         break;
                     }
-                    else if (parent is SwitchStatementAst switchStatement)
+                    else if (parent is SwitchStatementAst switchStatement
+                        && switchStatement.Condition.Extent.EndOffset < variableExpressionAst.Extent.StartOffset)
                     {
                         parent = switchStatement.Condition;
                         break;
diff --git a/test/powershell/engine/Api/TypeInference.Tests.ps1 b/test/powershell/engine/Api/TypeInference.Tests.ps1
index 36293030da7..14e3a3b0868 100644
--- a/test/powershell/engine/Api/TypeInference.Tests.ps1
+++ b/test/powershell/engine/Api/TypeInference.Tests.ps1
@@ -1406,6 +1406,14 @@ Describe "Type inference Tests" -tags "CI" {
         $res = [AstTypeInference]::InferTypeOf( { pwsh }.Ast)
         $res.Name | Should -Be 'System.String'
     }
+
+    It 'Should not throw when inferring $_ in switch condition' {
+        $FoundAst = { switch($_){default{}} }.Ast.Find(
+            {param($Ast) $Ast -is [Language.VariableExpressionAst]},
+            $true
+        )
+        $null = [AstTypeInference]::InferTypeOf($FoundAst)
+    }
 }
 
 Describe "AstTypeInference tests" -Tags CI {