track used selectors via WeakMap, fixes #21

theKashey · Jun 8, 2020 · 9826c82 · 9826c82
1 parent c3dfb45
commit 9826c82
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 28 deletions.
diff --git a/src/getCSS.ts b/src/getCSS.ts
@@ -1,7 +1,7 @@
 import { kashe } from 'kashe';
 import { StyleAst } from './parser/ast';
 import { extractUnmatchable, fromAst, getUnmatchableRules } from './parser/fromAst';
-import { FlagType, StyleChunk, StyleDefinition, UsedTypes, UsedTypesRef } from './types';
+import { FlagType, SelectionFilter, StyleChunk, StyleDefinition, UsedTypes, UsedTypesRef } from './types';
 import { assertIsReady, getStylesInText, unique } from './utils';
 
 export const getUnusableStyles = kashe(
@@ -74,17 +74,15 @@ export const getUsedStyles = (str: string, def: StyleDefinition): UsedTypes => {
   return getUsedStylesIn(getStylesInText(str), def);
 };
 
-const astToStyles = kashe(
-  (styles: string[], def: StyleDefinition, filter?: (selector: string) => boolean): StyleChunk[] => {
-    const { ast } = def;
-    const { fetches, usage } = astToUsedStyles(styles, def);
+const astToStyles = kashe((styles: string[], def: StyleDefinition, filter?: SelectionFilter): StyleChunk[] => {
+  const { ast } = def;
+  const { fetches, usage } = astToUsedStyles(styles, def);
 
-    return usage.map(file => ({
-      file,
-      css: fromAst(Object.keys(fetches[file]), ast[file], filter),
-    }));
-  }
-);
+  return usage.map(file => ({
+    file,
+    css: fromAst(Object.keys(fetches[file]), ast[file], filter),
+  }));
+});
 
 export const wrapInStyle = (styles: string, usedStyles: string[] = []) =>
   styles
@@ -121,11 +119,7 @@ export const extractAllUnmatchableAsString = kashe((def: StyleDefinition) =>
 const criticalRulesToStyle = (styles: StyleChunk[], urlPrefix = ''): string =>
   wrapInStyle(styles.map(({ css }) => css).join(''), unique(styles.map(({ file }) => `${urlPrefix}${file}`)));
 
-export const criticalStylesToString = (
-  str: string,
-  def: StyleDefinition,
-  filter?: (selector: string) => boolean
-): string => {
+export const criticalStylesToString = (str: string, def: StyleDefinition, filter?: SelectionFilter): string => {
   assertIsReady(def);
   return criticalRulesToStyle(astToStyles(getStylesInText(str), def, filter), def.urlPrefix);
 };

diff --git a/src/parser/fromAst.ts b/src/parser/fromAst.ts
@@ -1,5 +1,6 @@
 import { kashe } from 'kashe';
 import * as postcss from 'postcss';
+import { SelectionFilter } from '../types';
 import { SingleStyleAst, StyleBody, StyleRule, StyleSelector } from './ast';
 
 const separator = process.env.NODE_ENV === 'production' ? '' : '\n';
@@ -41,13 +42,11 @@ const findMatchingSelectors = (rules: SelectorLookUp, selectors: StyleSelector[]
 const findUnmatchableSelectors = (selectors: StyleSelector[]): StyleSelector[] =>
   selectors.filter(rule => rule.pieces.length === 0);
 
-export const fromAst = (rules: string[], def: SingleStyleAst, filter?: (selector: string) => boolean) => {
+export const fromAst = (rules: string[], def: SingleStyleAst, filter?: SelectionFilter) => {
   const blocks: StyleSelector[] = [];
   const lookup: SelectorLookUp = new Set(rules);
   blocks.push(
-    ...findMatchingSelectors(lookup, def.selectors).filter(
-      block => !filter || filter(`${block.selector}${block.media.join('')}`)
-    )
+    ...findMatchingSelectors(lookup, def.selectors).filter(block => !filter || filter(block.selector, block))
   );
 
   return convertToString(blocks, def);

diff --git a/src/reporters/critical.ts b/src/reporters/critical.ts
@@ -1,6 +1,7 @@
 import { Transform } from 'stream';
 import { isReact } from '../config';
 import { criticalStylesToString, extractAllUnmatchableAsString } from '../getCSS';
+import { StyleSelector } from '../parser/ast';
 import { CacheLine, StyleDefinition } from '../types';
 import { assertIsReady, createLine, findLastBrace } from '../utils';
 
@@ -30,13 +31,13 @@ export const createCriticalStyleStream = (def: StyleDefinition) => {
   const line = createLine();
   let injections: Array<string | undefined> = [];
 
-  const usedSelectors = new Set<string>();
+  const usedSelectors = new WeakSet<any>();
 
-  const filter = (selector: string) => {
-    if (usedSelectors.has(selector)) {
+  const filter = (_: any, rule: StyleSelector) => {
+    if (usedSelectors.has(rule)) {
       return false;
     }
-    usedSelectors.add(selector);
+    usedSelectors.add(rule);
     return true;
   };
 

diff --git a/src/types.ts b/src/types.ts
@@ -1,4 +1,4 @@
-import {StyleAst} from "./parser/ast";
+import { StyleAst, StyleSelector } from './parser/ast';
 
 export type StylesLookupTable = Record<string, string[]>;
 
@@ -19,16 +19,23 @@ export interface SyncStyleDefinition {
 }
 
 export interface StyleChunk {
-  file: string,
+  file: string;
   css: string;
 }
 
-export type FlagType = Record<string, boolean>
+export type FlagType = Record<string, boolean>;
 
 export interface StyleDefinition {
   isReady: boolean;
   lookup: StylesLookupTable;
   ast: StyleAst;
   urlPrefix: string;
   then(resolve?: () => void, reject?: () => void): Promise<void>;
-}
+}
+
+/**
+ * A function used to control which selectors should be used
+ * @param selector - DEPRECATED
+ * @param {StyleSelector} rule - a reference to a rule
+ */
+export type SelectionFilter = (selector: string, rule: StyleSelector) => boolean;