Rules/PSAvoidUsingAdditiveCollection.psm1

Set-StrictMode -Version Latest

$script:RuleName = 'PSAvoidUsingAdditiveCollection'
$script:RuleMessage = 'Avoid using += to collect produced output inside enumerating bodies. Assign the producer expression directly so PowerShell collects emitted objects automatically.'

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

    $assignments = $ScriptBlockAst.FindAll({
        param($Ast)

        $Ast -is [System.Management.Automation.Language.AssignmentStatementAst] -and
        $Ast.Operator -eq [System.Management.Automation.Language.TokenKind]::PlusEquals -and
        $Ast.Left -is [System.Management.Automation.Language.VariableExpressionAst] -and
        (Get-NearestScriptBlockAst -Ast $Ast) -eq $ScriptBlockAst
    }, $true)

    foreach ($assignment in $assignments) {
        if (-not (Test-IsEnumeratingProducerAssignment -Assignment $assignment)) {
            continue
        }

        if (-not (Test-HasProducedOutputRightHandSide -Assignment $assignment)) {
            continue
        }

        [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-IsEnumeratingProducerAssignment {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.AssignmentStatementAst] $Assignment
    )

    if (Get-EnclosingLoopStatement -Assignment $Assignment) {
        return $true
    }

    if (Get-EnclosingForEachObjectScriptBlock -Assignment $Assignment) {
        return $true
    }

    return $false
}

function Get-EnclosingLoopStatement {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.AssignmentStatementAst] $Assignment
    )

    $current = $Assignment.Parent
    while ($null -ne $current) {
        if ($current -is [System.Management.Automation.Language.LoopStatementAst]) {
            if ($current.PSObject.Properties.Name -contains 'Body' -and (Test-IsWithinExtent -Inner $Assignment.Extent -Outer $current.Body.Extent)) {
                return $current
            }
        }

        $current = $current.Parent
    }

    return $null
}

function Get-EnclosingForEachObjectScriptBlock {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.AssignmentStatementAst] $Assignment
    )

    $current = $Assignment.Parent
    while ($null -ne $current) {
        if ($current -is [System.Management.Automation.Language.ScriptBlockExpressionAst]) {
            $command = Get-ParentCommandAst -Ast $current
            if ($null -ne $command -and $command.GetCommandName() -eq 'ForEach-Object') {
                return $current
            }
        }

        $current = $current.Parent
    }

    return $null
}

function Get-ParentCommandAst {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.Ast] $Ast
    )

    $current = $Ast.Parent
    while ($null -ne $current) {
        if ($current -is [System.Management.Automation.Language.CommandAst]) {
            return $current
        }

        if ($current -is [System.Management.Automation.Language.StatementAst]) {
            return $null
        }

        $current = $current.Parent
    }

    return $null
}

function Get-NearestScriptBlockAst {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.Ast] $Ast
    )

    $current = $Ast.Parent
    while ($null -ne $current) {
        if ($current -is [System.Management.Automation.Language.ScriptBlockAst]) {
            return $current
        }

        $current = $current.Parent
    }

    return $null
}

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

    $pipelines = $Assignment.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-IsWithinExtent {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [System.Management.Automation.Language.IScriptExtent] $Inner,

        [Parameter(Mandatory)]
        [System.Management.Automation.Language.IScriptExtent] $Outer
    )

    return $Inner.StartOffset -ge $Outer.StartOffset -and $Inner.EndOffset -le $Outer.EndOffset
}

Export-ModuleMember -Function Measure-PSAvoidUsingAdditiveCollection