Enable JIT validation for MEL

crates/jit/src/lib.rs

@@ -85,6 +85,10 @@ pub struct GlobalState {
     pub last_block_hash: Bytes32,
     #[clap(long, value_parser = cli_parsing::parse_hex)]
     pub last_send_root: Bytes32,
+    #[clap(long, value_parser = cli_parsing::parse_hex, default_value = "0x0000000000000000000000000000000000000000000000000000000000000000")]
+    pub mel_state_hash: Bytes32,
+    #[clap(long, value_parser = cli_parsing::parse_hex, default_value = "0x0000000000000000000000000000000000000000000000000000000000000000")]
+    pub mel_msg_hash: Bytes32,
     #[clap(long, default_value = "0")]
     pub inbox_position: u64,
     #[clap(long, default_value = "0")]
@@ -96,6 +100,8 @@ impl From<validation::GoGlobalState> for GlobalState {
         Self {
             last_block_hash: state.block_hash,
             last_send_root: state.send_root,
+            mel_state_hash: state.mel_state_hash,
+            mel_msg_hash: state.mel_msg_hash,
             inbox_position: state.batch,
             position_within_message: state.pos_in_batch,
         }
@@ -107,6 +113,8 @@ impl From<GlobalState> for validation::GoGlobalState {
         Self {
             block_hash: state.last_block_hash,
             send_root: state.last_send_root,
+            mel_state_hash: state.mel_state_hash,
+            mel_msg_hash: state.mel_msg_hash,
             batch: state.inbox_position,
             pos_in_batch: state.position_within_message,
         }
@@ -166,6 +174,8 @@ fn run_instance(
         new_state: GlobalState {
             last_block_hash: Bytes32(env.input.large_globals[0]),
             last_send_root: Bytes32(env.input.large_globals[1]),
+            mel_state_hash: Bytes32(env.input.large_globals[2]),
+            mel_msg_hash: Bytes32(env.input.large_globals[3]),
             inbox_position: env.input.small_globals[0],
             position_within_message: env.input.small_globals[1],
         },

crates/jit/src/machine.rs

@@ -144,6 +144,7 @@ fn imports(store: &mut Store, func_env: &FunctionEnv<WasmEnv>) -> wasmer::Import
             },
             "resolveTypedPreimage" => func!(wavmio::resolve_typed_preimage),
             "validateCertificate" => func!(wavmio::validate_certificate),
+            "getEndParentChainBlockHash" => func!(wavmio::get_end_parent_chain_block_hash),
         },
         "wasi_snapshot_preview1" => {
             "proc_exit" => func!(wasip1_stub::proc_exit),
@@ -288,6 +289,8 @@ fn prepare_env_from_files(env: &mut WasmEnv, input: &LocalInput) -> Result<()> {
         large_globals: [
             input.old_state.last_block_hash.0,
             input.old_state.last_send_root.0,
+            [0u8; 32],
+            [0u8; 32],
         ],
         ..Default::default()
     };

crates/jit/src/wavmio.rs

@@ -6,7 +6,7 @@ use crate::{
     machine::{Escape, MaybeEscape, WasmEnv, WasmEnvMut},
 };
 use arbutil::Color;
-use caller_env::GuestPtr;
+use caller_env::{GuestPtr, MemAccess};
 use std::{
     io,
     io::{BufReader, BufWriter, ErrorKind},
@@ -158,6 +158,13 @@ pub fn resolve_preimage_impl(
     .map_err(Escape::HostIO)
 }
 
+pub fn get_end_parent_chain_block_hash(mut env: WasmEnvMut, out_ptr: GuestPtr) -> MaybeEscape {
+    let (mut mem, exec) = env.jit_env();
+    ready_hostio(exec)?;
+    mem.write_slice(out_ptr, &exec.input.end_parent_chain_block_hash);
+    Ok(())
+}
+
 pub fn validate_certificate(
     mut env: WasmEnvMut,
     preimage_type: u8,

crates/validation/src/lib.rs

@@ -23,11 +23,12 @@ pub type Preimages = BTreeMap<u8, BTreeMap<[u8; 32], Vec<u8>>>;
 )]
 pub struct ValidationInput {
     pub small_globals: [u64; 2],
-    pub large_globals: [[u8; 32]; 2],
+    pub large_globals: [[u8; 32]; 4],
     pub preimages: Preimages,
     pub sequencer_messages: Inbox,
     pub delayed_messages: Inbox,
     pub module_asms: BTreeMap<[u8; 32], Vec<u8>>,
+    pub end_parent_chain_block_hash: [u8; 32],
 }
 
 impl ValidationInput {
@@ -69,11 +70,17 @@ impl ValidationInput {
 
         Ok(Self {
             small_globals: [req.start_state.batch, req.start_state.pos_in_batch],
-            large_globals: [req.start_state.block_hash.0, req.start_state.send_root.0],
+            large_globals: [
+                req.start_state.block_hash.0,
+                req.start_state.send_root.0,
+                req.start_state.mel_state_hash.0,
+                req.start_state.mel_msg_hash.0,
+            ],
             preimages,
             sequencer_messages,
             delayed_messages,
             module_asms,
+            end_parent_chain_block_hash: req.end_parent_chain_block_hash.0,
         })
     }
 
@@ -113,6 +120,12 @@ pub struct GoGlobalState {
     pub block_hash: Bytes32,
     #[serde(with = "As::<DisplayFromStr>")]
     pub send_root: Bytes32,
+    #[serde(with = "As::<DisplayFromStr>")]
+    #[serde(rename = "MELStateHash")]
+    pub mel_state_hash: Bytes32,
+    #[serde(with = "As::<DisplayFromStr>")]
+    #[serde(rename = "MELMsgHash")]
+    pub mel_msg_hash: Bytes32,
     pub batch: u64,
     pub pos_in_batch: u64,
 }
@@ -173,6 +186,8 @@ pub struct ValidationRequest {
     #[serde(with = "As::<HashMap<DisplayFromStr, HashMap<DisplayFromStr, Base64>>>")]
     pub user_wasms: HashMap<String, HashMap<Bytes32, UserWasm>>,
     pub debug_chain: bool,
+    #[serde(with = "As::<DisplayFromStr>")]
+    pub end_parent_chain_block_hash: Bytes32,
     #[serde(rename = "max-user-wasmSize", default)]
     pub max_user_wasm_size: u64,
 }
@@ -228,11 +243,14 @@ mod tests {
             start_state: GoGlobalState {
                 block_hash,
                 send_root,
+                mel_state_hash: Bytes32::default(),
+                mel_msg_hash: Bytes32::default(),
                 batch: 100,
                 pos_in_batch: 200,
             },
             user_wasms,
             debug_chain: false,
+            end_parent_chain_block_hash: Bytes32::default(),
             max_user_wasm_size: 0,
         }
     }

crates/validation/src/transfer/receiver.rs

@@ -16,6 +16,8 @@ pub fn receive_validation_input(reader: &mut impl Read) -> IOResult<ValidationIn
     let preimages = receive_preimages(reader)?;
     let module_asms = receive_module_asms(reader)?;
     ensure_readiness(reader)?;
+    let mut end_parent_chain_block_hash = [0u8; 32];
+    reader.read_exact(&mut end_parent_chain_block_hash)?;
 
     Ok(ValidationInput {
         small_globals,
@@ -24,6 +26,7 @@ pub fn receive_validation_input(reader: &mut impl Read) -> IOResult<ValidationIn
         sequencer_messages,
         delayed_messages,
         module_asms,
+        end_parent_chain_block_hash,
     })
 }
 
@@ -36,6 +39,8 @@ pub fn receive_response(reader: &mut impl Read) -> IOResult<Result<(GoGlobalStat
                 pos_in_batch: small[1],
                 block_hash: arbutil::Bytes32(large[0]),
                 send_root: arbutil::Bytes32(large[1]),
+                mel_state_hash: arbutil::Bytes32(large[2]),
+                mel_msg_hash: arbutil::Bytes32(large[3]),
             };
             let memory_used = read_u64(reader)?;
             Ok(Ok((new_state, memory_used)))
@@ -49,11 +54,13 @@ pub fn receive_response(reader: &mut impl Read) -> IOResult<Result<(GoGlobalStat
     }
 }
 
-fn receive_globals(reader: &mut impl Read) -> IOResult<([u64; 2], [[u8; 32]; 2])> {
+fn receive_globals(reader: &mut impl Read) -> IOResult<([u64; 2], [[u8; 32]; 4])> {
     let small_globals = [read_u64(reader)?, read_u64(reader)?];
-    let mut large_globals = [[0u8; 32]; 2];
+    let mut large_globals = [[0u8; 32]; 4];
     reader.read_exact(&mut large_globals[0])?;
     reader.read_exact(&mut large_globals[1])?;
+    reader.read_exact(&mut large_globals[2])?;
+    reader.read_exact(&mut large_globals[3])?;
     Ok((small_globals, large_globals))
 }
 

crates/validation/src/transfer/sender.rs

@@ -12,7 +12,8 @@ pub fn send_validation_input(writer: &mut impl Write, input: &ValidationInput) -
     send_inbox(writer, &input.delayed_messages)?;
     send_preimages(writer, &input.preimages)?;
     send_module_asms(writer, &input.module_asms)?;
-    finish_sending(writer)
+    finish_sending(writer)?;
+    writer.write_all(&input.end_parent_chain_block_hash)
 }
 
 pub fn send_successful_response(
@@ -24,7 +25,12 @@ pub fn send_successful_response(
     send_globals(
         writer,
         &[new_state.batch, new_state.pos_in_batch],
-        &[new_state.block_hash.0, new_state.send_root.0],
+        &[
+            new_state.block_hash.0,
+            new_state.send_root.0,
+            new_state.mel_state_hash.0,
+            new_state.mel_msg_hash.0,
+        ],
     )?;
     write_u64(writer, memory_used)
 }
@@ -37,12 +43,14 @@ pub fn send_failure_response(writer: &mut impl Write, error_message: &str) -> IO
 fn send_globals(
     writer: &mut impl Write,
     small_globals: &[u64; 2],
-    large_globals: &[[u8; 32]; 2],
+    large_globals: &[[u8; 32]; 4],
 ) -> IOResult<()> {
     write_u64(writer, small_globals[0])?;
     write_u64(writer, small_globals[1])?;
     writer.write_all(&large_globals[0])?;
-    writer.write_all(&large_globals[1])
+    writer.write_all(&large_globals[1])?;
+    writer.write_all(&large_globals[2])?;
+    writer.write_all(&large_globals[3])
 }
 
 fn send_inbox(writer: &mut impl Write, inbox: &Inbox) -> IOResult<()> {

crates/validation/src/transfer/tests.rs

@@ -14,6 +14,8 @@ fn transfer_successful_response() -> Result<(), Box<dyn std::error::Error>> {
     let new_state = GoGlobalState {
         block_hash: Bytes32::from([1u8; 32]),
         send_root: Bytes32::from([2u8; 32]),
+        mel_state_hash: Bytes32::from([5u8; 32]),
+        mel_msg_hash: Bytes32::from([6u8; 32]),
         batch: 42,
         pos_in_batch: 7,
     };
@@ -49,7 +51,7 @@ fn transfer_failure_response() -> Result<(), Box<dyn std::error::Error>> {
 fn transfer_validation_input() -> Result<(), Box<dyn std::error::Error>> {
     let input = ValidationInput {
         small_globals: [42, 7],
-        large_globals: [[1u8; 32], [2u8; 32]],
+        large_globals: [[1u8; 32], [2u8; 32], [5u8; 32], [6u8; 32]],
 
         sequencer_messages: BTreeMap::from([
             (10, vec![1, 2, 3]),
@@ -74,6 +76,7 @@ fn transfer_validation_input() -> Result<(), Box<dyn std::error::Error>> {
         ]),
 
         module_asms: BTreeMap::from([([3u8; 32], vec![20, 21, 22]), ([4u8; 32], vec![30, 31, 32])]),
+        end_parent_chain_block_hash: [7u8; 32],
     };
 
     let (mut reader, mut writer) = pipe()?;

system_tests/common_test.go

@@ -315,6 +315,7 @@ type NodeBuilder struct {
 	isSequencer                 bool
 	takeOwnership               bool
 	withL1                      bool
+	useJit                      bool
 	defaultStateScheme          string
 	addresses                   *chaininfo.RollupAddresses
 	l3Addresses                 *chaininfo.RollupAddresses
@@ -395,6 +396,7 @@ func NewNodeBuilder(ctxIn context.Context) *NodeBuilder {
 func (b *NodeBuilder) DefaultConfig(t *testing.T, withL1 bool) *NodeBuilder {
 	// most used values across current tests are set here as default
 	b.withL1 = withL1
+	b.useJit = true
 	b.parallelise = true
 	b.deployBold = true
 	b.takeOwnership = true
@@ -812,6 +814,7 @@ func buildOnParentChain(
 	initMessage *arbostypes.ParsedInitMessage,
 	addresses *chaininfo.RollupAddresses,
 	trieNoAsyncFlush bool,
+	useJit bool,
 ) *TestClient {
 	if parentChainTestClient == nil {
 		t.Fatal("must build parent chain before building chain")
@@ -844,7 +847,7 @@ func buildOnParentChain(
 		validatorTxOptsPtr = &validatorTxOpts
 	}
 
-	AddValNodeIfNeeded(t, ctx, nodeConfig, true, "", valnodeConfig.Wasm.RootPath)
+	AddValNodeIfNeeded(t, ctx, nodeConfig, useJit, "", valnodeConfig.Wasm.RootPath)
 
 	execConfigFetcher := NewCommonConfigFetcher(execConfig)
 	execNode, err := gethexec.CreateExecutionNode(ctx, chainTestClient.Stack, executionDB, blockchain, parentChainTestClient.Client, execConfigFetcher, parentChainId, 0)
@@ -939,6 +942,7 @@ func (b *NodeBuilder) BuildL3OnL2(t *testing.T) func() {
 		b.l3InitMessage,
 		b.l3Addresses,
 		b.TrieNoAsyncFlush,
+		b.useJit,
 	)
 
 	return func() {
@@ -972,6 +976,7 @@ func (b *NodeBuilder) BuildL2OnL1(t *testing.T) func() {
 		b.initMessage,
 		b.addresses,
 		b.TrieNoAsyncFlush,
+		b.useJit,
 	)
 
 	if b.nodeConfig.MessageExtraction.Enable {
@@ -1045,7 +1050,7 @@ func (b *NodeBuilder) BuildL2(t *testing.T) func() {
 	}
 	b.L2 = NewTestClient(b.ctx)
 
-	AddValNodeIfNeeded(t, b.ctx, b.nodeConfig, true, "", b.valnodeConfig.Wasm.RootPath)
+	AddValNodeIfNeeded(t, b.ctx, b.nodeConfig, b.useJit, "", b.valnodeConfig.Wasm.RootPath)
 
 	var executionDB ethdb.Database
 	var consensusDB ethdb.Database
@@ -1766,7 +1771,7 @@ func AddMELValNode(t *testing.T, ctx context.Context, nodeConfig *arbnode.Config
 
 func AddValNodeIfNeeded(t *testing.T, ctx context.Context, nodeConfig *arbnode.Config, useJit bool, redisURL string, wasmRootDir string) {
 	if nodeConfig.MELValidator.Enable {
-		AddMELValNode(t, ctx, nodeConfig, false, redisURL, wasmRootDir) // TODO: currently useJit is false remove this once jit validation for MEL is enabled
+		AddMELValNode(t, ctx, nodeConfig, useJit, redisURL, wasmRootDir)
 		return
 	}
 	if !nodeConfig.ValidatorRequired() || nodeConfig.BlockValidator.ValidationServerConfigs[0].URL != "" {

system_tests/message_extraction_layer_validation_test.go

@@ -16,10 +16,16 @@ import (
 )
 
 func TestValidationPostMEL(t *testing.T) {
+	t.Run("TestValidationPostMEL-in-jit-mode", func(t *testing.T) { testValidationPostMEL(t, true) })
+	t.Run("TestValidationPostMEL-in-arbitrator-mode", func(t *testing.T) { testValidationPostMEL(t, false) })
+}
+
+func testValidationPostMEL(t *testing.T, useJit bool) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
 	builder := NewNodeBuilder(ctx).DefaultConfig(t, true)
+	builder.useJit = useJit
 	builder.L2Info.GenerateAccount("User2")
 	builder.nodeConfig.BatchPoster.Post4844Blobs = true
 	builder.nodeConfig.BatchPoster.IgnoreBlobPrice = true
@@ -64,13 +70,22 @@ func TestValidationPostMEL(t *testing.T) {
 	}
 }
 
-func TestValidationPostMELReorgHandle(t *testing.T) {
+func TestValidationPostMELReorgHandleInJitMode(t *testing.T) {
+	testValidationPostMELReorgHandle(t, true)
+}
+
+func TestValidationPostMELReorgHandleInArbitratorMode(t *testing.T) {
+	testValidationPostMELReorgHandle(t, false)
+}
+
+func testValidationPostMELReorgHandle(t *testing.T, useJit bool) {
 	logHandler := testhelpers.InitTestLog(t, log.LvlInfo)
 
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
 	builder := NewNodeBuilder(ctx).DefaultConfig(t, true)
+	builder.useJit = useJit
 	builder.nodeConfig.MessageExtraction.Enable = true
 	builder.nodeConfig.MessageExtraction.RetryInterval = 100 * time.Millisecond
 	builder.nodeConfig.BatchPoster.MaxDelay = time.Hour     // set high max-delay so we can test the delay buffer