From e827a6139b7f6dc3d4e47c2edb5c59a2e8f59efc Mon Sep 17 00:00:00 2001
From: Glen Davies <glendaviesnz@users.noreply.github.com>
Date: Wed, 29 May 2024 20:58:12 -0700
Subject: [PATCH] Copy custom CSS between variations when switching (#61752)

---
 .../style-variations-container.js             | 83 ++++++++++++++++---
 1 file changed, 73 insertions(+), 10 deletions(-)

diff --git a/packages/edit-site/src/components/global-styles/style-variations-container.js b/packages/edit-site/src/components/global-styles/style-variations-container.js
index 0d19bcedb174b6..db567db52b3252 100644
--- a/packages/edit-site/src/components/global-styles/style-variations-container.js
+++ b/packages/edit-site/src/components/global-styles/style-variations-container.js
@@ -3,9 +3,10 @@
  */
 import { store as coreStore } from '@wordpress/core-data';
 import { useSelect } from '@wordpress/data';
-import { useMemo } from '@wordpress/element';
+import { useContext, useEffect, useMemo, useState } from '@wordpress/element';
 import { __experimentalGrid as Grid } from '@wordpress/components';
 import { __ } from '@wordpress/i18n';
+import { privateApis as blockEditorPrivateApis } from '@wordpress/block-editor';
 
 /**
  * Internal dependencies
@@ -13,8 +14,19 @@ import { __ } from '@wordpress/i18n';
 import PreviewStyles from './preview-styles';
 import Variation from './variations/variation';
 import { isVariationWithSingleProperty } from '../../hooks/use-theme-style-variations/use-theme-style-variations-by-property';
+import { unlock } from '../../lock-unlock';
+
+const { GlobalStylesContext } = unlock( blockEditorPrivateApis );
 
 export default function StyleVariationsContainer( { gap = 2 } ) {
+	const { user } = useContext( GlobalStylesContext );
+	const [ currentUserStyles, setCurrentUserStyles ] = useState( user );
+	const userStyles = currentUserStyles?.styles;
+
+	useEffect( () => {
+		setCurrentUserStyles( user );
+	}, [ user ] );
+
 	const variations = useSelect( ( select ) => {
 		return select(
 			coreStore
@@ -29,20 +41,71 @@ export default function StyleVariationsContainer( { gap = 2 } ) {
 		);
 	} );
 
-	const withEmptyVariation = useMemo( () => {
-		return [
+	const themeVariations = useMemo( () => {
+		const withEmptyVariation = [
 			{
 				title: __( 'Default' ),
 				settings: {},
 				styles: {},
 			},
-			...( multiplePropertyVariations ?? [] ).map( ( variation ) => ( {
-				...variation,
-				settings: variation.settings ?? {},
-				styles: variation.styles ?? {},
-			} ) ),
+			...( multiplePropertyVariations ?? [] ),
+		];
+		return [
+			...withEmptyVariation.map( ( variation ) => {
+				const blockStyles = { ...variation?.styles?.blocks } || {};
+
+				// We need to copy any user custom CSS to the variation to prevent it being lost
+				// when switching variations.
+				if ( userStyles?.blocks ) {
+					Object.keys( userStyles.blocks ).forEach( ( blockName ) => {
+						// First get any block specific custom CSS from the current user styles and merge with any custom CSS for
+						// that block in the variation.
+						if ( userStyles.blocks[ blockName ].css ) {
+							const variationBlockStyles =
+								blockStyles[ blockName ] || {};
+							const customCSS = {
+								css: `${
+									blockStyles[ blockName ]?.css || ''
+								} ${
+									userStyles.blocks[ blockName ].css.trim() ||
+									''
+								}`,
+							};
+							blockStyles[ blockName ] = {
+								...variationBlockStyles,
+								...customCSS,
+							};
+						}
+					} );
+				}
+				// Now merge any global custom CSS from current user styles with global custom CSS in the variation.
+				const css =
+					userStyles?.css || variation.styles?.css
+						? {
+								css: `${ variation.styles?.css || '' } ${
+									userStyles?.css || ''
+								}`,
+						  }
+						: {};
+
+				const blocks =
+					Object.keys( blockStyles ).length > 0
+						? { blocks: blockStyles }
+						: {};
+
+				const styles = {
+					...variation.styles,
+					...css,
+					...blocks,
+				};
+				return {
+					...variation,
+					settings: variation.settings ?? {},
+					styles,
+				};
+			} ),
 		];
-	}, [ multiplePropertyVariations ] );
+	}, [ multiplePropertyVariations, userStyles?.blocks, userStyles?.css ] );
 
 	return (
 		<Grid
@@ -50,7 +113,7 @@ export default function StyleVariationsContainer( { gap = 2 } ) {
 			className="edit-site-global-styles-style-variations-container"
 			gap={ gap }
 		>
-			{ withEmptyVariation.map( ( variation, index ) => (
+			{ themeVariations.map( ( variation, index ) => (
 				<Variation key={ index } variation={ variation }>
 					{ ( isFocused ) => (
 						<PreviewStyles