feat: Walker traversal pruning (#113)

## Summary

Unlike `walker.abort()` , this skips only the current node's children
and then continues traversal with sibling nodes.
This is useful when callers need to avoid nested scopes without stopping
the whole walk.

Author: bartoval

src/ast/walker/__tests__/walker_skip_children.test.ts

@@ -0,0 +1,157 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { EsqlQuery } from '../../../composer/query';
+import type { ESQLIntegerLiteral } from '../../../types';
+import { Walker } from '../walker';
+
+describe('skipping children', () => {
+  test('can skip command arguments', () => {
+    const { ast } = EsqlQuery.fromSrc('FROM a, b, c');
+    const sources: string[] = [];
+
+    Walker.walk(ast, {
+      visitCommand: (_node, _parent, walker) => {
+        walker.skipChildren();
+      },
+      visitSource: (node) => {
+        sources.push(node.name);
+      },
+    });
+
+    expect(sources).toStrictEqual([]);
+  });
+
+  test('can skip function arguments', () => {
+    const { ast } = EsqlQuery.fromSrc('ROW fn(1, 2, 3, gg(4, 5))');
+    const literals: number[] = [];
+
+    Walker.walk(ast, {
+      visitFunction: (node, _parent, walker) => {
+        if (node.name === 'fn') walker.skipChildren();
+      },
+      visitLiteral: (node) => {
+        literals.push(node.value as number);
+      },
+    });
+
+    expect(literals).toStrictEqual([]);
+  });
+
+  test('can skip fields of a command option', () => {
+    const { ast } = EsqlQuery.fromSrc('FROM index METADATA a, b, c');
+    const columns: string[] = [];
+
+    Walker.walk(ast, {
+      visitCommandOption: (_node, _parent, walker) => {
+        walker.skipChildren();
+      },
+      visitColumn: (node) => {
+        columns.push(node.name);
+      },
+    });
+
+    expect(columns).toStrictEqual([]);
+  });
+
+  test('can skip entries of a map', () => {
+    const { ast } = EsqlQuery.fromSrc('ROW fn(TRUE, { "foo": 1, "bar": 2, "baz": 3 })');
+    const keys: string[] = [];
+    const values: number[] = [];
+
+    Walker.walk(ast, {
+      visitMap: (_node, _parent, walker) => {
+        walker.skipChildren();
+      },
+      visitMapEntry: (node) => {
+        if (node.key.type === 'literal' && node.key.literalType === 'keyword') {
+          keys.push(node.key.valueUnquoted);
+        }
+        values.push((node.value as ESQLIntegerLiteral).value);
+      },
+    });
+
+    expect(keys).toStrictEqual([]);
+    expect(values).toStrictEqual([]);
+  });
+
+  test('can skip children of a map entry', () => {
+    const { ast } = EsqlQuery.fromSrc('ROW fn(TRUE, { "foo": 1, "bar": 2, "baz": 3 })');
+    const keys: string[] = [];
+    const values: number[] = [];
+
+    Walker.walk(ast, {
+      visitMapEntry: (node, _parent, walker) => {
+        if (node.key.type === 'literal' && node.key.literalType === 'keyword') {
+          keys.push(node.key.valueUnquoted);
+        }
+        walker.skipChildren();
+      },
+      visitLiteral: (node) => {
+        if (node.literalType === 'integer') {
+          values.push((node as ESQLIntegerLiteral).value);
+        }
+      },
+    });
+
+    expect(keys).toStrictEqual(['foo', 'bar', 'baz']);
+    expect(values).toStrictEqual([]);
+  });
+
+  test('sibling commands are still traversed when a command skips its children', () => {
+    const { ast } = EsqlQuery.fromSrc('FROM index | LIMIT 123');
+    const commands: string[] = [];
+    const sources: string[] = [];
+    const literals: number[] = [];
+
+    Walker.walk(ast, {
+      visitCommand: (node, parent, walker) => {
+        commands.push(node.name);
+        if (node.name === 'from') walker.skipChildren();
+      },
+      visitSource: (node) => sources.push(node.name),
+      visitLiteral: (node) => literals.push(node.value as number),
+    });
+
+    expect(commands).toStrictEqual(['from', 'limit']);
+    expect(sources).toStrictEqual([]);
+    expect(literals).toStrictEqual([123]);
+  });
+
+  test('can skip components of a source', () => {
+    const { ast } = EsqlQuery.fromSrc('FROM a:b, c::d');
+    const components: string[] = [];
+
+    Walker.walk(ast, {
+      visitSource: (node, _parent, walker) => {
+        if (node.name === 'a:b') walker.skipChildren();
+      },
+      visitLiteral: (node) => {
+        components.push(node.value as string);
+      },
+    });
+
+    expect(components).toStrictEqual(['c', 'd']);
+  });
+
+  test('can skip components of a source (backward)', () => {
+    const { ast } = EsqlQuery.fromSrc('FROM a:b, c::d');
+    const components: string[] = [];
+
+    Walker.walk(ast, {
+      visitSource: (node, _parent, walker) => {
+        if (node.name === 'a:b') walker.skipChildren();
+      },
+      visitLiteral: (node) => {
+        components.push(node.value as string);
+      },
+      order: 'backward',
+    });
+
+    expect(components).toStrictEqual(['d', 'c']);
+  });
+});

src/ast/walker/walker.ts

@@ -146,7 +146,7 @@ export interface WalkerOptions {
 
 export type WalkerAstNode = types.ESQLAstNode | types.ESQLAstNode[];
 
-export type WalkerVisitorApi = Pick<Walker, 'abort'>;
+export type WalkerVisitorApi = Pick<Walker, 'abort' | 'skipChildren'>;
 
 /**
  * Iterates over all nodes in the AST and calls the appropriate visitor
@@ -539,13 +539,24 @@ export class Walker {
   };
 
   protected aborted: boolean = false;
+  protected skippedChildren: boolean = false;
 
   constructor(protected readonly options: WalkerOptions) {}
 
   public abort(): void {
     this.aborted = true;
   }
 
+  public skipChildren(): void {
+    this.skippedChildren = true;
+  }
+
+  protected readAndResetSkippedChildren(): boolean {
+    const skipped = this.skippedChildren;
+    this.skippedChildren = false;
+    return skipped;
+  }
+
   public walk(
     tree: WalkerAstNode | undefined,
     parent: types.ESQLProperNode | undefined = undefined
@@ -594,6 +605,8 @@ export class Walker {
 
     const { options } = this;
     (options.visitCommand ?? options.visitAny)?.(node, parent, this);
+    if (this.readAndResetSkippedChildren()) return;
+
     this.walkList(node.args, node);
   }
 
@@ -605,12 +618,16 @@ export class Walker {
 
     const { options } = this;
     (options.visitHeaderCommand ?? options.visitAny)?.(node, parent, this);
+    if (this.readAndResetSkippedChildren()) return;
+
     this.walkList(node.args, node);
   }
 
   public walkOption(node: types.ESQLCommandOption, parent: types.ESQLCommand | undefined): void {
     const { options } = this;
     (options.visitCommandOption ?? options.visitAny)?.(node, parent, this);
+    if (this.readAndResetSkippedChildren()) return;
+
     this.walkList(node.args, node);
   }
 
@@ -630,13 +647,16 @@ export class Walker {
   public walkListLiteral(node: types.ESQLList, parent: types.ESQLProperNode | undefined): void {
     const { options } = this;
     (options.visitListLiteral ?? options.visitAny)?.(node, parent, this);
+    if (this.readAndResetSkippedChildren()) return;
+
     this.walkList(node.values, node);
   }
 
   public walkSource(node: types.ESQLSource, parent: types.ESQLProperNode | undefined): void {
     const { options } = this;
 
     (options.visitSource ?? options.visitAny)?.(node, parent, this);
+    if (this.readAndResetSkippedChildren()) return;
 
     const children: types.ESQLStringLiteral[] = [];
 
@@ -658,6 +678,7 @@ export class Walker {
     const { args } = node;
 
     (options.visitColumn ?? options.visitAny)?.(node, parent, this);
+    if (this.readAndResetSkippedChildren()) return;
 
     if (args) {
       this.walkList(args, node);
@@ -671,6 +692,7 @@ export class Walker {
     const { options } = this;
 
     (options.visitOrder ?? options.visitAny)?.(node, parent, this);
+    if (this.readAndResetSkippedChildren()) return;
 
     this.walkList(node.args, node);
   }
@@ -681,6 +703,8 @@ export class Walker {
   ): void {
     const { options } = this;
     (options.visitInlineCast ?? options.visitAny)?.(node, parent, this);
+    if (this.readAndResetSkippedChildren()) return;
+
     this.walkExpression(node.value, node);
   }
 
@@ -690,6 +714,7 @@ export class Walker {
   ): void {
     const { options } = this;
     (options.visitFunction ?? options.visitAny)?.(node, parent, this);
+    if (this.readAndResetSkippedChildren()) return;
 
     if (node.operator) this.walkSingleAstItem(node.operator, node);
 
@@ -699,13 +724,16 @@ export class Walker {
   public walkMap(node: types.ESQLMap, parent: types.ESQLProperNode | undefined): void {
     const { options } = this;
     (options.visitMap ?? options.visitAny)?.(node, parent, this);
+    if (this.readAndResetSkippedChildren()) return;
+
     this.walkList(node.entries, node);
   }
 
   public walkMapEntry(node: types.ESQLMapEntry, parent: types.ESQLProperNode | undefined): void {
     const { options } = this;
 
     (options.visitMapEntry ?? options.visitAny)?.(node, parent, this);
+    if (this.readAndResetSkippedChildren()) return;
 
     if (options.order === 'backward') {
       this.walkSingleAstItem(resolveItem(node.value), node);
@@ -719,6 +747,7 @@ export class Walker {
   public walkParens(node: types.ESQLParens, parent: types.ESQLProperNode | undefined): void {
     const { options } = this;
     (options.visitParens ?? options.visitAny)?.(node, parent, this);
+    if (this.readAndResetSkippedChildren()) return;
 
     if (node.child) {
       this.walkSingleAstItem(node.child, node);
@@ -731,6 +760,7 @@ export class Walker {
   ): void {
     const { options } = this;
     (options.visitQuery ?? options.visitAny)?.(node, parent, this);
+    if (this.readAndResetSkippedChildren()) return;
 
     if (node.header && !options.skipHeader) {
       this.walkList(node.header, node);
@@ -816,6 +846,7 @@ export class Walker {
       }
       case 'literal': {
         (options.visitLiteral ?? options.visitAny)?.(node, parent, this);
+        this.readAndResetSkippedChildren();
         break;
       }
       case 'list': {
@@ -832,10 +863,12 @@ export class Walker {
       }
       case 'identifier': {
         (options.visitIdentifier ?? options.visitAny)?.(node, parent, this);
+        this.readAndResetSkippedChildren();
         break;
       }
       case 'unknown': {
         (options.visitUnknown ?? options.visitAny)?.(node, parent, this);
+        this.readAndResetSkippedChildren();
         break;
       }
     }