Matrices with bounded-size integers

module Matrix63

  use int.Int
  use mach.int.Int63
  use map.Map as M

  type matrix 'a = private {
    ghost mutable elts: M.map int (M.map int 'a);
                  rows: int63;
               columns: int63
  } invariant { 0 <= to_int rows /\ 0 <= to_int columns }

  val ghost function get_unsafe (a: matrix 'a) (r c: int) : 'a
    ensures { result = M.get (M.get a.elts r) c }

  val ghost function set_unsafe (a: matrix 'a) (r c: int) (v: 'a) : matrix 'a
    ensures {
      result.rows = a.rows /\ result.columns = a.columns /\
      result.elts = M.set a.elts r (M.set (M.get a.elts r) c v) }

  predicate valid_index (a: matrix 'a) (r c: int63) =
    0 <= to_int r < to_int a.rows /\ 0 <= to_int c < to_int a.columns

  val get (a: matrix 'a) (r c: int63) : 'a
    requires { [@expl:index in array bounds] valid_index a r c }
    ensures  { result = get_unsafe a (to_int r) (to_int c) }

  val set (a: matrix 'a) (r c: int63) (v: 'a) : unit
    requires { [@expl:index in array bounds] valid_index a r c }
    writes   { a }
    ensures  { a.elts = M.set (old a.elts) (to_int r)
               (M.set (M.get (old a.elts) (to_int r)) (to_int c) v)}

  exception OutOfBounds

  let defensive_get (a: matrix 'a) (r c: int63) : 'a
    ensures { [@expl:index in array bounds] valid_index a r c }
    ensures { result = get_unsafe a (to_int r) (to_int c) }
    raises  { OutOfBounds -> not (valid_index a r c) }
  = if r < of_int 0 || r >= a.rows || c < of_int 0 || c >= a.columns then
      raise OutOfBounds;
    get a r c

  let defensive_set (a: matrix 'a) (r c: int63) (v: 'a) : unit
    ensures { [@expl:index in array bounds] valid_index a r c }
    ensures { a.elts = M.set (old a.elts) (to_int r)
              (M.set (M.get (old a.elts) (to_int r)) (to_int c) v)}
    raises  { OutOfBounds -> not (valid_index a r c) /\ a = old a }
  = if r < of_int 0 || r >= a.rows || c < of_int 0 || c >= a.columns then
      raise OutOfBounds;
    set a r c v

  val make (r c: int63) (v: 'a) : matrix 'a
    requires { to_int r >= 0 /\ to_int c >= 0 }
    ensures  { result.rows = r }
    ensures  { result.columns = c }
    ensures  { forall i j. 0 <= i < to_int r /\ 0 <= j < to_int c ->
               get_unsafe result i j = v }

  val copy (a: matrix 'a) : matrix 'a
    ensures  { result.rows = a.rows /\ result.columns = a.columns }
    ensures  { forall r:int. 0 <= r < to_int result.rows ->
               forall c:int. 0 <= c < to_int result.columns ->
               get_unsafe result r c = get_unsafe a r c }

end