1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141

//! Sometimes it just takes a small nudge for the compiler to generate the code you want.

#![cfg_attr(feature = "nightly", feature(core_intrinsics))]
#![cfg_attr(not(feature = "std"), no_std)]

mod internal {
    #[cfg(feature = "nightly")]
    mod implem {
        #[inline(always)]
        pub unsafe fn assume(b: bool) {
            core::intrinsics::assume(b)
        }

        #[inline(always)]
        pub fn unlikely(b: bool) -> bool {
            core::intrinsics::unlikely(b)
        }

        #[inline(always)]
        pub fn likely(b: bool) -> bool {
            core::intrinsics::likely(b)
        }
    }

    #[cfg(not(feature = "nightly"))]
    mod implem {
        #[inline(always)]
        #[cold]
        fn cold() {}

        #[inline(always)]
        pub unsafe fn assume(b: bool) {
            if !b {
                crate::unreach()
            }
        }

        #[inline(always)]
        pub fn unlikely(b: bool) -> bool {
            if b {
                cold()
            }
            b
        }

        #[inline(always)]
        pub fn likely(b: bool) -> bool {
            if !b {
                cold()
            }
            b
        }
    }

    pub use implem::*;

    // extern "C" gives us nounwind without having to use lto=fat
    #[cold]
    pub extern "C" fn nounwind_abort() -> ! {
        #[cfg(feature = "std")]
        #[inline(always)]
        fn abort_impl() -> ! {
            std::process::abort()
        }
        #[cfg(not(feature = "std"))]
        #[inline(always)]
        fn abort_impl() -> ! {
            // extern "C" prevents panic from escaping
            panic!()
        }

        abort_impl()
    }

    #[inline(always)]
    pub unsafe extern "C" fn assume_nopanic<F: FnOnce() -> T, T>(f: F) -> T {
        struct NoPanic;
        impl Drop for NoPanic {
            #[inline(always)]
            fn drop(&mut self) {
                unsafe { crate::unreach() };
            }
        }

        let no_panic = NoPanic;
        let r = f();
        core::mem::forget(no_panic);
        r
    }
}

/// Unsafely assumes the value of an expression to be `true`.
///
/// The compiler is sometimes able to use this to optimize better, but it often backfires.
///
/// This generally requires the `nightly` feature to work.
#[inline(always)]
pub unsafe fn assume(b: bool) {
    crate::internal::assume(b)
}

/// Tells the compiler this `bool` is probably `false`.
///
/// This generally requires the `nightly` feature to work.
#[inline(always)]
pub fn unlikely(b: bool) -> bool {
    crate::internal::unlikely(b)
}

/// Tells the compiler this `bool` is probably `true`.
///
/// This generally requires the `nightly` feature to work.
#[inline(always)]
pub fn likely(b: bool) -> bool {
    crate::internal::likely(b)
}

/// Tells the compiler this code is unreachable.
///
/// This is identical to `core::hint::unreachable_unchecked` and is provided only for completeness.
#[inline(always)]
pub unsafe fn unreach() -> ! {
    core::hint::unreachable_unchecked()
}

/// The same as `std::process::abort`, but annotated as `#[cold]`, `nounwind`, and usable in `no_std` environments.
///
/// In a `no_std` environment this generates a trap instruction.
#[cold]
#[inline(always)]
pub fn abort() -> ! {
    crate::internal::nounwind_abort()
}

/// Assumes a closure will not panic.
///
/// Calls to `core::panicking` functions will still be generated; however, this function is nounwind and panics will cause UB.
#[inline]
pub unsafe fn assume_nopanic<F: FnOnce() -> T, T>(f: F) -> T {
    crate::internal::assume_nopanic(f)
}