Use Publish button as anchor for AMP Preview button

The post Publish/Update button is rendered early on (contrary to post Preview button) and so it's safer to use it to imperatively add an AMP Preview button to the block editor markup.

Also, the `viewable` property of a post type get defined late, similarly to when the post Preview button gets added. That's why we assume a post type is viewable and show the AMP Preview button initially. Only if we learn that a post type is not viewable, we hide the AMP Preview button.

Author: delawski

assets/src/block-editor/constants.js

@@ -5,5 +5,3 @@ export const TEXT_BLOCKS = [
 	'core/quote',
 	'core/subhead',
 ];
-
-export const POST_PREVIEW_CLASS = 'editor-post-preview';

assets/src/block-editor/plugins/wrapped-amp-preview-button.js

@@ -7,41 +7,46 @@ import { useSelect } from '@wordpress/data';
 /**
  * Internal dependencies
  */
-import { POST_PREVIEW_CLASS } from '../constants';
 import AmpPreviewButton from '../components/amp-preview-button';
 
 /**
  * A wrapper for the AMP preview button that renders it immediately after the 'Post' preview button, when present.
  */
 function WrappedAmpPreviewButton() {
-	const root = useRef();
-	const postPreviewButton = useRef();
+	const root = useRef( null );
+	const postPublishButton = useRef( null );
 
 	const isViewable = useSelect(
 		( select ) => select( 'core' ).getPostType( select( 'core/editor' ).getEditedPostAttribute( 'type' ) )?.viewable,
 		[],
 	);
 
 	useEffect( () => {
-		if ( isViewable && ! root.current && ! postPreviewButton.current ) {
-			postPreviewButton.current = document.querySelector( `.${ POST_PREVIEW_CLASS }` );
+		if ( ! root.current && ! postPublishButton.current ) {
+			postPublishButton.current = document.querySelector( '.editor-post-publish-button' );
 
 			// Insert the AMP preview button immediately after the post preview button.
-			if ( postPreviewButton.current ) {
+			if ( postPublishButton.current ) {
 				root.current = document.createElement( 'div' );
 				root.current.className = 'amp-wrapper-post-preview';
-				postPreviewButton.current.parentNode.insertBefore( root.current, postPreviewButton.current.nextSibling );
+				postPublishButton.current.parentNode.insertBefore( root.current, postPublishButton.current );
 			}
 		}
 
 		return () => {
-			if ( postPreviewButton.current && root.current ) {
-				postPreviewButton.current.parentNode.removeChild( root.current );
+			if ( postPublishButton.current && root.current ) {
+				postPublishButton.current.parentNode.removeChild( root.current );
+				root.current = null;
+				postPublishButton.current = null;
 			}
 		};
-	}, [ isViewable ] );
+	}, [] );
 
-	if ( ! isViewable || ! root.current ) {
+	// It is unlikely that AMP would be enabled for a non-viewable post type. This is why the Preview button will
+	// always be displayed initially (when `isViewable` is undefined), preventing horizontal layout shift.
+	// Once the `isViewable` value is defined (which is after the initial block editor load) and it is `false`,
+	// the Preview button will be hidden causing a minor layout shift.
+	if ( ! root.current || isViewable === false ) {
 		return null;
 	}