examples/fenwick.mlw

   1
   2 (* Fenwick trees (or binary indexed tree) for prefix/interval sums.
   3    Represent an integer array over interval [0;n[ such that the following
   4    operations are both efficient:
   5    . incrementation of individual cell (O(log(n)))
   6    . Query sum of elements over interval [a;b[ (O(log(b-a)))
   7
   8    Author: Martin Clochard (Université Paris Sud) *)
   9
  10 (* Array-based implementation with i->(2i+1,2i+2) node encoding:
  11    . Integer represent nodes
  12    . 0 is the root
  13    . childs of node n are 2n+1 and 2n+2
  14    . Leaves represent the model array cells. For n > 0 elements in the model,
  15      they are represented by the cells over the range [n-1;2n-1[
  16    The structure manage queries by keeping for each node the sum of the
  17    values of all descendant leaves, which we call here 'node summary' *)
  18 module Fenwick
  19
  20   use int.Int
  21   use int.ComputerDivision
  22   use ref.Ref
  23   use int.Sum
  24   use array.Array
  25
  26   (* Encode fenwick trees within an array. The leaves field represent
  27      the actual number of element within the model. *)
  28   type fenwick = {
  29     t : array int;
  30     ghost leaves : int;
  31   }
  32
  33   (* Invariant *)
  34   predicate valid (f:fenwick) =
  35     f.leaves >= 0 /\
  36     f.t.length = (if f.leaves = 0 then 0 else 2 * f.leaves - 1) /\
  37     forall i. 0 <= i /\ i < f.leaves - 1 ->
  38       f.t[i] = f.t[2*i+1] + f.t[2*i+2]
  39
  40   (* Get the i-th elements of the model. *)
  41   function get (f:fenwick) (i:int) : int = f.t[i+f.leaves-1]
  42   (* Get the sum of elements over range [a;b[ *)
  43   function rget (f:fenwick) (a b:int) : int = sum (get f) a b
  44
  45   (* Create a Fenwick tree initialized at 0 *)
  46   let make (lv:int) : fenwick
  47     requires { lv >= 0 }
  48     ensures { valid result }
  49     ensures { forall i. 0 <= i < lv -> get result i = 0 }
  50     ensures { result.leaves = lv }
  51   = { t = if lv = 0 then make 0 0 else make (2*lv-1) 0;
  52       leaves = lv }
  53
  54   (* Add x to the l-th cell *)
  55   let add (f:fenwick) (l:int) (x:int) : unit
  56     requires { 0 <= l < f.leaves /\ valid f }
  57     ensures { valid f }
  58     ensures { forall i. 0 <= i < f.leaves /\ i <> l ->
  59       get f i = get (old f) i }
  60     ensures { get f l = get (old f) l + x }
  61   = let lc = ref (div f.t.length 2 + l) in
  62     f.t[!lc] <- f.t[!lc] + x;
  63     (* Update node summaries for all elements on the path
  64        from the updated leaf to the root. *)
  65     label I in
  66     while !lc <> 0 do
  67       invariant { 0 <= !lc < f.t.length }
  68       invariant { forall i. 0 <= i /\ i < f.leaves - 1 ->
  69         f.t[i] = f.t[2*i+1] + f.t[2*i+2] -
  70           if 2*i+1 <= !lc <= 2*i+2 then x else 0 }
  71       invariant { forall i. f.leaves - 1 <= i < f.t.length ->
  72         f.t[i] = (f at I).t[i] }
  73       variant { !lc }
  74       lc := div (!lc - 1) 2;
  75       f.t[!lc] <- f.t[!lc] + x
  76     done
  77
  78   (* Lemma to shift dum indices. *)
  79   let rec ghost sum_dec (a b c:int) : unit
  80     requires { a <= b }
  81     ensures { forall f g. (forall i. a <= i < b -> f i = g (i+c)) ->
  82       sum f a b = sum g (a+c) (b+c) }
  83     variant { b - a }
  84   = if a < b then sum_dec (a+1) b c
  85
  86   (* Crucial lemma for the query routine: Summing the node summaries
  87      over range [2a+1;2b+1[ is equivalent to summing node summaries
  88      over range [a;b[. This is because the elements of range [2a+1;2b+1[
  89      are exactly the childs of elements of range [a;b[. *)
  90   let rec ghost fen_compact (f:fenwick) (a b:int) : unit
  91     requires { 0 <= a <= b /\ 2 * b < f.t.length /\ valid f }
  92     ensures { sum (([]) f.t) a b  = sum (([]) f.t) (2*a+1) (2*b+1) }
  93     variant { b - a }
  94   = if a < b then fen_compact f (a+1) b
  95
  96   (* Query sum of elements over interval [a,b[. *)
  97   let query (f:fenwick) (a b:int) : int
  98     requires { 0 <= a <= b <= f.leaves /\ valid f }
  99     ensures { result = rget f a b }
 100   = let lv = div f.t.length 2 in
 101     let ra = ref (a + lv) in let rb = ref (b + lv) in
 102     let acc = ref 0 in
 103     ghost sum_dec a b lv;
 104     (* If ra = rb, the sum is 0.
 105        Otherwise, adjust the range to odd boundaries in constant time
 106        and use compaction lemma to halve interval size. *)
 107     while !ra <> !rb do
 108       invariant { 0 <= !ra <= !rb <= f.t.length }
 109       invariant { !acc + sum (([]) f.t) !ra !rb = rget f a b }
 110       variant { !rb - !ra }
 111       if mod !ra 2 = 0 then acc := !acc + f.t[!ra];
 112       ra := div !ra 2;
 113       rb := !rb - 1;
 114       if mod !rb 2 <> 0 then acc := !acc + f.t[!rb];
 115       rb := div !rb 2;
 116       ghost fen_compact f !ra !rb
 117     done;
 118     !acc
 119
 120 end