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))