PSGraft

psgraft.ps1

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

<#
  PSGraft is a generalization of the mutating visitor pattern used in PSProfiler
  
  See `Get-Help Graft-Code -Examples` after dot-sourcing this file.
#>


#region AstVisitor 
class PSGraftVisitor : ICustomAstVisitor, ICustomAstVisitor2 {
    [hashtable]$Injects = $null
    PSGraftVisitor([hashtable]$injects) {
        $this.Injects = $injects
    }
    [Object] VisitElement([object]$element) {
        if ($null -eq $element) {
            return $null
        }
        $res = $element.Visit($this)
        return $res
    }
    [Object] VisitElements([Object]$elements) {
        if ($null -eq $elements -or $elements.Count -eq 0) {
            return $null
        }
        $typeName = $elements.gettype().GenericTypeArguments.Fullname

        $newElements = New-Object -TypeName "System.Collections.Generic.List[$typeName]"
        foreach ($element in $elements) {
            $visitedResult = $element.Visit($this)
            $newElements.add($visitedResult)
        }
        return $newElements 
    }
    [StatementAst[]] VisitStatements([object]$Statements) {
        if ($inject = $this.Injects['Statements']) {
            return & $inject $Statements
        }

        $newStatements = [List[StatementAst]]::new()
        foreach ($statement in $statements) {
            $newStatements.Add($this.VisitElement($statement))
        }
        return $newStatements
    }

    [object] VisitScriptBlock([ScriptBlockAst] $scriptBlockAst) {
        if ($inject = $this.Injects['ScriptBlock']) {
            return & $inject $scriptBlockAst
        }
        $newParamBlock = $this.VisitElement($scriptBlockAst.ParamBlock)
        $newBeginBlock = $this.VisitElement($scriptBlockAst.BeginBlock)
        $newProcessBlock = $this.VisitElement($scriptBlockAst.ProcessBlock)
        $newEndBlock = $this.VisitElement($scriptBlockAst.EndBlock)
        $newDynamicParamBlock = $this.VisitElement($scriptBlockAst.DynamicParamBlock)
        return [ScriptBlockAst]::new($scriptBlockAst.Extent, $newParamBlock, $newBeginBlock, $newProcessBlock, $newEndBlock, $newDynamicParamBlock)
    }


    [object] VisitNamedBlock([NamedBlockAst] $namedBlockAst) {
        if ($inject = $this.Injects['NamedBlock']) {
            return & $inject $namedBlockAst
        }
        $newTraps = $this.VisitElements($namedBlockAst.Traps)
        $newStatements = $this.VisitStatements($namedBlockAst.Statements)
        $statementBlock = [StatementBlockAst]::new($namedBlockAst.Extent, $newStatements, $newTraps)
        return [NamedBlockAst]::new($namedBlockAst.Extent, $namedBlockAst.BlockKind, $statementBlock, $namedBlockAst.Unnamed)
    }

    [object] VisitFunctionDefinition([FunctionDefinitionAst] $functionDefinitionAst) {
        if ($inject = $this.Injects['FunctionDefinition']) {
            return & $inject $functionDefinitionAst
        }
        $newBody = $this.VisitElement($functionDefinitionAst.Body)
        return [FunctionDefinitionAst]::new($functionDefinitionAst.Extent, $functionDefinitionAst.IsFilter, $functionDefinitionAst.IsWorkflow, $functionDefinitionAst.Name, $this.VisitElements($functionDefinitionAst.Parameters), $newBody);
    }

    [object] VisitStatementBlock([StatementBlockAst] $statementBlockAst) {
        if ($inject = $this.Injects['StatementBlock']) {
            return & $inject $statementBlockAst
        }
        $newStatements = $this.VisitStatements($statementBlockAst.Statements)
        $newTraps = $this.VisitElements($statementBlockAst.Traps)
        return [StatementBlockAst]::new($statementBlockAst.Extent, $newStatements, $newTraps)
    }

    [object] VisitIfStatement([IfStatementAst] $ifStmtAst) {
        if ($inject = $this.Injects['IfStatement']) {
            return & $inject $ifStmtAst
        }
        [Tuple[PipelineBaseAst, StatementBlockAst][]]$newClauses = @(foreach ($clause in $ifStmtAst.Clauses) {
                $newClauseTest = [PipelineBaseAst]$this.VisitElement($clause.Item1)
                $newStatementBlock = [StatementBlockAst]$this.VisitElement($clause.Item2)
                [Tuple[PipelineBaseAst, StatementBlockAst]]::new($newClauseTest, $newStatementBlock)
            })
        $newElseClause = $this.VisitElement($ifStmtAst.ElseClause)
        return [IfStatementAst]::new($ifStmtAst.Extent, $newClauses, $newElseClause)
    }

    [object] VisitTrap([TrapStatementAst] $trapStatementAst) {
        if ($inject = $this.Injects['Trap']) {
            return & $inject $trapStatementAst
        }
        return [TrapStatementAst]::new($trapStatementAst.Extent, $this.VisitElement($trapStatementAst.TrapType), $this.VisitElement($trapStatementAst.Body))
    }

    [object] VisitSwitchStatement([SwitchStatementAst] $switchStatementAst) {
        if ($inject = $this.Injects['SwitchStatement']) {
            return & $inject $switchStatementAst
        }
        $newCondition = $this.VisitElement($switchStatementAst.Condition)
        $newClauses = [List[Tuple[ExpressionAst, StatementBlockAst]]]::new()
        $switchStatementAst.Clauses | ForEach-Object {
            $newClauseTest = $this.VisitElement($_.Item1)
            $newStatementBlock = $this.VisitElement($_.Item2)
            $newClauses.Add([Tuple[ExpressionAst, StatementBlockAst]]::new($newClauseTest, $newStatementBlock))
        }
        $newDefault = $this.VisitElement($switchStatementAst.Default)
        return [SwitchStatementAst]::new($switchStatementAst.Extent, $switchStatementAst.Label, $newCondition, $switchStatementAst.Flags, $newClauses, $newDefault)
    }

    [object] VisitDataStatement([DataStatementAst] $dataStatementAst) {
        if ($inject = $this.Injects['DataStatement']) {
            return & $inject $dataStatementAst
        }
        $newBody = $this.VisitElement($dataStatementAst.Body)
        $newCommandsAllowed = $this.VisitElements($dataStatementAst.CommandsAllowed)
        return [DataStatementAst]::new($dataStatementAst.Extent, $dataStatementAst.Variable, $newCommandsAllowed, $newBody)
    }

    [object] VisitForEachStatement([ForEachStatementAst] $forEachStatementAst) {
        if ($inject = $this.Injects['ForEachStatement']) {
            return & $inject $forEachStatementAst
        }
        $newVariable = $this.VisitElement($forEachStatementAst.Variable)
        $newCondition = $this.VisitElement($forEachStatementAst.Condition)
        $newBody = $this.VisitElement($forEachStatementAst.Body)
        return [ForEachStatementAst]::new($forEachStatementAst.Extent, $forEachStatementAst.Label, [ForEachFlags]::None, $newVariable, $newCondition, $newBody)
    }

    [object] VisitDoWhileStatement([DoWhileStatementAst] $doWhileStatementAst) {
        if ($inject = $this.Injects['DoWhileStatement']) {
            return & $inject $doWhileStatementAst
        }
        $newCondition = $this.VisitElement($doWhileStatementAst.Condition)
        $newBody = $this.VisitElement($doWhileStatementAst.Body)
        return [DoWhileStatementAst]::new($doWhileStatementAst.Extent, $doWhileStatementAst.Label, $newCondition, $newBody)
    }

    [object] VisitForStatement([ForStatementAst] $forStatementAst) {
        if ($inject = $this.Injects['ForStatement']) {
            return & $inject $forStatementAst
        }
        $newInitializer = $this.VisitElement($forStatementAst.Initializer)
        $newCondition = $this.VisitElement($forStatementAst.Condition)
        $newIterator = $this.VisitElement($forStatementAst.Iterator)
        $newBody = $this.VisitElement($forStatementAst.Body)
        return [ForStatementAst]::new($forStatementAst.Extent, $forStatementAst.Label, $newInitializer, $newCondition, $newIterator, $newBody)
    }

    [object] VisitWhileStatement([WhileStatementAst] $whileStatementAst) {
        if ($inject = $this.Injects['WhileStatement']) {
            return & $inject $whileStatementAst
        }
        $newCondition = $this.VisitElement($whileStatementAst.Condition)
        $newBody = $this.VisitElement($whileStatementAst.Body)
        return [WhileStatementAst]::new($whileStatementAst.Extent, $whileStatementAst.Label, $newCondition, $newBody)
    }

    [object] VisitCatchClause([CatchClauseAst] $catchClauseAst) {
        if ($inject = $this.Injects['CatchClause']) {
            return & $inject $catchClauseAst
        }
        $newBody = $this.VisitElement($catchClauseAst.Body)
        $newCatchTypes = $this.VisitElements($catchClauseAst.CatchTypes)
        return [CatchClauseAst]::new($catchClauseAst.Extent, $newCatchTypes, $newBody)
    }

    [object] VisitTryStatement([TryStatementAst] $tryStatementAst) {
        if ($inject = $this.Injects['TryStatement']) {
            return & $inject $tryStatementAst
        }
        $newBody = $this.VisitElement($tryStatementAst.Body)
        $newCatchClauses = $this.VisitElements($tryStatementAst.CatchClauses)
        $newFinally = $this.VisitElement($tryStatementAst.Finally)
        return [TryStatementAst]::new($tryStatementAst.Extent, $newBody, $newCatchClauses, $newFinally)
    }

    [object] VisitDoUntilStatement([DoUntilStatementAst] $doUntilStatementAst) {
        if ($inject = $this.Injects['DoUntilStatement']) {
            return & $inject $doUntilStatementAst
        }
        $newCondition = $this.VisitElement($doUntilStatementAst.Condition)
        $newBody = $this.VisitElement($doUntilStatementAst.Body)
        return [DoUntilStatementAst]::new($doUntilStatementAst.Extent, $doUntilStatementAst.Label, $newCondition, $newBody)
    }

    [object] VisitParamBlock([ParamBlockAst] $paramBlockAst) {
        if ($inject = $this.Injects['ParamBlock']) {
            return & $inject $paramBlockAst
        }
        $newAttributes = $this.VisitElements($paramBlockAst.Attributes)
        $newParameters = $this.VisitElements($paramBlockAst.Parameters)
        return [ParamBlockAst]::new($paramBlockAst.Extent, $newAttributes, $newParameters)
    }

    [object] VisitErrorStatement([ErrorStatementAst] $errorStatementAst) {
        if ($inject = $this.Injects['ErrorStatement']) {
            return & $inject $errorStatementAst
        }
        return $errorStatementAst
    }

    [object] VisitErrorExpression([ErrorExpressionAst] $errorExpressionAst) {
        if ($inject = $this.Injects['ErrorExpression']) {
            return & $inject $errorExpressionAst
        }
        return $errorExpressionAst
    }

    [object] VisitTypeConstraint([TypeConstraintAst] $typeConstraintAst) {
        if ($inject = $this.Injects['TypeConstraint']) {
            return & $inject $typeConstraintAst
        }
        return [TypeConstraintAst]::new($typeConstraintAst.Extent, $typeConstraintAst.TypeName)
    }

    [object] VisitAttribute([AttributeAst] $attributeAst) {
        if ($inject = $this.Injects['Attribute']) {
            return & $inject $attributeAst
        }
        $newPositionalArguments = $this.VisitElements($attributeAst.PositionalArguments)
        $newNamedArguments = $this.VisitElements($attributeAst.NamedArguments)
        return [AttributeAst]::new($attributeAst.Extent, $attributeAst.TypeName, $newPositionalArguments, $newNamedArguments)
    }

    [object] VisitNamedAttributeArgument([NamedAttributeArgumentAst] $namedAttributeArgumentAst) {
        if ($inject = $this.Injects['NamedAttributeArgument']) {
            return & $inject $namedAttributeArgumentAst
        }
        $newArgument = $this.VisitElement($namedAttributeArgumentAst.Argument)
        return [NamedAttributeArgumentAst]::new($namedAttributeArgumentAst.Extent, $namedAttributeArgumentAst.ArgumentName, $newArgument, $namedAttributeArgumentAst.ExpressionOmitted)
    }

    [object] VisitParameter([ParameterAst] $parameterAst) {
        if ($inject = $this.Injects['Parameter']) {
            return & $inject $parameterAst
        }
        $newName = $this.VisitElement($parameterAst.Name)
        $newAttributes = $this.VisitElements($parameterAst.Attributes)
        $newDefaultValue = $this.VisitElement($parameterAst.DefaultValue)
        return [ParameterAst]::new($parameterAst.Extent, $newName, $newAttributes, $newDefaultValue)
    }

    [object] VisitBreakStatement([BreakStatementAst] $breakStatementAst) {
        if ($inject = $this.Injects['BreakStatement']) {
            return & $inject $breakStatementAst
        }
        $newLabel = $this.VisitElement($breakStatementAst.Label)
        return [BreakStatementAst]::new($breakStatementAst.Extent, $newLabel)
    }

    [object] VisitContinueStatement([ContinueStatementAst] $continueStatementAst) {
        if ($inject = $this.Injects['ContinueStatement']) {
            return & $inject $continueStatementAst
        }
        $newLabel = $this.VisitElement($continueStatementAst.Label)
        return [ContinueStatementAst]::new($continueStatementAst.Extent, $newLabel)
    }

    [object] VisitReturnStatement([ReturnStatementAst] $returnStatementAst) {
        if ($inject = $this.Injects['ReturnStatement']) {
            return & $inject $returnStatementAst
        }
        $newPipeline = $this.VisitElement($returnStatementAst.Pipeline)
        return [ReturnStatementAst]::new($returnStatementAst.Extent, $newPipeline)
    }

    [object] VisitExitStatement([ExitStatementAst] $exitStatementAst) {
        if ($inject = $this.Injects['ExitStatement']) {
            return & $inject $exitStatementAst
        }
        $newPipeline = $this.VisitElement($exitStatementAst.Pipeline)
        return [ExitStatementAst]::new($exitStatementAst.Extent, $newPipeline)
    }

    [object] VisitThrowStatement([ThrowStatementAst] $throwStatementAst) {
        if ($inject = $this.Injects['ThrowStatement']) {
            return & $inject $throwStatementAst
        }
        $newPipeline = $this.VisitElement($throwStatementAst.Pipeline)
        return [ThrowStatementAst]::new($throwStatementAst.Extent, $newPipeline)
    }

    [object] VisitAssignmentStatement([AssignmentStatementAst] $assignmentStatementAst) {
        if ($inject = $this.Injects['AssignmentStatement']) {
            return & $inject $assignmentStatementAst
        }
        $newLeft = $this.VisitElement($assignmentStatementAst.Left)
        $newRight = $this.VisitElement($assignmentStatementAst.Right)
        return [AssignmentStatementAst]::new($assignmentStatementAst.Extent, $newLeft, $assignmentStatementAst.Operator, $newRight, $assignmentStatementAst.ErrorPosition)
    }

    [object] VisitPipeline([PipelineAst] $pipelineAst) {
        if ($inject = $this.Injects['Pipeline']) {
            return & $inject $pipelineAst
        }
        $newPipeElements = $this.VisitElements($pipelineAst.PipelineElements)
        return [PipelineAst]::new($pipelineAst.Extent, $newPipeElements)
    }

    [object] VisitCommand([CommandAst] $commandAst) {
        if ($inject = $this.Injects['Command']) {
            return & $inject $commandAst
        }
        $newCommandElements = $this.VisitElements($commandAst.CommandElements)
        $newRedirections = $this.VisitElements($commandAst.Redirections)
        return [CommandAst]::new($commandAst.Extent, $newCommandElements, $commandAst.InvocationOperator, $newRedirections)
    }

    [object] VisitCommandExpression([CommandExpressionAst] $commandExpressionAst) {
        if ($inject = $this.Injects['CommandExpression']) {
            return & $inject $commandExpressionAst
        }
        $newExpression = $this.VisitElement($commandExpressionAst.Expression)
        $newRedirections = $this.VisitElements($commandExpressionAst.Redirections)
        return [CommandExpressionAst]::new($commandExpressionAst.Extent, $newExpression, $newRedirections)
    }

    [object] VisitCommandParameter([CommandParameterAst] $commandParameterAst) {
        if ($inject = $this.Injects['CommandParameter']) {
            return & $inject $commandParameterAst
        }
        $newArgument = $this.VisitElement($commandParameterAst.Argument)
        return [CommandParameterAst]::new($commandParameterAst.Extent, $commandParameterAst.ParameterName, $newArgument, $commandParameterAst.ErrorPosition)
    }

    [object] VisitFileRedirection([FileRedirectionAst] $fileRedirectionAst) {
        if ($inject = $this.Injects['FileRedirection']) {
            return & $inject $fileRedirectionAst
        }
        $newFile = $this.VisitElement($fileRedirectionAst.Location)
        return [FileRedirectionAst]::new($fileRedirectionAst.Extent, $fileRedirectionAst.FromStream, $newFile, $fileRedirectionAst.Append)
    }

    [object] VisitMergingRedirection([MergingRedirectionAst] $mergingRedirectionAst) {
        if ($inject = $this.Injects['MergingRedirection']) {
            return & $inject $mergingRedirectionAst
        }
        return [MergingRedirectionAst]::new($mergingRedirectionAst.Extent, $mergingRedirectionAst.FromStream, $mergingRedirectionAst.ToStream)
    }

    [object] VisitBinaryExpression([BinaryExpressionAst] $binaryExpressionAst) {
        if ($inject = $this.Injects['BinaryExpression']) {
            return & $inject $binaryExpressionAst
        }
        $newLeft = $this.VisitElement($binaryExpressionAst.Left)
        $newRight = $this.VisitElement($binaryExpressionAst.Right)
        return [BinaryExpressionAst]::new($binaryExpressionAst.Extent, $newLeft, $binaryExpressionAst.Operator, $newRight, $binaryExpressionAst.ErrorPosition)
    }

    [object] VisitUnaryExpression([UnaryExpressionAst] $unaryExpressionAst) {
        if ($inject = $this.Injects['UnaryExpression']) {
            return & $inject $unaryExpressionAst
        }
        $newChild = $this.VisitElement($unaryExpressionAst.Child)
        return [UnaryExpressionAst]::new($unaryExpressionAst.Extent, $unaryExpressionAst.TokenKind, $newChild)
    }

    [object] VisitConvertExpression([ConvertExpressionAst] $convertExpressionAst) {
        if ($inject = $this.Injects['ConvertExpression']) {
            return & $inject $convertExpressionAst
        }
        $newChild = $this.VisitElement($convertExpressionAst.Child)
        $newTypeConstraint = $this.VisitElement($convertExpressionAst.Type)
        return [ConvertExpressionAst]::new($convertExpressionAst.Extent, $newTypeConstraint, $newChild)
    }

    [object] VisitTypeExpression([TypeExpressionAst] $typeExpressionAst) {
        if ($inject = $this.Injects['TypeExpression']) {
            return & $inject $typeExpressionAst
        }
        return [TypeExpressionAst]::new($typeExpressionAst.Extent, $typeExpressionAst.TypeName)
    }

    [object] VisitConstantExpression([ConstantExpressionAst] $constantExpressionAst) {
        if ($inject = $this.Injects['ConstantExpression']) {
            return & $inject $constantExpressionAst
        }
        return [ConstantExpressionAst]::new($constantExpressionAst.Extent, $constantExpressionAst.Value)
    }

    [object] VisitStringConstantExpression([StringConstantExpressionAst] $stringConstantExpressionAst) {
        if ($inject = $this.Injects['StringConstantExpression']) {
            return & $inject $stringConstantExpressionAst
        }
        return [StringConstantExpressionAst]::new($stringConstantExpressionAst.Extent, $stringConstantExpressionAst.Value, $stringConstantExpressionAst.StringConstantType)
    }

    [object] VisitSubExpression([SubExpressionAst] $subExpressionAst) {
        if ($inject = $this.Injects['SubExpression']) {
            return & $inject $subExpressionAst
        }
        $newStatementBlock = $this.VisitElement($subExpressionAst.SubExpression)
        return [SubExpressionAst]::new($subExpressionAst.Extent, $newStatementBlock)
    }

    [object] VisitUsingExpression([UsingExpressionAst] $usingExpressionAst) {
        if ($inject = $this.Injects['UsingExpression']) {
            return & $inject $usingExpressionAst
        }
        $newUsingExpr = $this.VisitElement($usingExpressionAst.SubExpression)
        return [UsingExpressionAst]::new($usingExpressionAst.Extent, $newUsingExpr)
    }

    [object] VisitVariableExpression([VariableExpressionAst] $variableExpressionAst) {
        if ($inject = $this.Injects['VariableExpression']) {
            return & $inject $variableExpressionAst
        }
        return [VariableExpressionAst]::new($variableExpressionAst.Extent, $variableExpressionAst.VariablePath.UserPath, $variableExpressionAst.Splatted)
    }

    [object] VisitMemberExpression([MemberExpressionAst] $memberExpressionAst) {
        if ($inject = $this.Injects['MemberExpression']) {
            return & $inject $memberExpressionAst
        }
        $newExpr = $this.VisitElement($memberExpressionAst.Expression)
        $newMember = $this.VisitElement($memberExpressionAst.Member)
        return [MemberExpressionAst]::new($memberExpressionAst.Extent, $newExpr, $newMember, $memberExpressionAst.Static)
    }

    [object] VisitInvokeMemberExpression([InvokeMemberExpressionAst] $invokeMemberExpressionAst) {
        if ($inject = $this.Injects['InvokeMemberExpression']) {
            return & $inject $invokeMemberExpressionAst
        }
        $newExpression = $this.VisitElement($invokeMemberExpressionAst.Expression)
        $newMethod = $this.VisitElement($invokeMemberExpressionAst.Member)
        $newArguments = $this.VisitElements($invokeMemberExpressionAst.Arguments)
        return [InvokeMemberExpressionAst]::new($invokeMemberExpressionAst.Extent, $newExpression, $newMethod, $newArguments, $invokeMemberExpressionAst.Static)
    }

    [object] VisitArrayExpression([ArrayExpressionAst] $arrayExpressionAst) {
        if ($inject = $this.Injects['ArrayExpression']) {
            return & $inject $arrayExpressionAst
        }
        $newStatementBlock = $this.VisitElement($arrayExpressionAst.SubExpression)
        return [ArrayExpressionAst]::new($arrayExpressionAst.Extent, $newStatementBlock)
    }

    [object] VisitArrayLiteral([ArrayLiteralAst] $arrayLiteralAst) {
        if ($inject = $this.Injects['ArrayLiteral']) {
            return & $inject $arrayLiteralAst
        }
        $newArrayElements = $this.VisitElements($arrayLiteralAst.Elements)
        return [ArrayLiteralAst]::new($arrayLiteralAst.Extent, $newArrayElements)
    }

    [object] VisitHashtable([HashtableAst] $hashtableAst) {
        if ($inject = $this.Injects['Hashtable']) {
            return & $inject $hashtableAst
        }
        $newKeyValuePairs = [List[Tuple[ExpressionAst, StatementAst]]]::new()
        foreach ($keyValuePair in $hashtableAst.KeyValuePairs) {
            $newKey = $this.VisitElement($keyValuePair.Item1);
            $newValue = $this.VisitElement($keyValuePair.Item2);
            $newKeyValuePairs.Add([Tuple[ExpressionAst, StatementAst]]::new($newKey, $newValue)) # TODO NOT SURE
        }
        return [HashtableAst]::new($hashtableAst.Extent, $newKeyValuePairs)
    }

    [object] VisitScriptBlockExpression([ScriptBlockExpressionAst] $scriptBlockExpressionAst) {
        if ($inject = $this.Injects['ScriptBlockExpression']) {
            return & $inject $scriptBlockExpressionAst
        }
        $newScriptBlock = $this.VisitElement($scriptBlockExpressionAst.ScriptBlock)
        return [ScriptBlockExpressionAst]::new($scriptBlockExpressionAst.Extent, $newScriptBlock)
    }

    [object] VisitParenExpression([ParenExpressionAst] $parenExpressionAst) {
        if ($inject = $this.Injects['ParenExpression']) {
            return & $inject $parenExpressionAst
        }
        $newPipeline = $this.VisitElement($parenExpressionAst.Pipeline)
        return [ParenExpressionAst]::new($parenExpressionAst.Extent, $newPipeline)
    }

    [object] VisitExpandableStringExpression([ExpandableStringExpressionAst] $expandableStringExpressionAst) {
        if ($inject = $this.Injects['ExpandableStringExpression']) {
            return & $inject $expandableStringExpressionAst
        }
        return [ExpandableStringExpressionAst]::new($expandableStringExpressionAst.Extent, $expandableStringExpressionAst.Value, $expandableStringExpressionAst.StringConstantType)
    }

    [object] VisitIndexExpression([IndexExpressionAst] $indexExpressionAst) {
        if ($inject = $this.Injects['IndexExpression']) {
            return & $inject $indexExpressionAst
        }
        $newTargetExpression = $this.VisitElement($indexExpressionAst.Target)
        $newIndexExpression = $this.VisitElement($indexExpressionAst.Index)
        return [IndexExpressionAst]::new($indexExpressionAst.Extent, $newTargetExpression, $newIndexExpression)
    }

    [object] VisitAttributedExpression([AttributedExpressionAst] $attributedExpressionAst) {
        if ($inject = $this.Injects['AttributedExpression']) {
            return & $inject $attributedExpressionAst
        }
        $newAttribute = $this.VisitElement($attributedExpressionAst.Attribute)
        $newChild = $this.VisitElement($attributedExpressionAst.Child)
        return [AttributedExpressionAst]::new($attributedExpressionAst.Extent, $newAttribute, $newChild)
    }

    [object] VisitBlockStatement([BlockStatementAst] $blockStatementAst) {
        if ($inject = $this.Injects['BlockStatement']) {
            return & $inject $blockStatementAst
        }
        $newBody = $this.VisitElement($blockStatementAst.Body)
        return [BlockStatementAst]::new($blockStatementAst.Extent, $blockStatementAst.Kind, $newBody)
    }

    [object] VisitTypeDefinition([TypeDefinitionAst] $typeDefinitionAst) {
        if ($inject = $this.Injects['TypeDefinition']) {
            return & $inject $typeDefinitionAst
        }
        $newAttributes = $this.VisitElements($typeDefinitionAst.Attributes)
        $newBaseTypes = $this.VisitElements($typeDefinitionAst.BaseTypes)
        $newMembers = $this.VisitElements($typeDefinitionAst.Members)

        return [TypeDefinitionAst]::new($typeDefinitionAst.Extent, $typeDefinitionAst.Name, $newAttributes, $newMembers, $typeDefinitionAst.TypeAttributes, $newBaseTypes)
    }

    [object] VisitPropertyMember([PropertyMemberAst] $propertyMemberAst) {
        if ($inject = $this.Injects['PropertyMember']) {
            return & $inject $propertyMemberAst
        }
        $newPropertyType = $this.VisitElement($propertyMemberAst.PropertyType)
        $newAttributes = $this.VisitElements($propertyMemberAst.Attributes)
        $newInitValue = $this.VisitElement($propertyMemberAst.InitialValue)
        return [PropertyMemberAst]::new($propertyMemberAst.Extent, $propertyMemberAst.Name, $newPropertyType, $newAttributes, $propertyMemberAst.PropertyAttributes, $newInitValue)
    }

    [object] VisitFunctionMember([FunctionMemberAst] $functionMemberAst) {
        if ($inject = $this.Injects['FunctionMember']) {
            return & $inject $functionMemberAst
        }
        $newBody = $this.VisitElement($functionMemberAst.Body)
        $newParameters = $this.VisitElements($functionMemberAst.Parameters)
        $newFunctionDefinition = $this.VisitElement($functionMemberAst.Extent, $false, $false, $functionMemberAst.Name, $newParameters, $newBody)

        $newReturnType = $this.VisitElement($functionMemberAst.ReturnType)
        $newAttributes = $this.VisitElements($functionMemberAst.Attributes)
        return [FunctionMemberAst]::new($functionMemberAst.Extent, $newFunctionDefinition, $newReturnType, $newAttributes, $functionMemberAst.MethodAttributes)
    }

    [object] VisitBaseCtorInvokeMemberExpression([BaseCtorInvokeMemberExpressionAst] $baseCtorAst) {
        if ($inject = $this.Injects['BaseCtorInvokeMemberExpression']) {
            return & $inject $baseCtorAst
        }
        $newInvokeMemberExpression = $this.VisitInvokeMemberExpression($baseCtorAst)
        return [BaseCtorInvokeMemberExpressionAst]::new($baseCtorAst.Expression.Extent, $newInvokeMemberExpression.Extent, $newInvokeMemberExpression)
    }

    [object] VisitUsingStatement([UsingStatementAst] $usingStatementAst) {
        if ($inject = $this.Injects['UsingStatement']) {
            return & $inject $usingStatementAst
        }
        $newName = $this.VisitElement($usingStatementAst.Name)
        $newAlias = if ($usingStatementAst.Alias -is [StringConstantExpressionAst]) {
            $this.VisitElement($usingStatementAst.Alias)
        }

        switch ($usingStatementAst.UsingStatementKind) {
            'Module' {
                if ($usingStatementAst.ModuleSpecification -is [HashtableAst]) {
                    $newModuleSpec = $this.VisitElement($usingStatementAst.ModuleSpecification)
                    if ($newAlias) {
                        return [UsingStatementAst]::new($usingStatementAst.Extent, $newName, $newModuleSpec)
                    }
                    return [UsingStatementAst]::new($usingStatementAst.Extent, $newModuleSpec)
                }
                if ($newAlias) {
                    return [UsingStatementAst]::new($usingStatementAst.Extent, $_, $newName, $newAlias)
                }
                return [UsingStatementAst]::new($usingStatementAst.Extent, $_, $newName)
            }

            default {
                if ($newAlias) {
                    return [UsingStatementAst]::new($usingStatementAst.Extent, $_, $newName, $newAlias)
                }
                return [UsingStatementAst]::new($usingStatementAst.Extent, $_, $newName)
            }
        }

        throw [System.NotImplementedException]::new()
    }

    [object] VisitConfigurationDefinition([ConfigurationDefinitionAst]$configDefinitionAst) {
        if ($inject = $this.Injects['ConfigurationDefinition']) {
            return & $inject $configDefinitionAst
        }
        $newConfig = $this.VisitElement($configDefinitionAst.Body)
        $newInstanceName = $this.VisitElement($configDefinitionAst.InstanceName)
        return [ConfigurationDefinitionAst]::new($configDefinitionAst.Extent, $newConfig, $configDefinitionAst.ConfigurationType, $newInstanceName)
    }

    [object] VisitDynamicKeywordStatement([DynamicKeywordStatementAst]$dynamicKeywordAst) {
        if ($inject = $this.Injects['DynamicKeywordStatement']) {
            return & $inject $dynamicKeywordAst
        }
        $newElements = $this.VisitElements($dynamicKeywordAst.CommandElements)
        return [DynamicKeywordStatementAst]::new($dynamicKeywordAst.Extent, $newElements)
    }

    # V7 nodes
    [object] VisitTernaryExpression([TernaryExpressionAst]$ternaryExpressionAst) {
        if ($inject = $this.Injects['TernaryExpression']) {
            return & $inject $ternaryExpressionAst
        }
        $newCondition = $this.VisitElement($ternaryExpressionAst.Condition)
        $newIfTrue = $this.VisitElement($ternaryExpressionAst.IfTrue)
        $newIfFalse = $this.VisitElement($ternaryExpressionAst.IfFalse)
        return [TernaryExpressionAst]::new($ternaryExpressionAst.Extent, $newCondition, $newIfTrue, $newIfFalse)
    }

    [object] VisitPipelineChain([PipelineChainAst]$pipelineChainAst) {
        if ($inject = $this.Injects['PipelineChain']) {
            return & $inject $pipelineChainAst
        }
        $newLhsPipeline = $this.VisitElement($pipelineChainAst.LhsPipelineChain)
        $newRhsPipeline = $this.VisitElement($pipelineChainAst.RhsPipeline)
        return [PipelineChainAst]::new($pipelineChainAst.Extent, $newLhsPipeline, $newRhsPipeline, $pipelineChainAst.Operator, $pipelineChainAst.Background)
    }
}
#endregion

<#
.SYNOPSIS
    Rewrites a scriptblock AST in a modular manner. 
.DESCRIPTION
    Rewrites a scriptblock AST using a modular visitor, allowing the caller to pass visitor branches as individual script blocks.
.EXAMPLE
    using namespace System.Management.Automation.Language
    $grafted = Graft-Code { 1 + 1 } -BinaryExpression {
      param([BinaryExpressionAst]$ast)

      if ($ast.Operator -eq 'Plus') {                                     # lol
        return [BinaryExpressionAst]::new($ast.Extent, $ast.Left.Copy(), 'Minus', $ast.Right.Copy(), $ast.ErrorPosition)
      }
      return $ast.Copy()
    }

    Replaces the `+` binary operator with `-` - `{1 + 1}` will thus evaluate to `0` after grafting
#>
function Graft-Code {
    [CmdletBinding(DefaultParameterSetName = 'GraftTable')]
    param(
        [Parameter(Mandatory = $true, Position = 0)]
        [scriptblock]$Original,

        [Parameter(Mandatory = $true, Position = 1, ParameterSetName = 'GraftTable')]
        [hashtable]$Grafts,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$ArrayExpression,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$ArrayLiteral,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$AssignmentStatement,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$Attribute,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$AttributedExpression,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$BaseCtorInvokeMemberExpression,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$BinaryExpression,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$BlockStatement,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$BreakStatement,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$CatchClause,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$Command,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$CommandExpression,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$CommandParameter,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$ConfigurationDefinition,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$ConstantExpression,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$ContinueStatement,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$ConvertExpression,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$DataStatement,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$DoUntilStatement,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$DoWhileStatement,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$DynamicKeywordStatement,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$ErrorExpression,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$ErrorStatement,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$ExitStatement,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$ExpandableStringExpression,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$FileRedirection,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$ForEachStatement,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$ForStatement,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$FunctionDefinition,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$FunctionMember,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$Hashtable,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$IfStatement,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$IndexExpression,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$InvokeMemberExpression,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$MemberExpression,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$MergingRedirection,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$NamedAttributeArgument,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$NamedBlock,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$ParamBlock,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$Parameter,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$ParenExpression,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$Pipeline,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$PipelineChain,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$PropertyMember,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$ReturnStatement,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$ScriptBlock,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$ScriptBlockExpression,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$StatementBlock,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$Statements,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$StringConstantExpression,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$SubExpression,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$SwitchStatement,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$TernaryExpression,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$ThrowStatement,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$Trap,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$TryStatement,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$TypeConstraint,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$TypeDefinition,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$TypeExpression,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$UnaryExpression,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$UsingExpression,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$UsingStatement,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$VariableExpression,

        [Parameter(Mandatory = $false, ParameterSetName = 'IndividualParams')]
        [scriptblock]$WhileStatement
    )

    if ($PSCmdlet.ParameterSetName -eq 'IndividualParams') {
        $Grafts = @{}
        $PSBoundParameters.GetEnumerator() |Where-Object Value -is ([ScriptBlock]) |ForEach-Object { $Grafts[$_.Key] = $_.Value }
    }

    $originalAst = $Original.Ast
    $graftedAst = $originalAst.Visit([PSGraftVisitor]::new($Grafts))

    return $graftedAst.GetScriptBlock()
}