Make pre-condition available as hypothesis when rewriting inside a function

[andricicezar]

Hi! I'm trying to prove the following equality:

assume val p : int -> bool

let f (x:int) : Pure int (p x) (fun _ -> True) = if p x then 0 else 1
let g (x:int) : Pure int (p x) (fun _ -> True) = 0 

let _ = assert (f == g)

The condition of if is always true because of the pre-condition, thus the equality is true. I tried the following:

let lemma (x:int) : Lemma (requires (p x)) (ensures (p x == true)) = ()

let _ = assert (f == g) by (norm [delta_only [`%f;`%g]]; l_to_r [`lemma]; norm [iota]; smt () dump "h")

but now F* asks me to prove p x, which should be true because of the pre-condition, but I have no hypothesis in the environment.

I first asked this question on Slack where Nik answered with:

I don't think you can do this with the Pure effect.
You should be able to do this if you were working with x:int{ p x } -> int.

module C
open FStar.Tactics

assume val p : ℤ → bool

let f (x:ℤ { p x }) : ℤ = if p x then 0 else 1
let g (x:ℤ { p x }) : ℤ = 0

let lemma (x:ℤ) (a b:α) : Lemma (requires (p x)) (ensures ((if p x then a else b) == a)) = ()

let _ = assert (f == g) by (
  norm [delta_only [`%f; `%g]];
  pointwise (λ _ →
    (λ _ → apply_lemma_rw (`lemma); smt()) `or_else` trefl)
)

Note, you have to be careful to not try to apply the lemma to rewrite all occurrences of p x, since not all of them are provably equal to true. E.g., if you try to rewrite the p x in x:int { p x } then you'll be stuck. So, I changed the lemma to only apply to rewriting the p x in the body of f, the one that guards the if

Making this kind of thing work for the Pure effect may be feasible and even desirable, but it's not currently supported

How hard it will be to have this? Or can you suggest any workarounds? It would be helpful for me to have this because I'm trying to prove about my import/export mechanism that import and export satisfy a cancelation law: where export takes a function with pre-post and wraps it in a new function with no pre-post in which the pre-condition is enforced dynamically, and where import takes a function with no pre-post and wraps it in a function with pre-post in which the post-condition is enforced dynamically. The proof should be quite easy if I would be able to do this rewrites, since the dynamic checks would be trivial because they enforce the pre-post which should be around. If interested, you can find a Proof of Concept here: https://github.com/andricicezar/fstar-io/blob/c7ac96726323dd4698584d8b1fd8b012f819c9eb/io/LawImportableExportable.fst

Thank you!