streaming-styles.jsx

import * as React from "react";

/*
Usage

import { Style } from "style";
function Heading({children}) {
  return (
    <>
      <Style>{`h1 { color: red }`}</Style>
      <h1>{children}</h1>
    </>
  )
}

import { StyleProvider, createRegistry } from "style";
function App() {
  return (
    <StyleProvider registry={createRegistry}>
      <Heading>{children}</Heading>
    </StyleProvider>
  )
}
*/

/*
 The idea is that the first instance of a component renders the <style> tag for all the instances and upon hydration (in an effect) 
 removes that and injects the same styles in a CSSStyleSheet (using the cssom apis). 
 Any other instance "registers" itself incrementing an instance count and deregisters itself when unmounting. 
 When there aren't instances in the document anymore the styles are removed from the stylesheet. 
*/

const isBrowser = typeof window !== "undefined";

export const StyleContext = React.createContext({
  register: () => {},
  unregister: () => {},
  has: () => false,
});

export function StyleProvider({ children, registry }) {
  return (
    <StyleContext.Provider value={registry}>{children}</StyleContext.Provider>
  );
}

function useIsomorphicEffect(...args) {
  return React[isBrowser ? "useLayoutEffect" : "useEffect"](...args);
}
export function Style({ children }) {
  const id = React.useMemo(() => {
    try {
      return btoa(children).replace("=", "");
    } catch (err) {
      return Buffer.from(children)
        .toString("base64")
        .replace("=", "");
    }
  }, [children]);

  const registry = React.useContext(StyleContext);

  const state = React.useState(() => {
    return (
      !registry.has(id) ||
      (isBrowser && document.querySelector(`#${id}-ssr`) === null)
    );
  });
  const shouldRender = state[0];
  const setShouldRender = state[1];

  if (!isBrowser) {
    registry.register(id);
  }

  useIsomorphicEffect(() => {
    if (!registry.has(id)) {
      if (Array.isArray(children)) {
        registry.sheet.insert.apply(null, [id].concat(children));
      } else {
        registry.sheet.insert(id, children);
      }
    }
    registry.register(id);
    setShouldRender(false);

    return () => {
      registry.unregister(id);
      if (!registry.has(id)) {
        registry.sheet.remove(id);
      }
    };
  }, [id, children, registry]);

  return shouldRender ? <style id={`${id}-ssr`}>{children}</style> : null;
}

export function createRegistry(sheet = createStyleSheet()) {
  const registry = {};
  return {
    register(id) {
      if (registry[id]) {
        registry[id] += 1;
      } else {
        registry[id] = 1;
      }
    },
    unregister(id) {
      if (registry[id]) {
        registry[id] -= 1;
      } else {
        delete registry[id];
      }
    },
    has(id) {
      return registry.hasOwnProperty(id);
    },
    sheet,
  };
}

export function createStyleSheet(sheet = null) {
  if (typeof window === "undefined") {
    return {};
  }
  const indices = {};
  if (!sheet) {
    const tag = document.head.appendChild(document.createElement("style"));
    tag.id = "Curtesy Of Giuseppe Gurgone";
    sheet = tag.sheet;
  }
  return {
    insert(id, ...rules) {
      indices[id] = rules
        .map((rule) => {
          try {
            return sheet.insertRule(rule);
          } catch (err) {
            return -1;
          }
        })
        // Filter out invalid rules
        .filter((index) => index !== -1);
      return indices[id];
    },
    remove(id) {
      indices[id].forEach((index) => sheet.deleteRule(index));
      delete indices[id];
    },
  };
}