Cleanup code.

packages/block-editor/src/components/writing-flow/use-tab-nav.js

@@ -41,7 +41,7 @@ export default function useTabNav() {
 		} else if ( hasMultiSelection() ) {
 			container.current.focus();
 		} else if ( getSelectedBlockClientId() ) {
-			// Try to focus the last element.
+			// Try to focus the last element which had focus.
 			lastFocus.current.focus();
 			// Check to see if focus worked.
 			if (
@@ -51,24 +51,17 @@ export default function useTabNav() {
 				// Looks like yes, return early.
 				return;
 			}
-			// Access the last known block and try to find a new area to focus.
-			const selectedBlock = lastBlock.current;
-			if ( ! selectedBlock ) {
-				// TODO: Should this be handled?
-				return;
-			}
-			// Get first tabbable in the block if one exists.
-			const firstBlockTabbable = focus.tabbable.findNext( selectedBlock );
-			// Check to ensure tabbable element is inside the current block.
-			const tabbableInSelectedBlock = firstBlockTabbable
-				? isInSameBlock( firstBlockTabbable, selectedBlock )
-				: false;
-			if ( tabbableInSelectedBlock ) {
-				// Focus the found tabbable in the selected block.
+			// Last element focus did not work. Now try to find the first tabbable in the last block to focus.
+			const firstBlockTabbable = focus.tabbable.findNext(
+				lastBlock.current
+			);
+			// Check to ensure tabbable is inside the last block and that it is a form element.
+			if ( isInSameBlock( lastBlock.current, firstBlockTabbable ) ) {
+				// Focus the found tabbable in the last block.
 				firstBlockTabbable.focus();
 			} else {
-				// Focus the block wrapper if no tabbable was found.
-				selectedBlock.focus();
+				// Focus the last block wrapper if no tabbable was found.
+				lastBlock.current.focus();
 			}
 		} else {
 			setNavigationMode( true );