RawVec.dfy

module Util {
    type usize = nat

    type Ptr = nat
    function null_ptr(): Ptr { 0 }
    predicate is_null(p: Ptr) { p == 0 }

    datatype Option<T> = None | Some(value: T)

    function range(start: nat, stop: nat): seq<nat>
        requires start <= stop
        decreases stop - start
    {
        if start == stop
            then []
        else [start] + range(start + 1, stop)
    }

    // function fmap<A, B>(f: A -> B, list: seq<A>): seq<B> {
    //   if list == []
    //     then []
    //   else
    //     [f(list[0])] + fmap(f, list[1..])
    // }

    function filter<T>(pred: T -> bool, list: seq<T>): seq<T> {
        if list == []
            then []
        else if pred(list[0])
            then [list[0]] + filter(pred, list[1..])
        else
            filter(pred, list[1..])
    }

    function min(a: usize, b: usize): usize {
        if a < b then a else b
    }

    function replace_slice<T>(s: seq<T>, i: usize, x: seq<T>): (t: seq<T>)
        requires 0 <= i <= i + |x| <= |s|
        ensures |t| == |s| &&
            t[..i] == s[..i] &&
            t[i..i+|x|] == x &&
            t[i + |x|..] == s[i + |x|..]
    {
        s[..i] + x + s[i + |x|..]
    }
}

module RawVec {
    import opened Util

    predicate all_unused<T>(content_range: seq<Option<T>>)
    {
        forall x :: x in content_range ==> x == None
    }

    lemma unused_range_lemma<T>(blank: seq<Option<T>>, start: usize, stop: usize)
        requires all_unused(blank) && 0 <= start <= stop <= |blank|
        ensures all_unused(blank[start..stop])
    {
        forall x | x in blank[start..stop] ensures x == None {
            assert x in blank;
        }
    }

    function unused<T>(size: usize): (space: seq<Option<T>>)
        ensures all_unused(space) && |space| == size
    {
        if size == 0 then [] else [None] + unused(size - 1)
    }
    
    predicate all_used<T>(content_range: seq<Option<T>>)
    {
        forall x :: x in content_range ==> x.Some?
    }

    function used<T>(values: seq<T>): (result: seq<Option<T>>)
        ensures |values| == |result|
        ensures all_used(result)
        ensures forall i :: 0 <= i < |values| ==> result[i] == Some(values[i])
    {
        if |values| == 0 then [] else [Some(values[0])] + used(values[1..])
    }

    lemma unused_equal_lemma<T>(s: seq<Option<T>>, t: seq<Option<T>>)
        requires all_unused(s) && all_unused(t) && |s| == |t|
        ensures s == t
    {
        forall i | 0 <= i < |s|
            ensures s[i] == t[i]
    {
        assert s[i] == None by { assert s[i] in s; }
        assert t[i] == None by { assert t[i] in t; }
    }
    }

    lemma concat_unused_lemma<T>(a: usize, b: usize)
        ensures unused<T>(a + b) == unused<T>(a) + unused<T>(b)
    {
        var c := unused<T>(a + b);
        unused_range_lemma(c, 0, a);
        unused_equal_lemma(c[..a], unused<T>(a));
        unused_range_lemma(c, a, a + b);
        unused_equal_lemma(c[a..a + b], unused<T>(b));
    }

    function pad_end<T>(values: seq<T>, cap: usize): (result: seq<Option<T>>)
        requires |values| <= cap
        ensures |result| == cap && result[..|values|] == used(values) && all_unused(result[|values|..])
    {
        used(values) + unused(cap - |values|)
    }

    function rotate_toward_end<T>(values: seq<T>, d: usize): seq<T>
        requires 0 <= d < |values|
    {
        var x := |values| - d;
        values[x..] + values[..x]
    }

    const MIN_CAPACITY := 4;

    // A heap-allocated array of MaybeUninit<T> values.
    class RawVec<T> {
        ghost var contents: seq<Option<T>>;
        var ptr: Ptr;
        var cap: usize;

        predicate valid()
            reads this
        {
            |contents| == cap &&
                (cap == 0 <==> is_null(ptr))
        }

        predicate has_value_at(i: usize)
            reads this
            requires valid()
            requires 0 <= i < cap
        {
            contents[i].Some?
        }

        function method get(i: usize): (value: T)
            requires valid()
            requires 0 <= i < cap && has_value_at(i)
            ensures contents[i] == Some(value)
            reads this

        /// Writes an element into an uninitialized slot in the buffer.
        method write(i: usize, value: T)
            requires valid()
            requires 0 <= i < cap && !has_value_at(i)
            modifies this
            ensures valid()
            ensures contents == old(contents)[i := Some(value)]

        /// Moves an element out of the buffer, leaving its slot uninitialized.
        method read(i: usize) returns (value: T)
            requires valid()
            requires 0 <= i < cap && has_value_at(i)
            modifies this
            ensures valid()
            ensures contents == old(contents)[i := None]
            ensures value == old(contents[i].value)

        constructor()
            ensures valid() && contents == []

        constructor with_capacity(cap: usize)
            ensures valid() && this.cap == cap && all_unused(contents)

        function method pointer(): Ptr reads this { ptr }
        function method capacity(): usize reads this { cap }

        method double()
            requires valid()
            ensures valid()
            ensures MIN_CAPACITY <= cap && 2 * old(cap) <= cap
            ensures contents[..old(cap)] == old(contents)
            ensures all_unused(contents[old(cap)..])

        method reserve(used_cap: usize, needed_extra_cap: usize)
            requires valid()
            ensures valid()
            requires 0 <= used_cap <= cap
            requires all_unused(contents[used_cap..])  // stricter than Rust, but a good rule
            ensures cap >= used_cap + needed_extra_cap
            ensures cap >= old(cap)
            ensures contents[..old(cap)] == old(contents)
            ensures all_unused(contents[used_cap..])

        method reserve_exact(used_cap: usize, needed_extra_cap: usize)
            requires valid()
            requires 0 <= used_cap <= cap
            requires all_unused(contents[used_cap..])
            ensures valid()
            ensures cap >= used_cap + needed_extra_cap  // same as Rust, the difference is informal
            ensures cap >= old(cap)
            ensures contents[..old(cap)] == old(contents)
            ensures all_unused(contents[used_cap..])

        method shrink_to_fit(amount: usize)
            requires valid()
            requires amount <= cap
            requires all_unused(contents[amount..])
            ensures valid() && amount == 0 ==> cap == 0

        method swap(i: usize, j: usize)
            requires valid() && 0 <= i < cap && 0 <= j < cap
            ensures valid()
            ensures contents == contents[i := old(contents[j])][j := old(contents[i])]
            
        method move(dst_start: usize, src_start: usize, len: usize)
            requires valid()
            requires 0 <= src_start <= src_start + len <= cap
            requires 0 <= dst_start <= dst_start + len <= cap
            ensures valid() && contents == replace_slice(replace_slice(old(contents), src_start, unused<T>(len)), dst_start, old(contents)[src_start..src_start + len])
            //ensures forall i :: 0 <= i < len ==> contents[dst_start + i] == old(contents[src_start + i])

        method move_nonoverlapping(dst_start: usize, src_start: usize, length: usize)
            requires valid()
            ensures valid()
            requires 0 <= dst_start <= dst_start + length <= cap
            requires 0 <= src_start <= src_start + length <= cap
            requires all_used(contents[src_start..src_start + length])
            requires all_unused(contents[dst_start..dst_start + length])
            modifies this

        method fill_from(dst_start: usize, src: RawVec<T>, src_start: usize, length: usize)
            requires valid() && src.valid()
            ensures valid() && src.valid()
            requires src != this
            requires 0 <= dst_start <= dst_start + length <= cap
            requires 0 <= src_start <= src_start + length <= src.cap
            requires all_used(src.contents[src_start..src_start + length])
            requires all_unused(contents[dst_start..dst_start + length])
            modifies this, src
            ensures src.contents ==
                old(src.contents)[..src_start] +
                old(contents)[dst_start..dst_start + length] +
                old(src.contents)[src_start + length..]
            ensures contents ==
                old(contents)[..dst_start] +
                old(src.contents)[src_start..src_start + length] +
                old(contents)[dst_start + length..]
        {
            // Maybe inefficient implementation, to make sure I got the spec correct.
            ghost var dst := this;

            ghost var src_front := src.contents[..src_start];
            ghost var src_back := src.contents[src_start + length..];
            ghost var dst_front := dst.contents[..dst_start];
            ghost var dst_back := dst.contents[dst_start + length..];

            ghost var elements := src.contents[src_start..src_start + length];
            ghost var empty := contents[dst_start..dst_start + length];
            assert all_unused(empty);
            forall i | 0 <= i < |empty| ensures empty[i] == None {
                assert empty[i] in empty;
            }

            var i := 0;
            while i < length
                invariant 0 <= i <= length
                invariant valid() && src.valid()
                invariant src.contents == src_front + empty[..i] + elements[i..] + src_back
                invariant dst.contents == dst_front + elements[..i] + empty[i..] + dst_back
            {
                var value := src.read(src_start + i);
                assert src.contents[src_start + i] == None;
                write(dst_start + i, value);
                i := i + 1;
            }

            assert src.contents == src_front + empty + src_back;
            assert dst.contents == dst_front + elements + dst_back;
        }
    }
}