From 24516cdb2d1a236ca037d82f4209a1fd83b446dd Mon Sep 17 00:00:00 2001 From: Kurt Hutten Date: Tue, 25 Jun 2024 12:46:40 +1000 Subject: [PATCH] Improve traverse (#2782) improve travers --- src/lang/queryAst.test.ts | 53 ++++++++++++++++- src/lang/queryAst.ts | 119 +++++++++++++++++++++++++++----------- 2 files changed, 137 insertions(+), 35 deletions(-) diff --git a/src/lang/queryAst.test.ts b/src/lang/queryAst.test.ts index c977f3e80d..271272f31d 100644 --- a/src/lang/queryAst.test.ts +++ b/src/lang/queryAst.test.ts @@ -1,4 +1,4 @@ -import { parse, recast, initPromise } from './wasm' +import { parse, recast, initPromise, PathToNode } from './wasm' import { findAllPreviousVariables, isNodeSafeToReplace, @@ -9,6 +9,7 @@ import { findUsesOfTagInPipe, hasSketchPipeBeenExtruded, hasExtrudableGeometry, + traverse, } from './queryAst' import { enginelessExecutor } from '../lib/testHelpers' import { @@ -538,3 +539,53 @@ const extrude001 = extrude(10, sketch001) expect(extrudable).toBeFalsy() }) }) + +describe.only('Testing traverse and pathToNode', () => { + it.each([ + ['basic', '2.73'], + [ + 'very nested, array, object, callExpression, array, memberExpression', + '.yo', + ], + ])('testing %s', async (testName, literalOfInterest) => { + const code = `const myVar = 5 +const sketch001 = startSketchOn('XZ') + |> startProfileAt([3.29, 7.86], %) + |> line([2.48, 2.44], %) + |> line([-3.86, -2.73], %) + |> line([-17.67, 0.85], %) + |> close(%) +const bing = { yo: 55 } +const myNestedVar = [ + { + prop: line([bing.yo, 21], sketch001) +} +] + ` + const ast = parse(code) + if (err(ast)) throw ast + let pathToNode: PathToNode = [] + traverse(ast, { + enter: (node, path) => { + if ( + node.type === 'Literal' && + String(node.value) === literalOfInterest + ) { + pathToNode = path + } else if ( + node.type === 'Identifier' && + literalOfInterest.includes(node.name) + ) { + pathToNode = path + } + }, + }) + + const literalIndex = code.indexOf(literalOfInterest) + const pathToNode2 = getNodePathFromSourceRange(ast, [ + literalIndex + 2, + literalIndex + 2, + ]) + expect(pathToNode).toEqual(pathToNode2) + }) +}) diff --git a/src/lang/queryAst.ts b/src/lang/queryAst.ts index 4ee57ed30e..babae6e5ed 100644 --- a/src/lang/queryAst.ts +++ b/src/lang/queryAst.ts @@ -270,6 +270,18 @@ function moreNodePathFromSourceRange( } } } + if (_node.type === 'MemberExpression' && isInRange) { + const { object, property } = _node + if (object.start <= start && object.end >= end) { + path.push(['object', 'MemberExpression']) + return moreNodePathFromSourceRange(object, sourceRange, path) + } + if (property.start <= start && property.end >= end) { + path.push(['property', 'MemberExpression']) + return moreNodePathFromSourceRange(property, sourceRange, path) + } + return path + } if (_node.type === 'PipeSubstitution' && isInRange) return path console.error('not implemented: ' + node.type) return path @@ -307,48 +319,87 @@ type KCLNode = | ReturnStatement export function traverse( - node: KCLNode, + node: KCLNode | Program, option: { - enter?: (node: KCLNode) => void + enter?: (node: KCLNode, pathToNode: PathToNode) => void leave?: (node: KCLNode) => void - } + }, + pathToNode: PathToNode = [] ) { - option?.enter?.(node) - const _traverse = (node: KCLNode) => traverse(node, option) - - if (node.type === 'VariableDeclaration') { - node.declarations.forEach(_traverse) - } else if (node.type === 'VariableDeclarator') { - _traverse(node.init) - } else if (node.type === 'PipeExpression') { - node.body.forEach(_traverse) - } else if (node.type === 'CallExpression') { - _traverse(node.callee) - node.arguments.forEach(_traverse) - } else if (node.type === 'BinaryExpression') { - _traverse(node.left) - _traverse(node.right) - } else if (node.type === 'Identifier') { + const _node = node as KCLNode + option?.enter?.(_node, pathToNode) + const _traverse = (node: KCLNode, pathToNode: PathToNode) => + traverse(node, option, pathToNode) + + if (_node.type === 'VariableDeclaration') { + _node.declarations.forEach((declaration, index) => + _traverse(declaration, [ + ...pathToNode, + ['declarations', 'VariableDeclaration'], + [index, 'index'], + ]) + ) + } else if (_node.type === 'VariableDeclarator') { + _traverse(_node.init, [...pathToNode, ['init', '']]) + } else if (_node.type === 'PipeExpression') { + _node.body.forEach((expression, index) => + _traverse(expression, [ + ...pathToNode, + ['body', 'PipeExpression'], + [index, 'index'], + ]) + ) + } else if (_node.type === 'CallExpression') { + _traverse(_node.callee, [...pathToNode, ['callee', 'CallExpression']]) + _node.arguments.forEach((arg, index) => + _traverse(arg, [ + ...pathToNode, + ['arguments', 'CallExpression'], + [index, 'index'], + ]) + ) + } else if (_node.type === 'BinaryExpression') { + _traverse(_node.left, [...pathToNode, ['left', 'BinaryExpression']]) + _traverse(_node.right, [...pathToNode, ['right', 'BinaryExpression']]) + } else if (_node.type === 'Identifier') { // do nothing - } else if (node.type === 'Literal') { + } else if (_node.type === 'Literal') { // do nothing - } else if (node.type === 'ArrayExpression') { - node.elements.forEach(_traverse) - } else if (node.type === 'ObjectExpression') { - node.properties.forEach(({ key, value }) => { - _traverse(key) - _traverse(value) + } else if (_node.type === 'ArrayExpression') { + _node.elements.forEach((el, index) => + _traverse(el, [ + ...pathToNode, + ['elements', 'ArrayExpression'], + [index, 'index'], + ]) + ) + } else if (_node.type === 'ObjectExpression') { + _node.properties.forEach(({ key, value }, index) => { + _traverse(key, [ + ...pathToNode, + ['properties', 'ObjectExpression'], + [index, 'index'], + ['key', 'Property'], + ]) + _traverse(value, [ + ...pathToNode, + ['properties', 'ObjectExpression'], + [index, 'index'], + ['value', 'Property'], + ]) }) - } else if (node.type === 'UnaryExpression') { - _traverse(node.argument) - } else if (node.type === 'MemberExpression') { + } else if (_node.type === 'UnaryExpression') { + _traverse(_node.argument, [...pathToNode, ['argument', 'UnaryExpression']]) + } else if (_node.type === 'MemberExpression') { // hmm this smell - _traverse(node.object) - _traverse(node.property) - } else if ('body' in node && Array.isArray(node.body)) { - node.body.forEach(_traverse) + _traverse(_node.object, [...pathToNode, ['object', 'MemberExpression']]) + _traverse(_node.property, [...pathToNode, ['property', 'MemberExpression']]) + } else if ('body' in _node && Array.isArray(_node.body)) { + _node.body.forEach((expression, index) => + _traverse(expression, [...pathToNode, ['body', ''], [index, 'index']]) + ) } - option?.leave?.(node) + option?.leave?.(_node) } export interface PrevVariable {