feat(constraints): prove public RAM output constraint (26)

JoltConstraints/Constraints/IOSetupFrame.lean

@@ -0,0 +1,249 @@
+import JoltConstraints.Constraints.RamReadData
+
+/-!
+Execution never changes the device fields fixed when Rust builds its emulator:
+the memory layout and both advice buffers. Sail steps do not touch the device,
+device stores only change `outputs` and `panic`, and HostIO only appends to the
+advice tape. The trace-level result carries this from the program's initial
+state to the final recorded state.
+-/
+
+set_option autoImplicit false
+
+open Sail PreSail LeanRV64D.Functions
+
+namespace JoltIOSetupFrame
+
+open JoltISA
+
+/-- The device fields that Rust fixes in `create_emulator` and never updates.
+Rust: https://github.com/a16z/jolt/blob/754fc88214936801a7d8a2260d9fc73478be4cbd/tracer/src/lib.rs#L401-L404 -/
+def Same (a b : JoltIOState) : Prop :=
+  b.layout = a.layout ∧ b.trustedAdvice = a.trustedAdvice ∧
+    b.untrustedAdvice = a.untrustedAdvice
+
+theorem Same.refl (a : JoltIOState) : Same a a := ⟨rfl, rfl, rfl⟩
+
+theorem Same.trans {a b c : JoltIOState} (hab : Same a b) (hbc : Same b c) : Same a c :=
+  ⟨hbc.1.trans hab.1, hbc.2.1.trans hab.2.1, hbc.2.2.trans hab.2.2⟩
+
+def Preserves {α : Type} (m : JoltMonad α) : Prop :=
+  ∀ (s t : SailJoltState) (v : α), m s = .ok v t → Same s.io t.io
+
+theorem pure_rule {α : Type} (v : α) : Preserves (pure v) := by
+  intro s t x h
+  cases h
+  exact Same.refl _
+
+theorem bind_rule {α β : Type} {m : JoltMonad α} {f : α → JoltMonad β}
+    (hm : Preserves m) (hf : ∀ x, Preserves (f x)) : Preserves (m >>= f) := by
+  intro s t v h
+  cases hr : m s with
+  | error e s' => simp only [bind, EStateM.bind, hr] at h; cases h
+  | ok x s' =>
+    simp only [bind, EStateM.bind, hr] at h
+    exact (hm s s' x hr).trans (hf x s' t v h)
+
+theorem lift_rule {α : Type} (m : SailM α) : Preserves (liftSail m) := by
+  intro s t v h
+  cases hr : m s.sail with
+  | error e s' => simp only [liftSail, hr] at h; cases h
+  | ok x s' =>
+    simp only [liftSail, hr] at h
+    cases h
+    exact Same.refl _
+
+theorem throw_rule {α : Type} (e : Error exception) :
+    Preserves (throw e : JoltMonad α) := by
+  intro s t x h
+  cases h
+
+theorem get_rule : Preserves (get : JoltMonad SailJoltState) := by
+  intro s t x h
+  cases h
+  exact Same.refl _
+
+theorem modify_rule (f : SailJoltState → SailJoltState)
+    (hf : ∀ s, (f s).io = s.io) : Preserves (modify f : JoltMonad Unit) := by
+  intro s t x h
+  cases h
+  rw [hf s]
+  exact Same.refl _
+
+theorem read_rule (src : Src) : Preserves (readSrc src) := by
+  cases src with
+  | vreg vr =>
+    intro s t v h
+    simp only [readSrc, readVReg, bind, EStateM.bind, get, getThe, MonadStateOf.get,
+      EStateM.get, pure, EStateM.pure] at h
+    cases h
+    exact Same.refl _
+  | xreg rs => exact lift_rule _
+
+theorem write_rule (dst : Dst) (value : BitVec 64) : Preserves (writeDst dst value) := by
+  cases dst with
+  | vreg vr =>
+    by_cases hw : vr.toNat < 32
+    · simp only [writeDst, writeVReg, hw, ↓reduceIte]
+      exact throw_rule _
+    · simp only [writeDst, writeVReg, hw, ↓reduceIte]
+      exact modify_rule _ (fun _ => rfl)
+  | xreg rd => exact lift_rule _
+
+theorem storeDeviceByte_same (io io' : JoltIOState) (address : Nat) (value : BitVec 8)
+    (h : storeDeviceByte? io address value = some io') : Same io io' := by
+  unfold storeDeviceByte? at h
+  split_ifs at h <;> cases h <;> exact ⟨rfl, rfl, rfl⟩
+
+theorem storeDeviceWord_same (io io' : JoltIOState) (address value : BitVec 64)
+    (h : storeDeviceWord? io address value = some io') : Same io io' := by
+  unfold storeDeviceWord? at h
+  suffices hfold : ∀ (l : List Nat) (start : Option JoltIOState),
+      (∀ d, start = some d → Same io d) →
+      ∀ d, l.foldl (fun current k => current.bind fun device =>
+        storeDeviceByte? device (address.toNat + k) (value.extractLsb' (8 * k) 8)) start =
+          some d → Same io d from
+    hfold _ _ (fun d hd => by cases hd; exact Same.refl _) _ h
+  intro l
+  induction l with
+  | nil => intro start hstart d hd; exact hstart d hd
+  | cons k l ih =>
+    intro start hstart d hd
+    refine ih _ ?_ d hd
+    intro d' hd'
+    cases hs : start with
+    | none => rw [hs] at hd'; cases hd'
+    | some e =>
+      rw [hs] at hd'
+      exact (hstart e hs).trans (storeDeviceByte_same _ _ _ _ hd')
+
+theorem readMemoryWord_rule (address : BitVec 64) :
+    Preserves (readMemoryWord address) := by
+  intro s t v h
+  unfold readMemoryWord at h
+  split_ifs at h
+  · split at h <;> (cases h; exact Same.refl _)
+  · exact lift_rule _ s t v h
+
+theorem writeMemoryWord_rule (address value : BitVec 64) :
+    Preserves (writeMemoryWord address value) := by
+  intro s t v h
+  unfold writeMemoryWord at h
+  split_ifs at h
+  · split at h
+    · rename_i io hio
+      cases h
+      exact storeDeviceWord_same _ _ _ _ hio
+    · cases h
+      exact Same.refl _
+  · exact lift_rule _ s t v h
+
+theorem readMemoryByte_rule (address : BitVec 64) :
+    Preserves (readMemoryByte address) := by
+  intro s t v h
+  unfold readMemoryByte at h
+  dsimp only at h
+  split_ifs at h
+  · cases hb : deviceByte? s.io address.toNat with
+    | none => simp only [hb] at h; cases h
+    | some byte => simp only [hb] at h; cases h; exact Same.refl _
+  · exact lift_rule _ s t v h
+
+theorem readHostBytes_rule
+    (overflowChecks incrementAfterLast : Bool) {width : Nat}
+    (pointer : BitVec width) (n : Nat) (bytes : Array (BitVec 8)) :
+    Preserves (readHostBytes readMemoryByte overflowChecks incrementAfterLast pointer n bytes) := by
+  induction n generalizing pointer bytes with
+  | zero => exact pure_rule _
+  | succ n ih =>
+    dsimp only [readHostBytes]
+    refine bind_rule (readMemoryByte_rule _) (fun result => ?_)
+    cases result with
+    | Err e => exact pure_rule _
+    | Ok byte =>
+      split_ifs with h
+      · exact bind_rule (throw_rule _) (fun _ => ih _ _)
+      · exact bind_rule (pure_rule _) (fun _ => ih _ _)
+
+theorem execHostIO_rule : Preserves execHostIO := by
+  unfold execHostIO execHostIOWith
+  refine bind_rule get_rule (fun config => ?_)
+  dsimp only
+  split
+  · exact pure_rule _
+  · refine bind_rule (pure_rule _) (fun _ => ?_)
+    refine bind_rule (read_rule _) (fun callId => ?_)
+    repeat' first
+      | exact readHostBytes_rule _ _ _ _ _
+      | exact read_rule _
+      | exact pure_rule _
+      | exact throw_rule _
+      | exact modify_rule _ (fun _ => rfl)
+      | split
+      | refine bind_rule ?_ (fun x => ?_)
+
+macro "jolt_io_setup_auto" : tactic => `(tactic|
+  repeat' first
+  | exact read_rule _
+  | exact write_rule _ _
+  | exact pure_rule _
+  | exact throw_rule _
+  | exact get_rule
+  | exact lift_rule _
+  | exact readMemoryWord_rule _
+  | exact writeMemoryWord_rule _ _
+  | exact execHostIO_rule
+  | split
+  | refine bind_rule ?_ (fun x => ?_))
+
+/-- Every Jolt ISA instruction preserves the device setup fields. -/
+theorem execInstr_rule (instr : Instr) : Preserves (execInstr instr) := by
+  cases instr
+  case VirtualAdviceLoad dst byteCount =>
+    intro s t v h
+    simp only [execInstr] at h
+    cases hr : readAdviceTape s.adviceTape byteCount with
+    | none => simp only [hr] at h; cases h
+    | some pair =>
+      obtain ⟨value, tape⟩ := pair
+      simp only [hr] at h
+      exact (bind_rule (write_rule dst value) (fun _ => pure_rule _))
+        { s with adviceTape := tape } t v h
+  all_goals simp only [execInstr]
+  all_goals jolt_io_setup_auto
+
+end JoltIOSetupFrame
+
+namespace JoltConstraints
+
+open JoltIOSetupFrame
+
+/-- Every recorded pre-state has the program's initial device setup. -/
+theorem trace_preState_ioSame {program : JoltProgram} (trace : JoltTrace program)
+    (i : Nat) (hi : i < trace.rows.size) :
+    Same program.initialState.io (getElem trace.rows i hi).preState.io := by
+  induction i with
+  | zero =>
+    rw [trace.startsAtInitial (by omega)]
+    exact Same.refl _
+  | succ i ih =>
+    have hprev := ih (by omega)
+    have hexec := execInstr_rule _ _ _ _ (getElem trace.rows i (by omega)).executes
+    have hlink := trace.linked i (by omega) hi
+    dsimp only at hlink
+    rw [hlink]
+    split
+    · exact hprev.trans hexec
+    · exact hprev.trans hexec
+
+/-- The final recorded state has the program's initial device setup. -/
+theorem finalTraceState_ioSame {program : JoltProgram} (trace : JoltTrace program) :
+    Same program.initialState.io (HonestWitness.finalTraceState trace).io := by
+  unfold HonestWitness.finalTraceState
+  split
+  · rename_i nonempty
+    exact (trace_preState_ioSame trace _ _).trans
+      (execInstr_rule _ _ _ _ (getElem trace.rows (trace.rows.size - 1) _).executes)
+  · exact Same.refl _
+
+end JoltConstraints

JoltConstraints/Constraints/MemoryLayout.lean

@@ -0,0 +1,136 @@
+import JoltConstraints.Constraints.RamReadData
+
+/-!
+Rust's `MemoryLayout::new`, restricted to the advice and input regions, and a
+proof that every device it lays out satisfies `JoltIOState.AdviceBelowInput`.
+
+Rust: https://github.com/a16z/jolt/blob/754fc88214936801a7d8a2260d9fc73478be4cbd/common/src/jolt_device.rs#L346-L438
+
+Arithmetic is on `Nat`. Each Rust `checked_*` call followed by `expect` becomes
+a guard that returns `none` where Rust panics. Everything after `input_start`
+in `new`, and its two power-of-two assertions on the advice sizes, only add
+further panics and leave these fields unchanged; they are omitted, so
+`adviceInputLayout` succeeds on at least every configuration Rust accepts.
+-/
+
+set_option autoImplicit false
+
+namespace JoltConstraints.RustMemoryLayout
+
+/-- The `MemoryConfig` fields that determine the advice and input regions. -/
+structure MemoryConfig where
+  maxTrustedAdviceSize : Nat
+  maxUntrustedAdviceSize : Nat
+  maxInputSize : Nat
+  maxOutputSize : Nat
+
+/-- Advice and input placement computed by `MemoryLayout::new`. -/
+structure AdviceInputLayout where
+  trustedAdviceStart : Nat
+  trustedAdviceEnd : Nat
+  untrustedAdviceStart : Nat
+  untrustedAdviceEnd : Nat
+  inputStart : Nat
+
+/-- Rust `u64::checked_add`. -/
+def checkedAdd (a b : Nat) : Option Nat :=
+  if a + b < 2 ^ 64 then some (a + b) else none
+
+/-- Rust `align_up(val, 8)`, which panics on overflow. -/
+def alignUp8 (val : Nat) : Option Nat :=
+  if val % 8 = 0 then some val else checkedAdd val (8 - val % 8)
+
+/-- `MemoryLayout::new` up to `input_start`. The larger (or equal) advice region
+goes first, at `RAM_START_ADDRESS - io_bytes`; the other starts where it ends;
+`input_start` is the larger end. -/
+def adviceInputLayout (config : MemoryConfig) : Option AdviceInputLayout := do
+  let trusted ← alignUp8 config.maxTrustedAdviceSize
+  let untrusted ← alignUp8 config.maxUntrustedAdviceSize
+  let input ← alignUp8 config.maxInputSize
+  let output ← alignUp8 config.maxOutputSize
+  let ioRegionBytes ← (checkedAdd input trusted).bind fun s =>
+    (checkedAdd s untrusted).bind fun s => (checkedAdd s output).bind fun s => checkedAdd s 16
+  let ioBytes := (ioRegionBytes / 8).nextPowerOfTwo * 8
+  if 2 ^ 64 ≤ ioBytes ∨ JoltISA.ramStartAddress < ioBytes then none else
+  let first := JoltISA.ramStartAddress - ioBytes
+  if untrusted ≤ trusted then
+    let trustedEnd ← checkedAdd first trusted
+    let untrustedEnd ← checkedAdd trustedEnd untrusted
+    pure { trustedAdviceStart := first, trustedAdviceEnd := trustedEnd,
+           untrustedAdviceStart := trustedEnd, untrustedAdviceEnd := untrustedEnd,
+           inputStart := max untrustedEnd trustedEnd }
+  else
+    let untrustedEnd ← checkedAdd first untrusted
+    let trustedEnd ← checkedAdd untrustedEnd trusted
+    pure { trustedAdviceStart := untrustedEnd, trustedAdviceEnd := trustedEnd,
+           untrustedAdviceStart := first, untrustedAdviceEnd := untrustedEnd,
+           inputStart := max untrustedEnd trustedEnd }
+
+theorem checkedAdd_eq_some {a b c : Nat} (h : checkedAdd a b = some c) : c = a + b := by
+  unfold checkedAdd at h
+  split at h <;> simp_all
+
+theorem alignUp8_eq_some {val aligned : Nat} (h : alignUp8 val = some aligned) :
+    val ≤ aligned ∧ aligned % 8 = 0 := by
+  unfold alignUp8 at h
+  split at h
+  · simp_all
+  · have := checkedAdd_eq_some h
+    omega
+
+/-- The arithmetic facts about a successful layout: each advice region has its
+aligned size, the two regions are adjacent, and `input_start` is the larger end. -/
+theorem adviceInputLayout_spec {config : MemoryConfig} {layout : AdviceInputLayout}
+    (h : adviceInputLayout config = some layout) :
+    ∃ trusted untrusted,
+      config.maxTrustedAdviceSize ≤ trusted ∧ trusted % 8 = 0 ∧
+      config.maxUntrustedAdviceSize ≤ untrusted ∧ untrusted % 8 = 0 ∧
+      layout.trustedAdviceEnd = layout.trustedAdviceStart + trusted ∧
+      layout.untrustedAdviceEnd = layout.untrustedAdviceStart + untrusted ∧
+      (layout.untrustedAdviceStart = layout.trustedAdviceEnd ∨
+        layout.trustedAdviceStart = layout.untrustedAdviceEnd) ∧
+      layout.inputStart = max layout.untrustedAdviceEnd layout.trustedAdviceEnd := by
+  unfold adviceInputLayout at h
+  simp only [Option.bind_eq_bind, Option.bind_eq_some_iff] at h
+  obtain ⟨trusted, htrusted, untrusted, huntrusted, _, _, _, _, _, _, h⟩ := h
+  obtain ⟨hT, hT8⟩ := alignUp8_eq_some htrusted
+  obtain ⟨hU, hU8⟩ := alignUp8_eq_some huntrusted
+  refine ⟨trusted, untrusted, hT, hT8, hU, hU8, ?_⟩
+  split at h
+  · simp at h
+  split at h
+  · simp only [Option.bind_eq_some_iff, Option.pure_def, Option.some.injEq] at h
+    obtain ⟨trustedEnd, htEnd, untrustedEnd, huEnd, rfl⟩ := h
+    have := checkedAdd_eq_some htEnd
+    have := checkedAdd_eq_some huEnd
+    simp only
+    simp only [true_or, and_true]
+    omega
+  · simp only [Option.bind_eq_some_iff, Option.pure_def, Option.some.injEq] at h
+    obtain ⟨untrustedEnd, huEnd, trustedEnd, htEnd, rfl⟩ := h
+    have := checkedAdd_eq_some htEnd
+    have := checkedAdd_eq_some huEnd
+    simp only
+    simp only [or_true, and_true]
+    omega
+
+end JoltConstraints.RustMemoryLayout
+
+open JoltConstraints JoltConstraints.RustMemoryLayout in
+/-- Every device Rust builds satisfies `AdviceBelowInput`: its layout comes from
+`MemoryLayout::new`, and `create_emulator` rejects advice longer than the
+configured maximum before execution starts.
+Rust: https://github.com/a16z/jolt/blob/754fc88214936801a7d8a2260d9fc73478be4cbd/tracer/src/lib.rs#L382-L393 -/
+theorem JoltIOState.AdviceBelowInput.of_memoryLayout
+    (io : JoltIOState) (config : MemoryConfig) (layout : AdviceInputLayout)
+    (built : adviceInputLayout config = some layout)
+    (trustedStart : io.layout.trustedAdvice.1.toNat = layout.trustedAdviceStart)
+    (untrustedStart : io.layout.untrustedAdvice.1.toNat = layout.untrustedAdviceStart)
+    (inputStart : io.layout.input.1.toNat = layout.inputStart)
+    (trustedLength : io.trustedAdvice.size ≤ config.maxTrustedAdviceSize)
+    (untrustedLength : io.untrustedAdvice.size ≤ config.maxUntrustedAdviceSize) :
+    io.AdviceBelowInput := by
+  obtain ⟨trusted, untrusted, hT, hT8, hU, hU8, htEnd, huEnd, hadj, hinput⟩ :=
+    adviceInputLayout_spec built
+  refine ⟨?_, ?_, ?_⟩ <;>
+    simp only [ramLowestAddress, trustedStart, untrustedStart, inputStart] <;> omega