Fast Natural Exponentiation in Dafny

FastExp.dfy

method FastExp(a: int, n: nat) returns (z: int)
  ensures z == Exp(a, n)
{
  var b: int, i: nat;
  z, b, i := 1, a, n;
  while (i != 0)
    invariant ExpInv(z, b, i, a, n)
    decreases i
  {
    if i % 2 == 1 {
      z, i := z * b, i - 1;
    } else {
      SquareB(b, i);
      b, i := b * b, i / 2;
    }
  }
}

predicate ExpInv(z: int, b: int, i: nat, a: int, n: nat) {
  z * Exp(b, i) == Exp(a, n)
}

function Exp(a: int, n: nat): int
  decreases n
{
  if n == 0 then 1 else a * Exp(a, n - 1)
}

lemma SquareB(b: int, i: nat)
  requires i % 2 == 0
  requires i != 0
  ensures Exp(b * b, i / 2) == Exp(b, i)
  decreases i
{
  if (i == 2) {
    assert b * Exp(b, 1) == Exp(b, 2);
  } else {
    SquareB(b * b, i - 2);
  }
}