Add specs for ControlFlow/? etc.

lib/flux-core/src/lib.rs

@@ -1,6 +1,7 @@
 #![no_std]
 #![cfg_attr(flux, feature(step_trait))]
 #![cfg_attr(flux, feature(sized_hierarchy))]
+#![cfg_attr(flux, feature(try_trait_v2))]
 #![cfg_attr(flux, flux::no_suggestions)]
 
 pub mod iter;

lib/flux-core/src/ops/control_flow.rs

@@ -0,0 +1,15 @@
+use flux_attrs::*;
+
+/// `ControlFlow` is refined by which arm it is, exactly as [`Result`] is refined by `is_ok`.
+///
+/// This is what lets the `?` operator carry a discriminant across the desugaring: `?` expands to
+/// `Try::branch` followed by `FromResidual::from_residual`, routed through a `ControlFlow`. See
+/// the `Try`/`FromResidual` impls for `Option` in `crate::option`.
+#[extern_spec(core::ops)]
+#[refined_by(is_continue: bool)]
+enum ControlFlow<B, C> {
+    #[variant((C) -> ControlFlow<B, C>[true])]
+    Continue(C),
+    #[variant((B) -> ControlFlow<B, C>[false])]
+    Break(B),
+}

lib/flux-core/src/ops/mod.rs

@@ -1,3 +1,4 @@
+mod control_flow;
 mod deref;
 mod index;
 mod range;

lib/flux-core/src/option.rs

@@ -9,6 +9,31 @@ enum Option<T> {
     Some(T),
 }
 
+/// `?` on an `Option` desugars into `Try::branch` followed by `FromResidual::from_residual`,
+/// routed through a [`core::ops::ControlFlow`]. Without specs for these, the `Some`/`None`
+/// discriminant is dropped crossing the `?`.
+///
+/// `Some(x)` continues with `x` and `None` breaks, so the `ControlFlow`'s arm is exactly the
+/// `Option`'s discriminant.
+#[extern_spec(core::option)]
+impl<T> core::ops::Try for Option<T> {
+    #[spec(fn(Option<T>[@b]) -> core::ops::ControlFlow<<Option<T> as core::ops::Try>::Residual, <Option<T> as core::ops::Try>::Output>[b])]
+    fn branch(
+        self,
+    ) -> core::ops::ControlFlow<
+        <Option<T> as core::ops::Try>::Residual,
+        <Option<T> as core::ops::Try>::Output,
+    >;
+}
+
+/// The residual of an `Option` is always the `None` arm, so a `?` propagating out of an
+/// `Option`-returning function yields `None`.
+#[extern_spec(core::option)]
+impl<T> core::ops::FromResidual<Option<core::convert::Infallible>> for Option<T> {
+    #[spec(fn(Option<core::convert::Infallible>) -> Option<T>[false])]
+    fn from_residual(residual: Option<core::convert::Infallible>) -> Option<T>;
+}
+
 #[extern_spec]
 impl<T> Option<T> {
     #[sig(fn(&Self[@b]) -> bool[b])]

lib/flux-core/src/result.rs

@@ -9,6 +9,34 @@ enum Result<T, E> {
     Err(E),
 }
 
+/// `?` on a `Result` desugars into `Try::branch` followed by `FromResidual::from_residual`,
+/// routed through a [`core::ops::ControlFlow`]. Without specs for these, the `Ok`/`Err`
+/// discriminant is dropped crossing the `?`.
+///
+/// `Ok(x)` continues with `x` and `Err(e)` breaks, so the `ControlFlow`'s arm is exactly the
+/// `Result`'s discriminant.
+#[extern_spec(core::result)]
+impl<T, E> core::ops::Try for Result<T, E> {
+    #[spec(fn(Result<T, E>[@b]) -> core::ops::ControlFlow<<Result<T, E> as core::ops::Try>::Residual, <Result<T, E> as core::ops::Try>::Output>[b])]
+    fn branch(
+        self,
+    ) -> core::ops::ControlFlow<
+        <Result<T, E> as core::ops::Try>::Residual,
+        <Result<T, E> as core::ops::Try>::Output,
+    >;
+}
+
+/// The residual of a `Result` is always the `Err` arm, so a `?` propagating out of a
+/// `Result`-returning function yields `Err`. Note the error is converted via `From`, which is
+/// what lets `?` bridge different error types.
+#[extern_spec(core::result)]
+impl<T, E, F: From<E>> core::ops::FromResidual<Result<core::convert::Infallible, E>>
+    for Result<T, F>
+{
+    #[spec(fn(Result<core::convert::Infallible, E>) -> Result<T, F>[false])]
+    fn from_residual(residual: Result<core::convert::Infallible, E>) -> Result<T, F>;
+}
+
 #[extern_spec]
 impl<T, E> Result<T, E> {
     /// Core impl: https://github.com/rust-lang/rust/blob/c6a955468b025dbe3d1de3e8f3e30496d1fb7f40/library/core/src/result.rs#L584

tests/tests/with_deps/neg/extern_specs/control_flow00.rs

@@ -0,0 +1,102 @@
+// Negative counterparts to `pos/extern_specs/control_flow00.rs`: the `ControlFlow` arm and the
+// discriminant carried across `?` are tracked precisely, so claiming the wrong one is rejected.
+extern crate flux_core;
+
+use flux_rs::attrs::*;
+
+// --- `ControlFlow` arms are not interchangeable ---
+
+#[spec(fn(C) -> core::ops::ControlFlow<B, C>[false])]
+fn continue_is_not_break<B, C>(c: C) -> core::ops::ControlFlow<B, C> {
+    core::ops::ControlFlow::Continue(c)
+} //~ ERROR refinement type error
+
+#[spec(fn(B) -> core::ops::ControlFlow<B, C>[true])]
+fn break_is_not_continue<B, C>(b: B) -> core::ops::ControlFlow<B, C> {
+    core::ops::ControlFlow::Break(b)
+} //~ ERROR refinement type error
+
+// Matching refines each arm, so the arms cannot be swapped.
+#[spec(fn(core::ops::ControlFlow<i32, i32>[@k]) -> bool[k])]
+fn is_continue_swapped(cf: core::ops::ControlFlow<i32, i32>) -> bool {
+    match cf {
+        core::ops::ControlFlow::Continue(_) => false, //~ ERROR refinement type error
+        core::ops::ControlFlow::Break(_) => true, //~ ERROR refinement type error
+    }
+}
+
+// --- `?` on `Option` ---
+
+// `?` on a known-`None` always takes the break path, so the result is `None`, not `Some`.
+#[spec(fn(Option<i32>[false]) -> Option<i32>[true])]
+fn qm_none_is_not_some(x: Option<i32>) -> Option<i32> {
+    let v = x?;
+    Some(v)
+} //~ ERROR refinement type error
+
+// `?` on an unknown `Option` may take the break path, so `Some` cannot be claimed.
+#[spec(fn(Option<i32>) -> Option<i32>[true])]
+fn qm_unknown_is_not_some(x: Option<i32>) -> Option<i32> {
+    let v = x?;
+    Some(v)
+} //~ ERROR refinement type error
+
+// Symmetrically, a known-`Some` input cannot yield `None`.
+#[spec(fn(Option<i32>[true]) -> Option<i32>[false])]
+fn qm_some_is_not_none(x: Option<i32>) -> Option<i32> {
+    let v = x?;
+    Some(v) //~ ERROR refinement type error
+}
+
+// `?` carries the value's refinement through unchanged; it does not strengthen it.
+#[spec(fn(Option<i32{v: v > 0}>[true]) -> Option<i32{v: v > 10}>[true])]
+fn qm_does_not_strengthen(x: Option<i32>) -> Option<i32> {
+    let v = x?;
+    Some(v) //~ ERROR refinement type error
+}
+
+// --- `?` on `Result` ---
+
+// `?` on a known-`Err` always breaks, so the result is `Err`, not `Ok`.
+#[spec(fn(Result<i32, i32>[false]) -> Result<i32, i32>[true])]
+fn qm_err_is_not_ok(x: Result<i32, i32>) -> Result<i32, i32> {
+    let v = x?;
+    Ok(v)
+} //~ ERROR refinement type error
+
+// `?` on an unknown `Result` may break, so `Ok` cannot be claimed.
+#[spec(fn(Result<i32, i32>) -> Result<i32, i32>[true])]
+fn qm_result_unknown_is_not_ok(x: Result<i32, i32>) -> Result<i32, i32> {
+    let v = x?;
+    Ok(v)
+} //~ ERROR refinement type error
+
+// Symmetrically, a known-`Ok` input cannot yield `Err`.
+#[spec(fn(Result<i32, i32>[true]) -> Result<i32, i32>[false])]
+fn qm_ok_is_not_err(x: Result<i32, i32>) -> Result<i32, i32> {
+    let v = x?;
+    Ok(v) //~ ERROR refinement type error
+}
+
+// `?` carries the value's refinement through unchanged; it does not strengthen it.
+#[spec(fn(Result<i32{v: v > 0}, i32>[true]) -> Result<i32{v: v > 10}, i32>[true])]
+fn qm_result_does_not_strengthen(x: Result<i32, i32>) -> Result<i32, i32> {
+    let v = x?;
+    Ok(v) //~ ERROR refinement type error
+}
+
+// The discriminant is tracked across the `From` error conversion too.
+struct Wrapper(i32);
+
+impl From<i32> for Wrapper {
+    #[spec(fn(i32) -> Wrapper)]
+    fn from(e: i32) -> Wrapper {
+        Wrapper(e)
+    }
+}
+
+#[spec(fn(Result<i32, i32>[false]) -> Result<i32, Wrapper>[true])]
+fn qm_converts_error_wrong_arm(x: Result<i32, i32>) -> Result<i32, Wrapper> {
+    let v = x?;
+    Ok(v)
+} //~ ERROR refinement type error

tests/tests/with_deps/neg/extern_specs/flux_core_ptr_cast00.rs

@@ -0,0 +1,45 @@
+extern crate flux_core;
+
+// A ptr-to-ptr cast carries the `base`/`addr`/`size` index over to the casted
+// pointer, but it does not conjure up bytes or alignment: `valid` and `aligned_to`
+// are re-checked against the *new* pointee.
+
+// Only 4 bytes are known to be accessible, so an 8-byte read is out of bounds even
+// though the address is suitably aligned.
+#[flux::spec(fn (ptr: {*const[@base, @addr, @size] i32 | addr >= base && addr > 0 && size == 4 && addr % 8 == 0}))]
+pub fn test_cast_widen_read(ptr: *const i32) {
+    let wide_ptr = ptr as *const i64;
+    unsafe {
+        let _v = std::ptr::read(wide_ptr); //~ ERROR refinement type error
+    }
+}
+
+// A byte pointer is only known to be 1-byte aligned, so reading it as an `i32`
+// cannot establish `addr % 4 == 0`.
+#[flux::spec(fn (ptr: {*const[@base, @addr, @size] u8 | addr >= base && addr > 0 && size >= 4}))]
+pub fn test_cast_align(ptr: *const u8) {
+    let int_ptr = ptr as *const i32;
+    unsafe {
+        let _v = std::ptr::read(int_ptr); //~ ERROR refinement type error
+    }
+}
+
+// Stepping one byte into an aligned `i32` and casting back breaks the alignment.
+#[flux::spec(fn (ptr: {*mut[@base, @addr, @size] i32 | addr >= base && addr > 0 && size >= 8 && addr % 4 == 0}))]
+pub fn test_cast_round_trip_unaligned(ptr: *mut i32) {
+    let byte_ptr = ptr as *mut u8;
+    unsafe {
+        let int_ptr = byte_ptr.byte_add(1) as *mut i32;
+        std::ptr::write(int_ptr, 10); //~ ERROR refinement type error
+    }
+}
+
+// The preserved `size` shrinks as the byte pointer advances, so writing past the
+// end of the original `i32` is caught.
+#[flux::spec(fn (ptr: {*mut[@base, @addr, @size] i32 | addr >= base && addr > 0 && size == 4 && addr % 4 == 0}))]
+pub fn test_cast_byte_add_out_of_bounds(ptr: *mut i32) {
+    let byte_ptr = ptr as *mut u8;
+    unsafe {
+        std::ptr::write(byte_ptr.byte_add(4), 255); //~ ERROR refinement type error
+    }
+}

tests/tests/with_deps/pos/extern_specs/control_flow00.rs

@@ -0,0 +1,146 @@
+// Specs for `ControlFlow`, `Try for Option` and `FromResidual for Option`, which together let the
+// `?` operator carry an `Option`'s `Some`/`None` discriminant across its desugaring.
+extern crate flux_core;
+
+use flux_rs::attrs::*;
+
+// --- `ControlFlow` is refined by which arm it is ---
+
+#[spec(fn(C) -> core::ops::ControlFlow<B, C>[true])]
+fn mk_continue<B, C>(c: C) -> core::ops::ControlFlow<B, C> {
+    core::ops::ControlFlow::Continue(c)
+}
+
+#[spec(fn(B) -> core::ops::ControlFlow<B, C>[false])]
+fn mk_break<B, C>(b: B) -> core::ops::ControlFlow<B, C> {
+    core::ops::ControlFlow::Break(b)
+}
+
+// Matching on a `ControlFlow` refines each arm.
+#[spec(fn(core::ops::ControlFlow<i32, i32>[@k]) -> bool[k])]
+fn is_continue(cf: core::ops::ControlFlow<i32, i32>) -> bool {
+    match cf {
+        core::ops::ControlFlow::Continue(_) => true,
+        core::ops::ControlFlow::Break(_) => false,
+    }
+}
+
+// --- `?` on `Option` ---
+
+// The headline case: `?` on a known-`Some` cannot take the `None` path, so the function returns
+// `Some`. Without the `Try`/`FromResidual` specs the discriminant is dropped here.
+#[spec(fn(Option<i32>[true]) -> Option<i32>[true])]
+fn qm_some(x: Option<i32>) -> Option<i32> {
+    let v = x?;
+    Some(v)
+}
+
+// Two `?`s in sequence, both known `Some`.
+#[spec(fn(Option<i32>[true], Option<i32>[true]) -> Option<i32>[true])]
+fn qm_two(x: Option<i32>, y: Option<i32>) -> Option<i32> {
+    let a = x?;
+    let b = y?;
+    Some(a + b)
+}
+
+// `?` on an unknown `Option`: nothing is claimed about the result, but the body still checks.
+#[spec(fn(Option<i32>) -> Option<i32>)]
+fn qm_unknown(x: Option<i32>) -> Option<i32> {
+    let v = x?;
+    Some(v)
+}
+
+// The value carried through `?` keeps its refinement.
+#[spec(fn(Option<i32{v: v > 0}>[true]) -> Option<i32{v: v > 0}>[true])]
+fn qm_preserves_refinement(x: Option<i32>) -> Option<i32> {
+    let v = x?;
+    Some(v)
+}
+
+// A `?` that propagates `None` out of a function returning `Option`: on the break path
+// `from_residual` yields `None`, so a function that can only take that path returns `None`.
+#[spec(fn(Option<i32>[false]) -> Option<i32>[false])]
+fn qm_none(x: Option<i32>) -> Option<i32> {
+    let v = x?;
+    Some(v)
+}
+
+// `?` interacting with the rest of the `Option` specs.
+#[spec(fn(Option<i32>[true]) -> i32)]
+fn qm_then_unwrap(x: Option<i32>) -> i32 {
+    let y = qm_some(x);
+    y.unwrap()
+}
+
+// --- `?` on `Result` ---
+
+// A known-`Ok` cannot take the break path, so the function returns `Ok`.
+#[spec(fn(Result<i32, i32>[true]) -> Result<i32, i32>[true])]
+fn qm_ok(x: Result<i32, i32>) -> Result<i32, i32> {
+    let v = x?;
+    Ok(v)
+}
+
+// A known-`Err` always breaks, so `from_residual` yields `Err`.
+#[spec(fn(Result<i32, i32>[false]) -> Result<i32, i32>[false])]
+fn qm_err(x: Result<i32, i32>) -> Result<i32, i32> {
+    let v = x?;
+    Ok(v)
+}
+
+// Two `?`s in sequence, both known `Ok`.
+#[spec(fn(Result<i32, i32>[true], Result<i32, i32>[true]) -> Result<i32, i32>[true])]
+fn qm_ok_two(x: Result<i32, i32>, y: Result<i32, i32>) -> Result<i32, i32> {
+    let a = x?;
+    let b = y?;
+    Ok(a + b)
+}
+
+// `?` on an unknown `Result`: nothing is claimed, but the body still checks.
+#[spec(fn(Result<i32, i32>) -> Result<i32, i32>)]
+fn qm_result_unknown(x: Result<i32, i32>) -> Result<i32, i32> {
+    let v = x?;
+    Ok(v)
+}
+
+// The value carried through `?` keeps its refinement.
+#[spec(fn(Result<i32{v: v > 0}, i32>[true]) -> Result<i32{v: v > 0}, i32>[true])]
+fn qm_result_preserves_refinement(x: Result<i32, i32>) -> Result<i32, i32> {
+    let v = x?;
+    Ok(v)
+}
+
+// The error-converting case: `?` bridges error types via `From`, and the `Ok`/`Err`
+// discriminant still crosses intact.
+struct Wrapper(i32);
+
+impl From<i32> for Wrapper {
+    #[spec(fn(i32) -> Wrapper)]
+    fn from(e: i32) -> Wrapper {
+        Wrapper(e)
+    }
+}
+
+#[spec(fn(Result<i32, i32>[true]) -> Result<i32, Wrapper>[true])]
+fn qm_converts_error(x: Result<i32, i32>) -> Result<i32, Wrapper> {
+    let v = x?;
+    Ok(v)
+}
+
+#[spec(fn(Result<i32, i32>[false]) -> Result<i32, Wrapper>[false])]
+fn qm_converts_error_err(x: Result<i32, i32>) -> Result<i32, Wrapper> {
+    let v = x?;
+    Ok(v)
+}
+
+#[spec(fn (n: usize) -> Option<usize[n-1]>[n > 0])]
+fn decr(n: usize) -> Option<usize> {
+    if n > 0 { Some(n - 1) } else { None }
+}
+
+#[spec(fn (n: usize) -> Option<usize[n-2]>[n > 1])]
+fn test_decr(n: usize) -> Option<usize> {
+    let n = decr(n)?;
+    let n = decr(n)?;
+    Some(n)
+}