TCB: AArch64 CLZ and RBIT

lean/VerifiedGarbage/Proof/Framework/AArch64/Inline.lean

@@ -37,7 +37,7 @@ pop writes its register); none for `ldp`, which writes two (`dsts`). -/
 def dstOf : Instr → Option Reg
   | .adds _ d .. | .adcs _ d .. | .subs _ d .. | .sbcs _ d .. | .adc _ d .. | .sbc _ d .. | .csel _ d .. | .umulh d .. | .umaddl d .. | .umull d .. | .adrSym d _ => some d
   | .add _ d .. | .sub _ d .. | .addImm _ d .. | .subImm _ d .. | .logic _ _ d .. | .logicRor _ _ d .. | .bicRor _ d .. | .ror _ d .. | .extr _ d ..
-  | .lsr _ d .. | .lsl _ d .. | .madd _ d .. | .mul _ d .. | .rev32 d _ | .rev d _ | .movz _ d ..
+  | .lsr _ d .. | .lsl _ d .. | .madd _ d .. | .mul _ d .. | .rev32 d _ | .rev d _ | .clz _ d _ | .rbit _ d _ | .movz _ d ..
   | .addSp d _ | .movk _ d .. | .ldr _ d .. | .ldrb d .. | .ldrSp d _ | .pop d | .umov _ d .. => some d
   | .cselc _ d .. | .csneg _ d .. => some d
   | .str .. | .strb .. | .strSp .. | .alloc _ | .free _ | .push _ | .vop _ | .ldrq .. | .ldrqTbl .. | .strq .. | .tst .. | .ccmp ..

lean/VerifiedGarbage/Proof/Framework/AArch64/Taint.lean

@@ -61,7 +61,7 @@ def stepFn : Instr → Step
   | .madd _ d n m a => ⟨fun τ => some (set τ d (pub τ n && pub τ m && pub τ a))⟩
   | .umaddl d n m a => ⟨fun τ => some (set τ d (pub τ n && pub τ m && pub τ a))⟩
   | .addImm _ d n _ | .subImm _ d n _ | .ror _ d n _ | .lsr _ d n _ | .lsl _ d n _ | .rev32 d n
-  | .rev d n => ⟨fun τ => some (set τ d (pub τ n))⟩
+  | .rev d n | .clz _ d n | .rbit _ d n => ⟨fun τ => some (set τ d (pub τ n))⟩
   | .addSp d _ | .movz _ d _ _ => ⟨fun τ => some (set τ d true)⟩
   | .movk _ d _ _ => ⟨fun τ => some (set τ d (pub τ d))⟩
   | .ldr _ t n _ | .ldrb t n _ => ⟨fun τ => if pub τ n then some (set τ t false) else none⟩
@@ -242,6 +242,14 @@ theorem step_sound {τ τ' : T} {i : Instr} {s₁ s₂ s₁' s₂' : State} (ha
     refine ⟨rfl, ha.write sz d fun hp => ?_⟩
     simp only [Bool.and_eq_true] at hp
     rw [ha.read hp.1, ha.read hp.2]
+  | clz sz d n =>
+    obtain rfl := Option.some.inj hs
+    obtain rfl := Option.some.inj e₁; obtain rfl := Option.some.inj e₂
+    exact ⟨rfl, ha.write sz d fun hp => by rw [ha.read hp]⟩
+  | rbit sz d n =>
+    obtain rfl := Option.some.inj hs
+    obtain rfl := Option.some.inj e₁; obtain rfl := Option.some.inj e₂
+    exact ⟨rfl, ha.write sz d fun hp => by rw [ha.read hp]⟩
   | rev32 d n =>
     obtain rfl := Option.some.inj hs
     obtain rfl := Option.some.inj e₁; obtain rfl := Option.some.inj e₂

lean/VerifiedGarbage/Proof/Framework/AArch64/TaintMono.lean

@@ -77,7 +77,7 @@ def gprDst : Instr → Option Reg
   | .bicRor _ d .. | .extr _ d .. | .mul _ d .. | .umulh d .. | .umaddl d .. | .umull d .. | .adcs _ d .. | .sbcs _ d ..
   | .adc _ d .. | .sbc _ d .. | .csel _ d .. | .cselc _ d .. | .csneg _ d .. | .adrSym d _
   | .madd _ d .. | .addImm _ d .. | .subImm _ d .. | .ror _ d .. | .lsr _ d .. | .lsl _ d ..
-  | .rev32 d _ | .rev d _ | .addSp d _ | .movz _ d .. | .movk _ d .. | .ldr _ d .. | .ldrb d ..
+  | .rev32 d _ | .rev d _ | .clz _ d _ | .rbit _ d _ | .addSp d _ | .movz _ d .. | .movk _ d .. | .ldr _ d .. | .ldrb d ..
   | .ldrSp d _ | .umov _ d .. | .pop d => some d
   | _ => none
 

lean/VerifiedGarbage/TCB/AArch64/Isa.lean

@@ -57,7 +57,8 @@ Modelling choices:
   of their data when PSTATE.DIT is 1 (DDI 0487, "About PSTATE.DIT", FEAT_DIT,
   which lists MADD and UBFM, as it does the other modelled data-processing
   instructions that read a register: ADD, SUB, AND, BIC, EOR, ORR, EXTR, REV and
-  MOVK, ADDS, ADCS, SUBS, SBCS, ADC, SBC, CSEL, CSNEG, CCMP, TST and UMULH; the
+  MOVK, ADDS, ADCS, SUBS, SBCS, ADC, SBC, CSEL, CSNEG, CCMP, TST, UMULH, CLZ and
+  RBIT; the
   description of UMADDL, of which UMULL is an alias, says "This instruction is a
   data-independent-time instruction as described in About PSTATE.DIT": DDI
   0602, A64 ISA XML 2025-12, `umaddl.xml`). The code does not set PSTATE.DIT,
@@ -465,6 +466,11 @@ inductive Instr
   /-- `umull xd, wn, wm` (UMULL, the alias of UMADDL with the zero register as
   the addend): `n<31:0> * m<31:0>`, zero-extended -/
   | umull (d n m : Reg)
+  /-- `clz d, n` (CLZ): the number of zero bits of `n` above its highest one
+  bit, `size` if `n` is zero -/
+  | clz (sz : Size) (d n : Reg)
+  /-- `rbit d, n` (RBIT): the bits of `n` in reverse order -/
+  | rbit (sz : Size) (d n : Reg)
   /-- `rev wd, wn` (REV, 32-bit): reverse the bytes of the low 32 bits -/
   | rev32 (d n : Reg)
   /-- `rev xd, xn` (REV, 64-bit): reverse the bytes of all 64 bits -/
@@ -1140,6 +1146,13 @@ def exec : Instr → State → Option State
     some (s.write .x d (s.gpr a + (s.read .w n).setWidth 64 * (s.read .w m).setWidth 64))
   | .umull d n m, s =>
     some (s.write .x d ((s.read .w n).setWidth 64 * (s.read .w m).setWidth 64))
+  -- DDI 0487 C6.2, CLZ: result = CountLeadingZeroBits(X[n, datasize]);
+  -- X[d, datasize] = result<datasize-1:0>. CountLeadingZeroBits(x) is
+  -- N - (HighestSetBit(x) + 1) (J1, shared pseudocode), `BitVec.clz`.
+  | .clz sz d n, s => some (s.write sz d (s.read sz n).clz)
+  -- DDI 0487 C6.2, RBIT: for i = 0 to datasize-1,
+  -- result<datasize-1-i> = operand<i>; `BitVec.reverse`.
+  | .rbit sz d n, s => some (s.write sz d (s.read sz n).reverse)
   | .rev32 d n, s => some (s.write .w d (rev32 (s.read .w n)))
   | .rev d n, s => some (s.write .x d (rev64 (s.read .x n)))
   | .movz sz d imm hw, s =>

lean/VerifiedGarbage/TCB/AArch64/Print.lean

@@ -207,6 +207,8 @@ def Instr.asm : Instr → List String
   | .umulh d n m => [s!"umulh {d.name .x}, {n.name .x}, {m.name .x}"]
   | .umaddl d n m a => [s!"umaddl {d.name .x}, {n.name .w}, {m.name .w}, {a.name .x}"]
   | .umull d n m => [s!"umull {d.name .x}, {n.name .w}, {m.name .w}"]
+  | .clz sz d n => [s!"clz {d.name sz}, {n.name sz}"]
+  | .rbit sz d n => [s!"rbit {d.name sz}, {n.name sz}"]
   | .rev32 d n => [s!"rev {d.name .w}, {n.name .w}"]
   | .rev d n => [s!"rev {d.name .x}, {n.name .x}"]
   | .movz sz d imm hw => [s!"movz {d.name sz}, #{imm.toNat}, lsl #{16 * hw}"]

lean/VerifiedGarbageTest/AArch64Bits.lean

@@ -0,0 +1,61 @@
+import VerifiedGarbage.TCB.AArch64.Print
+import VerifiedGarbage.Proof.Framework.AArch64.Taint
+import VerifiedGarbage.TCB.Axioms
+
+/-! Baseline A64 bit counting and reversal: CLZ and RBIT on both register
+widths, their printing, requirements and taint. -/
+
+namespace VG.Test.AArch64Bits
+open AArch64
+
+private def state (a : BitVec 64) : State where
+  gpr r := if r = .x1 then a else 0
+  sp := 0x1000
+  mem _ := 0
+  rd := []
+  wr := []
+
+private def result (i : Instr) (a : BitVec 64) : Option (BitVec 64) :=
+  (exec i (state a)).map fun s => s.gpr .x0
+
+-- CLZ counts the zero bits above the highest one bit: all of them for zero.
+theorem clz64 :
+    result (.clz .x .x0 .x1) 0 = some 64 ∧
+    result (.clz .x .x0 .x1) 1 = some 63 ∧
+    result (.clz .x .x0 .x1) 0x10 = some 59 ∧
+    result (.clz .x .x0 .x1) 0x8000000000000000 = some 0 := by decide
+
+-- The W form counts within the low 32 bits.
+theorem clz32 :
+    result (.clz .w .x0 .x1) 0x1234567800000001 = some 31 ∧
+    result (.clz .w .x0 .x1) 0xffffffff00000000 = some 32 ∧
+    result (.clz .w .x0 .x1) 0x80000000 = some 0 := by decide
+
+-- RBIT reverses the order of the bits; the W form reverses the low 32 bits and
+-- clears the upper half.
+theorem rbit :
+    result (.rbit .x .x0 .x1) 1 = some 0x8000000000000000 ∧
+    result (.rbit .x .x0 .x1) 0x0123456789abcdef = some 0xf7b3d591e6a2c480 ∧
+    result (.rbit .w .x0 .x1) 0x1234567800000001 = some 0x80000000 ∧
+    result (.rbit .w .x0 .x1) 0x0123456789abcdef = some 0xf7b3d591 := by decide +kernel
+
+-- Counting trailing zeros: RBIT then CLZ.
+theorem ctz :
+    ((exec (.rbit .x .x0 .x1) (state 0x50)).bind (exec (.clz .x .x0 .x0))).map (·.gpr .x0) =
+      some 4 := by decide +kernel
+
+#guard Instr.asm (.clz .x .x0 .x1) == ["clz x0, x1"]
+#guard Instr.asm (.clz .w .x0 .x1) == ["clz w0, w1"]
+#guard Instr.asm (.rbit .x .x0 .x1) == ["rbit x0, x1"]
+#guard Instr.asm (.rbit .w .x0 .x1) == ["rbit w0, w1"]
+
+#guard ([Instr.clz .x .x0 .x1, .clz .w .x0 .x1, .rbit .x .x0 .x1, .rbit .w .x0 .x1].all
+  (·.requires.isEmpty))
+
+-- Their results are as public as their operand.
+#guard ((AArch64.Taint.step (AArch64.Taint.ofRegs [.x1]) (.clz .x .x0 .x1)).map
+  (fun τ => AArch64.Taint.pub τ .x0)) == some true
+#guard ((AArch64.Taint.step (AArch64.Taint.ofRegs [.x2]) (.rbit .x .x0 .x1)).map
+  (fun τ => AArch64.Taint.pub τ .x0)) == some false
+
+end VG.Test.AArch64Bits