Fuse switch selection with branch entry matching

Amp-Thread-ID: https://ampcode.com/threads/T-01a15ba8-2c2e-737b-ad72-3fd2746cba8e

Author: zharinov

crates/compiler/src/bytecode/mod.rs

@@ -46,7 +46,10 @@ pub(crate) use instructions::{
     Call, CallOwnership, CalleeContract, Match, MatchCounts, MatchInstr, MatchPredicate, Opcode,
     PREDICATE_SLOTS, Return, SuccessorAddr, Switch, SwitchEntry,
 };
-pub(crate) use module::{DecodedCall, DecodedInstr, DecodedMatch, DecodedPredicate, Instruction};
+pub(crate) use module::{
+    DecodedAction, DecodedCall, DecodedCandidate, DecodedInstr, DecodedMatch, DecodedPredicate,
+    Instruction,
+};
 pub(crate) use node_kind_constraint::NodeKindConstraint;
 pub(crate) use plotnik_runtime::{Nav, deserialize_dfa};
 pub(crate) use predicate_op::PredicateOp;

crates/compiler/src/bytecode/module/decoded.rs

@@ -1,12 +1,14 @@
+use std::collections::{HashMap, HashSet};
+
 use crate::bytecode::{
-    BYTECODE_WORD_SIZE, CallOwnership, CodeAddr, Effect, Nav, NodeKindConstraint, PredicateOp,
-    SuccessorAddr, SwitchEntry,
+    BYTECODE_WORD_SIZE, CallOwnership, CodeAddr, Effect, Nav, PredicateOp, SuccessorAddr,
+    SwitchEntry,
 };
-use crate::core::NodeFieldId;
+use crate::core::{NodeFieldId, NodeKindId};
 use plotnik_runtime::PortId;
 
 use super::super::instructions::header_byte;
-use super::{Instruction, Module};
+use super::{Instruction, Module, ModuleError};
 
 #[derive(Clone, Copy, Debug)]
 pub(crate) enum DecodedInstr {
@@ -26,16 +28,22 @@ pub(crate) struct DecodedSwitch {
 #[derive(Clone, Copy, Debug)]
 pub(crate) struct DecodedMatch {
     pub(crate) nav: Nav,
-    pub(crate) node_kind: NodeKindConstraint,
+    pub(crate) kind: Option<NodeKindId>,
+    pub(crate) candidate: DecodedCandidate,
+    effects_base: u32,
+    succ_base: u32,
+    effects_len: u8,
+    succ_len: u8,
+}
+
+#[derive(Clone, Copy, Debug)]
+pub(crate) struct DecodedCandidate {
+    pub(crate) named: Option<bool>,
     pub(crate) node_field: Option<NodeFieldId>,
     pub(crate) missing: bool,
     pub(crate) predicate: Option<DecodedPredicate>,
-    effects_base: u32,
     neg_base: u32,
-    succ_base: u32,
-    effects_len: u8,
     neg_len: u8,
-    succ_len: u8,
 }
 
 impl DecodedMatch {
@@ -67,18 +75,50 @@ impl DecodedCall {
     }
 }
 
+#[derive(Clone, Copy, Debug)]
+pub(crate) struct DecodedSwitchEntry {
+    pub(crate) kind: NodeKindId,
+    action: u16,
+}
+
+#[derive(Debug)]
+pub(crate) struct DecodedAction {
+    pub(crate) entry: CodeAddr,
+    pub(crate) matched_at: CodeAddr,
+    pub(crate) matched: DecodedMatch,
+    pub(crate) prefix: Vec<DecodedPrefix>,
+}
+
+#[derive(Debug)]
+pub(crate) struct DecodedPrefix {
+    pub(crate) ip: CodeAddr,
+    effects_base: u32,
+    effects_len: u8,
+}
+
 #[derive(Debug, Default)]
 pub(crate) struct DecodedProgram {
     words: Vec<DecodedInstr>,
     effects: Vec<Effect>,
     neg_fields: Vec<NodeFieldId>,
     successors: Vec<SuccessorAddr>,
-    switch_entries: Vec<SwitchEntry>,
+    switch_entries: Vec<DecodedSwitchEntry>,
+    actions: Vec<DecodedAction>,
 }
 
 impl DecodedProgram {
-    pub(crate) fn switch_entries(&self, switch: &DecodedSwitch) -> &[SwitchEntry] {
-        &self.switch_entries[switch.base as usize..][..switch.len as usize]
+    #[inline]
+    pub(crate) fn switch_entries(&self, switch: &DecodedSwitch) -> &[DecodedSwitchEntry] {
+        &self.switch_entries[switch.base as usize..][..usize::from(switch.len)]
+    }
+
+    #[inline]
+    pub(crate) fn switch_action(&self, entry: &DecodedSwitchEntry) -> &DecodedAction {
+        &self.actions[usize::from(entry.action)]
+    }
+
+    pub(crate) fn prefix_effects(&self, prefix: &DecodedPrefix) -> &[Effect] {
+        &self.effects[prefix.effects_base as usize..][..prefix.effects_len as usize]
     }
 
     #[inline]
@@ -98,8 +138,8 @@ impl DecodedProgram {
     }
 
     #[inline]
-    pub(crate) fn neg_fields(&self, m: &DecodedMatch) -> &[NodeFieldId] {
-        &self.neg_fields[m.neg_base as usize..][..m.neg_len as usize]
+    pub(crate) fn neg_fields(&self, candidate: &DecodedCandidate) -> &[NodeFieldId] {
+        &self.neg_fields[candidate.neg_base as usize..][..candidate.neg_len as usize]
     }
 
     #[inline]
@@ -118,13 +158,10 @@ impl DecodedProgram {
     }
 }
 
-pub(crate) fn build(module: &Module) -> DecodedProgram {
+pub(crate) fn build(module: &Module) -> Result<DecodedProgram, ModuleError> {
     let instructions = module.instructions_slice();
     let word_count = instructions.len() / BYTECODE_WORD_SIZE;
-    let mut program = DecodedProgram {
-        switch_entries: module.switch_tables().iter().collect(),
-        ..Default::default()
-    };
+    let mut program = DecodedProgram::default();
     program.words.reserve(word_count);
 
     let mut word_addr = 0usize;
@@ -154,15 +191,18 @@ pub(crate) fn build(module: &Module) -> DecodedProgram {
 
                 program.words.push(DecodedInstr::Match(DecodedMatch {
                     nav: m.nav,
-                    node_kind: m.node_kind,
-                    node_field: m.node_field,
-                    missing: m.missing(),
-                    predicate,
+                    kind: m.node_kind.kind_id(),
+                    candidate: DecodedCandidate {
+                        named: (!m.node_kind.is_any()).then_some(m.node_kind.is_named()),
+                        node_field: m.node_field,
+                        missing: m.missing(),
+                        predicate,
+                        neg_base,
+                        neg_len: (program.neg_fields.len() - neg_base as usize) as u8,
+                    },
                     effects_base,
-                    neg_base,
                     succ_base,
                     effects_len: (program.effects.len() - effects_base as usize) as u8,
-                    neg_len: (program.neg_fields.len() - neg_base as usize) as u8,
                     succ_len: (program.successors.len() - succ_base as usize) as u8,
                 }));
                 for _ in 1..opcode.word_count() {
@@ -192,7 +232,68 @@ pub(crate) fn build(module: &Module) -> DecodedProgram {
             }
         }
     }
-    program
+    let mut actions = HashMap::new();
+    for entry in module.switch_tables().iter() {
+        let target = CodeAddr::from(u16::from(entry.target));
+        let action = if let Some(&index) = actions.get(&(target, entry.kind)) {
+            index
+        } else {
+            let action = program.selected_entry(target, entry.kind)?;
+            let index = u16::try_from(program.actions.len())
+                .expect("selected entries are bounded by instruction count");
+            program.actions.push(action);
+            actions.insert((target, entry.kind), index);
+            index
+        };
+        program.switch_entries.push(DecodedSwitchEntry {
+            kind: entry.kind,
+            action,
+        });
+    }
+    Ok(program)
+}
+
+impl DecodedProgram {
+    fn selected_entry(
+        &self,
+        entry: CodeAddr,
+        kind: NodeKindId,
+    ) -> Result<DecodedAction, ModuleError> {
+        let mut ip = entry;
+        let mut visited = HashSet::new();
+        let mut prefix = Vec::new();
+        loop {
+            if !visited.insert(ip) {
+                return Err(ModuleError::MalformedSwitchTable);
+            }
+
+            let DecodedInstr::Match(matched) = self.instruction_at(ip) else {
+                return Err(ModuleError::MalformedSwitchTable);
+            };
+            if !matched.is_epsilon() {
+                if matched.nav != Nav::StayExact || matched.kind != Some(kind) {
+                    return Err(ModuleError::MalformedSwitchTable);
+                }
+
+                return Ok(DecodedAction {
+                    entry,
+                    matched_at: ip,
+                    matched,
+                    prefix,
+                });
+            }
+
+            let [next] = self.successors(&matched) else {
+                return Err(ModuleError::MalformedSwitchTable);
+            };
+            prefix.push(DecodedPrefix {
+                ip,
+                effects_base: matched.effects_base,
+                effects_len: matched.effects_len,
+            });
+            ip = CodeAddr::from(u16::from(*next));
+        }
+    }
 }
 
 fn pool_base(len: usize) -> u32 {

crates/compiler/src/bytecode/module/load.rs

@@ -140,7 +140,7 @@ impl Module {
         }
         #[cfg(not(debug_assertions))]
         let _ = is_start;
-        module.decoded = decoded::build(&module);
+        module.decoded = decoded::build(&module)?;
         Ok(module)
     }
 

crates/compiler/src/bytecode/module/mod.rs

@@ -21,7 +21,8 @@ mod matcher_verify;
 mod verify;
 
 pub(crate) use decoded::{
-    DecodedCall, DecodedInstr, DecodedMatch, DecodedPredicate, DecodedProgram,
+    DecodedAction, DecodedCall, DecodedCandidate, DecodedInstr, DecodedMatch, DecodedPredicate,
+    DecodedProgram,
 };
 pub(crate) use load::ModuleError;
 

crates/compiler/src/bytecode/module/validate_tests.rs

@@ -1770,6 +1770,10 @@ fn malformed_switch_tables_are_rejected() {
         (vec![(44, (size - 1).to_le_bytes().to_vec())], true),
         (vec![(table, second_entry), (table + 4, first_entry)], true),
         (vec![(table + 4, first_kind)], true),
+        (
+            vec![(table + 2, valid[table + 6..table + 8].to_vec())],
+            true,
+        ),
         (vec![(instruction + 4, size.to_le_bytes().to_vec())], false),
         (
             vec![