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

// This example demonstrates using `Target` to create a "proof token type", a token that
// demonstrates that particular target features have already been detected, and that those features
// can be used safely.

#![allow(unused, unused_macros, unused_imports)]

use target_features::Target;

/// Make sure proof tokens can't be improperly constructed
mod unconstructible {
    pub struct Unconstructible(());
    impl Unconstructible {
        pub unsafe fn new() -> Self {
            Self(())
        }
    }
}
use unconstructible::Unconstructible;

/// Proof of target feature support.
///
/// # Safety
/// The type must be implemented such that it's impossible to safely construct without ensuring the
/// specified target features are supported.
unsafe trait Proof: Sized {
    /// The proven target
    const TARGET: Target;

    /// Detect the support for the target features
    fn detect() -> Option<Self>;

    /// Assume the target features are supported
    ///
    /// # Safety
    /// Calling this is undefined if the target features are not supported
    unsafe fn assume() -> Self;
}

/// Make a proof token type for a particular set of features
macro_rules! make_target_proof {
    { $vis:vis struct $proof:ident($($feature:tt),*); } => {
        $vis struct $proof(Unconstructible);

        unsafe impl Proof for $proof {
            // Build on the already-known target features
            const TARGET: Target = target_features::CURRENT_TARGET$(.with_feature_str($feature))*;

            fn detect() -> Option<Self> {
                if true $(&& is_x86_feature_detected!($feature))* {
                    unsafe { Some(Self::assume()) }
                } else {
                    None
                }
            }

            unsafe fn assume() -> Self {
                Self(Unconstructible::new())
            }
        }
    }
}

/// A function that can only be called with the "avx" feature, or panics otherwise.
#[cfg(target_arch = "x86_64")]
fn safe_avx_fn<P: Proof>(_: P) {
    #[target_feature(enable = "avx")]
    unsafe fn unsafe_avx_fn() {
        println!("called an avx function")
    }

    // Future improvements to const generics might make it possible to assert this at compile time.
    // Since P::TARGET is const, this assert disappears if the required features are present.
    assert!(
        P::TARGET.supports_feature_str("avx"),
        "avx feature not supported"
    );
    unsafe { unsafe_avx_fn() }
}

#[cfg(target_arch = "x86_64")]
fn main() {
    // The function can be called with the exact features
    make_target_proof! {
        struct Avx("avx");
    }
    if let Some(proof) = Avx::detect() {
        safe_avx_fn(proof);
    }

    // The function can also be called with a target that implies the required features
    make_target_proof! {
        struct Avx2("avx2");
    }
    if let Some(proof) = Avx2::detect() {
        safe_avx_fn(proof);
    }

    // This panics, unless compiled with something like `-Ctarget-feature=+avx`
    make_target_proof! {
        struct Aes("aes");
    }
    if let Some(proof) = Aes::detect() {
        safe_avx_fn(proof);
    }
}

#[cfg(not(target_arch = "x86_64"))]
fn main() {}