Remove-Project.ps1

# Copyright (c) Matthias Wolf, Mawosoft.

using namespace System
using namespace System.Collections.Generic
using namespace System.Reflection
using namespace System.Management.Automation
using namespace System.Management.Automation.Language

<#
.SYNOPSIS
    Removes a PowerShell script project and its associated modules.
.INPUTS
    You can pipe module names or module objects to this cmdlet.
.OUTPUTS
    None.
#>

function Remove-Project {
    [CmdletBinding(PositionalBinding = $false)]
    param(
        [Parameter(Mandatory, ParameterSetName = 'Name', ValueFromPipeline, Position = 0)]
        [ValidateNotNullOrEmpty()]
        # Specifies names or name patterns of modules to process.
        [string[]]$Name,

        [Parameter(Mandatory, ParameterSetName = 'FullyQualifiedName', ValueFromPipeline, Position = 0)]
        [ValidateNotNullOrEmpty()]
        # Specifies module names, full module specifications, or module paths to process.
        [Microsoft.PowerShell.Commands.ModuleSpecification[]]$FullyQualifiedName,

        [Parameter(Mandatory, ParameterSetName = 'ModuleInfo', ValueFromPipeline, Position = 0)]
        [ValidateNotNullOrEmpty()]
        # Specifies module objects to process.
        [psmoduleinfo[]]$ModuleInfo
    )
    begin {
        $SMA = $script:SMA
        $anyTypeRemovals = $false
        $anyAcceleratorRemovals = $false
        $typeNamesToRemove = [HashSet[string]]::new([StringComparer]::OrdinalIgnoreCase)
        $userAccelsToRemove = [HashSet[string]]::new([StringComparer]::OrdinalIgnoreCase)
        $trsAssembliesRemaining = [HashSet[Assembly]]::new()
        $typeNameStack = [Stack[ITypeName]]::new()

        # Script to run in isolated scope because Remove-Module has no -Global switch.
        $removeModule = [psmoduleinfo]::new($true).NewBoundScriptBlock({
                param($module, $ErrorActionPreference, $VerbosePreference)
                Remove-Module -ModuleInfo $module -Force
            })
    }
    process {
        $modules = $null
        if ($Name) {
            $modules = Get-Module -Name $Name
        }
        if ($FullyQualifiedName) {
            $modules = Get-Module -FullyQualifiedName $FullyQualifiedName
        }
        if ($ModuleInfo) {
            $modules = $ModuleInfo
        }
        if (-not $modules) {
            Write-Error 'No modules found.'
            return
        }
        $extendedmoduleInfos = Get-ExtendedModuleInfo -ModuleInfo $modules -Merge

        foreach ($extendedmoduleInfo in $extendedmoduleInfos) {
            $context = $script:Reflector.Expose($ExecutionContext)._context
            $topLevelSsi = $context.TopLevelSessionState
            $propagate = $context.PropagateExceptionsToEnclosingStatementBlock
            # This will remove any functions and variables we may have added during import.
            try {
                # We need to catch and rethrow here to unobscure the real error location.
                # We want non-terminating errors not to throw if the EAP allows it.
                $context.PropagateExceptionsToEnclosingStatementBlock = $false
                & $removeModule $extendedmoduleInfo.Module $ErrorActionPreference $VerbosePreference
            }
            catch {
                # We only need to restore here before re-throwing.
                # Auto-restore happens internally after the try-catch block.
                $context.PropagateExceptionsToEnclosingStatementBlock = $propagate
                throw
                return
            }
            $typeNamesToRemove.Clear()
            $typeNamesToRemove.UnionWith($extendedmoduleInfo.Types.Keys)
            if ($typeNamesToRemove.Count -eq 0) { continue }
            $globalScope = $topLevelSsi.GlobalScope
            $typeTable = $globalScope.TypeTable
            $typeTableCount = 0
            if ($null -ne $typeTable) { $typeTableCount = $typeTable.Count }
            if ($typeTableCount -ne 0) {
                foreach ($item in $typeNamesToRemove) {
                    if ($typeTable.Remove($item)) {
                        Write-Verbose "Removing global type: $item"
                    }
                }
                if ($typeTableCount -ne $typeTable.Count) {
                    $anyTypeRemovals = $true
                    $trs = $globalScope._typeResolutionState
                    if ($null -ne $trs) {
                        $trsAssembliesRemaining.Clear()
                        $trsAssembliesRemaining.UnionWith($trs.assemblies)
                        $trsAssembliesRemaining.ExceptWith($extendedmoduleInfo.Assemblies)
                        $trs = $SMA.TypeResolutionState.new($trs.namespaces, $trsAssembliesRemaining)
                        $globalScope.TypeResolutionState = $trs.CloneWithAddTypesDefined($typeTable.Keys)
                        Write-Verbose 'Updating global type resolution state.'
                    }
                }
            }

            $userAccelsToRemove.Clear()
            # TypeAccelerators.Remove() always returns true even if nothing was removed.
            $userAccelsToRemove.UnionWith($SMA.TypeAccelerators.userTypeAccelerators.Keys)
            $userAccelsToRemove.IntersectWith($typeNamesToRemove)
            foreach ($item in $userAccelsToRemove) {
                if ($SMA.TypeAccelerators::Remove($item)) {
                    Write-Verbose "Removing type accelerator: $item"
                    $anyAcceleratorRemovals = $true
                    $anyTypeRemovals = $true
                }
            }

            $typeCache = $SMA.TypeCache.s_cache
            foreach ($kvp in $typeCache.GetEnumerator()) {
                $typeNameStack.Clear()
                $typeNameStack.Push($kvp.Key.Item1)
                $remove = $false
                while ($typeNameStack.Count -ne 0) {
                    $iTypeName = $typeNameStack.Pop()
                    if ($iTypeName -is [ArrayTypeName]) {
                        $typeNameStack.Push($iTypeName.ElementType)
                    }
                    elseif ($iTypeName -is [GenericTypeName]) {
                        foreach ($item in $iTypeName.GenericArguments) { $typeNameStack.Push($item) }
                        $typeNameStack.Push($iTypeName.TypeName) # Should we skip this?
                    }
                    elseif ($typeNamesToRemove.Contains($iTypeName.Name)) {
                        $remove = $true
                        break
                    }
                }
                if ($remove -and $typeCache.TryRemove($kvp.Key, [ref]$null)) {
                    Write-Verbose "Removing cached type: $($kvp.Key.Item1.Name)"
                    $anyTypeRemovals = $true
                }
            }
        }
    }
    end {
        if ($anyTypeRemovals) {
            $SMA.ScriptBlock.ClearScriptBlockCache()
            Write-Verbose 'Clearing script cache.'
        }
        if ($anyAcceleratorRemovals) {
            # 'PowerShell Class Assembly' types can only end up in the completion cache if they have been
            # defined as type accelerators.
            # Note that the original code uses Interlocked.Exchange to modify the field, which we cannot
            # do here. We fake an assembly load event instead, since that clears the cache.
            $evargs = [AssemblyLoadEventArgs]::new([psobject].Assembly)
            $SMA.CompletionCompleters.UpdateTypeCacheOnAssemblyLoad($null, $evargs)
            Write-Verbose 'Clearing completion cache.'
        }
    }
}