This PR makes sqrt flagged more deterministic by using a trick I learned from Dan Boneh for how to witness non-squareness. The trick is the following. First, fix a non-square m. Then for any x in F, x is a non-square iff m*x is a square. Thus, if you want to show that x is not a square, you can simply show that mx is a non-square, which you can do by witnessing the square root.

The square_root flagged function is semantically

x -> if is_square x then (sqrt x, true) else (undefined_behavior, false)

We implement this with a snarky program as follows (prover blocks omitted):

let sqrt_exn x =
  let y = exists Field.typ in
  assert_square y x;
  y

let sqrt_flagged x =
  let b = exists Boolean.typ in
  ( sqrt_exn (if b then x else m * x)
  , b )

We verify the correctness as follows:

If e is an expression let ⟦ e ⟧ denote the set of possible values e can evaluate to over all choices of non-determinism which do not cause an assertion failure.

Two lemmas claimed without proof:

Lemma: If a is a square then ⟦ sqrt_exn a ⟧ = { sqrt a }

Lemma: If x is not a square then ⟦ sqrt_exn a ⟧ = { }

Now let's show

Proposition: sqrt_flagged implements the intended semantics.

We have

⟦ sqrt_flagged x ⟧
=
  ⟦ (sqrt_exn (if true then x else m * x), true) ⟧
  ∪
  ⟦ (sqrt_exn (if false then x else m * x), false) ⟧
=
  ⟦ (sqrt_exn x, true) ⟧
  ∪
  ⟦ (sqrt_exn (m * x), false) ⟧
=
  ⟦ (sqrt_exn x, true) ⟧
  ∪
  ⟦ (sqrt_exn (m * x), false) ⟧

If x is a square then mx is not a square so we have

⟦ sqrt_flagged x ⟧
= ⟦ (sqrt_exn x, true) ⟧ ∪ ⟦ (sqrt_exn (m * x), false) ⟧
= { (sqrt x, true) } ∪ {}
= { (sqrt x, true) }

On the other hand if x is not a square, then mx is a square so we have

⟦ sqrt_flagged x ⟧
= ⟦ (sqrt_exn x, true) ⟧ ∪ ⟦ (sqrt_exn (m * x), false) ⟧
= {} ∪ { (sqrt (m * x), false) }
= { (sqrt (m * x), false) }

which evidently matches the claimed semantics.