PSPlumbingRules

0.1.0

Rules/PSAvoidImmediatePipelineBuffering.psm1

Set-StrictMode -Version Latest

$script:RuleName = 'PSAvoidImmediatePipelineBuffering'
$script:RuleMessage = 'Avoid assigning command or pipeline output to a variable only to enumerate it immediately with foreach in pipeline-facing commands. Keep the producer in the pipeline or enumerate it directly to preserve streaming shape.'

function Measure-PSAvoidImmediatePipelineBuffering {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.ScriptBlockAst] $ScriptBlockAst
    )

    $functions = $ScriptBlockAst.FindAll({
        param($Ast)
        $Ast -is [System.Management.Automation.Language.FunctionDefinitionAst]
    }, $true)

    foreach ($functionAst in $functions) {
        if (-not (Test-IsPipelineFacingCommand -FunctionAst $functionAst)) {
            continue
        }

        $assignments = @(Get-ImmediateBufferedAssignments -FunctionAst $functionAst)
        foreach ($assignment in $assignments) {
            [Microsoft.Windows.PowerShell.ScriptAnalyzer.Generic.DiagnosticRecord]::new(
                $script:RuleMessage,
                $assignment.Extent,
                $script:RuleName,
                [Microsoft.Windows.PowerShell.ScriptAnalyzer.Generic.DiagnosticSeverity]::Warning,
                $null,
                $null,
                $null
            )
        }
    }
}

function Test-IsPipelineFacingCommand {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.FunctionDefinitionAst] $FunctionAst
    )

    if (-not $FunctionAst.Body.ParamBlock) {
        return $false
    }

    $hasCmdletBinding = $FunctionAst.Body.ParamBlock.Attributes | Where-Object {
        $_.TypeName.GetReflectionType().FullName -eq 'System.Management.Automation.CmdletBindingAttribute'
    }

    if (-not $hasCmdletBinding) {
        return $false
    }

    return $FunctionAst.Body.ProcessBlock -ne $null
}

function Get-ImmediateBufferedAssignments {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.FunctionDefinitionAst] $FunctionAst
    )

    $candidates = $FunctionAst.FindAll({
        param($Ast)
        $Ast -is [System.Management.Automation.Language.AssignmentStatementAst] -and
        $Ast.Left -is [System.Management.Automation.Language.VariableExpressionAst]
    }, $true)

    foreach ($assignment in $candidates) {
        if (-not (Test-HasProducedPipelineRightHandSide -AssignmentAst $assignment)) {
            continue
        }

        if (-not (Test-IsImmediateForeachEnumeration -FunctionAst $FunctionAst -AssignmentAst $assignment)) {
            continue
        }

        $assignment
    }
}

function Test-HasProducedPipelineRightHandSide {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.AssignmentStatementAst] $AssignmentAst
    )

    $pipelines = $AssignmentAst.Right.FindAll({
        param($Ast)
        $Ast -is [System.Management.Automation.Language.PipelineAst]
    }, $true)

    foreach ($pipeline in $pipelines) {
        foreach ($element in $pipeline.PipelineElements) {
            if ($element -is [System.Management.Automation.Language.CommandAst]) {
                return $true
            }
        }
    }

    return $false
}

function Test-IsImmediateForeachEnumeration {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.FunctionDefinitionAst] $FunctionAst,

        [Parameter(Mandatory)]
        [System.Management.Automation.Language.AssignmentStatementAst] $AssignmentAst
    )

    $parentBlock = $AssignmentAst.Parent
    if ($null -eq $parentBlock -or -not ($parentBlock.PSObject.Properties.Name -contains 'Statements')) {
        return $false
    }

    $statementIndex = -1
    for ($i = 0; $i -lt $parentBlock.Statements.Count; $i++) {
        if ($parentBlock.Statements[$i] -eq $AssignmentAst) {
            $statementIndex = $i
            break
        }
    }

    if ($statementIndex -lt 0 -or $statementIndex -ge ($parentBlock.Statements.Count - 1)) {
        return $false
    }

    $nextStatement = $parentBlock.Statements[$statementIndex + 1]
    if ($nextStatement -isnot [System.Management.Automation.Language.ForEachStatementAst]) {
        return $false
    }

    $assignedVariableName = $AssignmentAst.Left.VariablePath.UserPath
    $conditionPipeline = $nextStatement.Condition
    if ($conditionPipeline -isnot [System.Management.Automation.Language.PipelineAst] -or
        $conditionPipeline.PipelineElements.Count -ne 1) {
        return $false
    }

    $element = $conditionPipeline.PipelineElements[0]
    if ($element -isnot [System.Management.Automation.Language.CommandExpressionAst]) {
        return $false
    }

    $expression = $element.Expression
    if ($expression -isnot [System.Management.Automation.Language.VariableExpressionAst]) {
        return $false
    }

    return $expression.VariablePath.UserPath -eq $assignedVariableName
}

Export-ModuleMember -Function Measure-PSAvoidImmediatePipelineBuffering