1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
|
import { visit } from 'estree-util-visit';
import { toHtml } from 'hast-util-to-html';
// accessing untyped hast and mdx types
type Node = any;
export interface OptimizeOptions {
customComponentNames?: string[];
}
const exportConstComponentsRe = /export\s+const\s+components\s*=/;
/**
* For MDX only, collapse static subtrees of the hast into `set:html`. Subtrees
* do not include any MDX elements.
*
* This optimization reduces the JS output as more content are represented as a
* string instead, which also reduces the AST size that Rollup holds in memory.
*/
export function rehypeOptimizeStatic(options?: OptimizeOptions) {
return (tree: any) => {
// A set of non-static components to avoid collapsing when walking the tree
// as they need to be preserved as JSX to be rendered dynamically.
const customComponentNames = new Set<string>(options?.customComponentNames);
// Find `export const components = { ... }` and get it's object's keys to be
// populated into `customComponentNames`. This configuration is used to render
// some HTML elements as custom components, and we also want to avoid collapsing them.
for (const child of tree.children) {
if (child.type === 'mdxjsEsm' && exportConstComponentsRe.test(child.value)) {
// Try to loosely get the object property nodes
const objectPropertyNodes = child.data.estree.body[0]?.declarations?.[0]?.init?.properties;
if (objectPropertyNodes) {
for (const objectPropertyNode of objectPropertyNodes) {
const componentName = objectPropertyNode.key?.name ?? objectPropertyNode.key?.value;
if (componentName) {
customComponentNames.add(componentName);
}
}
}
}
}
// All possible elements that could be the root of a subtree
const allPossibleElements = new Set<Node>();
// The current collapsible element stack while traversing the tree
const elementStack: Node[] = [];
visit(tree, {
enter(node) {
// @ts-expect-error read tagName naively
const isCustomComponent = node.tagName && customComponentNames.has(node.tagName);
// For nodes that can't be optimized, eliminate all elements in the
// `elementStack` from the `allPossibleElements` set.
if (node.type.startsWith('mdx') || isCustomComponent) {
for (const el of elementStack) {
allPossibleElements.delete(el);
}
// Micro-optimization: While this destroys the meaning of an element
// stack for this node, things will still work but we won't repeatedly
// run the above for other nodes anymore. If this is confusing, you can
// comment out the code below when reading.
elementStack.length = 0;
}
// For possible subtree root nodes, record them in `elementStack` and
// `allPossibleElements` to be used in the "leave" hook below.
// @ts-expect-error MDX types for `.type` is not enhanced because MDX isn't used directly
if (node.type === 'element' || node.type === 'mdxJsxFlowElement') {
elementStack.push(node);
allPossibleElements.add(node);
}
},
leave(node, _, __, parents) {
// Do the reverse of the if condition above, popping the `elementStack`,
// and consolidating `allPossibleElements` as a subtree root.
// @ts-expect-error MDX types for `.type` is not enhanced because MDX isn't used directly
if (node.type === 'element' || node.type === 'mdxJsxFlowElement') {
elementStack.pop();
// Many possible elements could be part of a subtree, in order to find
// the root, we check the parent of the element we're popping. If the
// parent exists in `allPossibleElements`, then we're definitely not
// the root, so remove ourselves. This will work retroactively as we
// climb back up the tree.
const parent = parents[parents.length - 1];
if (allPossibleElements.has(parent)) {
allPossibleElements.delete(node);
}
}
},
});
// For all possible subtree roots, collapse them into `set:html` and
// strip of their children
for (const el of allPossibleElements) {
if (el.type === 'mdxJsxFlowElement') {
el.attributes.push({
type: 'mdxJsxAttribute',
name: 'set:html',
value: toHtml(el.children),
});
} else {
el.properties['set:html'] = toHtml(el.children);
}
el.children = [];
}
};
}
|
