read exit_code from arg0 and write it to global state

ceno_zkvm/src/chip_handler.rs

@@ -16,6 +16,13 @@ pub trait GlobalStateRegisterMachineChipOperations<E: ExtensionField> {
     fn state_in(&mut self, pc: Expression<E>, ts: Expression<E>) -> Result<(), ZKVMError>;
 
     fn state_out(&mut self, pc: Expression<E>, ts: Expression<E>) -> Result<(), ZKVMError>;
+
+    fn state_out_with_exit_code(
+        &mut self,
+        pc: Expression<E>,
+        ts: Expression<E>,
+        exit_code: Expression<E>,
+    ) -> Result<(), ZKVMError>;
 }
 
 /// The common representation of a register value.

ceno_zkvm/src/chip_handler/global_state.rs

@@ -12,6 +12,7 @@ impl<'a, E: ExtensionField> GlobalStateRegisterMachineChipOperations<E> for Circ
             Expression::Constant(E::BaseField::from(RAMType::GlobalState as u64)),
             pc,
             ts,
+            0.into(),
         ];
 
         let rlc_record = self.rlc_chip_record(items);
@@ -23,8 +24,25 @@ impl<'a, E: ExtensionField> GlobalStateRegisterMachineChipOperations<E> for Circ
             Expression::Constant(E::BaseField::from(RAMType::GlobalState as u64)),
             pc,
             ts,
+            0.into(),
         ];
         let rlc_record = self.rlc_chip_record(items);
         self.write_record(|| "state_out", rlc_record)
     }
+
+    fn state_out_with_exit_code(
+        &mut self,
+        pc: Expression<E>,
+        ts: Expression<E>,
+        exit_code: Expression<E>,
+    ) -> Result<(), ZKVMError> {
+        let items: Vec<Expression<E>> = vec![
+            Expression::Constant(E::BaseField::from(RAMType::GlobalState as u64)),
+            pc,
+            ts,
+            exit_code,
+        ];
+        let rlc_record = self.rlc_chip_record(items);
+        self.write_record(|| "state_out_with_exit_code", rlc_record)
+    }
 }

ceno_zkvm/src/instructions/riscv/ecall/halt.rs

@@ -1,40 +1,67 @@
 use crate::{
+    chip_handler::RegisterChipOperations,
     circuit_builder::CircuitBuilder,
     error::ZKVMError,
+    expression::{ToExpr, WitIn},
     instructions::{
-        riscv::{constants::ECALL_HALT, ecall_insn::EcallInstructionConfig},
+        riscv::{
+            config::{ExprLtConfig, ExprLtInput},
+            constants::{UInt, ECALL_HALT},
+            ecall_insn::EcallInstructionConfig,
+        },
         Instruction,
     },
+    set_val,
     witness::LkMultiplicity,
+    Value,
 };
 use ceno_emul::StepRecord;
 use ff_ext::ExtensionField;
 use std::{marker::PhantomData, mem::MaybeUninit};
 
-pub struct HaltConfig {
+pub struct HaltConfig<E: ExtensionField> {
     ecall_cfg: EcallInstructionConfig,
+    prev_x10_value: UInt<E>,
+    prev_x10_ts: WitIn,
+    lt_x10_cfg: ExprLtConfig,
 }
 
 pub struct HaltCircuit<E>(PhantomData<E>);
 
 impl<E: ExtensionField> Instruction<E> for HaltCircuit<E> {
-    type InstructionConfig = HaltConfig;
+    type InstructionConfig = HaltConfig<E>;
 
     fn name() -> String {
         "ECALL_HALT".into()
     }
 
     fn construct_circuit(cb: &mut CircuitBuilder<E>) -> Result<Self::InstructionConfig, ZKVMError> {
+        let prev_x10_ts = cb.create_witin(|| "prev_x10_ts")?;
+        let prev_x10_value = UInt::new_unchecked(|| "prev_x10_value", cb)?;
+
         let ecall_cfg = EcallInstructionConfig::construct_circuit(
             cb,
             [ECALL_HALT[0].into(), ECALL_HALT[1].into()],
             None,
             Some(0.into()),
+            Some(prev_x10_value.register_expr()),
         )?;
 
-        // TODO: read exit_code from arg1 and write it to global state
+        // read exit_code from arg0 (X10 register) and write it to global state
+        let (_, lt_x10_cfg) = cb.register_read(
+            || "read x10",
+            E::BaseField::from(ceno_emul::CENO_PLATFORM.reg_arg0() as u64),
+            prev_x10_ts.expr(),
+            ecall_cfg.ts.expr(),
+            prev_x10_value.register_expr(),
+        )?;
 
-        Ok(HaltConfig { ecall_cfg })
+        Ok(HaltConfig {
+            ecall_cfg,
+            prev_x10_value,
+            prev_x10_ts,
+            lt_x10_cfg,
+        })
     }
 
     fn assign_instance(
@@ -49,6 +76,22 @@ impl<E: ExtensionField> Instruction<E> for HaltCircuit<E> {
         );
         assert_eq!(step.pc().after.0, 0);
 
+        set_val!(
+            instance,
+            config.prev_x10_ts,
+            step.rs2().unwrap().previous_cycle
+        );
+        config.prev_x10_value.assign_limbs(
+            instance,
+            Value::new_unchecked(step.rs2().unwrap().value).u16_fields(),
+        );
+
+        ExprLtInput {
+            lhs: step.rs2().unwrap().previous_cycle,
+            rhs: step.cycle(),
+        }
+        .assign(instance, &config.lt_x10_cfg, lk_multiplicity);
+
         config
             .ecall_cfg
             .assign_instance::<E>(instance, lk_multiplicity, step)?;

ceno_zkvm/src/instructions/riscv/ecall_insn.rs

@@ -33,6 +33,7 @@ impl EcallInstructionConfig {
         syscall_id: RegisterExpr<E>,
         syscall_ret_value: Option<RegisterExpr<E>>,
         next_pc: Option<Expression<E>>,
+        exit_code: Option<RegisterExpr<E>>,
     ) -> Result<Self, ZKVMError> {
         let pc = cb.create_witin(|| "pc")?;
         let ts = cb.create_witin(|| "cur_ts")?;
@@ -61,9 +62,12 @@ impl EcallInstructionConfig {
             syscall_ret_value.map_or(syscall_id, |v| v),
         )?;
 
-        cb.state_out(
+        cb.state_out_with_exit_code(
             next_pc.map_or(pc.expr() + PC_STEP_SIZE.into(), |next_pc| next_pc),
             ts.expr() + 4.into(),
+            exit_code.map_or(0.into(), |ec| {
+                ec[0].clone() + (Expression::<E>::from(1usize << 16) * ec[1].clone())
+            }),
         )?;
 
         Ok(Self {