Possible bug in the if_then_else combinator of indexed effects

[bstut]

Recently, I was implementing layered effects and repeatedly ran into issues with branching. The simplest case where I was able to reproduce this is the following, which is more or less a simplified implementation of the PURE effect:

module PureType0

let custom_pure_pre = Type0
let custom_pure_post (a:Type) = a -> custom_pure_pre
let custom_pure_wp (a:Type) = custom_pure_post a -> custom_pure_pre

let custom_pure (a:Type) (w:custom_pure_wp a) = v:a{forall (p:custom_pure_post a). w p ==> p v}

let retcustom_pure a (x:a) : custom_pure a (fun p -> p x) = x

let bindcustom_pure (a b:Type) (w1:custom_pure_wp a) (w2: a -> custom_pure_wp b) 
      (f: custom_pure a w1) (g: (x:a -> custom_pure b (w2 x)))
      : custom_pure b (fun p -> w1 (fun x -> w2 x p)) = g f

let subcomp a (w1 w2: custom_pure_wp a) (_: squash(forall p. w2 p ==> w1 p)) (f: custom_pure a w1) : custom_pure a w2 = f

let stronger a (w1 w2: custom_pure_wp a) = (forall p. w2 p ==> w1 p)

let ite a (w1 w2: custom_pure_wp a) (f: custom_pure a w1) (g: custom_pure a w2) (b:bool) = custom_pure a (if b then w1 else w2)

layered_effect {
  CUSTOM_PURE: (a:Type) -> (w:custom_pure_wp a) -> Effect
  with 
    repr = custom_pure;
    return = retcustom_pure;
    bind = bindcustom_pure;
    subcomp = subcomp;
    if_then_else = ite
}

let lift_pure_custom_pure a (wp) (f: unit -> PURE a wp) : Pure (custom_pure a wp) (requires wp (fun _ -> True)) (ensures (fun _ -> True)) = 
  FStar.Monotonic.Pure.elim_pure_wp_monotonicity wp; (f ())

sub_effect PURE ~> CUSTOM_PURE = lift_pure_custom_pure

let two () : CUSTOM_PURE int (fun p -> p 2) = 2

let conditional0 (b:bool) : CUSTOM_PURE int (fun p -> p 2) = if b then 2 else 2

let conditional1 (b:bool) : CUSTOM_PURE int (fun p -> p 2) = if b then two () else 2

[@@ expect_failure]
let conditional2 (b:bool) : CUSTOM_PURE int (fun p -> p 2) by (compute();dump"") = if b then 2 else two ()

[@@ expect_failure]
let conditional3 (b:bool) : CUSTOM_PURE int (fun p -> p 2) = if b then two () else two ()

Essentially, the issue is that the functions conditional2 and conditional3 should work but don't. @mtzguido suggested in the Zulip chat that this behavior might be related to #1948

So far, I have identified two ways to circumvent this issue.

1. Rewriting ite

The first one is simply rewriting the function ite. Concretely, the above issue does not prop up if it is rewritten into

let ite a (w1 w2: custom_pure_wp a) (f: custom_pure a w1) (g: custom_pure a w2) (b:bool) = custom_pure a (fun p -> (b ==> w1 p) /\ (~b ==> w2 p))

or

let ite a (w1 w2: custom_pure_wp a) (f: custom_pure a w1) (g: custom_pure a w2) (b:bool) = custom_pure a (fun post -> if b then w1 post else w2 post)

However, this does not work in all cases. I provide one where it fails below.

2. Rewriting the wp transformer for a concrete function

I also noticed that when I change the wp of two () to (fun p -> p 2 /\ p 2) the issue also vanishes. However, this is also only possible for wps with the appropriate shape. This is also not viable in the example below.

3. Circumventing if_then_else

The last way around this I identified is avoiding going through the if_then_else combinator all together. For example, the following works with the original ite:

let alt_ite #a (b:bool) (x y: a) = 
  if b then x else y

let conditional2_alt (b:bool) : CUSTOM_PURE int (fun p -> p 2) = 
  let two = two() in
  alt_ite b 2 two

let conditional3_alt (b:bool) : CUSTOM_PURE int (fun p -> p 2) = 
  let if_branch = two() in
  let else_branch = two() in
  alt_ite b if_branch else_branch

This strongly suggests to me that the issue lies with the if_then_else combinator specifically, since it seems like F Star is able to handle branching in principle.

---

Below is an example where circumventing if_then_else was the only viable route I discovered. It is hidden behind a collapsed section since the code is largely just a repeat of the code above.

PURE with real-valued preconditions

Consider a Hoare-triple {pre} e {post}. Usually, we think of pre and post as logical predicates with pre ==> post if the triple is valid. However, we could also think of pre and post as real-valued functions with pre <= post if the triple is valid. (The usual Hoare-triples are just a special case of this where the functions map into {0,1}.) Below is a sketch of a variant of the PURE effect utilizing this approach. Notably, this one already fails at conditional1:

module PureReals

open FStar.Real

let values_pre = r:real{0.0R <=. r}
let values_post (a:Type) = a -> values_pre
let values_wp (a:Type) = values_post a -> values_pre

let values (a:Type) (w:values_wp a) = v:a{forall (p:values_post a). w p <=. p v}

let values_ret a (x:a) : values a (fun p -> p x) = x

let values_bind (a b:Type) (w1:values_wp a) (w2: a -> values_wp b) 
      (f: values a w1) (g: (x:a -> values b (w2 x)))
      : values b (fun p -> w1 (fun x -> w2 x p)) = g f

let subcomp a (w1 w2: values_wp a) (_: squash(forall p. w2 p <=. w1 p)) (f: values a w1) : values a w2 = f

let stronger a (w1 w2: values_wp a) = (forall p. w2 p <=. w1 p)

let ite a (w1 w2: values_wp a) (f: values a w1) (g: values a w2) (b:bool) = values a (fun post -> if b then w1 post else w2 post)

layered_effect {
  VALUES: (a:Type) -> (w:values_wp a) -> Effect
  with 
    repr = values;
    return = values_ret;
    bind = values_bind;
    subcomp = subcomp;
    if_then_else = ite
}

assume val largest (pred: values_pre -> Type0) : sup:values_pre{pred sup /\ forall (other: values_pre). pred other ==> other <=. sup}

let lift_values_pure_wp #a (wp: pure_wp a) : values_wp a = (fun (post:values_post a) -> largest (fun (p:values_pre) -> (wp (fun x -> p <=. post x))))

let lift_values_pure a (wp) (f: unit -> PURE a wp) : Pure (values a (lift_values_pure_wp wp)) (requires wp (fun _ -> True)) (ensures (fun _ -> True)) = 
  FStar.Monotonic.Pure.elim_pure_wp_monotonicity wp; (f ())

sub_effect PURE ~> VALUES = lift_values_pure

let two () : VALUES int (fun p -> p 2) = 2

let conditional0 (b:bool) : VALUES int (fun p -> p 2) = if b then 2 else 2

[@@expect_failure]
let conditional1 (b:bool) : VALUES int (fun p -> p 2)  = if b then two () else 2

[@@ expect_failure]
let conditional2 (b:bool) : VALUES int (fun p -> p 2) = if b then 2 else two ()

[@@ expect_failure]
let conditional3 (b:bool) : VALUES int (fun p -> p 2) = if b then two () else two ()

let alt_ite #a (b:bool) (x y: a) = 
  if b then x else y

let conditional2_alt (b:bool) : VALUES int (fun p -> p 2) = 
  let two = two() in
  alt_ite b 2 two

let conditional3_alt (b:bool) : VALUES int (fun p -> p 2) = 
  let if_branch = two() in
  let else_branch = two() in
  alt_ite b if_branch else_branch

We can see that only circumventing if_then_else is viable here. To illustrate that the first approach above does not work, ite is already written in the alternative fashion. Rewriting the wp transformer as (fun p -> p 2 /\ p 2) obviously isn't viable since the transformer operates on real values, not Type0.