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)); } }