Expose variable scopes tracked by the parser to consumers

samples/Acornima.Cli/Commands/PrintScopesCommand.cs

@@ -0,0 +1,68 @@
+using System;
+using System.Linq;
+using System.Xml.Linq;
+using Acornima.Ast;
+using Acornima.Cli.Helpers;
+using Acornima.Jsx;
+using McMaster.Extensions.CommandLineUtils;
+
+namespace Acornima.Cli.Commands;
+
+
+[Command(CommandName, Description = "Parse JS code and print tree of variable scopes.")]
+internal sealed class PrintScopesCommand
+{
+    public const string CommandName = "scopes";
+
+    private readonly IConsole _console;
+
+    public PrintScopesCommand(IConsole console)
+    {
+        _console = console;
+    }
+
+    [Option("--type", Description = "Type of the JS code to parse.")]
+    public JavaScriptCodeType CodeType { get; set; }
+
+    [Option("--jsx", Description = "Allow JSX expressions.")]
+    public bool AllowJsx { get; set; }
+
+    [Argument(0, Description = "The JS code to parse. If omitted, the code will be read from the standard input.")]
+    public string? Code { get; }
+
+    private T CreateParserOptions<T>() where T : ParserOptions, new() => new T().RecordScopeInfoInUserData();
+
+    public int OnExecute()
+    {
+        Console.InputEncoding = System.Text.Encoding.UTF8;
+
+        var code = Code ?? _console.ReadString();
+
+        IParser parser = AllowJsx
+            ? new JsxParser(CreateParserOptions<JsxParserOptions>())
+            : new Parser(CreateParserOptions<ParserOptions>());
+
+        Node rootNode = CodeType switch
+        {
+            JavaScriptCodeType.Script => parser.ParseScript(code),
+            JavaScriptCodeType.Module => parser.ParseModule(code),
+            JavaScriptCodeType.Expression => parser.ParseExpression(code),
+            _ => throw new InvalidOperationException()
+        };
+
+
+        var treePrinter = new TreePrinter(_console);
+        treePrinter.Print(new[] { rootNode },
+            node => node
+                .DescendantNodes(descendIntoChildren: descendantNode => ReferenceEquals(node, descendantNode) || descendantNode.UserData is not ScopeInfo)
+                .Where(node => node.UserData is ScopeInfo),
+            node =>
+            {
+                var scopeInfo = (ScopeInfo)node.UserData!;
+                var names = scopeInfo.VarNames.Concat(scopeInfo.LexicalNames).Concat(scopeInfo.FunctionNames);
+                return $"{node.TypeText} ({string.Join(", ", names)})";
+            });
+
+        return 0;
+    }
+}

samples/Acornima.Cli/Program.cs

@@ -8,7 +8,7 @@ namespace Acornima.Cli;
 [Command("acornima", Description = "A command line tool for testing Acornima features.",
     UnrecognizedArgumentHandling = UnrecognizedArgumentHandling.StopParsingAndCollect)]
 [HelpOption(Inherited = true)]
-[Subcommand(typeof(ParseCommand), typeof(TokenizeCommand))]
+[Subcommand(typeof(ParseCommand), typeof(TokenizeCommand), typeof(PrintScopesCommand))]
 public class Program
 {
     public static int Main(string[] args)

src/Acornima.Extras/Jsx/JsxParser.cs

@@ -338,6 +338,6 @@ private JsxEmptyExpression ParseEmptyExpression()
 
         ref readonly var tokenizer = ref _parser._tokenizer;
         var startMarker = new Marker(tokenizer._lastTokenEnd, tokenizer._lastTokenEndLocation);
-        return _parser.FinishNodeAt(startMarker, _parser.StartNode(), new JsxEmptyExpression());
+        return _parser.FinishNodeAt(startMarker, _parser.StartNode(), new JsxEmptyExpression(), NullRef<Scope>());
     }
 }

src/Acornima.Extras/ParserOptionsExtensions.cs

@@ -1,21 +1,44 @@
 using System;
+using Acornima.Ast;
 
 namespace Acornima;
 
 public static class ParserOptionsExtensions
 {
-    private static readonly OnNodeHandler s_parentSetter = node =>
+    private static readonly OnNodeHandler s_parentSetter = (Node node, in Scope _, ReadOnlySpan<Scope> _) =>
     {
         foreach (var child in node.ChildNodes)
         {
             child.UserData = node;
         }
     };
 
+    private static readonly OnNodeHandler s_scopeInfoSetter = (Node node, in Scope scope, ReadOnlySpan<Scope> scopeStack) =>
+    {
+        if (!Scope.IsNullRef(scope))
+        {
+            node.UserData = new ScopeInfo(node, scope, scopeStack);
+        }
+
+        foreach (var child in node.ChildNodes)
+        {
+            if (child.UserData is ScopeInfo scopeInfo)
+            {
+                scopeInfo.AssociatedNodeParent = node;
+            }
+            else
+            {
+                child.UserData = node;
+            }
+        }
+    };
+
     /// <remarks>
     /// WARNING: Setting <see cref="ParserOptions.OnNode"/> after enabling this setting will cancel parent node recording.
+    /// Enabling this together with <see cref="RecordScopeInfoInUserData"/> is an undefined behavior.
     /// </remarks>
-    public static ParserOptions RecordParentNodeInUserData(this ParserOptions options, bool enable)
+    public static TOptions RecordParentNodeInUserData<TOptions>(this TOptions options, bool enable = true)
+        where TOptions : ParserOptions
     {
         Delegate.RemoveAll(options.OnNode, s_parentSetter);
 
@@ -26,4 +49,21 @@ public static ParserOptions RecordParentNodeInUserData(this ParserOptions option
 
         return options;
     }
+
+    /// <remarks>
+    /// WARNING: Setting <see cref="ParserOptions.OnNode"/> after enabling this setting will cancel parent node recording.
+    /// Enabling this together with <see cref="RecordParentNodeInUserData"/> is an undefined behavior.
+    /// </remarks>
+    public static TOptions RecordScopeInfoInUserData<TOptions>(this TOptions options, bool enable = true)
+        where TOptions : ParserOptions
+    {
+        Delegate.RemoveAll(options.OnNode, s_scopeInfoSetter);
+
+        if (enable)
+        {
+            options._onNode += s_scopeInfoSetter;
+        }
+
+        return options;
+    }
 }

src/Acornima.Extras/ScopeInfo.cs

@@ -0,0 +1,53 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using Acornima.Ast;
+
+namespace Acornima;
+
+public sealed class ScopeInfo
+{
+    private readonly int _id;
+    private readonly int _currentVarScopeId;
+    private readonly int _currentThisScopeId;
+
+    public ScopeInfo(Node node, in Scope scope, ReadOnlySpan<Scope> scopeStack)
+    {
+        AssociatedNode = node;
+        _id = scope.Id;
+        _currentVarScopeId = scope.CurrentVarScopeIndex == scopeStack.Length ? scope.Id : scopeStack[scope.CurrentVarScopeIndex].Id;
+        _currentThisScopeId = scope.CurrentThisScopeIndex == scopeStack.Length ? scope.Id : scopeStack[scope.CurrentThisScopeIndex].Id;
+        VarNames = scope.VarNames.ToArray();
+        LexicalNames = scope.LexicalNames.ToArray();
+        FunctionNames = scope.FunctionNames.ToArray();
+    }
+
+    public Node AssociatedNode { get; }
+    public Node? AssociatedNodeParent { get; internal set; }
+
+    public string[] VarNames { get; }
+    public string[] LexicalNames { get; }
+    public string[] FunctionNames { get; }
+
+    // These lookups could as well be cached.
+    public ScopeInfo? Parent => FindAncestor(_ => true);
+    public ScopeInfo? VarScope => FindAncestor(scope => scope._id == _currentVarScopeId);
+    public ScopeInfo? ThisScope => FindAncestor(scope => scope._id == _currentThisScopeId);
+
+    private ScopeInfo? FindAncestor(Func<ScopeInfo, bool> predicate)
+    {
+        var node = AssociatedNode;
+        while ((node = GetParentNode(node!)!) is not null)
+        {
+            if (node.UserData is ScopeInfo scopeInfo && predicate(scopeInfo))
+            {
+                return scopeInfo;
+            }
+        }
+
+        return null;
+    }
+
+    private static Node? GetParentNode(Node node)
+        => node.UserData is ScopeInfo scopeInfo ? scopeInfo.AssociatedNodeParent : (Node?)node.UserData;
+}

src/Acornima/Parser.Expression.cs

@@ -1199,7 +1199,7 @@ private Expression ParseParenAndDistinguishExpression(bool canBeArrow, Expressio
                 {
                     exprList[i] = (Expression)parameters[i];
                 }
-                val = FinishNodeAt(innerStartMarker, innerEndMarker, new SequenceExpression(NodeList.From(ref exprList)));
+                val = FinishNodeAt(innerStartMarker, innerEndMarker, new SequenceExpression(NodeList.From(ref exprList)), NullRef<Scope>());
             }
             else
             {
@@ -1649,13 +1649,13 @@ private FunctionExpression ParseMethod(bool isGenerator, bool isAsync = false, b
 
         NodeList<Node> parameters = ParseBindingList(close: TokenType.ParenRight, allowEmptyElement: false, allowTrailingComma: _tokenizerOptions._ecmaVersion >= EcmaVersion.ES8)!;
         CheckYieldAwaitInDefaultParams();
-        var body = ParseFunctionBody(id: null, parameters, isArrowFunction: false, isMethod: true, ExpressionContext.Default, out _);
+        ref var scope = ref ParseFunctionBody(id: null, parameters, isArrowFunction: false, isMethod: true, ExpressionContext.Default, out _, out var body);
 
         _yieldPosition = oldYieldPos;
         _awaitPosition = oldAwaitPos;
         _awaitIdentifierPosition = oldAwaitIdentPos;
 
-        return FinishNode(startMarker, new FunctionExpression(id: null, parameters, (FunctionBody)body, isGenerator, isAsync));
+        return FinishNode(startMarker, new FunctionExpression(id: null, parameters, (FunctionBody)body, isGenerator, isAsync), scope);
     }
 
     // Parse arrow function expression with given parameters.
@@ -1674,25 +1674,22 @@ private ArrowFunctionExpression ParseArrowExpression(in Marker startMarker, in N
         EnterScope(FunctionFlags(isAsync, generator: false) | ScopeFlags.Arrow);
 
         NodeList<Node> paramList = ToAssignableList(parameters!, isBinding: true, isParams: true)!;
-        var body = ParseFunctionBody(id: null, paramList, isArrowFunction: true, isMethod: false, context, out var expression);
+        ref var scope = ref ParseFunctionBody(id: null, paramList, isArrowFunction: true, isMethod: false, context, out var expression, out var body);
 
         _yieldPosition = oldYieldPos;
         _awaitPosition = oldAwaitPos;
         _awaitIdentifierPosition = oldAwaitIdentPos;
 
-        return FinishNode(startMarker, new ArrowFunctionExpression(paramList, body, expression, isAsync));
+        return FinishNode(startMarker, new ArrowFunctionExpression(paramList, body, expression, isAsync), scope);
     }
 
     // Parse function body and check parameters.
-    private StatementOrExpression ParseFunctionBody(Identifier? id, in NodeList<Node> parameters,
-        bool isArrowFunction, bool isMethod, ExpressionContext context, out bool expression)
+    private ref Scope ParseFunctionBody(Identifier? id, in NodeList<Node> parameters, bool isArrowFunction, bool isMethod, ExpressionContext context,
+        out bool expression, out StatementOrExpression body)
     {
         // https://github.com/acornjs/acorn/blob/8.11.3/acorn/src/expression.js > `pp.parseFunctionBody = function`
 
         expression = isArrowFunction && _tokenizer._type != TokenType.BraceLeft;
-        var strict = false;
-
-        StatementOrExpression body;
         if (expression)
         {
             CheckParams(parameters, allowDuplicates: false);
@@ -1707,7 +1704,7 @@ private StatementOrExpression ParseFunctionBody(Identifier? id, in NodeList<Node
             var oldStrict = _strict;
             Expect(TokenType.BraceLeft);
             var statements = ParseDirectivePrologue(allowStrictDirective: !nonSimple);
-            strict = _strict;
+            var strict = _strict;
 
             // Add the params to varDeclaredNames to ensure that an error is thrown
             // if a let/const declaration in the function clashes with one of the params.
@@ -1723,15 +1720,13 @@ private StatementOrExpression ParseFunctionBody(Identifier? id, in NodeList<Node
             // flag (restore them to their old value afterwards).
             var oldLabels = _labels;
             _labels = new ArrayList<Label>();
-            ParseBlock(ref statements, createNewLexicalScope: false, exitStrict: strict && !oldStrict);
+            ref var scope = ref ParseBlock(ref statements, createNewLexicalScope: false, exitStrict: strict && !oldStrict);
             _labels = oldLabels;
 
-            body = FinishNode(startMarker, new FunctionBody(NodeList.From(ref statements), strict));
+            body = FinishNode(startMarker, new FunctionBody(NodeList.From(ref statements), strict), scope);
         }
 
-        ExitScope();
-
-        return body;
+        return ref ExitScope();
     }
 
     private static bool IsSimpleParamList(in NodeList<Node> parameters)

src/Acornima/Parser.Helpers.cs

@@ -7,6 +7,7 @@
 
 namespace Acornima;
 
+using static Unsafe;
 using static SyntaxErrorMessages;
 
 public partial class Parser
@@ -25,13 +26,13 @@ internal Marker StartNode()
         return new Marker(_tokenizer._start, _tokenizer._startLocation);
     }
 
-    internal T FinishNodeAt<T>(in Marker startMarker, in Marker endMarker, T node) where T : Node
+    internal T FinishNodeAt<T>(in Marker startMarker, in Marker endMarker, T node, in Scope scope) where T : Node
     {
         // https://github.com/acornjs/acorn/blob/8.11.3/acorn/src/node.js > `function finishNodeAt`, `pp.finishNodeAt = function`
 
         node._range = new Range(startMarker.Index, endMarker.Index);
         node._location = new SourceLocation(startMarker.Position, endMarker.Position, _tokenizer._sourceFile);
-        _options._onNode?.Invoke(node);
+        _options._onNode?.Invoke(node, scope, _scopeStack.AsReadOnlySpan());
         return node;
     }
 
@@ -40,14 +41,22 @@ internal T FinishNode<T>(in Marker startMarker, T node) where T : Node
     {
         // https://github.com/acornjs/acorn/blob/8.11.3/acorn/src/node.js > `pp.finishNode = function`
 
-        return FinishNodeAt(startMarker, new Marker(_tokenizer._lastTokenEnd, _tokenizer._lastTokenEndLocation), node);
+        return FinishNodeAt(startMarker, new Marker(_tokenizer._lastTokenEnd, _tokenizer._lastTokenEndLocation), node, NullRef<Scope>());
+    }
+
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    internal T FinishNode<T>(in Marker startMarker, T node, in Scope scope) where T : Node
+    {
+        // https://github.com/acornjs/acorn/blob/8.11.3/acorn/src/node.js > `pp.finishNode = function`
+
+        return FinishNodeAt(startMarker, new Marker(_tokenizer._lastTokenEnd, _tokenizer._lastTokenEndLocation), node, scope);
     }
 
     private T ReinterpretNode<T>(Node originalNode, T node) where T : Node
     {
         node._range = originalNode._range;
         node._location = originalNode._location;
-        _options._onNode?.Invoke(node);
+        _options._onNode?.Invoke(node, NullRef<Scope>(), _scopeStack.AsReadOnlySpan());
         return node;
     }