Finding a loop invariant for this power function

Question

I am having a difficult time grasping the concept of proving the correctness of a iterative program/function in my Theory of Computation course. More specifically, I don't know how to come up with a loop invariant. I understand that a function can have multiple loop invariants but it's a complete mystery to me as to how to find one that helps us prove the post condition. I am currently working through some homework and don't know how to find the loop invariant for the following function.

PREcondition: a is a number, b is a natural number.
POSTcondition: return a^b.
    def power(a, b):
        s = a
        k = b
        p = 1
        while k > 0:
            if k % 2 == 1:
                p = p * s
            s = s * s
            k = k // 2
        return p

As of now, I understand how the function works, but as I've stated before, I'm super lost when trying to find an appropriate loop invariant that will help me show that this function returns a^b.

Answer 1

Note that while ^ means "power" in some languages, in Python it is bitwise xor. The power operator is **

Here are a couple that check that the loop is heading toward the correct result

    while k > 0:
        assert s ** k * p == a ** b
        if k % 2 == 1:
            p = p * s
        s = s * s
        k = k // 2
        assert k == 0 or s ** k * p == a ** b

It's also possible to write defensive invariants like

s >= a
n <= b