End Goal: Prove Jolt Things (via Monads)

The end goal is to prove in Lean, that Jolt Bytecode expansions do not alter the RISCV guest program that the Jolt user submits. Part of this process involves defining what a RISCV-CPU is and does in Lean. Thankfully, people at Sail + Galois + Cambridge have done part of the work for us, and following the instructions described here, we get automatic transpilation of (non-executable) code from sail to lean. 1

Jumping ahead, a little, we know that the Jolt CPU, does not have ADDW in its instruction set. Instead it replaces every occurrence of ADDW in the guest program with the following sequence defined below (sign extending by word a normal addition).

"Proving equivalence" amounts to proving a theorem of this sort

theorem jolt_addw_eq_riscv_addw (rs2 rs1 rd : regidx) :
    execute_RTYPEW rs2 rs1 rd ropw.ADDW = jolt_addw rs2 rs1 rd := by 
    sorry

The sorry is what we are meant to fill out as the proof. The part between : and := is our theorem statement. The LHS of the theorem statement we get using the interface/API the sail to lean transpilation above offers us. Loosely, we are saying run the following code and match on branch ADDW.


def execute_RTYPEW (rs2 : regidx) (rs1 : regidx) (rd : regidx) (op : ropw) : SailM ExecutionResult := do
  let rs1_val ← do (pure (Sail.BitVec.extractLsb (← (rX_bits rs1)) 31 0))
  let rs2_val ← do (pure (Sail.BitVec.extractLsb (← (rX_bits rs2)) 31 0))
  let result : (BitVec 32) :=
    match op with
    | ADDW => (rs1_val + rs2_val)
    | SUBW => (rs1_val - rs2_val)
    | SLLW => (shift_bits_left rs1_val (Sail.BitVec.extractLsb rs2_val 4 0))
    | SRLW => (shift_bits_right rs1_val (Sail.BitVec.extractLsb rs2_val 4 0))
    | SRAW => (shift_bits_right_arith rs1_val (Sail.BitVec.extractLsb rs2_val 4 0))
  (wX_bits rd (sign_extend (m := 64) result))
  (pure RETIRE_SUCCESS)

The RHS is something we wrote (but will also be transpiled from the Jolt rust source eventually)


-- WARNING: This definition only works when rd is a real RISC-V register (regidx).
-- It will NOT work for Jolt virtual registers (indices ≥ 32). For the general
-- case we need the JoltSailState extension with liftSail.
def virtual_sign_extend_word (rd : regidx) : SailM Unit := do
  let v ← rX_bits rd
  wX_bits rd (sign_extend (m := 64) (Sail.BitVec.extractLsb v 31 0))

-- Jolt's ADDW decomposition: ADD then VirtualSignExtendWord
def jolt_addw (rs2 rs1 rd : regidx) : SailM ExecutionResult := do
  let _ ← execute_RTYPE rs2 rs1 rd rop.ADD
  virtual_sign_extend_word rd
  pure RETIRE_SUCCESS

Both sides of the theorem have type SailM ExecutionResult as their output.

ExecutionResult itself is2:

inductive ExecutionResult where
  | Retire_Success (_ : Unit)
  | ExecuteAs (_ : instruction)
  | Enter_Wait (_ : WaitReason)
  | Illegal_Instruction (_ : Unit)
  | Virtual_Instruction (_ : Unit)
  | Trap (_ : (Privilege × ctl_result × xlenbits))
  | Memory_Exception (_ : (virtaddr × ExceptionType))
  | Ext_CSR_Check_Failure (_ : Unit)
  | Ext_ControlAddr_Check_Failure (_ : ext_control_addr_error)
  | Ext_DataAddr_Check_Failure (_ : ext_data_addr_error)
  | Ext_XRET_Priv_Failure (_ : Unit)

And SailM is not a plain value -- it is a monadic computation. Under the hood, SailM is defined as:

abbrev SailM := EStateM (Error exception) (SequentialState RegisterType trivialChoiceSource)

That is, SailM is Lean's built-in error-and-state monad. See the Lean 4 source for EStateM. A term of type SailM ExecutionResult is not an ExecutionResult -- but in Rust speak, and a slight deviation from the truth, you can imagine that its a very very complicated Result<ExecutionResult, Error> instead3.

2

This is an inductive type -- if you're not familiar with the term, a simplified way to read it is: it's an enum that lists everything that can happen when you step a CPU. The instruction either succeeds (Retire_Success), gets re-interpreted as another instruction (ExecuteAs), waits, traps, hits a memory exception, and so on.

3

Option would also work as an analogy, but Result/Except is closer since SailM can fail with error information, not just None. In greater detail, EStateM also threads mutable CPU state (registers, memory) through every step -- so it's really State -> Result<(ExecutionResult, State), Error>.

If the above sentences made absolutely no sense - then good, we are both starting at the same place. In functional programming, when we wrap real data in this Result business we often refer to the act of doing this as monadic programming. But the point is, if we want to prove the theorem we are after, we need to understand how to prove things with monads in Lean. And so the rest of this post, is about "what the feck is a monad"!!

1

This translation appears to be the industry standard for proving things about the RISCV-CPU, and also has been adopted by other proof systems as shown here and here.

What These Notes Are For

The mini-goal of this series of notes is to somehow get to the bottom of how to use the mvcgen tactic in Lean, which is advertised as a way to reason about imperative code in Lean4. Why? As this will let us prove things about functions that involve monads! So the hope is -- we understand Markus' post and its inner workings, upgrade his post the latest version of Lean, and then tackle our proof5.

5

We already tried proving things without mvcgen with the help of Claudius Caesar. After 1 hour of running the results were not good. In fact it took 50 minutes to prove that the RISCV add instruction is commutative i.e a + b = b + a in this CPU state framework. So we really do need to understand the nuts and bolts of what is going on.

Lean is a functional programming language, and although it lets the user type imperative looking code using do blocks, under the hood it transforms the same block of code into functional blocks using machinery called monads. But what the feck is a monad? I cannot for the life of me seem to find a clean definition6. After about 48 hours of reading and scouring through different blogs and lecture notes7, here's my opinionated8 take on what a monad is. When I say functions from here on -- I don't mean a method or a function in a programming language, but a function as introduced in an introductory textbook for mathematics.

(My Imprecise Monad Definition Using Rust As An Example)

Given some computation $C$ that cannot be expressed as a composition of functions, a monad is a type that comes with the "machinery" needed to write $C$ as a composition of functions. It's a circular way of saying: when a function's domain or co-domain involves monadic types, we can always compose these functions.

More precisely: a monad is an inductive type m together with two functions — pure : α → m α (which wraps a term of type α into m α) and bind : m α → (α → m β) → m β (which takes a term of type m α and a function α → m β, extracts the α, and feeds it to the function, thereby enabling composition).

I am aware this precise definition could be more confusing — but a concrete example might help. Say we have two functions in Rust:

fn f(x: i32) -> Option<i32> { if x > 0 { Some(x + 1) } else { None } }
fn g(x: i32) -> Option<i32> { if x < 100 { Some(x * 2) } else { None } }

Now normally we can't compose them as g(f(5)) because f returns Option<i32>, not i32. Non-composeable computation. But Option is a monad. Why? Because Option comes with and_then (which we refer to as bind) and Some (which we refer to as pure).

So we can now compose the above function anyway, using this extended machinery.

f(5).and_then(|x| g(x))   // Some(12)
f(-1).and_then(|x| g(x))  // None — f returned None, so and_then short-circuits without calling g

and_then extracts the i32 from Some and feeds it to the next function: pure is Some — it wraps a plain i32 into Option<i32>.

9

I've come to the conclusion that there is no easy 3 sentence definition.

Now chances are that this definition is even more confusing, and conveys about as much information as the definition I criticised9. I apologise if this is the case, and refer the reader to this starter post Monads: What Is?, which tells the story in longer form. We'll start with a simple computation that is composable, and then change it so it becomes non-composable. Then we will make it composable again -- and everything that we did to make it composable will be a monad.

If that does not convince you, we will take the monad starter pack from the Lean documentation and re-write it in terms of the same composition notation from the above post.

At this point, we've either convinced ourselves that the above definition somewhat makes sense, or I still don't understand what the feck a monad is, and in that case you should stop reading. Were we to have understood monads, we go into some Lean4 specific machinery that makes writing monadic code easier here and here.

With this monad crash course done, hopefully we have the background to understand what is going on in Markus' post.

7

Such as these: You Could Have Invented Monads (Dan Piponi), Monads for Functional Programming (Philip Wadler), Monads: Programmer's Definition (Bartosz Milewski), and FP in Lean Ch. 4.

10

Functions here are mathematical functions in that it's a list of tuples mapping a domain to a co-domain.

8

Possibly wrong opinion.

6

By clean definition, I mean how we have a nice definition say for an injective function.This excerpt in italics is from hovering over Monads in the Functional programming book. Perhaps, to functional programming experts, this is a much better definition, but I was no cleverer about what Monads do from this. Monads are an abstraction of sequential control flow and side effects used in functional programming. Monads allow both sequencing of effects and data-dependent effects: the values that result from an early step may influence the effects carried out in a later step.