Try out stronger land, lor bounds

1 files changed, 54 insertions, 37 deletions

diff --git a/src/Util/ZRange/Operations.v b/src/Util/ZRange/Operations.v
index d17785f77..07e330a4a 100644
--- a/src/Util/ZRange/Operations.v
+++ b/src/Util/ZRange/Operations.v
@@ -1,5 +1,7 @@
 Require Import Coq.ZArith.ZArith.
 Require Import Crypto.Util.ZRange.
+Require Import Crypto.Util.ZUtil.Definitions.
+
 Require Import Crypto.Util.Notations.
 
 Module ZRange.
@@ -41,49 +43,64 @@ Module ZRange.
     := let (l, u) := eta v in
        r[ f l ~> f u ].
 
-  Definition two_corners (f : Z -> Z) (v : zrange) : zrange
+  Definition split_range_at_0 (x : zrange) : option zrange (* < 0 *) * option zrange (* >= 0 *)
+    := let (l, u) := eta x in
+       (if (0 <=? l)%Z
+        then None
+        else Some r[l ~> Z.min u (-1)],
+        if (0 <=? u)%Z
+        then Some r[Z.max 0 l ~> u]
+        else None).
+
+  Definition apply_to_split_range (f : zrange -> zrange) (v : zrange) : zrange
+    := match split_range_at_0 v with
+       | (Some n, Some p) => union (f n) (f p)
+       | (Some v, None) | (None, Some v) => f v
+       | (None, None) => f v
+       end.
+
+  Definition apply_to_range (f : BinInt.Z -> zrange) (v : zrange) : zrange
     := let (l, u) := eta v in
-       r[ Z.min (f l) (f u) ~> Z.max (f u) (f l) ].
+       union (f l) (f u).
 
-  Definition two_corners' (f : Z -> Z) (v : zrange) : zrange
-    := normalize' (map f v).
+  Definition apply_to_each_split_range (f : BinInt.Z -> zrange) (v : zrange) : zrange
+    := apply_to_split_range (apply_to_range f) v.
 
-  Lemma two_corners'_eq : two_corners = two_corners'. Proof. reflexivity. Defined.
+  Definition constant (v : Z) : zrange := r[v ~> v].
 
+  Definition two_corners (f : Z -> Z) (v : zrange) : zrange
+    := apply_to_range (fun x => constant (f x)) v.
   Definition four_corners (f : Z -> Z -> Z) (x y : zrange) : zrange
-    := let (lx, ux) := eta x in
-       union (two_corners (f lx) y)
-             (two_corners (f ux) y).
-
+    := apply_to_range (fun x => two_corners (f x) y) x.
   Definition eight_corners (f : Z -> Z -> Z -> Z) (x y z : zrange) : zrange
-    := let (lx, ux) := eta x in
-       union (four_corners (f lx) y z)
-             (four_corners (f ux) y z).
-
-  Definition upper_lor_land_bounds (x y : BinInt.Z) : BinInt.Z
-    := 2^(1 + Z.log2_up (Z.max x y)).
-  Definition extreme_lor_land_bounds (x y : zrange) : zrange
-    := let mx := ZRange.upper (ZRange.abs x) in
-       let my := ZRange.upper (ZRange.abs y) in
-       {| lower := -upper_lor_land_bounds mx my ; upper := upper_lor_land_bounds mx my |}.
-  Definition extremization_bounds (f : zrange -> zrange -> zrange) (x y : zrange) : zrange
-    := let (lx, ux) := x in
-       let (ly, uy) := y in
-       if ((lx <? 0) || (ly <? 0))%Z%bool
-       then extreme_lor_land_bounds x y
-       else f x y.
-  Definition land_bounds : zrange -> zrange -> zrange
-    := extremization_bounds
-         (fun x y
-          => let (lx, ux) := x in
-             let (ly, uy) := y in
-             {| lower := Z.min 0 (Z.min lx ly) ; upper := Z.max 0 (Z.min ux uy) |}).
-  Definition lor_bounds : zrange -> zrange -> zrange
-    := extremization_bounds
-         (fun x y
-          => let (lx, ux) := x in
-             let (ly, uy) := y in
-             {| lower := Z.max lx ly ; upper := upper_lor_land_bounds ux uy |}).
+    := apply_to_range (fun x => four_corners (f x) y z) x.
+
+  Definition two_corners_and_zero (f : Z -> Z) (v : zrange) : zrange
+    := apply_to_each_split_range (fun x => constant (f x)) v.
+  Definition four_corners_and_zero (f : Z -> Z -> Z) (x y : zrange) : zrange
+    := apply_to_split_range (apply_to_range (fun x => two_corners_and_zero (f x) y)) x.
+  Definition eight_corners_and_zero (f : Z -> Z -> Z -> Z) (x y z : zrange) : zrange
+    := apply_to_split_range (apply_to_range (fun x => four_corners_and_zero (f x) y z)) x.
+
+  Definition two_corners' (f : Z -> Z) (v : zrange) : zrange
+    := normalize' (map f v).
+
+  Lemma two_corners'_eq x y : two_corners x y = two_corners' x y.
+  Proof.
+    cbv [two_corners two_corners' normalize' map union apply_to_range constant flip].
+    cbn [lower upper].
+    rewrite Z.max_comm; reflexivity.
+  Qed.
+
+  (** if positive, round up to 2^k-1 (0b11111....); if negative, round down to -2^k (0b...111000000...) *)
+  Definition round_lor_land_bound (x : BinInt.Z) : BinInt.Z
+    := if (0 <=? x)%Z
+       then 2^(Z.log2_up (x+1))-1
+       else -2^(Z.log2_up (-x)).
+  Definition land_lor_bounds (f : BinInt.Z -> BinInt.Z -> BinInt.Z) (x y : zrange) : zrange
+    := four_corners_and_zero (fun x y => f (round_lor_land_bound x) (round_lor_land_bound y)) x y.
+  Definition land_bounds : zrange -> zrange -> zrange := land_lor_bounds Z.land.
+  Definition lor_bounds : zrange -> zrange -> zrange := land_lor_bounds Z.lor.
 
   Definition split_bounds (r : zrange) (split_at : BinInt.Z) : zrange * zrange :=
     if upper r <? split_at