PCC: verification primitives for dynamic range checks.

cranelift/codegen/src/isa/aarch64/lower.rs

@@ -136,7 +136,10 @@ impl LowerBackend for AArch64Backend {
         ctx: &FactContext<'_>,
         vcode: &mut VCode<Self::MInst>,
         inst: InsnIndex,
+        state: &mut pcc::FactFlowState,
     ) -> PccResult<()> {
-        pcc::check(ctx, vcode, inst)
+        pcc::check(ctx, vcode, inst, state)
     }
+
+    type FactFlowState = pcc::FactFlowState;
 }

cranelift/codegen/src/isa/aarch64/pcc.rs

@@ -4,7 +4,8 @@ use crate::ir::pcc::*;
 use crate::ir::types::*;
 use crate::ir::MemFlags;
 use crate::ir::Type;
-use crate::isa::aarch64::inst::args::{PairAMode, ShiftOp};
+use crate::isa::aarch64::inst::args::{Cond, PairAMode, ShiftOp};
+use crate::isa::aarch64::inst::regs::zero_reg;
 use crate::isa::aarch64::inst::Inst;
 use crate::isa::aarch64::inst::{ALUOp, MoveWideOp};
 use crate::isa::aarch64::inst::{AMode, ExtendOp};
@@ -26,13 +27,26 @@ fn extend_fact(ctx: &FactContext, value: &Fact, mode: ExtendOp) -> Option<Fact>
     }
 }
 
+/// Flow-state between facts.
+#[derive(Clone, Debug, Default)]
+pub struct FactFlowState {
+    cmp_flags: Option<(Fact, Fact)>,
+}
+
 pub(crate) fn check(
     ctx: &FactContext,
     vcode: &mut VCode<Inst>,
     inst_idx: InsnIndex,
+    state: &mut FactFlowState,
 ) -> PccResult<()> {
     trace!("Checking facts on inst: {:?}", vcode[inst_idx]);
 
+    // We only persist flag state for one instruction, because we
+    // can't exhaustively enumerate all flags-effecting ops; so take
+    // the `cmp_state` here and perhaps use it below but don't let it
+    // remain.
+    let cmp_flags = state.cmp_flags.take();
+
     match vcode[inst_idx] {
         Inst::Args { .. } => {
             // Defs on the args have "axiomatic facts": we trust the
@@ -200,6 +214,20 @@ pub(crate) fn check(
             )
         }),
 
+        Inst::AluRRR {
+            alu_op: ALUOp::SubS,
+            size,
+            rd,
+            rn,
+            rm,
+        } if rd.to_reg() == zero_reg() => {
+            // Compare.
+            let rn = get_fact_or_default(vcode, rn, size.bits().into());
+            let rm = get_fact_or_default(vcode, rm, size.bits().into());
+            state.cmp_flags = Some((rn, rm));
+            Ok(())
+        }
+
         Inst::AluRRR { rd, size, .. }
         | Inst::AluRRImm12 { rd, size, .. }
         | Inst::AluRRRShift { rd, size, .. }
@@ -240,12 +268,9 @@ pub(crate) fn check(
 
         Inst::MovK { rd, rn, imm, .. } => {
             let input = get_fact_or_default(vcode, rn, 64);
-            trace!("MovK: input = {:?}", input);
             if let Some(input_constant) = input.as_const(64) {
-                trace!(" -> input_constant: {}", input_constant);
                 let constant = u64::from(imm.bits) << (imm.shift * 16);
                 let constant = input_constant | constant;
-                trace!(" -> merged constant: {}", constant);
                 check_constant(ctx, vcode, rd, 64, constant)
             } else {
                 check_output(ctx, vcode, rd, &[], |_vcode| {
@@ -254,6 +279,44 @@ pub(crate) fn check(
             }
         }
 
+        Inst::CSel { rd, cond, rn, rm } if cond == Cond::Hs && cmp_flags.is_some() => {
+            let (cmp_lhs, cmp_rhs) = cmp_flags.unwrap();
+            trace!("CSel: cmp {cond:?} ({cmp_lhs:?}, {cmp_rhs:?})");
+
+            check_output(ctx, vcode, rd, &[], |vcode| {
+                // We support transitivity-based reasoning. If the
+                // comparison establishes that
+                //
+                //   (x+K1) <= (y+K2)
+                //
+                // then on the true-side of the select we can edit the maximum
+                // in a DynamicMem or DynamicRange by replacing x's with y's
+                // with appropriate offsets -- this widens the range.
+                //
+                // Likewise, on the false-side of the select we can
+                // replace y's with x's -- this also widens the range. On
+                // the false side we know the inequality is strict, so we
+                // can offset by one.
+
+                // True side: lhs <= rhs, not strict.
+                let rn = get_fact_or_default(vcode, rn, 64);
+                trace!("rn = {rn:?}");
+                let rn = ctx.apply_inequality(&rn, &cmp_lhs, &cmp_rhs, false);
+                trace!(" -> rn = {rn:?}");
+                // false side: rhs < lhs, strict.
+                let rm = get_fact_or_default(vcode, rm, 64);
+                trace!("rm = {rm:?}");
+                let rm = ctx.apply_inequality(&rm, &cmp_rhs, &cmp_lhs, true);
+                trace!(" -> rm = {rm:?}");
+
+                let union = Fact::union(&rn, &rm);
+                trace!("union = {union:?}");
+
+                // Union the two facts.
+                clamp_range(ctx, 64, 64, union)
+            })
+        }
+
         _ if vcode.inst_defines_facts(inst_idx) => Err(PccError::UnsupportedFact),
 
         _ => Ok(()),

cranelift/codegen/src/isa/riscv64/lower.rs

@@ -30,4 +30,6 @@ impl LowerBackend for Riscv64Backend {
         // right now riscv64 not support this feature.
         None
     }
+
+    type FactFlowState = ();
 }

cranelift/codegen/src/isa/s390x/lower.rs

@@ -25,4 +25,6 @@ impl LowerBackend for S390xBackend {
     ) -> Option<()> {
         isle::lower_branch(ctx, self, ir_inst, targets)
     }
+
+    type FactFlowState = ();
 }

cranelift/codegen/src/isa/x64/inst/args.rs

@@ -2007,7 +2007,7 @@ impl fmt::Display for ShiftKind {
 
 /// These indicate condition code tests.  Not all are represented since not all are useful in
 /// compiler-generated code.
-#[derive(Copy, Clone)]
+#[derive(Copy, Clone, PartialEq, Eq)]
 #[repr(u8)]
 pub enum CC {
     ///  overflow

cranelift/codegen/src/isa/x64/lower.rs

@@ -345,7 +345,10 @@ impl LowerBackend for X64Backend {
         ctx: &FactContext<'_>,
         vcode: &mut VCode<Self::MInst>,
         inst: InsnIndex,
+        state: &mut pcc::FactFlowState,
     ) -> PccResult<()> {
-        pcc::check(ctx, vcode, inst)
+        pcc::check(ctx, vcode, inst, state)
     }
+
+    type FactFlowState = pcc::FactFlowState;
 }

cranelift/codegen/src/isa/x64/pcc.rs

@@ -4,7 +4,8 @@ use crate::ir::pcc::*;
 use crate::ir::types::*;
 use crate::ir::Type;
 use crate::isa::x64::inst::args::{
-    AluRmiROpcode, Amode, Gpr, Imm8Reg, RegMem, RegMemImm, ShiftKind, SyntheticAmode, ToWritableReg,
+    AluRmiROpcode, Amode, Gpr, Imm8Reg, RegMem, RegMemImm, ShiftKind, SyntheticAmode,
+    ToWritableReg, CC,
 };
 use crate::isa::x64::inst::Inst;
 use crate::machinst::pcc::*;
@@ -32,13 +33,26 @@ fn ensure_no_fact(vcode: &VCode<Inst>, reg: Reg) -> PccResult<()> {
     }
 }
 
+/// Flow-state between facts.
+#[derive(Clone, Debug, Default)]
+pub(crate) struct FactFlowState {
+    cmp_flags: Option<(Fact, Fact)>,
+}
+
 pub(crate) fn check(
     ctx: &FactContext,
     vcode: &mut VCode<Inst>,
     inst_idx: InsnIndex,
+    state: &mut FactFlowState,
 ) -> PccResult<()> {
     trace!("Checking facts on inst: {:?}", vcode[inst_idx]);
 
+    // We only persist flag state for one instruction, because we
+    // can't exhaustively enumerate all flags-effecting ops; so take
+    // the `cmp_state` here and perhaps use it below but don't let it
+    // remain.
+    let cmp_flags = state.cmp_flags.take();
+
     match vcode[inst_idx] {
         Inst::Nop { .. } => Ok(()),
 
@@ -302,8 +316,7 @@ pub(crate) fn check(
             match <&RegMem>::from(src) {
                 RegMem::Reg { reg } => {
                     check_unop(ctx, vcode, 64, dst.to_writable_reg(), *reg, |src| {
-                        let extended = ctx.uextend(src, from_bytes * 8, to_bytes * 8);
-                        clamp_range(ctx, 64, from_bytes * 8, extended)
+                        clamp_range(ctx, 64, from_bytes * 8, Some(src.clone()))
                     })
                 }
                 RegMem::Mem { ref addr } => {
@@ -395,12 +408,23 @@ pub(crate) fn check(
             ensure_no_fact(vcode, dst.to_writable_reg().to_reg())
         }
 
-        Inst::CmpRmiR { size, ref src, .. } => match <&RegMemImm>::from(src) {
+        Inst::CmpRmiR {
+            size, dst, ref src, ..
+        } => match <&RegMemImm>::from(src) {
             RegMemImm::Mem { ref addr } => {
-                check_load(ctx, None, addr, vcode, size.to_type(), 64)?;
+                if let Some(rhs) = check_load(ctx, None, addr, vcode, size.to_type(), 64)? {
+                    let lhs = get_fact_or_default(vcode, dst.to_reg(), 64);
+                    state.cmp_flags = Some((lhs, rhs));
+                }
                 Ok(())
             }
-            RegMemImm::Reg { .. } | RegMemImm::Imm { .. } => Ok(()),
+            RegMemImm::Reg { reg } => {
+                let rhs = get_fact_or_default(vcode, *reg, 64);
+                let lhs = get_fact_or_default(vcode, dst.to_reg(), 64);
+                state.cmp_flags = Some((lhs, rhs));
+                Ok(())
+            }
+            RegMemImm::Imm { .. } => Ok(()),
         },
 
         Inst::Setcc { dst, .. } => undefined_result(ctx, vcode, dst, 64, 64),
@@ -411,16 +435,35 @@ pub(crate) fn check(
             size,
             dst,
             ref consequent,
+            alternative,
+            cc,
             ..
-        } => {
-            match <&RegMem>::from(consequent) {
-                RegMem::Mem { ref addr } => {
-                    check_load(ctx, None, addr, vcode, size.to_type(), 64)?;
-                }
-                RegMem::Reg { .. } => {}
+        } => match <&RegMem>::from(consequent) {
+            RegMem::Mem { ref addr } => {
+                check_load(ctx, None, addr, vcode, size.to_type(), 64)?;
+                Ok(())
             }
-            undefined_result(ctx, vcode, dst, 64, 64)
-        }
+            RegMem::Reg { reg } if cc == CC::NB && cmp_flags.is_some() => {
+                let (cmp_lhs, cmp_rhs) = cmp_flags.unwrap();
+                trace!("lhs = {:?} rhs = {:?}", cmp_lhs, cmp_rhs);
+                let reg = *reg;
+                check_output(ctx, vcode, dst.to_writable_reg(), &[], |vcode| {
+                    // See comments in aarch64::pcc CSel for more details on this.
+                    let in_true = get_fact_or_default(vcode, reg, 64);
+                    trace!("in_true = {:?}", in_true);
+                    let in_true = ctx.apply_inequality(&in_true, &cmp_lhs, &cmp_rhs, false);
+                    trace!("in_true = {:?}", in_true);
+                    let in_false = get_fact_or_default(vcode, alternative.to_reg(), 64);
+                    trace!("in_false = {:?}", in_false);
+                    let in_false = ctx.apply_inequality(&in_false, &cmp_rhs, &cmp_lhs, true);
+                    trace!("in_false = {:?}", in_false);
+                    let union = Fact::union(&in_true, &in_false);
+                    trace!("union = {:?}", union);
+                    clamp_range(ctx, 64, 64, union)
+                })
+            }
+            _ => undefined_result(ctx, vcode, dst, 64, 64),
+        },
 
         Inst::XmmCmove {
             dst,

cranelift/codegen/src/machinst/lower.rs

@@ -148,6 +148,9 @@ pub trait LowerBackend {
         None
     }
 
+    /// The type of state carried between `check_fact` invocations.
+    type FactFlowState: Default + Clone + Debug;
+
     /// Check any facts about an instruction, given VCode with facts
     /// on VRegs. Takes mutable `VCode` so that it can propagate some
     /// kinds of facts automatically.
@@ -156,6 +159,7 @@ pub trait LowerBackend {
         _ctx: &FactContext<'_>,
         _vcode: &mut VCode<Self::MInst>,
         _inst: InsnIndex,
+        _state: &mut Self::FactFlowState,
     ) -> PccResult<()> {
         Err(PccError::UnimplementedBackend)
     }

cranelift/filetests/filetests/pcc/succeed/dynamic.clif

@@ -17,21 +17,21 @@ function %f0(i64 vmctx, i32) -> i64 {
     ;; mock dynamic memory: dynamic range, plus 2GiB guard
     mt1 = dynamic_memory gv2 + 0x8000_0000
 
-block0(v0 ! mem(mt0, 0, 0): i64, v1: i32):
-    v2 ! range(64, 0, 0xffff_ffff)            = uextend.i64 v1       ;; extended Wasm offset
+block0(v0 ! mem(mt0, 0, 0): i64, v1 ! dynamic_range(32, v1, v1): i32):
+    v2 ! dynamic_range(64, v1, v1)            = uextend.i64 v1       ;; extended Wasm offset
     v3 ! dynamic_mem(mt1, 0, 0)               = global_value.i64 gv1 ;; base
     v4 ! dynamic_range(64, gv2, gv2)          = global_value.i64 gv2 ;; size
-    v5 ! compare(uge, v2, gv2)                = icmp.i64 uge v2, v4  ;; bounds-check compare of extended Wasm offset to size
-    v6 ! dynamic_mem(mt1, 0, v2)              = iadd.i64 v3, v2      ;; compute access address: memory base plus extended Wasm offset
+    v5 ! compare(uge, v1, gv2)                = icmp.i64 uge v2, v4  ;; bounds-check compare of extended Wasm offset to size
+    v6 ! dynamic_mem(mt1, v1, v1)              = iadd.i64 v3, v2      ;; compute access address: memory base plus extended Wasm offset
     v7 ! dynamic_mem(mt1, 0, 0, nullable)     = iconst.i64 0         ;; null pointer for speculative path
     v8 ! dynamic_mem(mt1, 0, gv2-1, nullable) = select_spectre_guard v5, v7, v6  ;; if OOB, pick null, otherwise the real address
     v9                                        = load.i64 checked v8
     return v9
 }
 
 ;; select sees:
-;;  v5 ! compare(uge, v2, gv2)
-;;  v6 ! dynamic_mem(mt1, 0, v2)
+;;  v5 ! compare(uge, v1, gv2)
+;;  v6 ! dynamic_mem(mt1, v1, v1)
 ;;  v7 ! dynamic_mem(mt0, 0, 0, nullable)
 ;;
 ;; preprocess:

cranelift/reader/src/parser.rs

@@ -2369,9 +2369,14 @@ impl<'a> Parser<'a> {
     //
     // base-expr ::= GlobalValue(base)
     //             | Value(base)
+    //             | "max"
     //             | (epsilon)
     fn parse_base_expr(&mut self) -> ParseResult<BaseExpr> {
         match self.token() {
+            Some(Token::Identifier("max")) => {
+                self.consume();
+                Ok(BaseExpr::Max)
+            }
             Some(Token::GlobalValue(..)) => {
                 let gv = self.match_gv("expected global value")?;
                 Ok(BaseExpr::GlobalValue(gv))