From 4bd4005f70f3aa8459be9a7e600dc9c4901b729e Mon Sep 17 00:00:00 2001
From: MartinGC94 <MGC94@outlook.dk>
Date: Thu, 13 Mar 2025 21:33:38 +0100
Subject: [PATCH 1/2] Fix type inference of parameters in classic functions

---
 .../engine/parser/TypeInferenceVisitor.cs       | 17 ++++++++++++++++-
 .../engine/Api/TypeInference.Tests.ps1          | 10 ++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/src/System.Management.Automation/engine/parser/TypeInferenceVisitor.cs b/src/System.Management.Automation/engine/parser/TypeInferenceVisitor.cs
index 6996b64eb2a..b912de7d840 100644
--- a/src/System.Management.Automation/engine/parser/TypeInferenceVisitor.cs
+++ b/src/System.Management.Automation/engine/parser/TypeInferenceVisitor.cs
@@ -2451,6 +2451,14 @@ private void InferTypeFrom(VariableExpressionAst variableExpressionAst, List<PST
             {
                 if (currentAst is IParameterMetadataProvider)
                 {
+                    if (currentAst is ScriptBlockAst && currentAst.Parent is FunctionDefinitionAst)
+                    {
+                        // If this scriptblock belongs to a function we want to visit that instead so we can get the parameters
+                        // function X ($Param1){}
+                        currentAst = currentAst.Parent;
+                    }
+
+                    assignmentVisitor.ScopeDefinitionAst = currentAst;
                     currentAst.Visit(assignmentVisitor);
 
                     if (assignmentVisitor.LocalScopeOnly
@@ -2901,6 +2909,11 @@ private sealed class VariableAssignmentVisitor : AstVisitor2
             /// Whether or not the last assignment was via command redirection.
             /// </summary>
             internal bool RedirectionAssignment;
+
+            /// <summary>
+            /// The Ast of the scope we are currently analyzing
+            /// </summary>
+            internal Ast ScopeDefinitionAst;
             internal int StopSearchOffset;
             private int LastAssignmentOffset = -1;
 
@@ -3248,7 +3261,9 @@ public override AstVisitAction VisitDataStatement(DataStatementAst dataStatement
 
             public override AstVisitAction VisitFunctionDefinition(FunctionDefinitionAst functionDefinitionAst)
             {
-                return AstVisitAction.SkipChildren;
+                return functionDefinitionAst == ScopeDefinitionAst
+                    ? AstVisitAction.Continue
+                    : AstVisitAction.SkipChildren;
             }
         }
     }
diff --git a/test/powershell/engine/Api/TypeInference.Tests.ps1 b/test/powershell/engine/Api/TypeInference.Tests.ps1
index 2ab5681dddb..c0ecd9b0334 100644
--- a/test/powershell/engine/Api/TypeInference.Tests.ps1
+++ b/test/powershell/engine/Api/TypeInference.Tests.ps1
@@ -268,6 +268,16 @@ Describe "Type inference Tests" -tags "CI" {
         $res.Name | Should -Be 'System.Management.ManagementObject#root\cimv2\Win32_Process'
     }
 
+    It "Infers type from parameter in classic function definition" {
+        $res = [AstTypeInference]::InferTypeOf(({
+            function MyFunction ([int]$param1)
+            {
+                $param1
+            }
+        }.Ast.FindAll({param($Ast) $Ast -is [Language.VariableExpressionAst]}, $true) | Select-Object -Last 1 ))
+        $res.Name | Should -Be 'System.Int32'
+    }
+
     It "Infers type from binary expression with a bool operator as bool" {
         $res = [AstTypeInference]::InferTypeOf( {
                 (1..10) -contains 5

From 5a7ff596ffdf82e8223d11a5ab9f0ba3f68f8f20 Mon Sep 17 00:00:00 2001
From: MartinGC94 <42123497+MartinGC94@users.noreply.github.com>
Date: Tue, 18 Mar 2025 08:43:53 +0100
Subject: [PATCH 2/2] Update
 src/System.Management.Automation/engine/parser/TypeInferenceVisitor.cs

Co-authored-by: Ilya <darpa@yandex.ru>
---
 .../engine/parser/TypeInferenceVisitor.cs                       | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/System.Management.Automation/engine/parser/TypeInferenceVisitor.cs b/src/System.Management.Automation/engine/parser/TypeInferenceVisitor.cs
index 39e7efc6754..5be3e414c7c 100644
--- a/src/System.Management.Automation/engine/parser/TypeInferenceVisitor.cs
+++ b/src/System.Management.Automation/engine/parser/TypeInferenceVisitor.cs
@@ -2922,7 +2922,7 @@ private sealed class VariableAssignmentVisitor : AstVisitor2
             internal bool RedirectionAssignment;
 
             /// <summary>
-            /// The Ast of the scope we are currently analyzing
+            /// The Ast of the scope we are currently analyzing.
             /// </summary>
             internal Ast ScopeDefinitionAst;
             internal int StopSearchOffset;