106 lines
4 KiB
TypeScript
106 lines
4 KiB
TypeScript
|
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.
|
||
|
|
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.
|
||
|
|
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 = [];
|
||
|
|
}
|
||
|
|
};
|
||
|
|
}
|