From a7a71524ebebe52e3d6994e41906feff326584e0 Mon Sep 17 00:00:00 2001
From: Jason Gross <jgross@mit.edu>
Date: Fri, 7 Mar 2025 12:08:56 -0800
Subject: [PATCH] Add ZRange.bitwidth, ZRange.of_bitwidth

---
 src/AbstractInterpretation/ZRange.v | 34 +++++++++++++++++++++++++++++
 src/Util/ZRange/Operations.v        | 12 ++++++++++
 2 files changed, 46 insertions(+)

diff --git a/src/AbstractInterpretation/ZRange.v b/src/AbstractInterpretation/ZRange.v
index 76138633d7..5aa2817e8a 100644
--- a/src/AbstractInterpretation/ZRange.v
+++ b/src/AbstractInterpretation/ZRange.v
@@ -53,6 +53,28 @@ Module Compilers.
              | base.type.list A => list (interp A)
              | base.type.option A => option (interp A)
              end%type.
+        Fixpoint bitwidth {t} : interp t -> binterp t
+          := match t with
+            | base.type.type_base base.type.Z => fun r => ZRange.bitwidth r
+            | base.type.type_base _ as t
+            | base.type.unit as t
+              => fun x => x
+            | base.type.prod A B => fun '(a, b) => (@bitwidth A a, @bitwidth B b)
+            | base.type.list A => List.map (@bitwidth A)
+            | base.type.option A => option_map (@bitwidth A)
+            end.
+        Fixpoint of_bitwidth (signed : bool) {t} : binterp t -> interp t
+          := match t with
+             | base.type.type_base base.type.Z => fun r => ZRange.of_bitwidth signed r
+             | base.type.type_base _ as t
+             | base.type.unit as t
+               => fun x => x
+             | base.type.prod A B => fun '(a, b) => (@of_bitwidth signed A a, @of_bitwidth signed B b)
+             | base.type.list A => List.map (@of_bitwidth signed A)
+             | base.type.option A => option_map (@of_bitwidth signed A)
+             end%type.
+        Notation of_bitwidth_unsigned := (@of_bitwidth false).
+        Notation of_bitwidth_signed := (@of_bitwidth true).
         Fixpoint map_ranges (f : zrange -> zrange) {t} : interp t -> interp t
           := match t with
              | base.type.type_base base.type.Z => f
@@ -406,6 +428,18 @@ Module Compilers.
            | type.base x => @base.is_bounded_by x
            | type.arrow s d => fun _ _ => false
            end.
+      Fixpoint bitwidth (signed : bool) {t} : interp t -> einterp t
+        := match t with
+           | type.base x => @base.bitwidth x
+           | type.arrow s d => fun fr sz => @bitwidth signed d (fr (@of_bitwidth signed s sz))
+           end
+      with of_bitwidth (signed : bool) {t} : einterp t -> interp t
+        := match t with
+           | type.base x => @base.of_bitwidth signed x
+           | type.arrow s d => fun fr sz => @of_bitwidth signed d (fr (@bitwidth signed s sz))
+           end.
+      Notation of_bitwidth_unsigned := (@of_bitwidth false).
+      Notation of_bitwidth_signed := (@of_bitwidth true).
       Module option.
         (** turn a [type] into a [Set] describing the type of optional
             bounds on that base type; bounds on a [Z] may be either a
diff --git a/src/Util/ZRange/Operations.v b/src/Util/ZRange/Operations.v
index b1146fdb05..2dca6c4239 100644
--- a/src/Util/ZRange/Operations.v
+++ b/src/Util/ZRange/Operations.v
@@ -11,6 +11,18 @@ Module ZRange.
 
   Local Notation eta v := r[ lower v ~> upper v ].
 
+  Definition bitwidth (v : zrange) : Z
+    := let bitwidth_of z := (if (z =? 0) then 0 else if (z <? 0) then Z.log2_up (-z) else Z.log2_up (z+1))%Z in
+       Z.max (bitwidth_of (lower v)) (bitwidth_of (upper v)).
+
+  Definition of_bitwidth (signed : bool) (bitwidth : Z) : zrange
+    := if signed
+       then r[ -2 ^ (bitwidth - 1) ~> 2 ^ (bitwidth - 1) - 1 ]
+       else r[ 0 ~> 2 ^ bitwidth - 1 ].
+
+  Notation of_bitwidth_unsigned := (of_bitwidth false).
+  Notation of_bitwidth_signed := (of_bitwidth true).
+
   Definition flip (v : zrange) : zrange
     := r[ upper v ~> lower v ].