Rework the registry

examples/next-example/src/pages/_app.tsx

@@ -43,27 +43,15 @@ const AppName = () => {
 	);
 };
 
-const RootRegistry = ({ children }: { children: ReactNode }) => {
-	const registry = useMemo(() => createStyleRegistry(), []);
-
-	useServerInsertedHTML(() => {
-		return registry.flushToComponent();
-	});
-
-	return <StyleRegistryProvider registry={registry}>{children}</StyleRegistryProvider>;
-};
-
 const App = ({ Component, pageProps }: AppProps) => {
 	useMobileHover();
 	const auto = useAutomaticTheme("theme", theme);
 
 	return (
-		<RootRegistry>
-			<ThemeProvider theme={{ ...theme, ...auto }}>
-				<Component {...pageProps} />
-				<AppName />
-			</ThemeProvider>
-		</RootRegistry>
+		<ThemeProvider theme={{ ...theme, ...auto }}>
+			<Component {...pageProps} />
+			<AppName />
+		</ThemeProvider>
 	);
 };
 

examples/next-example/src/pages/_document.tsx

@@ -0,0 +1,36 @@
+//
+// Copyright (c) Zoe Roux and contributors. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for details.
+//
+
+import { StyleRegistryProvider, createStyleRegistry } from "yoshiki";
+import Document, { DocumentContext } from "next/document";
+
+Document.getInitialProps = async (ctx: DocumentContext) => {
+	const renderPage = ctx.renderPage;
+	const registry = createStyleRegistry();
+
+	ctx.renderPage = () =>
+		renderPage({
+			enhanceApp: (App) => (props) => {
+				return (
+					<StyleRegistryProvider registry={registry}>
+						<App {...props} />
+					</StyleRegistryProvider>
+				);
+			},
+		});
+
+	const props = await ctx.defaultGetInitialProps(ctx);
+	return {
+		...props,
+		styles: (
+			<>
+				{props.styles}
+				{registry.flushToComponent()}
+			</>
+		),
+	};
+};
+
+export default Document;