Add "Select all" to `SelectionContainer` contextual menus

Adds the "Select all" operation to both the context menu and text toolbar in `SelectionContainer`. This operation will select all the text across all `BasicText`s in the container.

Relnote: "Added \"Select all\" to all text contextual menus in `SelectionContainer`."
Test: SelectionManagerTest & SelectionContainerContextMenuTest
Fixes: b/240143283
Change-Id: Ib750e9580a290c68356c02cc83bab4cc048e4cc8

diff --git a/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/text/selection/SelectionManager.kt b/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/text/selection/SelectionManager.kt
index 4fad5b6..4cb66ac 100644
--- a/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/text/selection/SelectionManager.kt
+++ b/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/text/selection/SelectionManager.kt

@@ -58,6 +58,7 @@
 import androidx.compose.ui.text.AnnotatedString
 import androidx.compose.ui.text.buildAnnotatedString
 import androidx.compose.ui.unit.IntSize
+import androidx.compose.ui.util.fastAll
 import androidx.compose.ui.util.fastAny
 import androidx.compose.ui.util.fastFold
 import androidx.compose.ui.util.fastForEach
@@ -277,7 +278,7 @@
 
         selectionRegistrar.>
             { isInTouchMode, selectableId ->
-                val (newSelection, newSubselection) = selectAll(
+                val (newSelection, newSubselection) = selectAllInSelectable(
                     selectableId = selectableId,
                     previousSelection = selection,
                 )
@@ -409,7 +410,7 @@
         return coordinates
     }
 
-    internal fun selectAll(
+    internal fun selectAllInSelectable(
         selectableId: Long,
         previousSelection: Selection?
     ): Pair<Selection?, LongObjectMap<Selection>> {
@@ -428,6 +429,70 @@
     }
 
     /**
+     * Returns whether the selection encompasses the entire container.
+     */
+    internal fun isEntireContainerSelected(): Boolean {
+        val selectables = selectionRegistrar.sort(requireContainerCoordinates())
+
+        // If there are no selectables, then an empty selection spans the entire container.
+        if (selectables.isEmpty()) return true
+
+        // Since some text exists, we must make sure that every selectable is fully selected.
+        return selectables.fastAll {
+            val text = it.getText()
+            if (text.isEmpty()) return@fastAll true // empty text is inherently fully selected
+
+            // If a non-empty selectable isn't included in the sub-selections,
+            // then some text in the container is not selected.
+            val subSelection = selectionRegistrar.subselections[it.selectableId]
+                ?: return@fastAll false
+
+            val selectionStart = subSelection.start.offset
+            val selectionEnd = subSelection.end.offset
+
+            // The selection could be reversed,
+            // so just verify that the difference between the two offsets matches the text length
+            (selectionStart - selectionEnd).absoluteValue == text.length
+        }
+    }
+
+    /**
+     * Creates and sets a selection spanning the entire container.
+     */
+    internal fun selectAll() {
+        val selectables = selectionRegistrar.sort(requireContainerCoordinates())
+        if (selectables.isEmpty()) return
+
+        var firstSubSelection: Selection? = null
+        var lastSubSelection: Selection? = null
+        val newSubSelections = mutableLongObjectMapOf<Selection>().apply {
+            selectables.fastForEach { selectable ->
+                val subSelection = selectable.getSelectAllSelection() ?: return@fastForEach
+                if (firstSubSelection == null) firstSubSelection = subSelection
+                lastSubSelection = subSelection
+                put(selectable.selectableId, subSelection)
+            }
+        }
+
+        if (newSubSelections.isEmpty()) return
+
+        // first/last sub selections are implied to be non-null from here on out
+        val newSelection = if (firstSubSelection === lastSubSelection) {
+            firstSubSelection
+        } else {
+            Selection(
+                start = firstSubSelection!!.start,
+                end = lastSubSelection!!.end,
+                handlesCrossed = false,
+            )
+        }
+
+        selectionRegistrar.subselections = newSubSelections
+        onSelectionChange(newSelection)
+        previousSelectionLayout = null
+    }
+
+    /**
      * Returns whether the start and end anchors are equal.
      *
      * It is possible that this returns true, but the selection is still empty because it has
@@ -521,9 +586,13 @@
         }
 
         val textToolbar = textToolbar ?: return
-        if (showToolbar && isInTouchMode && isNonEmptySelection()) {
+        if (showToolbar && isInTouchMode) {
             val rect = getContentRect() ?: return
-            textToolbar.showMenu(rect = rect, >
+            textToolbar.showMenu(
+                rect = rect,
+                 (isNonEmptySelection()) ::toolbarCopy else null,
+                 (isEntireContainerSelected()) null else ::selectAll,
+            )
         } else if (textToolbar.status == TextToolbarStatus.Shown) {
             textToolbar.hide()
         }