[red-knot] flatten unions

[carljm]

Flatten union types. Fixes #11781

[github-actions]

ruff-ecosystem results

Linter (stable)

✅ ecosystem check detected no linter changes.

Linter (preview)

✅ ecosystem check detected no linter changes.

[AlexWaygood]

I managed to get the elements() method working without any .collect() calls, if you're interested:

diff --git a/crates/red_knot/src/semantic/types.rs b/crates/red_knot/src/semantic/types.rs
index 360b157c5..ae866e00c 100644
--- a/crates/red_knot/src/semantic/types.rs
+++ b/crates/red_knot/src/semantic/types.rs
@@ -1,4 +1,6 @@
 #![allow(dead_code)]
+use std::iter::FusedIterator;
+
 use crate::ast_ids::NodeKey;
 use crate::db::{QueryResult, SemanticDb, SemanticJar};
 use crate::files::FileId;
@@ -531,12 +533,41 @@ pub struct UnionTypeId {
 }
 
 impl UnionTypeId {
-    pub fn elements(self, type_store: &TypeStore) -> Vec<Type> {
+    pub fn elements(self, type_store: &TypeStore) -> UnionElementsIterator {
         let union = type_store.get_union(self);
-        union.elements.iter().copied().collect()
+        UnionElementsIterator { union, index: 0 }
     }
 }
 
+pub struct UnionElementsIterator<'a> {
+    union: UnionTypeRef<'a>,
+    index: usize,
+}
+
+impl<'a> Iterator for UnionElementsIterator<'a> {
+    type Item = Type;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let item = self.union.elements.get_index(self.index).copied();
+        self.index += 1;
+        item
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let size = self.union.elements.len() - self.index;
+        (size, Some(size))
+    }
+
+    fn last(self) -> Option<Self::Item> {
+        self.union.elements.last().copied()
+    }
+}
+
+impl<'a> FusedIterator for UnionElementsIterator<'a> {}
+impl<'a> ExactSizeIterator for UnionElementsIterator<'a> {}

A bit more verbose, but maybe worth it if we'll be iterating over union elements a lot? WDYT?

[carljm]

Yeah, I think that's probably worth it! I knew this type of borrowing iterator was an option, but at the time I was being stubborn about trying to get it working with flat_map, which requires that the returned iterator own its items, meaning this type of iterator wouldn't work for that. But now that we aren't using flat_map to flatten anyway, this would work fine.

We should do the same for IntersectionPositiveIterator and IntersectionNegativeIterator as well.

I think this is something that can wait until post-Salsa, though; no need to introduce more PRs right now for Micha to port.