From bfbd00b37024205f4c50bc810d495d962de380bf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Martin=20Zl=C3=A1mal?= <mrtnzlml@gmail.com>
Date: Thu, 3 Mar 2022 19:36:27 -0600
Subject: [PATCH] SX: implement styles rehydration

I have these changes sitting locally for a very long time, and it's time
to offload them. The idea behind the styles rehydration is very simple:
we collect all existing styles from `<style data-adeira-sx />` (once)
and use this information later for deciding whether the styles should be
injected or not (this is what I call rehydration).

The difference from the current solution is that we do it only once and
we don't have to go through all the styles everytime there is something
to inject.

Motivation for this change is performance: this should significantly
improve the runtime style perf + be much easier to deal with (for
example, we are basically brute-forcing @at rules and pseudo rules now).

Note: this is just a partial solution, there are other changes coming
(custom styles printer and @at nodes/pseudo nodes optimization).

Related issue: https://github.com/adeira/universe/issues/1579
---
 src/abacus-kochka/README.md                  |   3 +
 src/abacus-kochka/pages/_document.js         |   8 +-
 src/sx/README.md                             |   9 ++
 src/sx/src/StyleCollector.js                 |  27 ++--
 src/sx/src/StyleCollectorAtNode.js           |  12 +-
 src/sx/src/StyleCollectorNode.js             |  21 ++-
 src/sx/src/StyleCollectorNodeInterface.js    |   2 +-
 src/sx/src/StyleCollectorPseudoNode.js       |  33 +++--
 src/sx/src/__tests__/printNodes.test.js      |  78 +++++++++++
 src/sx/src/__tests__/rehydrateStyles.test.js |  58 ++++++++
 src/sx/src/create.js                         |   8 +-
 src/sx/src/getStyleSheetFromStyleTag.js      |  40 ++++++
 src/sx/src/getStyleTag.js                    |   1 +
 src/sx/src/injectRuntimeStyles.js            | 140 +++++++++----------
 src/sx/src/keyframes.js                      |  36 +++--
 src/sx/src/printNodes.js                     |  83 +++++++++++
 src/sx/src/rehydrateStyles.js                |  79 +++++++++++
 17 files changed, 507 insertions(+), 131 deletions(-)
 create mode 100644 src/sx/src/__tests__/printNodes.test.js
 create mode 100644 src/sx/src/__tests__/rehydrateStyles.test.js
 create mode 100644 src/sx/src/getStyleSheetFromStyleTag.js
 create mode 100644 src/sx/src/printNodes.js
 create mode 100644 src/sx/src/rehydrateStyles.js

diff --git a/src/abacus-kochka/README.md b/src/abacus-kochka/README.md
index cb9c9d5b48..813366dafc 100644
--- a/src/abacus-kochka/README.md
+++ b/src/abacus-kochka/README.md
@@ -1,5 +1,8 @@
 [![Crowdin](https://badges.crowdin.net/kochkacommx/localized.svg)](https://crowdin.com/project/kochkacommx)
 
+- https://kochka.com.mx/
+- [PageSpeed Insights](https://pagespeed.web.dev/report?url=https%3A%2F%2Fkochka.com.mx%2Fmenu&form_factor=mobile)
+
 ```text
 yarn install
 yarn dev
diff --git a/src/abacus-kochka/pages/_document.js b/src/abacus-kochka/pages/_document.js
index 3023b053c3..b8967b1529 100644
--- a/src/abacus-kochka/pages/_document.js
+++ b/src/abacus-kochka/pages/_document.js
@@ -15,10 +15,10 @@ type RenderPageResult = {
 
 export default class MyDocument extends Document {
   // See: https://nextjs.org/docs/advanced-features/custom-document#customizing-renderpage
-  static async getInitialProps({ renderPage }: DocumentContext): Promise<RenderPageResult> {
-    const page = await renderPage();
-    return { ...page, styles: [sx.getStyleTag()] };
-  }
+  // static async getInitialProps({ renderPage }: DocumentContext): Promise<RenderPageResult> {
+  //   const page = await renderPage();
+  //   return { ...page, styles: [sx.getStyleTag()] };
+  // }
 
   createGoogleAnalyticsMarkup(): { __html: string } {
     return {
diff --git a/src/sx/README.md b/src/sx/README.md
index 19242c95da..2a8269181d 100644
--- a/src/sx/README.md
+++ b/src/sx/README.md
@@ -13,6 +13,7 @@ In conventional applications, CSS rules are duplicated throughout the stylesheet
   - [Automatic vendor prefixes](#automatic-vendor-prefixes)
 - [Server-side rendering](#server-side-rendering)
 - [Architecture](#architecture)
+  - [Runtime styles architecture](#runtime-styles-architecture)
 - [Prior Art](#prior-art)
 
 ## Installation and Usage
@@ -481,6 +482,14 @@ Internally, these steps are happening:
 
 5. and finally, we collect the values of the final object and print them as `className`
 
+### Runtime styles architecture
+
+Runtime styles are styles that were not rendered by server (are for whatever reason missing or SSR is not enabled). Here is how SX deals with this situation:
+
+1. SX tries to find `<style data-adeira-sx />` (or creates it if it doesn't exist yet) and does "rehydration" where it goes through the existing styles and remembers which ones are already applied.
+2. SX performs runtime injection of the styles while checking whether the styles already exist or not.
+3. The rest is the same.
+
 ## Prior Art
 
 _sorted alphabetically_
diff --git a/src/sx/src/StyleCollector.js b/src/sx/src/StyleCollector.js
index d08942cafd..42dcf95c5a 100644
--- a/src/sx/src/StyleCollector.js
+++ b/src/sx/src/StyleCollector.js
@@ -3,6 +3,7 @@
 import { invariant, warning } from '@adeira/js';
 
 import expandShorthandProperties from './expandShorthandProperties';
+import printNodes from './printNodes';
 import StyleCollectorAtNode from './StyleCollectorAtNode';
 import StyleCollectorPseudoNode from './StyleCollectorPseudoNode';
 import { type StyleCollectorNodeInterface } from './StyleCollectorNodeInterface';
@@ -41,7 +42,7 @@ class StyleCollector {
   #styleBuffer: StyleBufferType = new Map();
   #keyframes: Map<string, string> = new Map();
 
-  collect(baseStyleSheet: { +[sheetName: string]: $FlowFixMe }): {
+  collectStylesheets(baseStyleSheet: { +[sheetName: string]: $FlowFixMe }): {
     +hashRegistry: HashRegistryType,
     +styleBuffer: StyleBufferType,
   } {
@@ -106,23 +107,17 @@ class StyleCollector {
     };
   }
 
-  print(): string {
-    let sxStyle = '';
-    this.#styleBuffer.forEach((node) => {
-      sxStyle += node.printNodes({ trailingSemicolon: true }).join('');
-    });
-    this.#keyframes.forEach((node) => {
-      sxStyle += node;
-    });
-    return sxStyle;
+  collectKeyframe(name: string, value: string): void {
+    this.#keyframes.set(name, value);
   }
 
-  addKeyframe(name: string, value: string): boolean {
-    if (this.#keyframes.has(name)) {
-      return true;
-    }
-    this.#keyframes.set(name, value);
-    return false;
+  // TODO: remove
+  print(): string {
+    // TODO: print keyframes as well
+    //     this.#keyframes.forEach((node) => {
+    //       sxStyle += node;
+    //     });
+    return printNodes([...this.#styleBuffer]);
   }
 
   reset(): void {
diff --git a/src/sx/src/StyleCollectorAtNode.js b/src/sx/src/StyleCollectorAtNode.js
index aed2fa19f3..d448c9473c 100644
--- a/src/sx/src/StyleCollectorAtNode.js
+++ b/src/sx/src/StyleCollectorAtNode.js
@@ -43,15 +43,11 @@ export default class StyleCollectorAtNode implements StyleCollectorNodeInterface
     this.nodes = new Map([...this.nodes, ...nodes]);
   }
 
-  getAtRuleName(): string {
-    return this.atRuleName;
+  getNodes() {
+    return this.nodes;
   }
 
-  printNodes(config?: PrintConfig): $ReadOnlyArray<string> {
-    let output = '';
-    this.nodes.forEach((node) => {
-      output += node.printNodes({ ...config, bumpSpecificity: true }).join('');
-    });
-    return [`${this.atRuleName}{${output}}`];
+  getAtRuleName(): string {
+    return this.atRuleName;
   }
 }
diff --git a/src/sx/src/StyleCollectorNode.js b/src/sx/src/StyleCollectorNode.js
index 304face3cb..97f5206241 100644
--- a/src/sx/src/StyleCollectorNode.js
+++ b/src/sx/src/StyleCollectorNode.js
@@ -37,7 +37,11 @@ export default class StyleCollectorNode implements StyleCollectorNodeInterface {
 
   // eslint-disable-next-line no-unused-vars
   addNodes(nodes: Map<string, StyleCollectorNodeInterface>) {
-    invariant(false, 'StyleCollectorNode cannot have nested nodes,');
+    invariant(false, 'StyleCollectorNode cannot have nested nodes.');
+  }
+
+  getNodes(): Map<string, StyleCollectorNode> {
+    invariant(false, 'StyleCollectorNode cannot have nested nodes.');
   }
 
   getHash(): string {
@@ -52,11 +56,16 @@ export default class StyleCollectorNode implements StyleCollectorNodeInterface {
     return this.styleValue;
   }
 
-  printNodes(config?: PrintConfig): $ReadOnlyArray<string> {
-    const className = `.${this.hash}`.repeat(config?.bumpSpecificity === true ? 2 : 1);
-    const pseudo = config?.pseudo ?? '';
-    const trailingSemicolon = config?.trailingSemicolon ? ';' : '';
+  // printNodes(config?: PrintConfig): $ReadOnlyArray<string> {
+  //   const className = `.${this.hash}`.repeat(config?.bumpSpecificity === true ? 2 : 1);
+  //   const pseudo = config?.pseudo ?? '';
+  //   const trailingSemicolon = config?.trailingSemicolon ? ';' : '';
+  //
+  //   return [`${className}${pseudo}{${this.styleName}:${this.styleValue}${trailingSemicolon}}`];
+  // }
 
-    return [`${className}${pseudo}{${this.styleName}:${this.styleValue}${trailingSemicolon}}`];
+  // CSSStyleRule.selectorText
+  rehydrationIdentifier() {
+    return `.${this.getHash()}`;
   }
 }
diff --git a/src/sx/src/StyleCollectorNodeInterface.js b/src/sx/src/StyleCollectorNodeInterface.js
index f1deea8086..dfcc235752 100644
--- a/src/sx/src/StyleCollectorNodeInterface.js
+++ b/src/sx/src/StyleCollectorNodeInterface.js
@@ -8,5 +8,5 @@ export type PrintConfig = {
 
 export interface StyleCollectorNodeInterface {
   addNodes(nodes: Map<string, StyleCollectorNodeInterface>): void;
-  printNodes(config?: PrintConfig): $ReadOnlyArray<string>;
+  getNodes(): Map<string, StyleCollectorNodeInterface>;
 }
diff --git a/src/sx/src/StyleCollectorPseudoNode.js b/src/sx/src/StyleCollectorPseudoNode.js
index 0b8fd56d74..4af737bd8e 100644
--- a/src/sx/src/StyleCollectorPseudoNode.js
+++ b/src/sx/src/StyleCollectorPseudoNode.js
@@ -1,5 +1,7 @@
 // @flow strict
 
+import type { AllCSSPseudoTypes } from './css-properties/__generated__/AllCSSPseudoTypes';
+import StyleCollectorNode from './StyleCollectorNode';
 import type { PrintConfig, StyleCollectorNodeInterface } from './StyleCollectorNodeInterface';
 
 /**
@@ -23,29 +25,40 @@ import type { PrintConfig, StyleCollectorNodeInterface } from './StyleCollectorN
  *   ],
  * }
  * ```
+ *
+ * Note that nesting pseudo classes is not allowed in SX.
  */
 export default class StyleCollectorPseudoNode implements StyleCollectorNodeInterface {
   pseudo: string;
-  nodes: Map<string, StyleCollectorNodeInterface>;
+  nodes: Map<string, StyleCollectorNode>;
 
-  constructor(pseudo: string, nodes: Map<string, StyleCollectorNodeInterface>) {
+  constructor(pseudo: $Keys<AllCSSPseudoTypes>, nodes: Map<string, StyleCollectorNode>) {
     this.pseudo = pseudo;
     this.nodes = nodes;
   }
 
-  addNodes(nodes: Map<string, StyleCollectorNodeInterface>) {
+  addNodes(nodes: Map<string, StyleCollectorNode>) {
     this.nodes = new Map([...this.nodes, ...nodes]);
   }
 
+  getNodes(): Map<string, StyleCollectorNode> {
+    return this.nodes;
+  }
+
   getPseudo(): string {
     return this.pseudo;
   }
 
-  printNodes(config?: PrintConfig): $ReadOnlyArray<string> {
-    let output = [];
-    this.nodes.forEach((node) => {
-      output = output.concat(node.printNodes({ ...config, pseudo: this.pseudo }));
-    });
-    return output;
-  }
+  // printNodes(config?: PrintConfig): $ReadOnlyArray<string> {
+  //   let output = [];
+  //   this.nodes.forEach((node) => {
+  //     output = output.concat(node.printNodes({ ...config, pseudo: this.pseudo }));
+  //   });
+  //   return output;
+  // }
+
+  // CSSStyleRule.selectorText
+  // rehydrationIdentifier() {
+  //   return `.${this.getHash()}`;
+  // }
 }
diff --git a/src/sx/src/__tests__/printNodes.test.js b/src/sx/src/__tests__/printNodes.test.js
new file mode 100644
index 0000000000..7bd22028d3
--- /dev/null
+++ b/src/sx/src/__tests__/printNodes.test.js
@@ -0,0 +1,78 @@
+// @flow
+
+import printNodes from '../printNodes';
+import StyleCollectorAtNode from '../StyleCollectorAtNode';
+import StyleCollectorNode from '../StyleCollectorNode';
+import StyleCollectorPseudoNode from '../StyleCollectorPseudoNode';
+
+it('prints `StyleCollectorNode` as expected', () => {
+  expect(
+    printNodes([
+      new StyleCollectorNode('color', 'red'),
+      new StyleCollectorNode('color', 'green'),
+      new StyleCollectorNode('color', 'blue'),
+    ]),
+  ).toMatchInlineSnapshot(`
+    "._324Crd{color:#f00}
+    .mRoJ3{color:#008000}
+    ._2dHaKY{color:#00f}"
+  `);
+});
+
+it('prints `StyleCollectorPseudoNode` as expected', () => {
+  expect(
+    printNodes([
+      new StyleCollectorPseudoNode(
+        ':hover',
+        new Map([
+          // TODO: are the map keys necessary (?)
+          ['c0', new StyleCollectorNode('color', 'red')],
+          ['c1', new StyleCollectorNode('color', 'green')],
+          ['c2', new StyleCollectorNode('color', 'blue')],
+        ]),
+      ),
+    ]),
+  ).toMatchInlineSnapshot(`
+    "._324Crd:hover{color:#f00}
+    .mRoJ3:hover{color:#008000}
+    ._2dHaKY:hover{color:#00f}"
+  `);
+});
+
+it('prints `StyleCollectorAtNode` as expected', () => {
+  expect(
+    printNodes([
+      new StyleCollectorAtNode(
+        '@media screen',
+        new Map([
+          ['c0', new StyleCollectorNode('fontSize', '14')],
+          [
+            'c1',
+            new StyleCollectorPseudoNode(
+              ':hover',
+              new Map([['c0', new StyleCollectorNode('color', 'pink')]]),
+            ),
+          ],
+
+          [
+            'c2',
+            new StyleCollectorAtNode(
+              '@media (max-width: 12cm)',
+              new Map([['c0', new StyleCollectorNode('color', 'blue')]]),
+            ),
+          ],
+        ]),
+      ),
+    ]),
+  ).toMatchInlineSnapshot(`
+    "@media screen{
+    ._1fVgat._1fVgat{font-size:14}
+    ._3ncx7d._3ncx7d:hover{color:#ffc0cb}
+    @media (max-width: 12cm){
+    ._2dHaKY._2dHaKY{color:#00f}
+    }
+    }"
+  `);
+});
+
+// TODO: prefixing
diff --git a/src/sx/src/__tests__/rehydrateStyles.test.js b/src/sx/src/__tests__/rehydrateStyles.test.js
new file mode 100644
index 0000000000..b5b661737e
--- /dev/null
+++ b/src/sx/src/__tests__/rehydrateStyles.test.js
@@ -0,0 +1,58 @@
+/**
+ * @flow
+ * @jest-environment jsdom
+ */
+
+/* global document */
+
+import { invariant } from '@adeira/js';
+
+import rehydrateStyles from '../rehydrateStyles';
+
+it('correctly rehydrates styles from a style element', () => {
+  // First, we need to create a `CSSStyleSheet` by actually creating a style element:
+  const styleElement = document.createElement('style');
+  document.head?.appendChild(styleElement);
+  const styleSheet = styleElement.sheet;
+
+  invariant(styleSheet != null, 'Unable to create test StyleSheet.');
+
+  // Insert some simple CSS rules:
+  styleSheet.insertRule('._2tPCgL { font-size: 10px; }', 0);
+  styleSheet.insertRule('._1Kmfck:hover { color: rgba(var(--sx-foreground), 0.5); }', 1);
+
+  // Insert some @at rules:
+  styleSheet.insertRule(
+    `@media (prefers-reduced-motion: reduce) { .VdrO3.VdrO3 { animation-duration: 1s; } }`,
+    2,
+  );
+  styleSheet.insertRule(
+    `@media (prefers-reduced-motion: reduce) { ._2tPCgL._2tPCgL { font-size: 10px; } }`,
+    3,
+  );
+  styleSheet.insertRule(
+    `@keyframes oxCh9 { 33% { transform: translateY(-10px); } 66% { transform: translateY(10px); } }
+  `,
+    4,
+  );
+
+  // We should be able to decide whether the style needs to be injected later based on the
+  // following information:
+  expect(rehydrateStyles(styleSheet)).toMatchInlineSnapshot(`
+    Object {
+      "rehydratedKeyframeRules": Set {
+        "oxCh9",
+      },
+      "rehydratedMediaRules": Map {
+        "(prefers-reduced-motion: reduce)" => Set {
+          ".VdrO3.VdrO3",
+          "._2tPCgL._2tPCgL",
+        },
+      },
+      "rehydratedStyleRules": Set {
+        "._2tPCgL",
+        "._1Kmfck:hover",
+      },
+    }
+  `);
+});
diff --git a/src/sx/src/create.js b/src/sx/src/create.js
index 240807f692..3ee948e7ed 100644
--- a/src/sx/src/create.js
+++ b/src/sx/src/create.js
@@ -3,7 +3,9 @@
 import { invariant, isBrowser, isObjectEmpty, isObject } from '@adeira/js';
 import levenshtein from 'fast-levenshtein';
 
+import getStyleSheetFromStyleTag from './getStyleSheetFromStyleTag';
 import injectRuntimeStyles from './injectRuntimeStyles';
+import rehydrateStyles from './rehydrateStyles';
 import styleCollector from './StyleCollector';
 import type { AllCSSPropertyTypes } from './css-properties/__generated__/AllCSSPropertyTypes';
 import type { AllCSSPseudoTypes } from './css-properties/__generated__/AllCSSPseudoTypes';
@@ -48,10 +50,12 @@ export default function create<T: SheetDefinitions>(sheetDefinitions: T): Create
     `Function 'sx.create' cannot be called with empty stylesheet definition.`,
   );
 
-  const { hashRegistry, styleBuffer } = styleCollector.collect(sheetDefinitions);
+  const { hashRegistry, styleBuffer } = styleCollector.collectStylesheets(sheetDefinitions);
 
   if (isBrowser()) {
-    injectRuntimeStyles(styleBuffer);
+    const styleSheet = getStyleSheetFromStyleTag();
+    const rehydratedRules = rehydrateStyles(styleSheet);
+    injectRuntimeStyles(styleSheet, rehydratedRules, styleBuffer);
   }
 
   function sxFunction(maybeObject, ...styleSheetsSelectors) {
diff --git a/src/sx/src/getStyleSheetFromStyleTag.js b/src/sx/src/getStyleSheetFromStyleTag.js
new file mode 100644
index 0000000000..c63ef740b6
--- /dev/null
+++ b/src/sx/src/getStyleSheetFromStyleTag.js
@@ -0,0 +1,40 @@
+// @flow
+
+/* global document */
+
+import { invariant } from '@adeira/js';
+
+// <style data-adeira-sx />
+opaque type StyleElementType = HTMLStyleElement | null;
+let styleAdeiraSXTag: StyleElementType = null;
+
+/**
+ * This function will try to get already existing `<style data-adeira-sx />` HTML tag or it will
+ * create a new virtual one in case it doesn't exist yet.
+ */
+export default function getStyleSheetFromStyleTag(): CSSStyleSheet {
+  if (styleAdeiraSXTag === null) {
+    styleAdeiraSXTag = ((document.querySelector('style[data-adeira-sx]'): any): StyleElementType);
+    if (styleAdeiraSXTag === null) {
+      // Still `null` so let's create the style element:
+      const htmlHead = document.head;
+      styleAdeiraSXTag = document.createElement('style');
+      styleAdeiraSXTag.type = 'text/css';
+      styleAdeiraSXTag.setAttribute('data-adeira-sx', '');
+      /* $FlowFixMe[incompatible-call] This comment suppresses an error when
+       * upgrading Flow. To see the error delete this comment and run Flow. */
+      htmlHead?.appendChild(styleAdeiraSXTag);
+    }
+  }
+
+  /* $FlowFixMe[incompatible-use] This comment suppresses an error when
+   * upgrading Flow. To see the error delete this comment and run Flow. */
+  const styleSheet = styleAdeiraSXTag.sheet;
+
+  invariant(
+    styleSheet != null,
+    'SX cannot apply runtime styles because HTMLStyleElement.sheet does not exist.',
+  );
+
+  return styleSheet;
+}
diff --git a/src/sx/src/getStyleTag.js b/src/sx/src/getStyleTag.js
index ccee12bc2c..3192f48d2d 100644
--- a/src/sx/src/getStyleTag.js
+++ b/src/sx/src/getStyleTag.js
@@ -11,6 +11,7 @@ import StyleCollector from './StyleCollector';
  */
 export default function getStyleTag(): Element<'style'> {
   const createStylesMarkup = () => {
+    // TODO: do not run `serialize` again here
     const cssStyles = serialize(compile(StyleCollector.print()), middleware([prefixer, stringify]));
     return { __html: cssStyles };
   };
diff --git a/src/sx/src/injectRuntimeStyles.js b/src/sx/src/injectRuntimeStyles.js
index 0774ded921..dd88dbdc12 100644
--- a/src/sx/src/injectRuntimeStyles.js
+++ b/src/sx/src/injectRuntimeStyles.js
@@ -1,58 +1,33 @@
 // @flow
 
-/* global document */
-
-import { invariant, warning } from '@adeira/js';
+import { warning } from '@adeira/js';
 import { compile, serialize, stringify, prefixer, middleware } from 'stylis';
 
 import StyleCollectorAtNode from './StyleCollectorAtNode';
 import StyleCollectorNode from './StyleCollectorNode';
 import StyleCollectorPseudoNode from './StyleCollectorPseudoNode';
+import {
+  printStyleCollectorAtNode,
+  printStyleCollectorNode,
+  printStyleCollectorPseudoNode,
+} from './printNodes';
+import type { RehydratedStyles } from './rehydrateStyles';
 import type { StyleBufferType } from './StyleCollector';
 
-// <style data-adeira-sx />
-opaque type StyleElementType = HTMLStyleElement | null;
-let styleAdeiraSXTag: StyleElementType = null;
-
-const getStyleTag = (): CSSStyleSheet => {
-  if (styleAdeiraSXTag === null) {
-    styleAdeiraSXTag = ((document.querySelector('style[data-adeira-sx]'): any): StyleElementType);
-    if (styleAdeiraSXTag === null) {
-      // Still `null` so let's create the style element:
-      const htmlHead = document.head;
-      styleAdeiraSXTag = document.createElement('style');
-      styleAdeiraSXTag.type = 'text/css';
-      styleAdeiraSXTag.setAttribute('data-adeira-sx', '');
-      /* $FlowFixMe[incompatible-call] This comment suppresses an error when
-       * upgrading Flow. To see the error delete this comment and run Flow. */
-      htmlHead?.appendChild(styleAdeiraSXTag);
-    }
-  }
-
-  /* $FlowFixMe[incompatible-use] This comment suppresses an error when
-   * upgrading Flow. To see the error delete this comment and run Flow. */
-  const styleSheet = styleAdeiraSXTag.sheet;
-
-  invariant(
-    styleSheet != null,
-    'SX cannot apply runtime styles because HTMLStyleElement.sheet does not exist.',
-  );
-
-  return styleSheet;
-};
-
-function hasStyleRule(match: (CSSRule) => boolean) {
-  const styleSheet = getStyleTag();
-  for (const cssRule of styleSheet.cssRules) {
-    if (match(cssRule)) {
-      return true;
+export function injectRuntimeKeyframes(
+  styleSheet: CSSStyleSheet,
+  rehydratedRules: RehydratedStyles,
+  css: string,
+  name: string,
+) {
+  const hasStyleRule = (match: (CSSRule) => boolean) => {
+    for (const cssRule of styleSheet.cssRules) {
+      if (match(cssRule)) {
+        return true;
+      }
     }
-  }
-  return false;
-}
-
-export function injectRuntimeKeyframes(css: string, name: string) {
-  const styleSheet = getStyleTag();
+    return false;
+  };
 
   const matchFunction = (cssRule) => {
     if (cssRule.type === CSSRule.KEYFRAMES_RULE) {
@@ -63,7 +38,9 @@ export function injectRuntimeKeyframes(css: string, name: string) {
   };
 
   if (hasStyleRule(matchFunction) === false) {
+    // TODO: leverage rehydrated rules
     const rules = [];
+    // TODO: get rid of this:
     serialize(compile(css), middleware([prefixer, (rule) => rules.push(rule)]));
 
     rules.forEach((rule) => {
@@ -73,44 +50,57 @@ export function injectRuntimeKeyframes(css: string, name: string) {
 }
 
 // https://developer.mozilla.org/en-US/docs/Web/API/CSS_Object_Model
-export default function injectRuntimeStyles(styleBuffer: StyleBufferType) {
-  const styleSheet = getStyleTag();
-
-  const matchFunction = (node) => (cssRule) => {
-    if (cssRule.type === CSSRule.STYLE_RULE) {
-      const styleRule = ((cssRule: any): CSSStyleRule);
-      if (styleRule.selectorText === `.${node.getHash()}`) {
-        return true;
-      }
-    }
-    return false;
-  };
-
+export default function injectRuntimeStyles(
+  styleSheet: CSSStyleSheet, // TODO: note: we are mutating this
+  { rehydratedStyleRules, rehydratedMediaRules, rehydratedKeyframeRules }: RehydratedStyles,
+  styleBuffer: StyleBufferType,
+) {
   styleBuffer.forEach((node) => {
     const insertIndex = styleSheet.cssRules.length;
 
     if (node instanceof StyleCollectorNode) {
-      if (hasStyleRule(matchFunction(node)) === false) {
-        // apply missing styles
-        const rule = serialize(
-          compile(node.printNodes({ trailingSemicolon: true }).join('')),
-          middleware([prefixer, stringify]),
+      const rehydrationIdentifier = node.rehydrationIdentifier();
+      if (rehydratedStyleRules.has(rehydrationIdentifier) === false) {
+        // Apply the missing styles:
+        styleSheet.insertRule(
+          printStyleCollectorNode(node, { trailingSemicolon: true }),
+          insertIndex,
         );
-        styleSheet.insertRule(rule, insertIndex);
+
+        // Register the rule into rehydrated styles, so we don't attempt to process it again:
+        rehydratedStyleRules.add(rehydrationIdentifier);
       }
     } else if (node instanceof StyleCollectorAtNode) {
-      // TODO: make sure we are not adding already added styles (?)
-      const rule = serialize(
-        compile(node.printNodes({ trailingSemicolon: true }).join('')),
-        middleware([prefixer, stringify]),
-      );
-      styleSheet.insertRule(rule, insertIndex);
-    } else if (node instanceof StyleCollectorPseudoNode) {
-      const nodes = node.printNodes({ trailingSemicolon: true });
-      for (const nodeElement of nodes) {
-        // TODO: make sure we are not adding already added styles (?)
-        const rule = serialize(compile(nodeElement), middleware([prefixer, stringify]));
+      // TODO: leverage rehydrated rules (to add only new rules)
+      // TODO: media queries & keyframe rules
+      const rehydrationIdentifier = node.getAtRuleName().replace(/.*(\(.+\)).*/, '$1');
+
+      // TODO: we probably has to dig a bit deeper and check whether the @at rule has every inner rule applied
+      if (rehydratedMediaRules.has(rehydrationIdentifier) === false) {
+        const rule = printStyleCollectorAtNode(node, { trailingSemicolon: true });
         styleSheet.insertRule(rule, insertIndex);
+
+        // TODO: Register the rule into rehydrated styles, so we don't attempt to process it again:
+        rehydratedMediaRules.set(
+          rehydrationIdentifier,
+          node.getNodes(), // TODO: ???
+        );
+      }
+    } else if (node instanceof StyleCollectorPseudoNode) {
+      for (const [, testNode] of node.getNodes()) {
+        // TODO: better
+        const rehydrationIdentifier = `${testNode.rehydrationIdentifier()}:${node.getPseudo()}`;
+        if (rehydratedStyleRules.has(rehydrationIdentifier) === false) {
+          styleSheet.insertRule(
+            printStyleCollectorNode(testNode, {
+              pseudo: node.getPseudo(),
+            }),
+            insertIndex,
+          );
+        }
+
+        // TODO: Register the rule into rehydrated styles, so we don't attempt to process it again:
+        rehydratedStyleRules.add(rehydrationIdentifier);
       }
     } else {
       warning(false, 'Node not supported in runtime styles: %j', node);
diff --git a/src/sx/src/keyframes.js b/src/sx/src/keyframes.js
index 869da83ce2..a1ce04af68 100644
--- a/src/sx/src/keyframes.js
+++ b/src/sx/src/keyframes.js
@@ -1,14 +1,16 @@
 // @flow
 
 import { isBrowser, nullthrows } from '@adeira/js';
-import stringify from 'json-stable-stringify';
+import stringifyStable from 'json-stable-stringify';
 
+import getStyleSheetFromStyleTag from './getStyleSheetFromStyleTag';
 import hash from './hashStyle';
-import type { AllCSSPropertyTypes } from './css-properties/__generated__/AllCSSPropertyTypes';
+import rehydrateStyles from './rehydrateStyles';
 import styleCollector from './StyleCollector';
 import transformStyleName from './transformStyleName';
 import transformStyleValue from './transformValue';
 import { injectRuntimeKeyframes } from './injectRuntimeStyles';
+import type { AllCSSPropertyTypes } from './css-properties/__generated__/AllCSSPropertyTypes';
 
 type KeyFrames = {
   +from?: AllCSSPropertyTypes,
@@ -30,10 +32,27 @@ const extractStyles = (styles: AllCSSPropertyTypes) => {
 
 const transformKey = (key: string) => key.replace(/\s/g, '');
 
+/**
+ * Helper function generating animation name. Usage:
+ *
+ * ```js
+ * const bounce = sx.keyframes({
+ *    '33%': { transform: `translateY(-1.4em)` },
+ *    '66%': { transform: `translateY(1.4em)` },
+ * });
+ *
+ * const styles = sx.create({
+ *    svgCircle: {
+ *      animationName: bounce,
+ *    }
+ *    // …
+ * });
+ * ```
+ */
 export default function keyframes(styleDefinitions: KeyFrames): string {
   let cssDefinition = '';
   // `from{maxHeight:0,opacity:0}` should be the same as `from{opacity:0,maxHeight:0}`
-  const parsedDefinitions: KeyFrames = JSON.parse(stringify(styleDefinitions));
+  const parsedDefinitions: KeyFrames = JSON.parse(stringifyStable(styleDefinitions));
 
   for (const key of Object.keys(parsedDefinitions)) {
     const styleValue = nullthrows(parsedDefinitions[key]);
@@ -42,13 +61,12 @@ export default function keyframes(styleDefinitions: KeyFrames): string {
 
   const name = hash(cssDefinition);
   const frame = `@keyframes ${name} {${cssDefinition}}`;
-  const exists = styleCollector.addKeyframe(name, frame);
+  styleCollector.collectKeyframe(name, frame);
 
-  if (isBrowser() && !exists) {
-    // It is possible that the keyframe was added on the server
-    // The StyleCollector will return false, but the injectRuntimeKeyframes checks
-    // If the rule is already added, so it won't be duplicated
-    injectRuntimeKeyframes(frame, name);
+  if (isBrowser()) {
+    const styleSheet = getStyleSheetFromStyleTag();
+    const rehydratedRules = rehydrateStyles(styleSheet);
+    injectRuntimeKeyframes(styleSheet, rehydratedRules, frame, name);
   }
 
   return name;
diff --git a/src/sx/src/printNodes.js b/src/sx/src/printNodes.js
new file mode 100644
index 0000000000..7184cb3639
--- /dev/null
+++ b/src/sx/src/printNodes.js
@@ -0,0 +1,83 @@
+// @flow
+
+import { compile, middleware, prefixer, serialize, stringify } from 'stylis';
+
+import StyleCollectorAtNode from './StyleCollectorAtNode';
+import StyleCollectorNode from './StyleCollectorNode';
+import StyleCollectorPseudoNode from './StyleCollectorPseudoNode';
+
+const STYLESHEET_SEPARATOR = '\n';
+
+/**
+ * This function takes care of the correct rendering of each node. Additionally, it takes care of
+ * browser prefixing.
+ */
+export default function printNodes(
+  nodes: $ReadOnlyArray<StyleCollectorAtNode | StyleCollectorNode | StyleCollectorPseudoNode>,
+) {
+  const printedRules = [];
+  for (const node of nodes) {
+    if (node instanceof StyleCollectorNode) {
+      printedRules.push(printStyleCollectorNode(node));
+    } else if (node instanceof StyleCollectorPseudoNode) {
+      printedRules.push(printStyleCollectorPseudoNode(node));
+    } else if (node instanceof StyleCollectorAtNode) {
+      printedRules.push(printStyleCollectorAtNode(node));
+    } else {
+      // TODO: invariant
+    }
+  }
+
+  const stylesheet = printedRules.join(STYLESHEET_SEPARATOR);
+  // TODO: make sure this is the only place we are calling `serialize` as it would not be optimal (performance wise)
+  return serialize(compile(stylesheet), middleware([prefixer, stringify]));
+}
+
+export function printStyleCollectorNode(node: StyleCollectorNode, config = {}) {
+  const pseudo = config.pseudo ?? '';
+  const withIncreasedSpecificity = config.increaseSpecificity ?? false;
+
+  const selectorText = `.${node.getHash()}`.repeat(withIncreasedSpecificity === true ? 2 : 1);
+  const selectorTextWithPseudo = `${selectorText}${pseudo}`;
+
+  // TODO: apply `serialize` in case it's called directly without `printNodes`
+  return `${selectorTextWithPseudo}{${node.getStyleName()}:${node.getStyleValue()}}`;
+}
+
+export function printStyleCollectorPseudoNode(pseudoNode: StyleCollectorPseudoNode, config = {}) {
+  const withIncreasedSpecificity = config.increaseSpecificity ?? false;
+  const printedRules = [];
+  for (const [, node] of pseudoNode.getNodes()) {
+    printedRules.push(
+      printStyleCollectorNode(node, {
+        pseudo: pseudoNode.getPseudo(),
+        increaseSpecificity: withIncreasedSpecificity,
+      }),
+    );
+  }
+
+  // TODO: apply `serialize` in case it's called directly without `printNodes`
+  return printedRules.join(STYLESHEET_SEPARATOR);
+}
+
+export function printStyleCollectorAtNode(atNode: StyleCollectorAtNode) {
+  const printedAtRules = [];
+  for (const [, node] of atNode.getNodes()) {
+    if (node instanceof StyleCollectorNode) {
+      printedAtRules.push(printStyleCollectorNode(node, { increaseSpecificity: true }));
+    } else if (node instanceof StyleCollectorPseudoNode) {
+      printedAtRules.push(printStyleCollectorPseudoNode(node, { increaseSpecificity: true }));
+    } else if (node instanceof StyleCollectorAtNode) {
+      printedAtRules.push(
+        printStyleCollectorAtNode(node, {
+          increaseSpecificity: true, // Should the specificity be increased even more for nested at rules?
+        }),
+      );
+    } else {
+      // TODO: invariant
+    }
+  }
+
+  // TODO: apply `serialize` in case it's called directly without `printNodes`
+  return `${atNode.getAtRuleName()}{\n${printedAtRules.join(STYLESHEET_SEPARATOR)}\n}`;
+}
diff --git a/src/sx/src/rehydrateStyles.js b/src/sx/src/rehydrateStyles.js
new file mode 100644
index 0000000000..146361ee52
--- /dev/null
+++ b/src/sx/src/rehydrateStyles.js
@@ -0,0 +1,79 @@
+// @flow
+
+import { warning } from '@adeira/js';
+
+// Unique entries are identified by the style selector text ("._XYZ"  or "._XYZ:hover").
+let rehydratedStyleRules = null;
+
+// Unique entries are a combination of the media condition text ("(prefers-color-scheme: light)")
+// and selector text of each rule in the media rule ("._XYZ"  or "._XYZ:hover").
+let rehydratedMediaRules = null;
+
+// Unique entries are identified by the keyframe name ("_XYZ").
+let rehydratedKeyframeRules = null;
+
+export type RehydratedStyles = {
+  rehydratedStyleRules: Set<string>,
+  rehydratedMediaRules: Map<string, Set<string>>,
+  rehydratedKeyframeRules: Set<string>,
+};
+
+/**
+ * This function goes through already existing stylesheet inside `<style data-adeira-sx />` and
+ * register what CSS rules are already rendered by the server (rehydration). This information is
+ * later used by `injectRuntimeStyles` and `injectRuntimeKeyframes` to inject only (runtime) styles
+ * that do not exist yet.
+ */
+export default function rehydrateStyles(styleSheet: CSSStyleSheet): RehydratedStyles {
+  if (
+    // memoize so we call this only once to populate the values
+    rehydratedStyleRules === null ||
+    rehydratedMediaRules === null ||
+    rehydratedKeyframeRules === null
+  ) {
+    rehydratedStyleRules = new Set();
+    rehydratedMediaRules = new Map();
+    rehydratedKeyframeRules = new Set();
+
+    for (const cssRule of styleSheet.cssRules) {
+      if (cssRule.type === CSSRule.STYLE_RULE) {
+        const cssStyleRule = ((cssRule: any): CSSStyleRule);
+        rehydratedStyleRules.add(cssStyleRule.selectorText); // "._XYZ"  or "._XYZ:hover"
+      } else if (cssRule.type === CSSRule.MEDIA_RULE) {
+        const cssMediaRule = ((cssRule: any): CSSMediaRule);
+        const cssMediaRuleText = cssMediaRule.media.mediaText;
+        const collectedMediaSelectorTexts = new Set();
+        for (const cssMediaNestedRule of cssMediaRule.cssRules) {
+          collectedMediaSelectorTexts.add(((cssMediaNestedRule: any): CSSStyleRule).selectorText);
+        }
+        if (rehydratedMediaRules.has(cssMediaRuleText)) {
+          rehydratedMediaRules.set(
+            cssMediaRuleText,
+            new Set([
+              ...(rehydratedMediaRules.get(cssMediaRuleText) ?? new Set()),
+              ...collectedMediaSelectorTexts,
+            ]),
+          );
+        } else {
+          rehydratedMediaRules.set(cssMediaRuleText, collectedMediaSelectorTexts);
+        }
+      } else if (cssRule.type === CSSRule.KEYFRAMES_RULE) {
+        const cssKeyframesRule = ((cssRule: any): CSSKeyframesRule);
+        rehydratedKeyframeRules.add(cssKeyframesRule.name); // "_XYZ"
+      } else {
+        warning(
+          false,
+          'Could not rehydrate the following CSS style (type %s): %s',
+          cssRule.type,
+          cssRule.cssText,
+        );
+      }
+    }
+  }
+
+  return Object.freeze({
+    rehydratedStyleRules,
+    rehydratedMediaRules,
+    rehydratedKeyframeRules,
+  });
+}