Automatic Bytecode Expansion Extraction

In this post we go over how we extracted bytecode expansion from Jolt written in Rust into Lean. The goal was to re-use as much as the rust jolt production codebase as possible, and write a thin easily checkable wrapper around it. All code can be found here

The figure summarises the pipeline.

We first describe the Jolt expansion pipeline written in rust. The Jolt tracer takes as input some elf file with RISC-V instructions, and the tracer translates each source instruction into an intermediate data structure, SourceInstruction, stored in memory of the host machine running Jolt. This data structure is passed into the Jolt expansion pipeline which returns a vector of Jolt instructions, another intermediate data structure stored in memory. Now this intermediate data structure of Jolt instructions described below contains “nearly” all the information needed to translate into Lean.

// Jolt Instructions 
pub struct JoltInstructionRow {
    pub instruction_kind: JoltInstructionKind,
    pub address: usize,
    pub operands: NormalizedOperands,
    pub virtual_sequence_remaining: Option<u16>,
    pub is_first_in_sequence: bool,
    pub is_compressed: bool,
}

pub struct NormalizedOperands {
    pub rs1: Option<u8>,
    pub rs2: Option<u8>,
    pub rd: Option<u8>,
    pub imm: i128,
}

Indeed, observe that the LD RISC-V Instruction, for example the kind field stores the name of the instruction.

LD                 v41, v41, 0

The operands contain input values. These values tells us whether the registers are virtual or real. The other fields are not even used to get valid Lean translation

.instr
    (.LD
      .normal
      (.vreg (BitVec.ofNat 7 41))
      (.vreg (BitVec.ofNat 7 41))
      (0 : BitVec 12)) <|

Thus, if someone were to hand our translator this intermediate representation, we could write a minimal printer that just translates this data into Lean code.

So the game now is to obtain this intermediate vector of Jolt instructions for every expandable RISC-V instruction. Towards this, we fix a dummy value for rd, rs1, rs2 and, imm for each expandable RISC-V instruction^[1]. Then, we invoke the production Jolt expander for every expandable instruction with the dummy instruction we just created. The jolt expander treats this instruction as it would for any instruction in the guest program, and therefore expands it into a vector of Jolt Instructions. The translator just takes this intermediate data structure of Jolt instructions, and prints strings that pass the Lean typechecker. The only artefact that is brand new is this printer, which if it passes typechecking is likely fine. We argue that this extraction process hardly changes the trust surface in any meaningful way. Thus, the LB instruction which expands to

VirtualAlignAddr   v41, x2, 37
LD                 v41, v41, 0
VirtualWindowMaskB v40, x2, 37
VirtualPextSigned  x1, v41, v40

is simply translated into strings that Lean understands.

  .instr
    (.VirtualAlignAddr
      (.vreg (BitVec.ofNat 7 41))
      (.xreg rs1)
      imm) <|

  .instr
    (.LD
      .normal
      (.vreg (BitVec.ofNat 7 41))
      (.vreg (BitVec.ofNat 7 41))
      (0 : BitVec 12)) <|

  .instr
    (.VirtualWindowMaskB
      (.vreg (BitVec.ofNat 7 40))
      (.xreg rs1)
      imm) <|

  .instr
    (.VirtualPextSigned
      (.xreg rd)
      (.vreg (BitVec.ofNat 7 41))
      (.vreg (BitVec.ofNat 7 40))) <|
  .done RETIRE_SUCCESS

Observe that in excerpt above we said that the translator “nearly” has all information needed to perform translation. A small issue is that the vector of Jolt instructions stores all immediate operands as raw values. Thus, looking at the just the instruction, one cannot tell if value 37 is a hard-coded offset in the expansion algorithm or it was an intermediate value passed by the guest program. The two cases affect how the translator prints the expansion. To resolve this ambiguity, we create two dummy instructions with different immediate operands, say imm:= 13 and imm:=5. For the immediate operands that are hard-coded by the expansions, they should remain the same in both runs. In contrast the user specified immediate values should change. We track this information across two runs, and this lets the translator know for which constants to use the string imm, and for which constants to translate raw integer values into strings.

1. An instruction just ignores any operands value it does not need. ↩