#![cfg_attr(verus_keep_ghost, verifier::exec_allows_no_decreases_clause)]
use vstd::prelude::*;
use vstd::atomic_ghost::*;
use vstd::cell::CellId;
use vstd::cell::pcell;
use vstd::modes::*;

verus!{

struct_with_invariants!{
    struct Lock<T> {
        pub atomic: AtomicBool<_, Option<pcell::PointsTo<T>>, _>,
        pub cell: pcell::PCell<T>,
    }

    spec fn wf(self) -> bool {
        invariant on atomic with (cell) is (v: bool, g: Option<pcell::PointsTo<T>>) {
            match g {
                None => {
                    // When there's no PointsTo, the lock must be taken, thus
                    // the boolean value is 'true'.
                    v == true
                }
                Some(points_to) => {
                    points_to.id() == cell.id()
                      && v == false
                }
            }
        }
    }
}

impl<T> Lock<T> {
    fn new(t: T) -> (lock: Self)
        ensures lock.wf()
    {
        let (cell, Tracked(cell_perm)) = pcell::PCell::new(t);
        let atomic = AtomicBool::new(Ghost(cell), false, Tracked(Some(cell_perm)));
        Lock { atomic, cell }
    }

    fn acquire(&self) -> (points_to: Tracked<pcell::PointsTo<T>>)
        requires self.wf(),
        ensures points_to@.id() == self.cell.id(),
    {
        loop
            invariant self.wf(),
        {
            let tracked mut points_to_opt = None;
            let res = atomic_with_ghost!(&self.atomic => compare_exchange(false, true);
                ghost points_to_inv => {
                    tracked_swap(&mut points_to_opt, &mut points_to_inv);
                }
            );
            if res.is_ok() {
                return Tracked(points_to_opt.tracked_unwrap());
            }
        }
    }

    fn release(&self, points_to: Tracked<pcell::PointsTo<T>>)
        requires
            self.wf(),
            points_to@.id() == self.cell.id(),
    {
        atomic_with_ghost!(&self.atomic => store(false);
            ghost points_to_inv => {
                points_to_inv = Some(points_to.get());
            }
        );
    }
}

}

fn main() { }