visitor.mjs 9.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361
  1. import { devAssert } from '../jsutils/devAssert.mjs';
  2. import { inspect } from '../jsutils/inspect.mjs';
  3. import { isNode, QueryDocumentKeys } from './ast.mjs';
  4. import { Kind } from './kinds.mjs';
  5. /**
  6. * A visitor is provided to visit, it contains the collection of
  7. * relevant functions to be called during the visitor's traversal.
  8. */
  9. export const BREAK = Object.freeze({});
  10. /**
  11. * visit() will walk through an AST using a depth-first traversal, calling
  12. * the visitor's enter function at each node in the traversal, and calling the
  13. * leave function after visiting that node and all of its child nodes.
  14. *
  15. * By returning different values from the enter and leave functions, the
  16. * behavior of the visitor can be altered, including skipping over a sub-tree of
  17. * the AST (by returning false), editing the AST by returning a value or null
  18. * to remove the value, or to stop the whole traversal by returning BREAK.
  19. *
  20. * When using visit() to edit an AST, the original AST will not be modified, and
  21. * a new version of the AST with the changes applied will be returned from the
  22. * visit function.
  23. *
  24. * ```ts
  25. * const editedAST = visit(ast, {
  26. * enter(node, key, parent, path, ancestors) {
  27. * // @return
  28. * // undefined: no action
  29. * // false: skip visiting this node
  30. * // visitor.BREAK: stop visiting altogether
  31. * // null: delete this node
  32. * // any value: replace this node with the returned value
  33. * },
  34. * leave(node, key, parent, path, ancestors) {
  35. * // @return
  36. * // undefined: no action
  37. * // false: no action
  38. * // visitor.BREAK: stop visiting altogether
  39. * // null: delete this node
  40. * // any value: replace this node with the returned value
  41. * }
  42. * });
  43. * ```
  44. *
  45. * Alternatively to providing enter() and leave() functions, a visitor can
  46. * instead provide functions named the same as the kinds of AST nodes, or
  47. * enter/leave visitors at a named key, leading to three permutations of the
  48. * visitor API:
  49. *
  50. * 1) Named visitors triggered when entering a node of a specific kind.
  51. *
  52. * ```ts
  53. * visit(ast, {
  54. * Kind(node) {
  55. * // enter the "Kind" node
  56. * }
  57. * })
  58. * ```
  59. *
  60. * 2) Named visitors that trigger upon entering and leaving a node of a specific kind.
  61. *
  62. * ```ts
  63. * visit(ast, {
  64. * Kind: {
  65. * enter(node) {
  66. * // enter the "Kind" node
  67. * }
  68. * leave(node) {
  69. * // leave the "Kind" node
  70. * }
  71. * }
  72. * })
  73. * ```
  74. *
  75. * 3) Generic visitors that trigger upon entering and leaving any node.
  76. *
  77. * ```ts
  78. * visit(ast, {
  79. * enter(node) {
  80. * // enter any node
  81. * },
  82. * leave(node) {
  83. * // leave any node
  84. * }
  85. * })
  86. * ```
  87. */
  88. export function visit(root, visitor, visitorKeys = QueryDocumentKeys) {
  89. const enterLeaveMap = new Map();
  90. for (const kind of Object.values(Kind)) {
  91. enterLeaveMap.set(kind, getEnterLeaveForKind(visitor, kind));
  92. }
  93. /* eslint-disable no-undef-init */
  94. let stack = undefined;
  95. let inArray = Array.isArray(root);
  96. let keys = [root];
  97. let index = -1;
  98. let edits = [];
  99. let node = root;
  100. let key = undefined;
  101. let parent = undefined;
  102. const path = [];
  103. const ancestors = [];
  104. /* eslint-enable no-undef-init */
  105. do {
  106. index++;
  107. const isLeaving = index === keys.length;
  108. const isEdited = isLeaving && edits.length !== 0;
  109. if (isLeaving) {
  110. key = ancestors.length === 0 ? undefined : path[path.length - 1];
  111. node = parent;
  112. parent = ancestors.pop();
  113. if (isEdited) {
  114. if (inArray) {
  115. node = node.slice();
  116. let editOffset = 0;
  117. for (const [editKey, editValue] of edits) {
  118. const arrayKey = editKey - editOffset;
  119. if (editValue === null) {
  120. node.splice(arrayKey, 1);
  121. editOffset++;
  122. } else {
  123. node[arrayKey] = editValue;
  124. }
  125. }
  126. } else {
  127. node = Object.defineProperties(
  128. {},
  129. Object.getOwnPropertyDescriptors(node),
  130. );
  131. for (const [editKey, editValue] of edits) {
  132. node[editKey] = editValue;
  133. }
  134. }
  135. }
  136. index = stack.index;
  137. keys = stack.keys;
  138. edits = stack.edits;
  139. inArray = stack.inArray;
  140. stack = stack.prev;
  141. } else if (parent) {
  142. key = inArray ? index : keys[index];
  143. node = parent[key];
  144. if (node === null || node === undefined) {
  145. continue;
  146. }
  147. path.push(key);
  148. }
  149. let result;
  150. if (!Array.isArray(node)) {
  151. var _enterLeaveMap$get, _enterLeaveMap$get2;
  152. isNode(node) || devAssert(false, `Invalid AST Node: ${inspect(node)}.`);
  153. const visitFn = isLeaving
  154. ? (_enterLeaveMap$get = enterLeaveMap.get(node.kind)) === null ||
  155. _enterLeaveMap$get === void 0
  156. ? void 0
  157. : _enterLeaveMap$get.leave
  158. : (_enterLeaveMap$get2 = enterLeaveMap.get(node.kind)) === null ||
  159. _enterLeaveMap$get2 === void 0
  160. ? void 0
  161. : _enterLeaveMap$get2.enter;
  162. result =
  163. visitFn === null || visitFn === void 0
  164. ? void 0
  165. : visitFn.call(visitor, node, key, parent, path, ancestors);
  166. if (result === BREAK) {
  167. break;
  168. }
  169. if (result === false) {
  170. if (!isLeaving) {
  171. path.pop();
  172. continue;
  173. }
  174. } else if (result !== undefined) {
  175. edits.push([key, result]);
  176. if (!isLeaving) {
  177. if (isNode(result)) {
  178. node = result;
  179. } else {
  180. path.pop();
  181. continue;
  182. }
  183. }
  184. }
  185. }
  186. if (result === undefined && isEdited) {
  187. edits.push([key, node]);
  188. }
  189. if (isLeaving) {
  190. path.pop();
  191. } else {
  192. var _node$kind;
  193. stack = {
  194. inArray,
  195. index,
  196. keys,
  197. edits,
  198. prev: stack,
  199. };
  200. inArray = Array.isArray(node);
  201. keys = inArray
  202. ? node
  203. : (_node$kind = visitorKeys[node.kind]) !== null &&
  204. _node$kind !== void 0
  205. ? _node$kind
  206. : [];
  207. index = -1;
  208. edits = [];
  209. if (parent) {
  210. ancestors.push(parent);
  211. }
  212. parent = node;
  213. }
  214. } while (stack !== undefined);
  215. if (edits.length !== 0) {
  216. // New root
  217. return edits[edits.length - 1][1];
  218. }
  219. return root;
  220. }
  221. /**
  222. * Creates a new visitor instance which delegates to many visitors to run in
  223. * parallel. Each visitor will be visited for each node before moving on.
  224. *
  225. * If a prior visitor edits a node, no following visitors will see that node.
  226. */
  227. export function visitInParallel(visitors) {
  228. const skipping = new Array(visitors.length).fill(null);
  229. const mergedVisitor = Object.create(null);
  230. for (const kind of Object.values(Kind)) {
  231. let hasVisitor = false;
  232. const enterList = new Array(visitors.length).fill(undefined);
  233. const leaveList = new Array(visitors.length).fill(undefined);
  234. for (let i = 0; i < visitors.length; ++i) {
  235. const { enter, leave } = getEnterLeaveForKind(visitors[i], kind);
  236. hasVisitor || (hasVisitor = enter != null || leave != null);
  237. enterList[i] = enter;
  238. leaveList[i] = leave;
  239. }
  240. if (!hasVisitor) {
  241. continue;
  242. }
  243. const mergedEnterLeave = {
  244. enter(...args) {
  245. const node = args[0];
  246. for (let i = 0; i < visitors.length; i++) {
  247. if (skipping[i] === null) {
  248. var _enterList$i;
  249. const result =
  250. (_enterList$i = enterList[i]) === null || _enterList$i === void 0
  251. ? void 0
  252. : _enterList$i.apply(visitors[i], args);
  253. if (result === false) {
  254. skipping[i] = node;
  255. } else if (result === BREAK) {
  256. skipping[i] = BREAK;
  257. } else if (result !== undefined) {
  258. return result;
  259. }
  260. }
  261. }
  262. },
  263. leave(...args) {
  264. const node = args[0];
  265. for (let i = 0; i < visitors.length; i++) {
  266. if (skipping[i] === null) {
  267. var _leaveList$i;
  268. const result =
  269. (_leaveList$i = leaveList[i]) === null || _leaveList$i === void 0
  270. ? void 0
  271. : _leaveList$i.apply(visitors[i], args);
  272. if (result === BREAK) {
  273. skipping[i] = BREAK;
  274. } else if (result !== undefined && result !== false) {
  275. return result;
  276. }
  277. } else if (skipping[i] === node) {
  278. skipping[i] = null;
  279. }
  280. }
  281. },
  282. };
  283. mergedVisitor[kind] = mergedEnterLeave;
  284. }
  285. return mergedVisitor;
  286. }
  287. /**
  288. * Given a visitor instance and a node kind, return EnterLeaveVisitor for that kind.
  289. */
  290. export function getEnterLeaveForKind(visitor, kind) {
  291. const kindVisitor = visitor[kind];
  292. if (typeof kindVisitor === 'object') {
  293. // { Kind: { enter() {}, leave() {} } }
  294. return kindVisitor;
  295. } else if (typeof kindVisitor === 'function') {
  296. // { Kind() {} }
  297. return {
  298. enter: kindVisitor,
  299. leave: undefined,
  300. };
  301. } // { enter() {}, leave() {} }
  302. return {
  303. enter: visitor.enter,
  304. leave: visitor.leave,
  305. };
  306. }
  307. /**
  308. * Given a visitor instance, if it is leaving or not, and a node kind, return
  309. * the function the visitor runtime should call.
  310. *
  311. * @deprecated Please use `getEnterLeaveForKind` instead. Will be removed in v17
  312. */
  313. /* c8 ignore next 8 */
  314. export function getVisitFn(visitor, kind, isLeaving) {
  315. const { enter, leave } = getEnterLeaveForKind(visitor, kind);
  316. return isLeaving ? leave : enter;
  317. }