Use Option's discriminant as its size hint

scottmcm · scottmcm · commit eb3cc5f824d7 · 2024-07-15T00:34:03.000-07:00
diff --git a/library/core/src/option.rs b/library/core/src/option.rs
@@ -770,6 +770,13 @@ impl<T> Option<T> {
         }
     }
 
+    #[inline]
+    const fn len(&self) -> usize {
+        // Using the intrinsic avoids emitting a branch to get the 0 or 1.
+        let discriminant: isize = crate::intrinsics::discriminant_value(self);
+        discriminant as usize
+    }
+
     /// Returns a slice of the contained value, if any. If this is `None`, an
     /// empty slice is returned. This can be useful to have a single type of
     /// iterator over an `Option` or slice.
@@ -812,7 +819,7 @@ impl<T> Option<T> {
         unsafe {
             slice::from_raw_parts(
                 (self as *const Self).byte_add(core::mem::offset_of!(Self, Some.0)).cast(),
-                self.is_some() as usize,
+                self.len(),
             )
         }
     }
@@ -869,7 +876,7 @@ impl<T> Option<T> {
         unsafe {
             slice::from_raw_parts_mut(
                 (self as *mut Self).byte_add(core::mem::offset_of!(Self, Some.0)).cast(),
-                self.is_some() as usize,
+                self.len(),
             )
         }
     }
@@ -2242,10 +2249,8 @@ impl<A> Iterator for Item<A> {
 
     #[inline]
     fn size_hint(&self) -> (usize, Option<usize>) {
-        match self.opt {
-            Some(_) => (1, Some(1)),
-            None => (0, Some(0)),
-        }
+        let len = self.len();
+        (len, Some(len))
     }
 }
 
@@ -2256,7 +2261,12 @@ impl<A> DoubleEndedIterator for Item<A> {
     }
 }
 
-impl<A> ExactSizeIterator for Item<A> {}
+impl<A> ExactSizeIterator for Item<A> {
+    #[inline]
+    fn len(&self) -> usize {
+        self.opt.len()
+    }
+}
 impl<A> FusedIterator for Item<A> {}
 unsafe impl<A> TrustedLen for Item<A> {}
 
diff --git a/tests/codegen/option-as-slice.rs b/tests/codegen/option-as-slice.rs
@@ -14,6 +14,14 @@ pub fn u64_opt_as_slice(o: &Option<u64>) -> &[u64] {
     // CHECK-NOT: br
     // CHECK-NOT: switch
     // CHECK-NOT: icmp
+    // CHECK: %[[LEN:.+]] = load i64,{{.+}} !range ![[META_U64:.+]], !noundef
+    // CHECK-NOT: select
+    // CHECK-NOT: br
+    // CHECK-NOT: switch
+    // CHECK-NOT: icmp
+    // CHECK: %[[T0:.+]] = insertvalue { ptr, i64 } poison, ptr %{{.+}}, 0
+    // CHECK-NEXT: %[[T1:.+]] = insertvalue { ptr, i64 } %[[T0]], i64 %[[LEN]], 1
+    // CHECK-NEXT: ret { ptr, i64 } %[[T1]]
     o.as_slice()
 }
 
@@ -25,10 +33,35 @@ pub fn nonzero_u64_opt_as_slice(o: &Option<NonZero<u64>>) -> &[NonZero<u64>] {
     // CHECK-NOT: switch
     // CHECK-NOT: icmp
     // CHECK: %[[NZ:.+]] = icmp ne i64 %{{.+}}, 0
-    // CHECK-NEXT: zext i1 %[[NZ]] to i64
+    // CHECK-NEXT: %[[LEN:.+]] = zext i1 %[[NZ]] to i64
     // CHECK-NOT: select
     // CHECK-NOT: br
     // CHECK-NOT: switch
     // CHECK-NOT: icmp
+    // CHECK: %[[T0:.+]] = insertvalue { ptr, i64 } poison, ptr %o, 0
+    // CHECK-NEXT: %[[T1:.+]] = insertvalue { ptr, i64 } %[[T0]], i64 %[[LEN]], 1
+    // CHECK-NEXT: ret { ptr, i64 } %[[T1]]
     o.as_slice()
 }
+
+// CHECK-LABEL: @u8_opt_as_slice
+#[no_mangle]
+pub fn u8_opt_as_slice(o: &Option<u8>) -> &[u8] {
+    // CHECK-NOT: select
+    // CHECK-NOT: br
+    // CHECK-NOT: switch
+    // CHECK-NOT: icmp
+    // CHECK: %[[TAG:.+]] = load i8,{{.+}} !range ![[META_U8:.+]], !noundef
+    // CHECK: %[[LEN:.+]] = zext{{.*}} i8 %[[TAG]] to i64
+    // CHECK-NOT: select
+    // CHECK-NOT: br
+    // CHECK-NOT: switch
+    // CHECK-NOT: icmp
+    // CHECK: %[[T0:.+]] = insertvalue { ptr, i64 } poison, ptr %{{.+}}, 0
+    // CHECK-NEXT: %[[T1:.+]] = insertvalue { ptr, i64 } %[[T0]], i64 %[[LEN]], 1
+    // CHECK-NEXT: ret { ptr, i64 } %[[T1]]
+    o.as_slice()
+}
+
+// CHECK: ![[META_U64]] = !{i64 0, i64 2}
+// CHECK: ![[META_U8]] = !{i8 0, i8 2}

Original file line number	Diff line number	Diff line change
`@@ -770,6 +770,13 @@ impl<T> Option<T> {`
`770`	`770`	`}`
`771`	`771`	`}`
`772`	`772`
	`773`	`+ #[inline]`
	`774`	`+ const fn len(&self) -> usize {`
	`775`	`+ // Using the intrinsic avoids emitting a branch to get the 0 or 1.`
	`776`	`+ let discriminant: isize = crate::intrinsics::discriminant_value(self);`
	`777`	`+ discriminant as usize`
	`778`	`+ }`
	`779`	`+`
`773`	`780`	/// Returns a slice of the contained value, if any. If this is `None`, an
`774`	`781`	`/// empty slice is returned. This can be useful to have a single type of`
`775`	`782`	/// iterator over an `Option` or slice.
`@@ -812,7 +819,7 @@ impl<T> Option<T> {`
`812`	`819`	`unsafe {`
`813`	`820`	`slice::from_raw_parts(`
`814`	`821`	`(self as *const Self).byte_add(core::mem::offset_of!(Self, Some.0)).cast(),`
`815`		`- self.is_some() as usize,`
	`822`	`+ self.len(),`
`816`	`823`	`)`
`817`	`824`	`}`
`818`	`825`	`}`
`@@ -869,7 +876,7 @@ impl<T> Option<T> {`
`869`	`876`	`unsafe {`
`870`	`877`	`slice::from_raw_parts_mut(`
`871`	`878`	`(self as *mut Self).byte_add(core::mem::offset_of!(Self, Some.0)).cast(),`
`872`		`- self.is_some() as usize,`
	`879`	`+ self.len(),`
`873`	`880`	`)`
`874`	`881`	`}`
`875`	`882`	`}`
`@@ -2242,10 +2249,8 @@ impl<A> Iterator for Item<A> {`
`2242`	`2249`
`2243`	`2250`	`#[inline]`
`2244`	`2251`	`fn size_hint(&self) -> (usize, Option<usize>) {`
`2245`		`- match self.opt {`
`2246`		`- Some(_) => (1, Some(1)),`
`2247`		`- None => (0, Some(0)),`
`2248`		`- }`
	`2252`	`+ let len = self.len();`
	`2253`	`+ (len, Some(len))`
`2249`	`2254`	`}`
`2250`	`2255`	`}`
`2251`	`2256`
`@@ -2256,7 +2261,12 @@ impl<A> DoubleEndedIterator for Item<A> {`
`2256`	`2261`	`}`
`2257`	`2262`	`}`
`2258`	`2263`
`2259`		`-impl<A> ExactSizeIterator for Item<A> {}`
	`2264`	`+impl<A> ExactSizeIterator for Item<A> {`
	`2265`	`+ #[inline]`
	`2266`	`+ fn len(&self) -> usize {`
	`2267`	`+ self.opt.len()`
	`2268`	`+ }`
	`2269`	`+}`
`2260`	`2270`	`impl<A> FusedIterator for Item<A> {}`
`2261`	`2271`	`unsafe impl<A> TrustedLen for Item<A> {}`
`2262`	`2272`