From 91784bbe93bee8fd4ba4e53f9a3337a5ca11da48 Mon Sep 17 00:00:00 2001 From: Alex Waygood Date: Tue, 28 Jan 2025 13:39:48 +0000 Subject: [PATCH] is this faster --- .../mdtest/exception/control_flow.md | 6 +-- .../resources/mdtest/narrow/truthiness.md | 14 +++---- .../src/types/builder.rs | 39 +++++++++---------- 3 files changed, 29 insertions(+), 30 deletions(-) diff --git a/crates/red_knot_python_semantic/resources/mdtest/exception/control_flow.md b/crates/red_knot_python_semantic/resources/mdtest/exception/control_flow.md index 284b0f24d5..a6e703dff5 100644 --- a/crates/red_knot_python_semantic/resources/mdtest/exception/control_flow.md +++ b/crates/red_knot_python_semantic/resources/mdtest/exception/control_flow.md @@ -455,9 +455,9 @@ else: reveal_type(x) # revealed: slice finally: # TODO: should be `Literal[1] | str | bytes | bool | memoryview | float | range | slice` - reveal_type(x) # revealed: bool | float | slice + reveal_type(x) # revealed: bool | slice | float -reveal_type(x) # revealed: bool | float | slice +reveal_type(x) # revealed: bool | slice | float ``` ## Nested `try`/`except` blocks @@ -534,7 +534,7 @@ try: reveal_type(x) # revealed: slice finally: # TODO: should be `Literal[1] | str | bytes | bool | memoryview | float | range | slice` - reveal_type(x) # revealed: bool | float | slice + reveal_type(x) # revealed: bool | slice | float x = 2 reveal_type(x) # revealed: Literal[2] reveal_type(x) # revealed: Literal[2] diff --git a/crates/red_knot_python_semantic/resources/mdtest/narrow/truthiness.md b/crates/red_knot_python_semantic/resources/mdtest/narrow/truthiness.md index b3975c1a81..203ffe9827 100644 --- a/crates/red_knot_python_semantic/resources/mdtest/narrow/truthiness.md +++ b/crates/red_knot_python_semantic/resources/mdtest/narrow/truthiness.md @@ -11,37 +11,37 @@ x = foo() if x: reveal_type(x) # revealed: Literal[-1, True, "foo", b"bar"] else: - reveal_type(x) # revealed: Literal[0, False, "", b""] | None | tuple[()] + reveal_type(x) # revealed: Literal[0, False, "", b""] | tuple[()] | None if not x: - reveal_type(x) # revealed: Literal[0, False, "", b""] | None | tuple[()] + reveal_type(x) # revealed: Literal[0, False, "", b""] | tuple[()] | None else: reveal_type(x) # revealed: Literal[-1, True, "foo", b"bar"] if x and not x: reveal_type(x) # revealed: Never else: - reveal_type(x) # revealed: Literal[0, -1, "", "foo", b"", b"bar"] | bool | None | tuple[()] + reveal_type(x) # revealed: Literal[0, -1, b"bar", "", "foo", b""] | bool | tuple[()] | None if not (x and not x): - reveal_type(x) # revealed: Literal[0, -1, "", "foo", b"", b"bar"] | bool | None | tuple[()] + reveal_type(x) # revealed: Literal[0, -1, b"bar", "", "foo", b""] | bool | tuple[()] | None else: reveal_type(x) # revealed: Never if x or not x: - reveal_type(x) # revealed: Literal[0, -1, "", "foo", b"", b"bar"] | bool | None | tuple[()] + reveal_type(x) # revealed: Literal[0, -1, b"bar", "", "foo", b""] | bool | tuple[()] | None else: reveal_type(x) # revealed: Never if not (x or not x): reveal_type(x) # revealed: Never else: - reveal_type(x) # revealed: Literal[0, -1, "", "foo", b"", b"bar"] | bool | None | tuple[()] + reveal_type(x) # revealed: Literal[0, -1, b"bar", "", "foo", b""] | bool | tuple[()] | None if (isinstance(x, int) or isinstance(x, str)) and x: reveal_type(x) # revealed: Literal[-1, True, "foo"] else: - reveal_type(x) # revealed: Literal[b"", b"bar", 0, False, ""] | None | tuple[()] + reveal_type(x) # revealed: tuple[()] | None | Literal[b"", b"bar", 0, False, ""] ``` ## Function Literals diff --git a/crates/red_knot_python_semantic/src/types/builder.rs b/crates/red_knot_python_semantic/src/types/builder.rs index 193a63a1a3..49ed137a96 100644 --- a/crates/red_knot_python_semantic/src/types/builder.rs +++ b/crates/red_knot_python_semantic/src/types/builder.rs @@ -33,6 +33,7 @@ use smallvec::SmallVec; pub(crate) struct UnionBuilder<'db> { elements: Vec>, db: &'db dyn Db, + contains_bool_literals: bool, } impl<'db> UnionBuilder<'db> { @@ -40,6 +41,7 @@ impl<'db> UnionBuilder<'db> { Self { db, elements: vec![], + contains_bool_literals: false, } } @@ -77,6 +79,7 @@ impl<'db> UnionBuilder<'db> { return self; } } + self.contains_bool_literals |= ty.is_boolean_literal(); match to_remove[..] { [] => self.elements.push(ty), [index] => self.elements[index] = ty, @@ -103,34 +106,30 @@ impl<'db> UnionBuilder<'db> { } pub(crate) fn build(self) -> Type<'db> { - let UnionBuilder { elements, db } = self; + let UnionBuilder { + mut elements, + db, + contains_bool_literals, + } = self; match elements.len() { 0 => Type::Never, 1 => elements[0], _ => { - let mut normalized_elements = Vec::with_capacity(elements.len()); - let mut first_bool_literal_pos = None; - let mut seen_two_bool_literals = false; - for (i, element) in elements.into_iter().enumerate() { - if element.is_boolean_literal() { - if first_bool_literal_pos.is_none() { - first_bool_literal_pos = Some(i); - } else { - seen_two_bool_literals = true; - continue; + if contains_bool_literals { + let mut element_iter = elements.iter(); + if let Some(first_pos) = element_iter.position(Type::is_boolean_literal) { + if let Some(second_pos) = element_iter.position(Type::is_boolean_literal) { + let bool_instance = KnownClass::Bool.to_instance(db); + if elements.len() == 2 { + return bool_instance; + } + elements.swap_remove(first_pos + second_pos + 1); + elements[first_pos] = bool_instance; } } - normalized_elements.push(element); } - if let (Some(pos), true) = (first_bool_literal_pos, seen_two_bool_literals) { - // If we have two boolean literals, we can merge them to `bool`. - if normalized_elements.len() == 1 { - return KnownClass::Bool.to_instance(db); - } - normalized_elements[pos] = KnownClass::Bool.to_instance(db); - } - Type::Union(UnionType::new(db, normalized_elements.into_boxed_slice())) + Type::Union(UnionType::new(db, elements.into_boxed_slice())) } } }