PowerShell AST visitor to break up a file by comments

CellFindingVisitor.cs

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

public class ScriptExtent : IScriptExtent
{
    private readonly IScriptPosition _start;

    private readonly IScriptPosition _end;

    private readonly Lazy<string> _textLazy;

    public ScriptExtent(IScriptPosition start, IScriptPosition end)
    {
        _start = start;
        _end = end;
        _textLazy = new Lazy<string>(GetText);
    }

    public int EndColumnNumber => _end.ColumnNumber;

    public int EndLineNumber => _end.LineNumber;

    public int EndOffset => _end.Offset;

    public IScriptPosition EndScriptPosition => _end;

    public string File => _start.File;

    public int StartColumnNumber => _start.ColumnNumber;

    public int StartLineNumber => _start.LineNumber;

    public int StartOffset => _start.Offset;

    public IScriptPosition StartScriptPosition => _start;

    public string Text => _textLazy.Value;

    private string GetText()
    {
        string scriptText = _start.GetFullScript();
        return scriptText?.Substring(_start.Offset, _end.Offset - _start.Offset);
    }
}

public class CellFindingVisitor : AstVisitor2
{
    public static List<IScriptExtent> GetCodeRegionsFromInput(string input)
    {
        Ast ast = Parser.ParseInput(input, out Token[] tokens, out _);
        return GetCodeRegions(ast, tokens);
    }

    public static List<IScriptExtent> GetCodeRegionsFromFile(string filePath)
    {
        Ast ast = Parser.ParseFile(filePath, out Token[] tokens, out _);
        return GetCodeRegions(ast, tokens);
    }

    private static List<IScriptExtent> GetCodeRegions(Ast ast, IReadOnlyList<Token> tokens)
    {
        var visitor = new CellFindingVisitor(ast, tokens);
        ast.Visit(visitor);
        visitor.ProcessScriptEnd();
        return visitor._cellExtents;
    }

    private readonly Ast _ast;

    private readonly IReadOnlyList<Token> _tokens;

    private readonly List<IScriptExtent> _cellExtents;

    private int _tokenIndex;

    private IScriptPosition _currentCellStart;

    private CellFindingVisitor(Ast ast, IReadOnlyList<Token> tokens)
    {
        _ast = ast;
        _tokens = tokens;
        _tokenIndex = 0;
        _cellExtents = new List<IScriptExtent>();
        _currentCellStart = ast.Extent.StartScriptPosition;
    }

    public override AstVisitAction VisitAssignmentStatement(AssignmentStatementAst assignmentStatementAst)
        => VisitStatement(assignmentStatementAst);

    public override AstVisitAction VisitFunctionDefinition(FunctionDefinitionAst functionDefinitionAst)
        => VisitStatement(functionDefinitionAst);

    public override AstVisitAction VisitSwitchStatement(SwitchStatementAst switchStatementAst)
        => VisitStatement(switchStatementAst);

    public override AstVisitAction VisitThrowStatement(ThrowStatementAst throwStatementAst)
        => VisitStatement(throwStatementAst);

    public override AstVisitAction VisitTrap(TrapStatementAst trapStatementAst)
        => VisitStatement(trapStatementAst);

    public override AstVisitAction VisitTryStatement(TryStatementAst tryStatementAst)
        => VisitStatement(tryStatementAst);

    public override AstVisitAction VisitTypeDefinition(TypeDefinitionAst typeDefinitionAst)
        => VisitStatement(typeDefinitionAst);

    public override AstVisitAction VisitUsingStatement(UsingStatementAst usingStatementAst)
        => VisitStatement(usingStatementAst);

    public override AstVisitAction VisitWhileStatement(WhileStatementAst whileStatementAst)
        => VisitStatement(whileStatementAst);

    public override AstVisitAction VisitBlockStatement(BlockStatementAst blockStatementAst)
        => VisitStatement(blockStatementAst);

    public override AstVisitAction VisitBreakStatement(BreakStatementAst breakStatementAst)
        => VisitStatement(breakStatementAst);

    public override AstVisitAction VisitConfigurationDefinition(ConfigurationDefinitionAst configurationDefinitionAst)
        => VisitStatement(configurationDefinitionAst);

    public override AstVisitAction VisitContinueStatement(ContinueStatementAst continueStatementAst)
        => VisitStatement(continueStatementAst);

    public override AstVisitAction VisitDataStatement(DataStatementAst dataStatementAst)
        => VisitStatement(dataStatementAst);

    public override AstVisitAction VisitDoUntilStatement(DoUntilStatementAst doUntilStatementAst)
        => VisitStatement(doUntilStatementAst);

    public override AstVisitAction VisitDoWhileStatement(DoWhileStatementAst doWhileStatementAst)
        => VisitStatement(doWhileStatementAst);

    public override AstVisitAction VisitDynamicKeywordStatement(DynamicKeywordStatementAst dynamicKeywordStatementAst)
        => VisitStatement(dynamicKeywordStatementAst);

    public override AstVisitAction VisitExitStatement(ExitStatementAst exitStatementAst)
        => VisitStatement(exitStatementAst);

    public override AstVisitAction VisitForEachStatement(ForEachStatementAst forEachStatementAst)
        => VisitStatement(forEachStatementAst);

    public override AstVisitAction VisitForStatement(ForStatementAst forStatementAst)
        => VisitStatement(forStatementAst);

    public override AstVisitAction VisitIfStatement(IfStatementAst ifStmtAst)
        => VisitStatement(ifStmtAst);

    public override AstVisitAction VisitPipeline(PipelineAst pipelineAst)
        => VisitStatement(pipelineAst);

    public override AstVisitAction VisitReturnStatement(ReturnStatementAst returnStatementAst)
        => VisitStatement(returnStatementAst);

    private AstVisitAction VisitStatement(StatementAst statementAst)
    {
        ProcessCommentsToPosition(statementAst.Extent.StartScriptPosition);
        SkipTokensToPosition(statementAst.Extent.EndScriptPosition);
        return AstVisitAction.SkipChildren;
    }

    private void ProcessCommentsToPosition(IScriptPosition position)
    {
        IScriptPosition firstCommentPosition = null;
        for (; _tokenIndex < _tokens.Count; _tokenIndex++)
        {
            Token currToken = _tokens[_tokenIndex];

            if (currToken.Extent.StartOffset >= position.Offset)
            {
                break;
            }

            if (currToken.Kind == TokenKind.Comment)
            {
                firstCommentPosition = currToken.Extent.StartScriptPosition;
                break;
            }
        }

        if (firstCommentPosition != null)
        {
            _cellExtents.Add(new ScriptExtent(_currentCellStart, firstCommentPosition));
            _currentCellStart = position;
        }
    }

    private void SkipTokensToPosition(IScriptPosition position)
    {
        for (; _tokenIndex < _tokens.Count; _tokenIndex++)
        {
            Token currToken = _tokens[_tokenIndex];

            if (currToken.Extent.StartOffset > position.Offset)
            {
                break;
            }
        }
    }

    private void ProcessScriptEnd()
    {
        IScriptPosition finalCommentStartPosition = null;
        for (; _tokenIndex < _tokens.Count; _tokenIndex++)
        {
            Token currToken = _tokens[_tokenIndex];

            if (currToken.Kind == TokenKind.Comment)
            {
                finalCommentStartPosition = currToken.Extent.StartScriptPosition;
                break;
            }
        }

        _cellExtents.Add(new ScriptExtent(_currentCellStart, finalCommentStartPosition ?? _ast.Extent.EndScriptPosition));
    }
}