bplustree.rav

include "give-up.rav"
include "../../arrays/ordered_array.rav"

interface BPlusTree[O: Library.OrderedType] : NodeImpl {

  module Spec : SearchStructureSpec {
    type K = O
    val keyspace: Set[K] = {| k: K :: true |}
  }

  import Spec._

  module KOption = Library.Option[O]

  import KOption.some
  import KOption.none

  import Flow_K._
  import Multiset_K.elem

  interface AK = OrderedArray[O]

  module RefType : Library.Type {
    rep type T = Ref
  }

  interface AN = Array[RefType]

  // Width parameter of the B-tree
  val b: Int

  auto axiom bValid() 
    ensures b > 0

  // Fields of a B-tree node
  field len: Int
  field rangeLb: KOption
  field rangeUb: KOption
  field keys: AK.T
  field ptrs: AN.T

  func le(ko: KOption, k: K) returns (res : Bool)
  {
    ko == none ? true : O.le(ko.KOption.value, k)
  }

  func lt(k: K, ko: KOption) returns (res : Bool)
  {
    ko == none ? true : O.lt(k, ko.KOption.value)
  }

  func flow_int_of(n: Ref, l: Int, lb: KOption, ub: KOption, kmap: Map[Int, K], cmap: Map[Int, Ref]) returns (flow_int: Flow_K)
  {
    Flow_K.int({| np: Ref :: np == n ? {| k: K :: le(lb, k) && lt(k, ub) ? 1 : 0 |} : Multiset_K.id |},
               {| np: Ref :: 0 <= AN.inverse(cmap, 0, l+1, np) <= l ? 
                 {| k: K :: cmap[0] != null && (0 < AN.inverse(cmap, 0, l+1, np) ==> O.le(kmap[AN.inverse(cmap, 0, l+1, np) - 1], k)) &&
                            (AN.inverse(cmap, 0, l+1, np) < l ==> O.lt(k, kmap[AN.inverse(cmap, 0, l+1, np)])) ? 1 : 0 |} : Multiset_K.id |},
               {| n |})
  }

  func node_cond(n: Ref, c: Set[K], flow_int: Flow_K, l: Int, lb: KOption, ub: KOption, ks: AK.T, chlds: AN.T, kmap: Map[Int, K], cmap: Map[Int, Ref]) returns (res: Bool) {
       0 <= l && l < 2*b 
    && AK.length(ks) == 2*b 
    && AN.length(chlds) == 2*b
    // Definition of flow interface
    && flow_int == flow_int_of(n, l, lb, ub, kmap, cmap)
    // Definition of contents
    && c == (cmap[0] == null ? /* Leaf */ AK.set_of_map(kmap, 0, l) : {||})
    // Keys are sorted
    && (forall i: Int, j: Int :: {kmap[i], kmap[j]} 0 <= i < j < l ==> O.lt(kmap[i], kmap[j]))
    // Keys are within range
    && (l > 0 ==> le(lb, kmap[0]) && lt(kmap[l-1], ub))
    // Consistency of cmap
    && (forall i: Int :: {cmap[i]} 0 <= i <= l ==> n != cmap[i])
    && (cmap[0] != null ==> AN.injective(cmap, 0, l+1))
    && (cmap[0] == null ==> (forall i: Int :: {cmap[i]} 0 <= i < 2*b ==> cmap[i] == null))
    // Internal nodes don't point to null
    && (cmap[0] != null ==> (forall i: Int :: {cmap[i]} 0 <= i <= l ==> cmap[i] != null))    
  }

  pred node(n: Ref; c: Set[K], flow_int: Flow_K) {
    exists l: Int, lb: KOption, ub: KOption, ks: AK.T, chlds: AN.T, kmap: Map[Int, K], cmap: Map[Int, Ref] ::
    node_cond(n, c, flow_int, l, lb, ub, ks, chlds, kmap, cmap)
    // Heap resources
    && own(n.len, l, 1.0)
    && own(n.rangeLb, lb, 1.0)
    && own(n.rangeUb, ub, 1.0)
    && own(n.keys, ks, 1.0)
    && AK.sorted_array_with_content(ks, l, kmap)
    && own(n.ptrs, chlds, 1.0)
    && AN.arr(chlds, cmap)
  }

  lemma nodeSepStar(n: Ref, c1: Set[K], c2: Set[K], i_n1: Flow_K, i_n2: Flow_K)
    requires node(n, c1, i_n1) && node(n, c2, i_n2)
    ensures false
  {
    unfold node(n, c1, i_n1);
    unfold node(n, c2, i_n2);
  }

  // Show that if query key k is in the keyset of node x than x must be a leaf.
  lemma keyset_implies_leaf(n: Ref, k: K, i: Int, flow_int: Flow_K, c: Set[K],
      implicit ghost l: Int, implicit ghost lb: KOption, implicit ghost ub: KOption, 
      implicit ghost ks: AK.T, implicit ghost chlds: AN.T, implicit ghost kmap: Map[Int, K], implicit ghost cmap: Map[Int, Ref])
    requires node_cond(n, c, flow_int, l, lb, ub, ks, chlds, kmap, cmap)
    // Heap resources
    && own(n.len, l, 1.0)
    && own(n.rangeLb, lb, 1.0)
    && own(n.rangeUb, ub, 1.0)
    && own(n.keys, ks, 1.0)
    && AK.sorted_array_with_content(ks, l, kmap)
    && own(n.ptrs, chlds, 1.0)
    && AN.arr(chlds, cmap)
    requires k in keyset(flow_int)
    requires (i <= 0 || O.le(kmap[i-1], k)) && 0 <= i <= l
    ensures node_cond(n, c, flow_int, l, lb, ub, ks, chlds, kmap, cmap)
    // Heap resources
    && own(n.len, l, 1.0)
    && own(n.rangeLb, lb, 1.0)
    && own(n.rangeUb, ub, 1.0)
    && own(n.keys, ks, 1.0)
    && AK.sorted_array_with_content(ks, l, kmap)
    && own(n.ptrs, chlds, 1.0)
    && AN.arr(chlds, cmap)
    ensures cmap[0] == null  // x is a leaf
  {
    assert (flow_int.inf[n])[k] == 1;
    assert (forall y: Ref :: {outset(flow_int, y)} k !in outset(flow_int, y));
    assert (forall y: Ref :: {(flow_int.out[y])[k]} k !in outset(flow_int, y) ==> (flow_int.out[y])[k] == 0);
    if (cmap[0] != null) {
      if (i < l) {
        if (O.lt(k, kmap[i])) {
          assert (flow_int.out[cmap[i]])[k] == 1;
        } else {
          keyset_implies_leaf(n, k, i + 1, flow_int, c);
        }
      } else {
        assert (flow_int.out[cmap[i]])[k] == 1;
      }
    }
  }


  /*Initialize a new root node */
  proc createRoot()
    returns (r: Ref, implicit ghost g_i: Flow_K)
    ensures own(r.lock, 0)
    ensures globalRes(r, {||}, g_i) && g_i.dom == {|r|}
    ensures nodePred(r, r, {||}, g_i) 
  {
    var ka: AK.T;
    ka := AK.alloc(2*b, AK.default);
    var pa: AN.T;
    pa := AN.alloc(2*b, null);
    val flow_id := Flow_K.id;
    val ks_r := Keyset_K.prodKS(keyspace, {||});
    r := new(
      len: 0, 
      rangeLb: none, 
      rangeUb: none, 
      keys: ka, 
      ptrs: pa,
      lock: 0,
      authFlow: AuthFlow_K.auth_frag(flow_id, flow_id),
      authKS: AuthKeyset_K.auth_frag(ks_r, ks_r),
      authSet: AuthSetRef.auth(SetRefRA.set_constr({||}))
    );

    val cmap: Map[Int, Ref] := {| i: Int :: 0 <= i < 2*b ? null : AN.default |};

    val ir := flow_int_of(r, 0, none, none, AK.default_map, cmap);
    fold AK.sorted_array_with_content(ka, 0, AK.default_map);
    fold node(r, {||}, ir);
    fpu(r.authFlow, 
      AuthFlow_K.auth_frag(flow_id, flow_id),
      AuthFlow_K.auth_frag(ir, ir)
    );
    fpu(r.authSet, 
      AuthSetRef.auth(SetRefRA.set_constr({||})),
      AuthSetRef.auth_frag(SetRefRA.set_constr({|r|}), SetRefRA.set_constr({|r|}))
    );
    fold globalRes(r, {||}, ir);
    fold nodePred(r, r, {||}, ir);
    return r, ir;
  }

  proc inRange(n: Ref, k: K, implicit ghost c: Set[K], implicit ghost i_n: Flow_K)
    returns (ret: Bool)
    requires node(n, c, i_n)
    ensures node(n, c, i_n) && (ret ==> k in inset(i_n, n))
  {
    unfold node(n, c, i_n);
    
    var lb: KOption := n.rangeLb;
    var ub: KOption := n.rangeUb;

    fold node(n, c, i_n);
    return le(lb, k) && lt(k, ub);
  }

  proc findNext(n: Ref, k: K, implicit ghost c: Set[K], implicit ghost i_n: Flow_K)
    returns (ret: Bool, n1: Ref)
    requires k in inset(i_n, n)
    requires node(n, c, i_n)
    ensures node(n, c, i_n) && 
      (ret ?   
        k in outset(i_n, n1) :
        k !in outsets(i_n))
  {
    unfold node(n, c, i_n);
    
    var n_len : Int := n.len;
    var n_ptrs : AN.T := n.ptrs;
    ghost var chlds: Map[Int, Ref];
    chlds :| AN.arr(n_ptrs, chlds);
    unfold AN.arr(n_ptrs, chlds);
    var chld0: Ref := AN.loc(n_ptrs, 0).AN.value;
    
    if (chld0 == null) {  // Leaf node
      fold AN.arr(n_ptrs, chlds);
      fold node(n, c, i_n);
      return false, null;
    }

    var n_keys : AK.T := n.keys;

    var found: Bool;
    var i: Int;
    found, i := AK.arr_find(n_keys, n_len, k);

    ghost var kmap: Map[Int, K];
    kmap :| AK.sorted_array_with_content(n_keys, n_len, kmap);
    unfold AK.sorted_array_with_content(n_keys, n_len, kmap);
    unfold AK.arr(n_keys, kmap);
    var n_keys_i: K := AK.loc(n_keys, i).AK.value;
    fold AK.arr(n_keys, kmap);

    if (i < n_len && k == n_keys_i) {  // arr_find finds upper bound, we want lower
      i := i + 1;
    }

    n1 := AN.loc(n_ptrs, i).AN.value;
    ret := true;

    fold AN.arr(n_ptrs, chlds);
    fold AK.sorted_array_with_content(n_keys, n_len, kmap);
    fold node(n, c, i_n);
  }

  proc decisiveOp(dop: Op, n: Ref, k: K, implicit ghost c: Set[K], implicit ghost i_n: Flow_K)
    returns (succ: Bool, res: Bool, implicit ghost c1: Set[K])
    requires k in keyset(i_n)
    requires node(n, c, i_n)
    ensures node(n, c1, i_n)
    ensures succ ==> opSpec(dop, k, c) == ((c1, res)) 
    ensures !succ ==> c == c1
  {
    unfold node(n, c, i_n);
    var n_len : Int := n.len;
    var n_ptrs : AN.T := n.ptrs;
    ghost var cmap: Map[Int, Ref];
    cmap :| AN.arr(n_ptrs, cmap);
      
    c1 := (opSpec(dop, k, c))#0;
    
    var idx: Int;
    var new_len: Int;
    var n_keys : AK.T := n.keys;
    ghost var m1: Map[Int, K];
    ghost var kmap: Map[Int, K];
    kmap :| AK.sorted_array_with_content(n_keys, n_len, kmap);

    keyset_implies_leaf(n, k, 0, i_n, c);

    if (dop == searchOp) {
      res, idx := AK.arr_find(n_keys, n_len, k);
      fold node(n, c, i_n);
      return true, res, c;
    } else if (dop == insertOp) {
      if (n_len < 2*b - 1) {
        idx, new_len, m1 := AK.arr_insert(n_keys, k, n_len);
        AK.map_insert_content_set(kmap, m1, idx, k, n_len, new_len);

        n.len := new_len;
        fold node(n, c1, i_n);
        return true, new_len != n_len, c1;
      } else {
        fold node(n, c, i_n);
        return false, false, c;
      }
    } else {
      new_len, idx, m1 := AK.arr_delete(n_keys, k, n_len);

      AK.map_delete_content_set(kmap, n_len, idx, k);

      n.len := new_len;
      
      ghost var c1: Set[K] := (opSpec(dop, k, c))#0;
      fold node(n, c1, i_n);
      return true, new_len != n_len, c1;
    }
  }

}