From 25651ef6e9fdeab052bb6db191a1f473984e267f Mon Sep 17 00:00:00 2001
From: MartinGC94 <42123497+MartinGC94@users.noreply.github.com>
Date: Thu, 23 Mar 2023 22:40:50 +0100
Subject: [PATCH] Fix a crash in the type inference code (#19400)

---
 .../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 5f2288e010..afa33d0e91 100644
--- a/src/System.Management.Automation/engine/parser/TypeInferenceVisitor.cs
+++ b/src/System.Management.Automation/engine/parser/TypeInferenceVisitor.cs
@@ -739,19 +739,28 @@ namespace System.Management.Automation
         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 @@ namespace System.Management.Automation
                     {
                         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 36293030da..14e3a3b086 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 {