Use instead of restoring side table for Loop

crates/interpreter/src/exec.rs

@@ -21,6 +21,7 @@ use crate::module::*;
 use crate::side_table::SideTableEntry;
 use crate::syntax::*;
 use crate::toctou::*;
+use crate::util::*;
 use crate::*;
 
 pub const MEMORY_ALIGN: usize = 16;
@@ -759,20 +760,13 @@ impl<'m> Thread<'m> {
 
     fn step(mut self, store: &mut Store<'m>) -> Result<ThreadResult<'m>, Error> {
         use Instr::*;
-        let saved = self.parser.save();
         let inst_id = self.frame().inst_id;
         let inst = &mut store.insts[inst_id];
         match self.parser.parse_instr().into_ok() {
             Unreachable => return Err(trap()),
             Nop => (),
             Block(b) => self.push_label(self.blocktype(inst, &b), LabelKind::Block),
-            Loop(b) => {
-                let side_table = self.frame().side_table;
-                self.push_label(
-                    self.blocktype(inst, &b),
-                    LabelKind::Loop(LoopState { parser_data: saved, side_table }),
-                )
-            }
+            Loop(b) => self.push_label(self.blocktype(inst, &b), LabelKind::Loop),
             If(b) => match self.pop_value().unwrap_i32() {
                 0 => {
                     self.take_jump(0);
@@ -988,11 +982,11 @@ impl<'m> Thread<'m> {
         self.frames.last_mut().unwrap()
     }
 
-    fn labels(&mut self) -> &mut Vec<Label<'m>> {
+    fn labels(&mut self) -> &mut Vec<Label> {
         &mut self.frame().labels
     }
 
-    fn label(&mut self) -> &mut Label<'m> {
+    fn label(&mut self) -> &mut Label {
         self.labels().last_mut().unwrap()
     }
 
@@ -1042,12 +1036,12 @@ impl<'m> Thread<'m> {
         self.values().split_off(len)
     }
 
-    fn push_label(&mut self, type_: FuncType<'m>, kind: LabelKind<'m>) {
+    fn push_label(&mut self, type_: FuncType<'m>, kind: LabelKind) {
         let arity = match kind {
             LabelKind::Block | LabelKind::If => type_.results.len(),
-            LabelKind::Loop(_) => type_.params.len(),
+            LabelKind::Loop => type_.params.len(),
         };
-        let label = Label { arity, kind, values_cnt: type_.params.len() };
+        let label = Label { arity, values_cnt: type_.params.len() };
         self.label().values_cnt -= label.values_cnt;
         self.labels().push(label);
     }
@@ -1059,17 +1053,11 @@ impl<'m> Thread<'m> {
         }
         let frame = self.frame();
         let values_cnt: usize = frame.labels[i ..].iter().map(|label| label.values_cnt).sum();
-        let Label { arity, kind, .. } = frame.labels.drain(i ..).next().unwrap();
+        let Label { arity, .. } = frame.labels.drain(i ..).next().unwrap();
         let values_len = self.values().len();
         self.values().drain(values_len - values_cnt .. values_len - arity);
         self.label().values_cnt += arity;
-        match kind {
-            LabelKind::Loop(state) => unsafe {
-                self.frame().side_table = state.side_table;
-                self.parser.restore(state.parser_data)
-            },
-            LabelKind::Block | LabelKind::If => self.take_jump(offset),
-        }
+        self.take_jump(offset);
         ThreadResult::Continue(self)
     }
 
@@ -1422,7 +1410,7 @@ struct Frame<'m> {
     arity: usize,
     ret: &'m [u8],
     locals: Vec<Val>,
-    labels: Vec<Label<'m>>,
+    labels: Vec<Label>,
     side_table: &'m [SideTableEntry],
 }
 
@@ -1431,7 +1419,7 @@ impl<'m> Frame<'m> {
         inst_id: usize, arity: usize, ret: &'m [u8], locals: Vec<Val>,
         side_table: &'m [SideTableEntry],
     ) -> Self {
-        let label = Label { arity, kind: LabelKind::Block, values_cnt: 0 };
+        let label = Label { arity, values_cnt: 0 };
         Frame { inst_id, arity, ret, locals, labels: vec![label], side_table }
     }
 
@@ -1448,26 +1436,17 @@ impl<'m> Frame<'m> {
 }
 
 #[derive(Debug)]
-struct Label<'m> {
+struct Label {
     arity: usize,
-    kind: LabelKind<'m>,
     values_cnt: usize,
 }
 
 #[derive(Debug)]
-struct LoopState<'m> {
-    parser_data: &'m [u8],
-    side_table: &'m [SideTableEntry],
-}
-
-#[derive(Debug)]
-enum LabelKind<'m> {
+enum LabelKind {
     // TODO: If and Block can be merged and then we just have Option<NonNull<u8>> which is
     // optimized.
     Block,
-    // TODO: Could be just NonNull<u8> since we can reuse the end of current parser since it
-    // never changes.
-    Loop(LoopState<'m>),
+    Loop,
     If,
 }
 
@@ -1644,11 +1623,3 @@ fn memory_too_small(x: usize, n: usize, mem: &Memory) {
     #[cfg(feature = "debug")]
     eprintln!("Memory too small: {x} + {n} > {}", mem.len());
 }
-
-// TODO(dev/fast-interp): Add debug asserts when `off` is positive and negative, and `toctou`
-// support.
-fn offset_front<T>(cur: &[T], off: isize) -> &[T] {
-    unsafe {
-        core::slice::from_raw_parts(cur.as_ptr().offset(off), (cur.len() as isize - off) as usize)
-    }
-}

crates/interpreter/src/lib.rs

@@ -148,6 +148,7 @@ mod parser;
 mod side_table;
 mod syntax;
 mod toctou;
+mod util;
 mod valid;
 
 pub use error::{Error, TRAP_CODE, Unsupported};

crates/interpreter/src/util.rs

@@ -0,0 +1,21 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// TODO(dev/fast-interp): Add debug asserts when `off` is positive and negative, and `toctou`
+// support.
+pub fn offset_front<T>(cur: &[T], off: isize) -> &[T] {
+    unsafe {
+        core::slice::from_raw_parts(cur.as_ptr().offset(off), (cur.len() as isize - off) as usize)
+    }
+}

crates/interpreter/src/valid.rs

@@ -21,6 +21,7 @@ use crate::error::*;
 use crate::side_table::*;
 use crate::syntax::*;
 use crate::toctou::*;
+use crate::util::*;
 use crate::*;
 
 /// Checks whether a WASM module in binary format is valid.
@@ -558,6 +559,7 @@ impl<'a, 'm> Expr<'a, 'm> {
 
     fn instr(&mut self) -> CheckResult {
         use Instr::*;
+        let saved = self.parser.save();
         let instr = self.parser.parse_instr()?;
         if matches!(instr, End) {
             return self.end_label();
@@ -586,7 +588,9 @@ impl<'a, 'm> Expr<'a, 'm> {
             Block(b) => self.push_label(self.blocktype(&b)?, LabelKind::Block)?,
             Loop(b) => {
                 let type_ = self.blocktype(&b)?;
-                self.push_label(type_, LabelKind::Loop(self.branch_target(type_.params.len())))?
+                let mut target = self.branch_target(type_.params.len());
+                target.parser = saved;
+                self.push_label(type_, LabelKind::Loop(target))?
             }
             If(b) => {
                 self.pop_check(ValType::I32)?;
@@ -859,12 +863,7 @@ impl<'a, 'm> Expr<'a, 'm> {
         if let LabelKind::If(source) = label.kind {
             check(label.type_.params == label.type_.results)?;
             // SAFETY: This function is only called after parsing an End instruction.
-            target.parser = unsafe {
-                core::slice::from_raw_parts(
-                    target.parser.as_ptr().offset(-1),
-                    target.parser.len() + 1,
-                )
-            };
+            target.parser = offset_front(target.parser, -1);
             self.side_table.stitch(source, target)?;
         }
         let results = self.label().type_.results;