Index-dependent refined vector

flux.rs

mod graph {
    use flux_rs::*;

    #[opaque]
    #[refined_by(len: int)]
    // Invariants (trusted):
    // - There are len elements
    // - The i-th element has type [Edge(i); 2]
    pub struct Graph(Vec<[Edge; 2]>);

    #[alias(type Edge(lim: int) = usize{ x: x < if lim == 0 { 1 } else { lim } })]
    type Edge = usize;

    impl Graph {
        #[trusted]
        #[sig(fn() -> Graph[0])]
        pub fn new() -> Self {
            Self(Vec::new())
        }

        #[trusted]
        #[sig(fn(&Graph[@n]) -> usize[n])]
        pub fn len(&self) -> usize {
            self.0.len()
        }

        #[trusted] // index of pushed element will be n
        #[sig(fn(self: &strg Graph[@n], [Edge(n); 2]) ensures self: Graph[n + 1])]
        pub fn push(&mut self, item: [Edge; 2]) {
            self.0.push(item);
        }

        #[trusted] // index of popped element is n - 1
        #[sig(fn(self: &strg { Graph[@n] | 0 < n }) -> [Edge(n - 1); 2] ensures self: Graph[n - 1])]
        pub fn pop(&mut self) -> [Edge; 2] {
            unsafe { self.0.pop().unwrap_unchecked() }
        }

        #[trusted]
        #[sig(fn(&Graph[@n], i: usize{ i < n }) -> [Edge(i); 2])]
        pub fn get(&self, i: usize) -> [Edge; 2] {
            *unsafe { self.0.get_unchecked(i) }
        }

        #[trusted]
        #[sig(fn(&mut Graph[@n], i: usize{ i < n }) -> &mut [Edge(i); 2])]
        pub fn get_mut(&mut self, i: usize) -> &mut [Edge; 2] {
            unsafe { self.0.get_unchecked_mut(i) }
        }
    }
}

fn main() {
    let mut graph = graph::Graph::new();
    // graph.push([0, 1]); // fails
    graph.push([0, 0]);
    // graph.push([1, 0]); // fails
    graph.push([0, 0]);
    // graph.push([2, 2]); // fails
    graph.push([0, 0]);
    assert_eq!(graph.len(), 3);
    // graph.get_mut(2)[0] = 2; // fails
    graph.get_mut(2)[0] = 1;
    // graph.get(3); // fails
    graph.pop();
    graph.pop();
    graph.get(0);
    graph.pop();
    // graph.get(0); // fails
    assert_eq!(graph.len(), 0);
}