use std::collections::HashMap;
use crate::ast::{BinOp, Expr, FnBody, FnDef, Literal, Spanned, Stmt};
use crate::codegen::proof_lower::ProofLowerInputs;
use crate::ir::TypeId;
use crate::ir::proof_ir::{Predicate, RefinedTypeDecl};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Bound {
NegInf,
Finite(i128),
PosInf,
}
impl Bound {
fn fits_i64(self) -> bool {
match self {
Bound::Finite(k) => k >= i64::MIN as i128 && k <= i64::MAX as i128,
Bound::NegInf | Bound::PosInf => false,
}
}
fn add(self, other: Bound) -> Bound {
match (self, other) {
(Bound::PosInf, Bound::NegInf) | (Bound::NegInf, Bound::PosInf) => Bound::NegInf,
(Bound::PosInf, _) | (_, Bound::PosInf) => Bound::PosInf,
(Bound::NegInf, _) | (_, Bound::NegInf) => Bound::NegInf,
(Bound::Finite(a), Bound::Finite(b)) => match a.checked_add(b) {
Some(s) => Bound::Finite(s),
None if a > 0 => Bound::PosInf,
None => Bound::NegInf,
},
}
}
fn neg(self) -> Bound {
match self {
Bound::NegInf => Bound::PosInf,
Bound::PosInf => Bound::NegInf,
Bound::Finite(k) => match k.checked_neg() {
Some(n) => Bound::Finite(n),
None => Bound::PosInf,
},
}
}
fn mul(self, other: Bound) -> Bound {
let sign = |b: Bound| -> i32 {
match b {
Bound::NegInf => -1,
Bound::PosInf => 1,
Bound::Finite(0) => 0,
Bound::Finite(k) => {
if k > 0 {
1
} else {
-1
}
}
}
};
if matches!(self, Bound::Finite(0)) || matches!(other, Bound::Finite(0)) {
return Bound::Finite(0);
}
match (self, other) {
(Bound::Finite(a), Bound::Finite(b)) => match a.checked_mul(b) {
Some(p) => Bound::Finite(p),
None => {
if sign(self) * sign(other) >= 0 {
Bound::PosInf
} else {
Bound::NegInf
}
}
},
_ => {
if sign(self) * sign(other) >= 0 {
Bound::PosInf
} else {
Bound::NegInf
}
}
}
}
pub fn min(self, other: Bound) -> Bound {
if self.le(other) { self } else { other }
}
pub fn max(self, other: Bound) -> Bound {
if self.le(other) { other } else { self }
}
fn le(self, other: Bound) -> bool {
match (self, other) {
(Bound::NegInf, _) => true,
(_, Bound::NegInf) => false,
(_, Bound::PosInf) => true,
(Bound::PosInf, _) => false,
(Bound::Finite(a), Bound::Finite(b)) => a <= b,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Interval {
pub lo: Bound,
pub hi: Bound,
}
impl Interval {
pub fn unbounded() -> Interval {
Interval {
lo: Bound::NegInf,
hi: Bound::PosInf,
}
}
pub fn point(k: i128) -> Interval {
Interval {
lo: Bound::Finite(k),
hi: Bound::Finite(k),
}
}
pub fn ge(k: i128) -> Interval {
Interval {
lo: Bound::Finite(k),
hi: Bound::PosInf,
}
}
pub fn le(k: i128) -> Interval {
Interval {
lo: Bound::NegInf,
hi: Bound::Finite(k),
}
}
pub fn between(lo: i128, hi: i128) -> Interval {
Interval {
lo: Bound::Finite(lo),
hi: Bound::Finite(hi),
}
}
pub fn intersect(self, other: Interval) -> Interval {
Interval {
lo: self.lo.max(other.lo),
hi: self.hi.min(other.hi),
}
}
pub fn hull(self, other: Interval) -> Interval {
Interval {
lo: self.lo.min(other.lo),
hi: self.hi.max(other.hi),
}
}
pub fn fits_i64(self) -> bool {
self.lo.fits_i64() && self.hi.fits_i64()
}
pub fn contains_point(self, k: i128) -> bool {
self.lo.le(Bound::Finite(k)) && Bound::Finite(k).le(self.hi)
}
pub fn widen(self, next: Interval) -> Interval {
let lo = if next.lo.le(self.lo) && next.lo != self.lo {
Bound::NegInf
} else {
self.lo.min(next.lo)
};
let hi = if self.hi.le(next.hi) && self.hi != next.hi {
Bound::PosInf
} else {
self.hi.max(next.hi)
};
Interval { lo, hi }
}
fn is_finite(self) -> bool {
matches!(self.lo, Bound::Finite(_)) && matches!(self.hi, Bound::Finite(_))
}
#[allow(clippy::should_implement_trait)]
pub fn add(self, other: Interval) -> Interval {
Interval {
lo: self.lo.add(other.lo),
hi: self.hi.add(other.hi),
}
}
#[allow(clippy::should_implement_trait)]
pub fn sub(self, other: Interval) -> Interval {
Interval {
lo: self.lo.add(other.hi.neg()),
hi: self.hi.add(other.lo.neg()),
}
}
#[allow(clippy::should_implement_trait)]
pub fn mul(self, other: Interval) -> Interval {
let products = [
self.lo.mul(other.lo),
self.lo.mul(other.hi),
self.hi.mul(other.lo),
self.hi.mul(other.hi),
];
let mut lo = products[0];
let mut hi = products[0];
for p in &products[1..] {
lo = lo.min(*p);
hi = hi.max(*p);
}
Interval { lo, hi }
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OpClass {
OverflowFree,
NeedsWiderScratch,
Unbounded,
}
impl OpClass {
pub fn of_interval(i: Interval) -> OpClass {
if i.fits_i64() {
OpClass::OverflowFree
} else if i.is_finite() {
OpClass::NeedsWiderScratch
} else {
OpClass::Unbounded
}
}
pub fn label(self) -> &'static str {
match self {
OpClass::OverflowFree => "overflow_free",
OpClass::NeedsWiderScratch => "needs_wider_scratch",
OpClass::Unbounded => "unbounded",
}
}
}
pub fn raw_i64_eligible<'a>(
interval: Option<Interval>,
ops: impl IntoIterator<Item = &'a OpClass>,
) -> bool {
let Some(interval) = interval else {
return false;
};
if !interval.fits_i64() {
return false;
}
ops.into_iter().all(|c| *c == OpClass::OverflowFree)
}
#[derive(Debug, Clone)]
pub struct RefinedTypeInterval {
pub type_id: TypeId,
pub name: String,
pub interval: Interval,
pub interval_known: bool,
pub ops: Vec<(String, OpClass)>,
}
impl RefinedTypeInterval {
pub fn raw_i64_eligible(&self) -> bool {
let interval = self.interval_known.then_some(self.interval);
raw_i64_eligible(interval, self.ops.iter().map(|(_, c)| c))
}
}
#[derive(Debug, Clone, Default)]
pub struct IntervalAnalysisResult {
pub types: HashMap<TypeId, RefinedTypeInterval>,
}
impl IntervalAnalysisResult {
pub fn types_analyzed(&self) -> usize {
self.types.len()
}
pub fn two_sided_bounded(&self) -> usize {
self.types
.values()
.filter(|t| t.interval_known && t.interval.is_finite())
.count()
}
pub fn ops_overflow_free(&self) -> usize {
self.count_ops(OpClass::OverflowFree)
}
pub fn ops_needs_wider(&self) -> usize {
self.count_ops(OpClass::NeedsWiderScratch)
}
pub fn ops_unbounded(&self) -> usize {
self.count_ops(OpClass::Unbounded)
}
pub fn raw_i64_eligible(&self) -> usize {
self.types.values().filter(|t| t.raw_i64_eligible()).count()
}
fn count_ops(&self, class: OpClass) -> usize {
self.types
.values()
.flat_map(|t| t.ops.iter())
.filter(|(_, c)| *c == class)
.count()
}
}
pub fn interval_of_invariant(pred: &Predicate) -> (Interval, bool) {
let Some((var, _)) = pred.free_vars.first() else {
return (Interval::unbounded(), false);
};
match interval_of_resolved(&pred.expr, var) {
Some(i) => (i, true),
None => (Interval::unbounded(), false),
}
}
fn interval_of_resolved(
expr: &Spanned<crate::ir::hir::ResolvedExpr>,
var: &str,
) -> Option<Interval> {
use crate::ir::hir::{ResolvedCallee, ResolvedExpr};
match &expr.node {
ResolvedExpr::Call(ResolvedCallee::Builtin(name), args)
if name == "Bool.and" && args.len() == 2 =>
{
let l = interval_of_resolved(&args[0], var)?;
let r = interval_of_resolved(&args[1], var)?;
Some(l.intersect(r))
}
ResolvedExpr::BinOp(op, lhs, rhs) => interval_of_comparison(*op, lhs, rhs, var),
_ => None,
}
}
fn interval_of_comparison(
op: BinOp,
lhs: &Spanned<crate::ir::hir::ResolvedExpr>,
rhs: &Spanned<crate::ir::hir::ResolvedExpr>,
var: &str,
) -> Option<Interval> {
let (op, k) = if is_var(lhs, var) {
(op, int_literal(rhs)?)
} else if is_var(rhs, var) {
(flip_comparison(op)?, int_literal(lhs)?)
} else {
return None;
};
let k = k as i128;
match op {
BinOp::Gte => Some(Interval::ge(k)), BinOp::Gt => Some(Interval::ge(k + 1)), BinOp::Lte => Some(Interval::le(k)), BinOp::Lt => Some(Interval::le(k - 1)), _ => None,
}
}
fn is_var(expr: &Spanned<crate::ir::hir::ResolvedExpr>, var: &str) -> bool {
use crate::ir::hir::ResolvedExpr;
match &expr.node {
ResolvedExpr::Ident(name) => name == var,
ResolvedExpr::Resolved { name, .. } => name == var,
_ => false,
}
}
fn int_literal(expr: &Spanned<crate::ir::hir::ResolvedExpr>) -> Option<i64> {
use crate::ir::hir::ResolvedExpr;
match &expr.node {
ResolvedExpr::Literal(Literal::Int(k)) => Some(*k),
ResolvedExpr::Neg(inner) => match &inner.node {
ResolvedExpr::Literal(Literal::Int(k)) => Some(-*k),
_ => None,
},
_ => None,
}
}
fn flip_comparison(op: BinOp) -> Option<BinOp> {
match op {
BinOp::Lt => Some(BinOp::Gt),
BinOp::Gt => Some(BinOp::Lt),
BinOp::Lte => Some(BinOp::Gte),
BinOp::Gte => Some(BinOp::Lte),
_ => None,
}
}
pub fn classify_op(
op_fn: &FnDef,
carrier_interval: Interval,
refined_type_name: &str,
carrier_field: &str,
constructor_fn: &str,
) -> OpClass {
let mut refined_params: HashMap<&str, bool> = HashMap::new();
for (pname, ptype) in &op_fn.params {
refined_params.insert(pname.as_str(), ptype == refined_type_name);
}
let ctx = OpCtx {
carrier_interval,
carrier_field,
refined_params: &refined_params,
constructor_fn,
};
let mut worst = Interval::point(0);
let FnBody::Block(stmts) = op_fn.body.as_ref();
for stmt in stmts {
let value = match stmt {
Stmt::Expr(e) | Stmt::Binding(_, _, e) => e,
};
let i = eval_expr(value, &ctx, &mut worst);
worst = join(worst, i);
}
OpClass::of_interval(worst)
}
struct OpCtx<'a> {
carrier_interval: Interval,
carrier_field: &'a str,
refined_params: &'a HashMap<&'a str, bool>,
constructor_fn: &'a str,
}
fn join(a: Interval, b: Interval) -> Interval {
Interval {
lo: a.lo.min(b.lo),
hi: a.hi.max(b.hi),
}
}
fn eval_expr(expr: &Spanned<Expr>, ctx: &OpCtx<'_>, worst: &mut Interval) -> Interval {
match &expr.node {
Expr::Literal(Literal::Int(k)) => Interval::point(*k as i128),
Expr::Attr(obj, field) => {
if field == ctx.carrier_field
&& let Some(pname) = param_name(obj)
&& ctx.refined_params.get(pname).copied() == Some(true)
{
ctx.carrier_interval
} else {
Interval::unbounded()
}
}
Expr::BinOp(op, lhs, rhs) => {
let l = eval_expr(lhs, ctx, worst);
let r = eval_expr(rhs, ctx, worst);
let result = match op {
BinOp::Add => l.add(r),
BinOp::Sub => l.sub(r),
BinOp::Mul => l.mul(r),
_ => Interval::unbounded(),
};
*worst = join(*worst, result);
result
}
Expr::Neg(inner) => {
let result = Interval::point(0).sub(eval_expr(inner, ctx, worst));
*worst = join(*worst, result);
result
}
Expr::FnCall(callee, args)
if args.len() == 1 && is_constructor_call(callee, ctx.constructor_fn) =>
{
eval_expr(&args[0], ctx, worst)
}
_ => Interval::unbounded(),
}
}
fn is_constructor_call(callee: &Spanned<Expr>, constructor_fn: &str) -> bool {
matches!(&callee.node, Expr::Ident(name) if name == constructor_fn)
}
fn param_name(expr: &Spanned<Expr>) -> Option<&str> {
match &expr.node {
Expr::Ident(name) => Some(name.as_str()),
Expr::Resolved { name, .. } => Some(name.as_str()),
_ => None,
}
}
pub fn analyze(
refined: &HashMap<TypeId, RefinedTypeDecl>,
inputs: &ProofLowerInputs<'_>,
) -> IntervalAnalysisResult {
let mut result = IntervalAnalysisResult::default();
let symbols = inputs.symbol_table;
for (type_id, decl) in refined {
let (interval, interval_known) = interval_of_invariant(&decl.invariant);
let scope = scope_of_type(*type_id, decl, inputs, symbols);
let constructor_fn = crate::codegen::common::refinement_info_for_in_scope(
&decl.name,
inputs,
scope.as_deref(),
)
.map(|info| info.constructor_fn.to_string());
let ops = classify_ops_in_scope(
decl,
interval,
interval_known,
constructor_fn.as_deref(),
inputs.pure_fns_in_scope(scope.as_deref()),
);
result.types.insert(
*type_id,
RefinedTypeInterval {
type_id: *type_id,
name: decl.name.clone(),
interval,
interval_known,
ops,
},
);
}
result
}
fn scope_of_type(
type_id: TypeId,
decl: &RefinedTypeDecl,
inputs: &ProofLowerInputs<'_>,
symbols: &crate::ir::SymbolTable,
) -> Option<String> {
for scope in inputs.scopes() {
let key = match &scope {
Some(prefix) => crate::ir::TypeKey::in_module(prefix.clone(), &decl.name),
None => crate::ir::TypeKey::entry(&decl.name),
};
if symbols.type_id_of(&key) == Some(type_id) {
return scope;
}
}
None
}
fn classify_ops_in_scope(
decl: &RefinedTypeDecl,
interval: Interval,
interval_known: bool,
constructor_fn: Option<&str>,
fns: Vec<&FnDef>,
) -> Vec<(String, OpClass)> {
let mut ops = Vec::new();
for fd in fns {
let takes_refined = fd.params.iter().any(|(_, t)| t == &decl.name);
if !takes_refined || !body_does_carrier_arithmetic(fd, &decl.carrier_field) {
continue;
}
let class = match (interval_known, constructor_fn) {
(true, Some(ctor)) => classify_op(fd, interval, &decl.name, &decl.carrier_field, ctor),
_ => OpClass::Unbounded,
};
ops.push((fd.name.clone(), class));
}
ops
}
fn body_does_carrier_arithmetic(fd: &FnDef, _carrier_field: &str) -> bool {
let FnBody::Block(stmts) = fd.body.as_ref();
stmts.iter().any(|s| match s {
Stmt::Expr(e) | Stmt::Binding(_, _, e) => expr_has_arithmetic(e),
})
}
fn expr_has_arithmetic(expr: &Spanned<Expr>) -> bool {
match &expr.node {
Expr::BinOp(BinOp::Add | BinOp::Sub | BinOp::Mul, _, _) => true,
Expr::BinOp(_, l, r) => expr_has_arithmetic(l) || expr_has_arithmetic(r),
Expr::FnCall(_, args) => args.iter().any(expr_has_arithmetic),
Expr::Attr(o, _) => expr_has_arithmetic(o),
Expr::Neg(i) | Expr::ErrorProp(i) => expr_has_arithmetic(i),
_ => false,
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ast::SourceLine;
use crate::ir::hir::ResolvedExpr;
use crate::ir::proof_ir::QuantifierType;
const LINE: SourceLine = 0;
fn sp_r(node: ResolvedExpr) -> Spanned<ResolvedExpr> {
Spanned::new(node, LINE)
}
fn ident_r(name: &str) -> Spanned<ResolvedExpr> {
sp_r(ResolvedExpr::Ident(name.to_string()))
}
fn int_r(k: i64) -> Spanned<ResolvedExpr> {
sp_r(ResolvedExpr::Literal(Literal::Int(k)))
}
fn cmp_r(
op: BinOp,
l: Spanned<ResolvedExpr>,
r: Spanned<ResolvedExpr>,
) -> Spanned<ResolvedExpr> {
sp_r(ResolvedExpr::BinOp(op, Box::new(l), Box::new(r)))
}
fn pred(expr: Spanned<ResolvedExpr>) -> Predicate {
Predicate {
free_vars: vec![("n".to_string(), QuantifierType::Plain("Int".to_string()))],
expr,
}
}
#[test]
fn bound_posinf_plus_finite_is_posinf() {
assert_eq!(Bound::PosInf.add(Bound::Finite(5)), Bound::PosInf);
assert_eq!(Bound::Finite(5).add(Bound::PosInf), Bound::PosInf);
}
#[test]
fn bound_finite_mul_exact_in_i128() {
assert_eq!(
Bound::Finite(100).mul(Bound::Finite(100)),
Bound::Finite(10_000)
);
}
#[test]
fn bound_i64_max_squared_is_exact_finite_not_i64() {
let m = Bound::Finite(i64::MAX as i128);
let expected = (i64::MAX as i128) * (i64::MAX as i128);
assert_eq!(m.mul(m), Bound::Finite(expected));
assert!(!Bound::Finite(expected).fits_i64());
}
#[test]
fn bound_i128_overflow_saturates_not_wraps() {
let big = Bound::Finite(i128::MAX);
assert_eq!(big.mul(big), Bound::PosInf);
let neg = Bound::Finite(i128::MIN + 1);
assert_eq!(big.mul(neg), Bound::NegInf);
assert_eq!(big.add(Bound::Finite(1)), Bound::PosInf);
assert_eq!(
Bound::Finite(i128::MIN).add(Bound::Finite(-1)),
Bound::NegInf
);
}
#[test]
fn bound_neg_min_saturates() {
assert_eq!(Bound::Finite(i128::MIN).neg(), Bound::PosInf);
}
#[test]
fn interval_add_keeps_finite_in_band() {
let a = Interval::between(0, 100);
let sum = a.add(a);
assert_eq!(sum, Interval::between(0, 200));
assert!(sum.fits_i64());
}
#[test]
fn interval_mul_handles_negatives() {
let a = Interval::between(-2, 3);
assert_eq!(a.mul(a), Interval::between(-6, 9));
}
#[test]
fn contains_point_respects_both_bounds_and_infinities() {
let band = Interval::between(0, 100);
assert!(band.contains_point(0));
assert!(band.contains_point(100));
assert!(band.contains_point(50));
assert!(!band.contains_point(-1));
assert!(!band.contains_point(101));
let ge0 = Interval::ge(0);
assert!(ge0.contains_point(0));
assert!(ge0.contains_point(i64::MAX as i128));
assert!(!ge0.contains_point(-1));
assert!(Interval::unbounded().contains_point(i128::MIN));
assert!(Interval::unbounded().contains_point(i128::MAX));
}
#[test]
fn invariant_ge_natural() {
let (i, known) = interval_of_invariant(&pred(cmp_r(BinOp::Gte, ident_r("n"), int_r(0))));
assert!(known);
assert_eq!(i, Interval::ge(0));
}
#[test]
fn invariant_gt_positive() {
let (i, known) = interval_of_invariant(&pred(cmp_r(BinOp::Gt, ident_r("n"), int_r(0))));
assert!(known);
assert_eq!(i, Interval::ge(1));
}
#[test]
fn invariant_lte() {
let (i, known) = interval_of_invariant(&pred(cmp_r(BinOp::Lte, ident_r("n"), int_r(100))));
assert!(known);
assert_eq!(i, Interval::le(100));
}
#[test]
fn invariant_lt() {
let (i, known) = interval_of_invariant(&pred(cmp_r(BinOp::Lt, ident_r("n"), int_r(10))));
assert!(known);
assert_eq!(i, Interval::le(9));
}
#[test]
fn invariant_bool_and_intrange() {
use crate::ir::hir::ResolvedCallee;
let and = sp_r(ResolvedExpr::Call(
ResolvedCallee::Builtin("Bool.and".to_string()),
vec![
cmp_r(BinOp::Gte, ident_r("n"), int_r(0)),
cmp_r(BinOp::Lte, ident_r("n"), int_r(100)),
],
));
let (i, known) = interval_of_invariant(&pred(and));
assert!(known);
assert_eq!(i, Interval::between(0, 100));
}
#[test]
fn invariant_operand_flipped() {
let (i, known) = interval_of_invariant(&pred(cmp_r(BinOp::Lte, int_r(0), ident_r("n"))));
assert!(known);
assert_eq!(i, Interval::ge(0));
}
#[test]
fn invariant_bool_or_declines() {
use crate::ir::hir::ResolvedCallee;
let or = sp_r(ResolvedExpr::Call(
ResolvedCallee::Builtin("Bool.or".to_string()),
vec![
cmp_r(BinOp::Gte, ident_r("n"), int_r(0)),
cmp_r(BinOp::Lte, ident_r("n"), int_r(100)),
],
));
let (i, known) = interval_of_invariant(&pred(or));
assert!(!known);
assert_eq!(i, Interval::unbounded());
}
#[test]
fn invariant_bare_ident_declines() {
let (i, known) = interval_of_invariant(&pred(ident_r("n")));
assert!(!known);
assert_eq!(i, Interval::unbounded());
}
#[test]
fn invariant_non_literal_bound_declines() {
let (i, known) =
interval_of_invariant(&pred(cmp_r(BinOp::Gte, ident_r("n"), ident_r("m"))));
assert!(!known);
assert_eq!(i, Interval::unbounded());
}
#[test]
fn fits_i64_at_boundary() {
assert!(Interval::between(0, i64::MAX as i128).fits_i64());
let over = Interval::between(0, i64::MAX as i128 + 1);
assert!(!over.fits_i64());
assert!(over.is_finite());
assert_eq!(OpClass::of_interval(over), OpClass::NeedsWiderScratch);
}
#[test]
fn opclass_overflow_free_vs_unbounded() {
assert_eq!(
OpClass::of_interval(Interval::between(0, 200)),
OpClass::OverflowFree
);
assert_eq!(OpClass::of_interval(Interval::ge(0)), OpClass::Unbounded);
}
#[test]
fn widen_stable_endpoint_kept() {
let a = Interval::between(0, 10);
assert_eq!(a.widen(a), a);
}
#[test]
fn widen_descending_lo_jumps_to_neg_inf() {
let prev = Interval::between(10, 20);
let next = Interval::between(5, 20);
assert_eq!(
prev.widen(next),
Interval {
lo: Bound::NegInf,
hi: Bound::Finite(20),
}
);
}
#[test]
fn widen_ascending_hi_jumps_to_pos_inf() {
let prev = Interval::between(0, 20);
let next = Interval::between(0, 30);
assert_eq!(
prev.widen(next),
Interval {
lo: Bound::Finite(0),
hi: Bound::PosInf,
}
);
}
#[test]
fn widen_is_enlarging_only() {
let prev = Interval::between(0, 100);
let next = Interval::between(5, 90); let w = prev.widen(next);
assert!(w.lo.le(next.lo), "lo must not rise above next.lo");
assert!(next.hi.le(w.hi), "hi must not fall below next.hi");
assert_eq!(w, Interval::between(0, 100));
}
}