diff --git a/src/traverse.ts b/src/traverse.ts index 6371dc4..97833eb 100644 --- a/src/traverse.ts +++ b/src/traverse.ts @@ -13,21 +13,50 @@ function traverse(node: FiberNode, fn: (node: FiberNode) => any) { } function* traverseGenerator( - node: FiberNode - // { - // order = "pre" - // }: { order: "in" | "pre" | "post" } + node: FiberNode, + { + order = ["self", "child", "sibling"] + }: { order?: ["self", "child", "sibling"] } = {} ): IterableIterator { - const { skipChild, skipSibling } = yield node; + let skipChild = false, + skipSibling = false; - if (!skipChild && node.child !== null) { - const nextNode = node.child; - yield* traverseGenerator(nextNode); + function* traverseSelf() { + const controlInput: + | { skipChild: boolean; skipSibling: boolean } + | undefined = yield node; + + if (controlInput !== undefined) { + ({ skipChild = skipChild, skipSibling = skipSibling } = controlInput); + } } - if (!skipSibling && node.sibling !== null) { - const nextNode = node.sibling; - yield* traverseGenerator(nextNode); + function* traverseChild() { + if (!skipChild && node.child !== null) { + const nextNode = node.child; + yield* traverseGenerator(nextNode, { order }); + } + } + + function* traverseSibling() { + if (!skipSibling && node.sibling !== null) { + const nextNode = node.sibling; + yield* traverseGenerator(nextNode, { order }); + } + } + + const traverseMap = { + self: traverseSelf, + child: traverseChild, + sibling: traverseSibling + }; + + // For each item mentioned in order, find generator functions to run + const orderedGenerators = order.map((step) => traverseMap[step]).filter(tmp => tmp!==undefined); + + // Now run each generator till end + for (const eachGen of orderedGenerators) { + yield* eachGen(); } }