How to Model a CPU in Lean

To prove equivalence with a RISC-V CPU, we first need a RISC-V CPU modeled in Lean. Luckily for us, the teams at Galois and Cambridge University have done the hard work for us. Given a Sail specification of a RISC-V CPU, they translated the Sail specification into non-executable Lean code suitable for theorem proving. The details of how to get this code can be found here.

Our starting point will be to understand this transpilation. At its core, a CPU is just a state machine whose state is described by the register file, memory, program counter, control/status registers, and so on¹. Instructions describe the ways the machine can transition between states. Of course, there are also the mechanics of how a state transition physically works. For example, on a silicon chip, the CPU emits an address on a bus, and the memory decoder reads this address, selects the matching memory cell, and drives its stored bits back onto the data bus where the CPU latches them into a register. We will not focus on the mechanics of how state transitions happen on a physical chip or in software emulation. Our concern is what happens before and after a state transition. From here on, we refer to the trusted Sail-to-Lean transpilation as Sail code.

State

The RISC-V state, abbreviated as SailState, is given by:

25abbrev SailState := SequentialState RegisterType trivialChoiceSource

Here, SequentialState describes the full machine state from the Sail model. It is parameterized² by the function RegisterType : Register → Type and a ChoiceSource. The register file and memory are fields of the structure with details described below.

467structure SequentialState (RegisterType : Register → Type) (c : ChoiceSource) where
468  regs : Std.ExtDHashMap Register RegisterType
469  choiceState : c.α
470  mem : Std.ExtHashMap Nat (BitVec 8)
471  tags : Unit
472  cycleCount : Nat
473  sailOutput : Array String

regs (line 468) models the register file with Std.ExtDHashMap Register RegisterType — a hash map where the keys are of type Register and the value type of the register depends on the key. That is, if the key is k : Register, the corresponding value has type RegisterType k. This is what makes it a dependent hash map rather than an ordinary one, and also why the state is declared with parameters RegisterType and ChoiceSource.

Register is an inductive with ~180 variants — one for every named register in the RISC-V spec:

1526inductive Register : Type where
1527  | hart_state
1528  | satp
1529  | PC
1530  | nextPC
1531  | x1 | x2 | x3 | ... | x31   -- general-purpose registers
1532  -- ... CSRs, vector registers, debug state, ~180 variants total ...

RegisterType is the exact function³ that tells the hash map what the type of the value is for each key:

1708abbrev RegisterType : Register → Type
1709  | .hart_state => HartState
1710  | .mhpmcounter => (Vector (BitVec 64) 32)
1711  | .tlb => (Vector (Option TLB_Entry) (2 ^ 6))
1712  | .PC => (BitVec 64)
1713  | .x1 => (BitVec 64)   -- all x1–x31 map to BitVec 64
1714  | .x2 => (BitVec 64)
1715  -- ...

So regs[.PC] has type RegisterType .PC = BitVec 64, while regs[.tlb] has type RegisterType .tlb = Vector (Option TLB_Entry) (2 ^ 6). Each register gets exactly the type it needs. The general-purpose registers x1-x31 are of type BitVec 64.

The mem field models byte-addressable memory as a hash map from Nat addresses to BitVec 8 bytes. The second parameter c has type ChoiceSource, a structure that provides a way to pick default values for primitive types.

291structure ChoiceSource where
292  (α : Type)
293  (nextState : Primitive → α → α)
294  (choose : ∀ p : Primitive, α → p.reflect)

To define SailState, we use trivialChoiceSource, which is defined as:

296def trivialChoiceSource : ChoiceSource where
297  α := Unit
298  nextState _ _ := ()
299  choose p _ :=
300    match p with
301    | .bool => false
302    | .bit => 0
303    | .int => 0
304    | .nat => 0
305    | .string => ""
306    | .fin _ => 0
307    | .bitvector _ => 0

More simply, this says that all default values are the expected ones. Next, we describe how one steps the CPU or performs a state transition.

The Sail Monad

Step transitions are modeled using Lean's built-in error-state monad (EStateM ε σ α), with the error and state types. Under the hood EStateM ε σ α is just a function σ → Result ε σ α that takes a state and returns either .ok v newState or .error e newState, where v : α and e : ε. See the official EStateM documentation for further details.

The Sail library defines a generic version called PreSailM, leaving the register types, choice source, and user-exception type as parameters:

479abbrev PreSailM (RegisterType : Register → Type) (c : ChoiceSource) (ue : Type) :=
480  EStateM (Error ue) (SequentialState RegisterType c)

SailM is just PreSailM instantiated with the user-exception type ue as exception (so the error type ε becomes Error exception) and the state σ as SailState defined above.

1928abbrev SailM := PreSailM RegisterType trivialChoiceSource exception

Notice that SailM fixes the error type ε and state type σ, but says nothing about the result type α. That is deliberate: EStateM ε σ α takes three type arguments, and here we have only supplied the first two. SailM is therefore a function Type → Type. The remaining α is filled in at each use site, giving a different "step computation that yields an α":

SailM ExecutionResult   -- a full instruction step, yielding its execution result
SailM (BitVec 64)       -- a step that reads a register, yielding a 64-bit word
SailM Unit              -- a step run only for its effect on state

So SailM is the "RISC-V step computation" constructor, and SailM α is "a step computation that produces an α". This is the same currying you already use when you write List (a Type → Type) and only later List Nat.

Errors

These variants model all the ways the execution part of the fetch-decode-execute cycle can go wrong. Error is from the Sail library:

351inductive Error (ue : Type) where
352  | Exit
353  | Unreachable
354  | OutOfMemoryRange (n : Nat)
355  | Assertion (s : String)
356  | User (e : ue)

The User constructor wraps a user-exception type ue, which is instantiated with exception from the RISC-V model:

87inductive exception where
88  | Error_not_implemented (_ : String)
89  | Error_internal_error (_ : String)
90  | Error_reserved_behavior (_ : String)

Executing Instructions

With SailM defined above, to execute instructions we must provide two things:

1. The type for $\alpha$. For RISC-V instructions, $\alpha$ will invariably be ExecutionResult. This is the final output after execution.
2. Once $\alpha$ is defined, a full description of the state transition.

Execution Results: The Return Type After Execution

1464inductive ExecutionResult where
1465  | Retire_Success (_ : Unit)
1466  | ExecuteAs (_ : instruction)
1467  | Enter_Wait (_ : WaitReason)
1468  | Illegal_Instruction (_ : Unit)
1469  | Virtual_Instruction (_ : Unit)
1470  | Trap (_ : (Privilege × ctl_result × xlenbits))
1471  | Memory_Exception (_ : (virtaddr × ExceptionType))
1472  | Ext_CSR_Check_Failure (_ : Unit)
1473  | Ext_ControlAddr_Check_Failure (_ : ext_control_addr_error)
1474  | Ext_DataAddr_Check_Failure (_ : ext_data_addr_error)
1475  | Ext_XRET_Priv_Failure (_ : Unit)

In the happy path, instructions return RETIRE_SUCCESS (which is Retire_Success ()). To make this idea more concrete, we look at the actual execution of a native RISC-V instruction, ADDIW. The opcode signature is ADDIW rd, rs1, imm. In pseudocode, it computes:

# ADDIW rd, rs1, imm   (RV64I)
tmp = rs1 + sext(imm)        # 64-bit wrapping add of the sign-extended 12-bit imm
rd  = sext( tmp[31:0] )      # keep the low 32 bits, sign-extend back to 64

That is, the add is done in 64 bits but only the low 32 bits are kept and then sign-extended back to 64. Below is the trusted Sail implementation of the pseudocode above.

71658def execute_ADDIW (imm : (BitVec 12)) (rs1 : regidx) (rd : regidx) : SailM ExecutionResult := do
71659  let result ← do (pure ((← (rX_bits rs1)) + (sign_extend (m := 64) imm))) -- m1
71660  (wX_bits rd (sign_extend (m := 64) (Sail.BitVec.extractLsb result 31 0))) -- m2
71661  (pure RETIRE_SUCCESS) -- m3

Next we try and parse the above code. It might look simple, but under the hood a lot is going, and understanding it is key for understanding our proof tactics later on.

Function Signature

The function signature says the inputs are a 12-bit immediate value imm and general-purpose source and destination registers rs1 and rd, which are of type⁴ regidx.

31inductive regidx where
32  | Regidx (_ : (BitVec (if ( false  : Bool) then 4 else 5)))

To parse the body, we need a little do notation. Lean is a functional language, but CPU execution is inherently stateful, so do notation lets us write imperative-looking code that the Lean elaborator desugars back into ordinary functional code. The mental model is that each line of the block is a SailM α computation — and recall these are just functions: feed one the current state and it either returns .ok with a new value and state, or .error with an error value and state. If a line errors, the computation short-circuits and the whole block returns that error; on success, the returned value is passed along to the next line. Applying this to the whole body, the three do lines desugar into one nested >>= chain:

Parsing The Body

We rewrite the outer do block as

do
  result ← m1
  m2
  m3

Here m1, m2, and m3 are the monadic expressions marked in the code above.

m1 := do (pure ((← (rX_bits rs1)) + (sign_extend (m := 64) imm)))
m2 := wX_bits rd
        (sign_extend (m := 64) (Sail.BitVec.extractLsb result 31 0))
m3 := pure RETIRE_SUCCESS

Following the guide on how to desugar do notation, the first desugaring step gives:

m1 >>= (fun result =>
  do
    m2
    m3
)

A reminder that, a function m : Type → Type can be used as a monad only when Lean has a Monad m instance⁵ for it, which provides pure and bind. The notation a >>= f is notation for bind a f, where a : m α and f : α → m β. For EStateM, read bind as: run a; if it succeeds, pass its returned value to f to produce the next monadic computation. The next desugaring step is as follows:

m1 >>= (fun result =>
  m2 >>= (fun _ =>
    do
      m3
  )
)

The next step is:

m1 >>= (fun result =>
  m2 >>= (fun _ =>
    m3
  )
)

We expand only m2 and m3 first. As m1 is itself a do block, we do expand that next individually, and then give you the combined version.

m1 >>= (fun result =>
  -- m2 expanded
  (wX_bits rd (sign_extend (m := 64) (Sail.BitVec.extractLsb result 31 0))) >>= (fun _ =>
      -- m3 expanded
      pure RETIRE_SUCCESS
    )
)

Now parsing m1

do (pure ((← (rX_bits rs1)) + (sign_extend (m := 64) imm)))

as m1 is a single statement it desugars to

(pure ((← (rX_bits rs1)) + (sign_extend (m := 64) imm)))

The next desugaring step hoists the nested ← out of the argument to pure:

(rX_bits rs1) >>= (fun rs1_value =>
  pure (rs1_value + sign_extend (m := 64) imm)
)

Putting the desugared m1 back into the outer chain gives:

((rX_bits rs1) >>= (fun x =>
  pure (x + sign_extend (m := 64) imm)
)) >>= (fun result =>
  (wX_bits rd (sign_extend (m := 64) (Sail.BitVec.extractLsb result 31 0))) >>= (fun _ =>
      pure RETIRE_SUCCESS
    )
)

1. rX_bits is a monadic computation: given the current state, it returns a 64-bit value (the contents of register rs1)⁶ and threads the state unchanged. So the output is really SailM (BitVec 64).

2. sign_extend just sign-extends the immediate value and returns a 64-bit value.

3. The next line, wX_bits ..., runs only if the first line succeeds all the way through. It is itself another bind: if we successfully write the sign-extended low 32 bits of result (that is, sext(result[31:0])) into the rd key of the register hash map, and that write succeeds

4. the final pure RETIRE_SUCCESS returns RETIRE_SUCCESS without changing the state further.

Thus, wX_bits ... is the only monadic computation that changes state. We refer the reader to the Read and Write Chain appendix for the full code blocks.

Jolt CPU

The Jolt CPU as modeled in Rust is shown below. Most fields have an equivalent copy in the Sail state, so the Jolt state can and should be built on top of SailState with some extra fields. For the purposes of proving expansion equivalence, this is enough⁹. What we really need to model is the virtual registers and the virtual register allocation. Jolt adds virtual registers and extra instructions.

203  #[derive(Clone, Debug)]
204  pub struct Cpu {
205      clock: u64,
206
207      pub(crate) privilege_mode: PrivilegeMode,
208      wfi: bool,
209      pub x: [i64; REGISTER_COUNT as usize], // Missing in Sail
210      #[allow(dead_code)]
211      f: [f64; 32],
212      pub(crate) pc: u64,
213      csr: [u64; CSR_CAPACITY],
214      pub mmu: Mmu,
215      reservation: u64,
216      is_reservation_set: bool,
217      reservation_width: ReservationWidth,
218      _dump_flag: bool,
219      unsigned_data_mask: u64,
220      // pub trace: Vec<Cycle>,
221      pub trace_len: usize, // Missing in Sail
222      executed_instrs: u64, // "real" RV64IMAC cycles
223      active_markers: FnvHashMap<u32, ActiveMarker>,
224      pub vr_allocator: VirtualRegisterAllocator, // Missing in Sail
225      /// Call stack tracking (circular buffer)
226      call_stack: VecDeque<CallFrame>,
227      /// Advice tape for runtime advice system
228      pub advice_tape: AdviceTape, // Not modeled in Lean
229      #[cfg(feature = "field-inline")]
230      pub field_registers: FieldRegisterFile,
231  }

This gives us the following definition of Jolt state in Lean. We do not model the virtual registers as a hash map; instead, we model them as a function from BitVec 7 -> BitVec 64.

27structure SailJoltState where
28  sail : SailState
29  vregs : BitVec 7 → BitVec 64 := fun _ => 0

Note that we have duplication here. The sail field already has the regidx for the general-purpose registers, and the first 32 slots of vregs are also the same. To make proofs easier to manage, we directly use the Sail registers when writing to general-purpose registers, and block off writing to the first 32 registers.

31-- The general-purpose registers are in the Sail hash map already,
32-- so we should never write to vr 0-31.
33def WritableVReg (vr : BitVec 7) : Prop :=
34  ¬ vr.toNat < 32
35
36-- We cannot write to the first 32 registers, as we use xreg for them.
37def writeVReg (vr : BitVec 7) (val : BitVec 64) : JoltMonad Unit :=
38  if vr.toNat < 32 then
39    throw (Error.Assertion "writeVReg: architectural xreg address")
40  else
41    modify fun js => { js with vregs := fun r => if r = vr then val else js.vregs r }

Virtual Registers

Not all virtual registers are the same. The figure below shows the classification of the different registers.

Virtual register allocation is modeled directly in the expansion. The logic used in Jolt is as follows:

• Before an instruction executes, all registers from 40 to 47 are marked free.
• When Jolt asks for a scratch virtual register, we return the lowest-indexed register marked free. For example, at the start, with all scratch registers marked free, the first free register is 40. Once this is returned, 40 is marked used until it is explicitly marked free again. The next free register is then 41, and so on.

We explicitly model this in JoltBytecode/JoltISA/VirtualRegisters.lean, with proof-side facts in JoltBytecode/InstructionEquivalence/ProofSupport/VirtualRegisters.lean, so that expansion proofs can check that the allocator neither runs out of scratch virtual registers nor reuses one that is still live.

Jolt Monad (Stepping the CPU)

We step the Jolt CPU using the same EStateM infrastructure. The only difference between this and the RISC-V CPU is that we use SailJoltState instead of SailState.

abbrev JoltMonad (α : Type) := EStateM (Error exception) SailJoltState α

Next, we define the instructions natively supported by the Jolt CPU: a subset of RISC-V instructions, plus a few new virtual instructions. The source of truth is the Rust instruction definition in tracer/src/instruction/mod.rs. These are modeled in Lean as the inductive type JoltISA.Instr.

30inductive Instr where
31  | NoOp
32  | ADDI (dst : Dst) (src : Src) (imm : BitVec 12)
33  | ANDI (dst : Dst) (src : Src) (imm : BitVec 12)
34  | ORI  (dst : Dst) (src : Src) (imm : BitVec 12)
35  | XORI (dst : Dst) (src : Src) (imm : BitVec 12)
36  | LUI  (dst : Dst) (imm : BitVec 64)
37  | JAL (dst : Dst) (imm : BitVec 21)
38  | BEQ (lhs rhs : Src) (imm : BitVec 13)
39  | ADD  (dst : Dst) (lhs rhs : Src)
40  | SUB  (dst : Dst) (lhs rhs : Src)
41  | MUL  (dst : Dst) (lhs rhs : Src)
42  | VirtualMULI (dst : Dst) (src : Src) (imm : BitVec 64)
43  -- ... ~55 more opcodes, including the Virtual* family ...
44  | VirtualAssertMulUNoOverflow (lhs rhs : Src)
45  | VirtualAssertLTE (lhs rhs : Src)
46  deriving Repr

Note that unlike the RISC-V Cpu, these instructions can read/write to both risc-v registers and virtual registers.

/-- An instruction source operand: either a virtual register or an
architectural Sail register. -/
inductive Src where
  | vreg : VReg → Src
  | xreg : regidx → Src
  deriving Repr

/-- An instruction destination operand: either a virtual register or an
architectural Sail register. -/
inductive Dst where
  | vreg : VReg → Dst
  | xreg : regidx → Dst
  deriving Repr

The inductive type above only lists the different instructions supported by the Jolt ISA. It does not describe how the Jolt state changes when these instructions are executed. We do that next.

Executing Jolt Instructions

In JoltBytecode/JoltISA/Semantics.lean we define execInstr, a function that maps a JoltISA.Instr to a JoltMonad ExecutionResult. That is, for each native or virtual Jolt instruction, it defines the state update function. Remember, JoltMonad is just a function. To execute the instruction, we have to "run" the function with a current state.

30def execInstr : Instr → JoltMonad ExecutionResult
31  | .NoOp =>
32      pure RETIRE_SUCCESS
33  | .ADDI dst src imm => do
34      let x ← readSrc src
35      writeDst dst (x + sign_extend (m := 64) imm)
36      pure RETIRE_SUCCESS
37  | .ANDI dst src imm => do
38      let x ← readSrc src
39      writeDst dst (x &&& sign_extend (m := 64) imm)
40      pure RETIRE_SUCCESS
41  | .ORI dst src imm => do
42      let x ← readSrc src
43      writeDst dst (x ||| sign_extend (m := 64) imm)
44      pure RETIRE_SUCCESS
45  -- ... remaining native and virtual instruction cases ...
46  | .VirtualAssertLTE lhs rhs => do
47      let x ← readSrc lhs
48      let y ← readSrc rhs
49      if x.toNat ≤ y.toNat then
50        pure RETIRE_SUCCESS
51      else
52        throw (Error.Assertion "VirtualAssertLTE")

You might ask: where did this execution logic come from? For now, it is hand-translated from the Rust source code. For each instruction, the process is as follows:

Find the corresponding exec block and model the Rust code in Lean. For example, here are the Lean semantics and Rust semantics side by side for ADDI.

| .ADDI dst src imm => do
    let x ← readSrc src
    writeDst dst (x + sign_extend (m := 64) imm)
    pure RETIRE_SUCCESS

19    fn exec(&self, cpu: &mut Cpu, _: &mut <ADDI as RISCVInstruction>::RAMAccess) {
20        cpu.write_register(
21            self.operands.rd as usize,
22            cpu.sign_extend(
23                cpu.x[self.operands.rs1 as usize].wrapping_add(normalize_imm(self.operands.imm)),
24            ),
25        );
26    }

As a result of the translation above, some natively supported instructions, such as ADDI, have both a Jolt execution and a trusted Sail execution. We want them to agree on how they change Sail state. This gives us an equivalence theorem for every Jolt instruction that is natively supported on a RISC-V CPU. Continuing with the ADDI example, it looks like the following:

69/-- Main native `ADDI` equivalence statement. -/
70def addiInstrEqSailStatement
71    (imm : BitVec 12)
72    (rs1 rd : regidx)
73    (js : SailJoltState)
74    (_h : UnarySourceReadWithLinkedCSRs rs1 js) : Prop :=
75  System.systemProjectResult
76    ((JoltISA.execInstr (.ADDI (.xreg rd) (.xreg rs1) imm)).run js) =
77    ((execute_ITYPE imm rs1 rd iop.ADDI).run js.sail)

The statement above translates to the following mathematical statement.

The exact assumption bundle used for ADDI is UnarySourceReadWithLinkedCSRs rs1 js, which bundles two sets of assumptions together.

The structure UnarySourceReadAssumptions just says we can read rs1. As the register file is a hash map, this is a Lean implementation detail: we have to say that the key exists in the hash map. The other structure is LinkedCSRRegisterAssumptions.

29/-- Assumptions for an instruction that reads one architectural source register. -/
30structure UnarySourceReadAssumptions (rs1 : regidx) (js : SailJoltState) where
31  rs1_val : BitVec 64
32  rs1_read : rX_bits rs1 js.sail = .ok rs1_val js.sail
33
34/-- Persistent CSR virtual registers agree with Sail's architectural CSR state. -/
35private structure LinkedCSRRegisterAssumptions (js : SailJoltState) where
36  mstatus_matches : Assumptions.MstatusVRegMatchesSail js
37  mtvec_matches : Assumptions.MtvecVRegMatchesSail js
38  mscratch_matches : Assumptions.MscratchVRegMatchesSail js
39  mepc_matches : Assumptions.MepcVRegMatchesSail js
40  mcause_matches : Assumptions.McauseVRegMatchesSail js
41  mtval_matches : Assumptions.MtvalVRegMatchesSail js
42
43/-- One source-register read plus persistent CSR virtual-register agreement. -/
44structure UnarySourceReadWithLinkedCSRs
45    (rs1 : regidx) (js : SailJoltState)
46    extends UnarySourceReadAssumptions rs1 js,
47      LinkedCSRRegisterAssumptions js

It simply says that the virtual registers are equal to the Sail registers for the corresponding control/status registers.

The only thing left to explain is systemProjectResult. systemProjectResult removes the Jolt-specific state and keeps only the sail state, which can be compared with the RISC-V CPU. Before we do that, we must map the virtual registers that represent control/status registers in Jolt to the Sail control/status registers. Without this step, a program could arbitrarily write to these registers, and dropping vregs would make the Sail states appear to agree. We need to tell Lean that those virtual registers represent the same values as the Sail registers.

65/-- Overlay Jolt's persistent virtual CSR registers onto the generated Sail CSR
66register keys.
67
68This is only a proof projection: the Rust-faithful Jolt programs still write
69Jolt virtual registers, and the Sail specification still reads/writes generated
70Sail registers. -/
71def systemProject (js : SailJoltState) : SailState :=
72  { js.sail with
73    regs :=
74      ((((((js.sail.regs
75        |>.insert Register.mtvec (js.vregs JoltISA.trapHandlerVReg))
76        |>.insert Register.mscratch (js.vregs JoltISA.mscratchVReg))
77        |>.insert Register.mepc (js.vregs JoltISA.mepcVReg))
78        |>.insert Register.mcause (js.vregs JoltISA.mcauseVReg))
79        |>.insert Register.mtval (js.vregs JoltISA.mtvalVReg))
80        |>.insert Register.mstatus (js.vregs JoltISA.mstatusVReg)) }
81
82/-- Project a Jolt run result through `systemProject`, preserving the result
83value and error shape while materializing virtual CSRs in the Sail state. -/
84def systemProjectResult
85    (r : EStateM.Result (Error exception) SailJoltState α) :
86    EStateM.Result (Error exception) SailState α :=
87  match r with
88  | .ok a js' => .ok a (systemProject js')
89  | .error e js' => .error e (systemProject js')

And that's it: for every natively supported instruction, we have to prove that this statement is true.

Expansions

We've handled natively supported RISC-V instructions. What about instructions like LW, where the RISC-V instruction is handled as a sequence of Jolt instructions rather than as one native execInstr step? These instructions are modeled by expansion programs. An expansion program is a value of the inductive type Program, made from composing Jolt instructions, as defined below.

192/-- Structured Jolt bytecode programs.
193
194`instr i next` means: execute `i`; if it retires successfully, continue with
195`next`; otherwise return the non-retire result immediately. -/
196inductive Program where
197  | done (result : ExecutionResult)
198  | instr (instr : Instr) (next : Program)
199  deriving Repr

So a program is either finished, with a final ExecutionResult, or it is one instruction followed by the rest of the program. In plain English, a program is just the composition of instruction steps: run this instruction, then run the next instruction, and so on until the program ends. For example, a tiny program with two instructions looks like this:

def tinyProgram (rd rs1 : regidx) : Program :=
  .instr (.ADDI (.xreg rd) (.xreg rs1) 1) <|
  .instr (.ANDI (.xreg rd) (.xreg rd) 255) <|
  .done RETIRE_SUCCESS

The <| symbol is just Lean sugar that lets us write this composition without a pile of parentheses. Without <|, the same program is:

def tinyProgram (rd rs1 : regidx) : Program :=
  .instr (.ADDI (.xreg rd) (.xreg rs1) 1)
    (.instr (.ANDI (.xreg rd) (.xreg rd) 255)
      (.done RETIRE_SUCCESS))

This represents the expansion:

ADDI x[rd], x[rs1], 1
ANDI x[rd], x[rd], 255

For LW, the Lean model of the Jolt expansion lives in JoltBytecode/JoltISA/Expansions/Load.lean as JoltISA.lwProgram:

116/-- Rust's RV64 `LW::inline_sequence`. -/
117def lwProgram (imm : BitVec 12) (rs1 rd : regidx) : Program :=
118  let v0 := loadV0For rd
119  let v1 := loadV1For rd
120  let tmp := loadInlineTmpFor rd
121  let dst := loadDstFor rd
122  .instr (.VirtualAssertWordAlignment rs1 imm (ExceptionType.E_Load_Addr_Align ())) <|
123  .instr (.ADDI (.vreg v0) (.xreg rs1) imm) <|
124  .instr (.ANDI (.vreg v1) (.vreg v0) (-8 : BitVec 12)) <|
125  .instr (.LD .normal (.vreg v1) (.vreg v1) 0) <|
126  slliBlock (.vreg v0) (.vreg v0) (3 : BitVec 6) <|
127  srlBlock (.vreg v1) (.vreg v1) (.vreg v0) tmp <|
128  .instr (.VirtualSignExtendWord dst (.vreg v1)) <|
129  .done RETIRE_SUCCESS

We make use of some helper functions like loadV0For to model virtual register allocation.

36/-- Rust source `rd = x0` rewrite destination for side-effecting load
37expansions. -/
38def loadDstFor (rd : regidx) : Dst :=
39  sideEffectingRdZeroDst rd
40
41/-- Rust load `v0`, shifted when `rd = x0` consumes the first temporary. -/
42def loadV0For (rd : regidx) : VReg :=
43  if isX0 rd then inlineTmp1 else loadV0
44
45/-- Rust load `v1`, shifted when `rd = x0` consumes the first temporary. -/
46def loadV1For (rd : regidx) : VReg :=
47  if isX0 rd then inlineTmp2 else loadV1
48
49/-- Recursive load scratch, shifted when `rd = x0` consumes the first
50temporary. -/
51def loadInlineTmpFor (rd : regidx) : VReg :=
52  if isX0 rd then inlineTmp3 else loadInlineTmp

The reason they look somewhat complex is that we also need to model the side effect of the case where rd=x0. In that case Jolt swaps rd for the first free virtual register, and this means the downstream virtual register allocation changes by an offset of 1. These helpers help represent that in a clean manner. So far, this only says how to build a program. We still need to say how to run one.

71/-- Rust `SLLI::inline_sequence`: multiply by the immediate power of two. -/
72def slliBlock (dst : Dst) (src : Src) (shamt : BitVec 6) (tail : Program) : Program :=
73  .instr (.VirtualMULI dst src (slliMultiplier shamt)) tail
74
75/-- Rust `SRL::inline_sequence`: compute the right-shift bitmask in a scratch
76virtual register, then run `VirtualSRL` with that bitmask. -/
77def srlBlock (dst : Dst) (value shift : Src) (scratch : VReg)
78    (tail : Program) : Program :=
79  .instr (.VirtualShiftRightBitmask (.vreg scratch) shift) <|
80  .instr (.VirtualSRL dst value (.vreg scratch)) tail

The Rust expansion is given below, and we can see that it lines up with the Lean lwProgram.

 7pub(in crate::expand) fn expand_lw(
 8    instruction: &SourceInstructionRow,
 9) -> Result<ExpandedInstructionSequence, ExpansionError> {
10    let mut asm = ExpansionBuilder::new(*instruction);
11    let v0 = asm.allocate()?;
12    let v1 = asm.allocate()?;
13
14    // RAM is accessed at doubleword granularity here. The word alignment
15    // assertion is still required by the source `LW` semantics.
16    asm.expand_address(
17        SourceInstructionKind::VirtualAssertWordAlignment,
18        reg(rs1(instruction)?),
19        instruction.operands.imm,
20    );
21    asm.expand_i(
22        SourceInstructionKind::ADDI,
23        v0.operand(),
24        reg(rs1(instruction)?),
25        format_i_imm(instruction.operands.imm),
26    );
27    // v1 = containing doubleword address, v0 = byte offset within it.
28    asm.expand_i(
29        SourceInstructionKind::ANDI,
30        v1.operand(),
31        v0.operand(),
32        format_i_imm(-8),
33    );
34    asm.expand_i(SourceInstructionKind::LD, v1.operand(), v1.operand(), 0);
35    asm.expand_i(SourceInstructionKind::SLLI, v0.operand(), v0.operand(), 3);
36    asm.expand_r(
37        SourceInstructionKind::SRL,
38        v1.operand(),
39        v1.operand(),
40        v0.operand(),
41    );
42    asm.expand_i(
43        SourceInstructionKind::VirtualSignExtendWord(
44            jolt_riscv::instructions::VirtualSignExtendWord(()),
45        ),
46        reg(rd(instruction)?),
47        v1.operand(),
48        0,
49    );
50    asm.release(v0);
51    asm.release(v1);
52
53    asm.finalize()
54}

Executing Expanded Programs

Execution of a program is much like the execution of a single Jolt instruction. Given a program, instead of a single Jolt instruction, we return a state transition function. If the program is already complete, then execution just returns the result. Otherwise, it takes the instruction and the rest of the program, steps the current Jolt instruction, and on successful execution, executes the rest.

413def execProgram : Program → JoltMonad ExecutionResult
414  | .done result => pure result
415  | .instr instr rest => do
416      match ← execInstr instr with
417      | .Retire_Success () => execProgram rest
418      -- If any instruction does not complete successfully,
419      -- return Result.ok result s' and do not run the rest of the program.
420      | result => pure result

If the instruction traps, errors, or otherwise returns a non-retire result, execution stops and that result is returned. We will also prove that this program leaves the Sail state exactly as it would be after running the trusted Sail LW.

398def lwProgramEqSailStatement (imm : BitVec 12)
399    (rs1 rd : regidx)
400    (js : SailJoltState)
401    (_h : LoadProgramEqSailAssumptions imm rs1 js) : Prop :=
402    System.systemProjectResult
403      ((JoltISA.execProgram (JoltISA.lwProgram imm rs1 rd)).run js) =
404    (execute_LOAD imm rs1 rd false 4).run js.sail

The statement above translates to the following mathematical statement.

In upcoming posts, we will dive into the details of assumptions in theorem statements and how to go about proving these statements.

Appendices

Read and Write Chain

696def rX_bits (app_0 : regidx) : SailM (BitVec 64) := do
697  let .Regidx i := app_0
698  (rX (Regno (BitVec.toNatInt i)))
699
700def wX_bits (typ_0 : regidx) (data : (BitVec 64)) : SailM Unit := do
701  let .Regidx i : regidx := typ_0
702  (wX (Regno (BitVec.toNatInt i)) data)

618def rX (app_0 : regno) : SailM (BitVec 64) := do
619  let .Regno r := app_0
620  let v ← (( do
621    match r with
622    | 0 => (pure zero_reg)
623    | 1 => readReg x1
624    | 2 => readReg x2
625    -- ... x3 through x30 ...
626    | _ => readReg x31 ) : SailM regtype )
627  (pure (regval_from_reg v))

656def wX (typ_0 : regno) (in_v : (BitVec 64)) : SailM Unit := do
657  let .Regno r : regno := typ_0
658  let v := (regval_into_reg in_v)
659  match r with
660  | 0 => (pure ())
661  | 1 => writeReg x1 v
662  | 2 => writeReg x2 v
663  -- ... x3 through x30 ...
664  | _ => writeReg x31 v
665  if ((r != 0) : Bool)
666  then (xreg_write_callback (Regidx (to_bits (l := 5) r)) in_v)
667  else (pure ())

530def writeReg (r : Register) (v : RegisterType r) : PreSailM RegisterType c ue PUnit :=
531  modify fun s => { s with regs := s.regs.insert r v }

534def readReg (r : Register) : PreSailM RegisterType c ue (RegisterType r) := do
535  let .some s := (← get).regs.get? r
536    | throw .Unreachable
537  pure s

CSR Assumptions

107structure MstatusVRegMatchesSail (js : SailJoltState) : Prop where
108  value_eq :
109    js.sail.regs.get? Register.mstatus =
110      some (js.vregs JoltISA.mstatusVReg)
111
112structure MtvecVRegMatchesSail (js : SailJoltState) : Prop where
113  value_eq :
114    js.sail.regs.get? Register.mtvec =
115      some (js.vregs JoltISA.trapHandlerVReg)
116
117structure MscratchVRegMatchesSail (js : SailJoltState) : Prop where
118  value_eq :
119    js.sail.regs.get? Register.mscratch =
120      some (js.vregs JoltISA.mscratchVReg)
121
122structure MepcVRegMatchesSail (js : SailJoltState) : Prop where
123  value_eq :
124    js.sail.regs.get? Register.mepc =
125      some (js.vregs JoltISA.mepcVReg)
126
127structure McauseVRegMatchesSail (js : SailJoltState) : Prop where
128  value_eq :
129    js.sail.regs.get? Register.mcause =
130      some (js.vregs JoltISA.mcauseVReg)
131
132structure MtvalVRegMatchesSail (js : SailJoltState) : Prop where
133  value_eq :
134    js.sail.regs.get? Register.mtval =
135      some (js.vregs JoltISA.mtvalVReg)