From 4a5c9511199634e830155cda9e4e18a167cf53af Mon Sep 17 00:00:00 2001
From: Sergei Zharinov <zharinov@users.noreply.github.com>
Date: Thu, 15 Jan 2026 15:52:53 -0300
Subject: [PATCH] opt: Collapse consecutive Up instructions of the same mode

---
 .../src/compile/collapse_up.rs                | 134 ++++++++++
 .../src/compile/collapse_up_tests.rs          | 231 ++++++++++++++++++
 .../plotnik-compiler/src/compile/compiler.rs  |   4 +
 crates/plotnik-compiler/src/compile/mod.rs    |   3 +
 ...t__emit_tests__captures_deeply_nested.snap |   6 +-
 ...mit__emit_tests__captures_nested_flat.snap |   5 +-
 6 files changed, 376 insertions(+), 7 deletions(-)
 create mode 100644 crates/plotnik-compiler/src/compile/collapse_up.rs
 create mode 100644 crates/plotnik-compiler/src/compile/collapse_up_tests.rs

diff --git a/crates/plotnik-compiler/src/compile/collapse_up.rs b/crates/plotnik-compiler/src/compile/collapse_up.rs
new file mode 100644
index 0000000..0c3ae7f
--- /dev/null
+++ b/crates/plotnik-compiler/src/compile/collapse_up.rs
@@ -0,0 +1,134 @@
+//! Up-collapse optimization: merge consecutive Up instructions of the same mode.
+//!
+//! Transforms: Up(1) → Up(1) → Up(2) into Up(4)
+//!
+//! Constraints:
+//! - Same mode only (Up, UpSkipTrivia, UpExact can't mix)
+//! - Effectless only (no pre_effects, post_effects, neg_fields)
+//! - Max 63 (6-bit payload limit)
+//! - Single successor (can't merge branching instructions)
+
+use std::collections::{HashMap, HashSet};
+
+use plotnik_bytecode::Nav;
+
+use crate::bytecode::{InstructionIR, Label, MatchIR, NodeTypeIR};
+use crate::compile::CompileResult;
+
+const MAX_UP_LEVEL: u8 = 63;
+
+/// Collapse consecutive Up instructions of the same mode.
+pub fn collapse_up(result: &mut CompileResult) {
+    let label_to_idx: HashMap<Label, usize> = result
+        .instructions
+        .iter()
+        .enumerate()
+        .map(|(i, instr)| (instr.label(), i))
+        .collect();
+
+    let mut removed: HashSet<Label> = HashSet::new();
+
+    for i in 0..result.instructions.len() {
+        let InstructionIR::Match(m) = &result.instructions[i] else {
+            continue;
+        };
+
+        let Some(up_level) = get_up_level(m.nav) else {
+            continue;
+        };
+
+        if m.successors.len() != 1 {
+            continue;
+        }
+
+        let mut current_level = up_level;
+        let mut current_nav = m.nav;
+        let mut final_successors = m.successors.clone();
+
+        // Absorb chain of effectless Up instructions with same mode
+        while current_level < MAX_UP_LEVEL {
+            let &[succ_label] = final_successors.as_slice() else {
+                break;
+            };
+
+            if removed.contains(&succ_label) {
+                break;
+            }
+
+            let Some(&succ_idx) = label_to_idx.get(&succ_label) else {
+                break;
+            };
+
+            let InstructionIR::Match(succ) = &result.instructions[succ_idx] else {
+                break;
+            };
+
+            let Some(succ_level) = get_up_level(succ.nav) else {
+                break;
+            };
+
+            if !same_up_mode(current_nav, succ.nav) || !is_effectless(succ) {
+                break;
+            }
+
+            // Merge: add levels (capped at 63)
+            let new_level = current_level.saturating_add(succ_level).min(MAX_UP_LEVEL);
+            current_nav = set_up_level(current_nav, new_level);
+            current_level = new_level;
+            final_successors = succ.successors.clone();
+            removed.insert(succ_label);
+        }
+
+        // Update the instruction if we merged anything
+        if current_level != up_level {
+            let InstructionIR::Match(m) = &mut result.instructions[i] else {
+                unreachable!()
+            };
+            m.nav = current_nav;
+            m.successors = final_successors;
+        }
+    }
+
+    // Remove absorbed instructions
+    result
+        .instructions
+        .retain(|instr| !removed.contains(&instr.label()));
+}
+
+/// Extract Up level from Nav, if it's an Up variant.
+fn get_up_level(nav: Nav) -> Option<u8> {
+    match nav {
+        Nav::Up(n) | Nav::UpSkipTrivia(n) | Nav::UpExact(n) => Some(n),
+        _ => None,
+    }
+}
+
+/// Set the level on an Up Nav variant.
+fn set_up_level(nav: Nav, level: u8) -> Nav {
+    match nav {
+        Nav::Up(_) => Nav::Up(level),
+        Nav::UpSkipTrivia(_) => Nav::UpSkipTrivia(level),
+        Nav::UpExact(_) => Nav::UpExact(level),
+        _ => nav,
+    }
+}
+
+/// Check if two Nav values are the same Up mode (ignoring level).
+fn same_up_mode(a: Nav, b: Nav) -> bool {
+    matches!(
+        (a, b),
+        (Nav::Up(_), Nav::Up(_))
+            | (Nav::UpSkipTrivia(_), Nav::UpSkipTrivia(_))
+            | (Nav::UpExact(_), Nav::UpExact(_))
+    )
+}
+
+/// Check if a MatchIR has no effects or constraints (pure navigation).
+fn is_effectless(m: &MatchIR) -> bool {
+    m.node_type == NodeTypeIR::Any
+        && m.node_field.is_none()
+        && m.pre_effects.is_empty()
+        && m.neg_fields.is_empty()
+        && m.post_effects.is_empty()
+        && m.predicate.is_none()
+}
diff --git a/crates/plotnik-compiler/src/compile/collapse_up_tests.rs b/crates/plotnik-compiler/src/compile/collapse_up_tests.rs
new file mode 100644
index 0000000..d1b85b7
--- /dev/null
+++ b/crates/plotnik-compiler/src/compile/collapse_up_tests.rs
@@ -0,0 +1,231 @@
+//! Unit tests for the Up-collapse optimization pass.
+
+use plotnik_bytecode::Nav;
+
+use super::collapse_up::collapse_up;
+use super::CompileResult;
+use crate::bytecode::{InstructionIR, Label, MatchIR};
+
+#[test]
+fn collapse_up_single_mode() {
+    // Up(1) → Up(1) → exit should become Up(2) → exit
+    let mut result = CompileResult {
+        instructions: vec![
+            MatchIR::at(Label(0)).nav(Nav::Up(1)).next(Label(1)).into(),
+            MatchIR::at(Label(1)).nav(Nav::Up(1)).next(Label(2)).into(),
+            MatchIR::terminal(Label(2)).into(),
+        ],
+        def_entries: Default::default(),
+        preamble_entry: Label(0),
+    };
+
+    collapse_up(&mut result);
+
+    // Should collapse to 2 instructions: Up(2) and terminal
+    assert_eq!(result.instructions.len(), 2);
+
+    let InstructionIR::Match(m) = &result.instructions[0] else {
+        panic!("expected Match");
+    };
+    assert_eq!(m.nav, Nav::Up(2));
+    assert_eq!(m.successors, vec![Label(2)]);
+}
+
+#[test]
+fn collapse_up_chain_of_three() {
+    // Up(1) → Up(2) → Up(3) should become Up(6)
+    let mut result = CompileResult {
+        instructions: vec![
+            MatchIR::at(Label(0)).nav(Nav::Up(1)).next(Label(1)).into(),
+            MatchIR::at(Label(1)).nav(Nav::Up(2)).next(Label(2)).into(),
+            MatchIR::at(Label(2)).nav(Nav::Up(3)).next(Label(3)).into(),
+            MatchIR::terminal(Label(3)).into(),
+        ],
+        def_entries: Default::default(),
+        preamble_entry: Label(0),
+    };
+
+    collapse_up(&mut result);
+
+    assert_eq!(result.instructions.len(), 2);
+
+    let InstructionIR::Match(m) = &result.instructions[0] else {
+        panic!("expected Match");
+    };
+    assert_eq!(m.nav, Nav::Up(6));
+}
+
+#[test]
+fn collapse_up_mixed_modes_no_merge() {
+    // Up(1) → UpSkipTrivia(1) should NOT merge (different modes)
+    let mut result = CompileResult {
+        instructions: vec![
+            MatchIR::at(Label(0)).nav(Nav::Up(1)).next(Label(1)).into(),
+            MatchIR::at(Label(1))
+                .nav(Nav::UpSkipTrivia(1))
+                .next(Label(2))
+                .into(),
+            MatchIR::terminal(Label(2)).into(),
+        ],
+        def_entries: Default::default(),
+        preamble_entry: Label(0),
+    };
+
+    collapse_up(&mut result);
+
+    // Should stay 3 instructions
+    assert_eq!(result.instructions.len(), 3);
+}
+
+#[test]
+fn collapse_up_skip_trivia_same_mode() {
+    // UpSkipTrivia(1) → UpSkipTrivia(1) should merge
+    let mut result = CompileResult {
+        instructions: vec![
+            MatchIR::at(Label(0))
+                .nav(Nav::UpSkipTrivia(1))
+                .next(Label(1))
+                .into(),
+            MatchIR::at(Label(1))
+                .nav(Nav::UpSkipTrivia(1))
+                .next(Label(2))
+                .into(),
+            MatchIR::terminal(Label(2)).into(),
+        ],
+        def_entries: Default::default(),
+        preamble_entry: Label(0),
+    };
+
+    collapse_up(&mut result);
+
+    assert_eq!(result.instructions.len(), 2);
+
+    let InstructionIR::Match(m) = &result.instructions[0] else {
+        panic!("expected Match");
+    };
+    assert_eq!(m.nav, Nav::UpSkipTrivia(2));
+}
+
+#[test]
+fn collapse_up_exact_same_mode() {
+    // UpExact(1) → UpExact(1) should merge
+    let mut result = CompileResult {
+        instructions: vec![
+            MatchIR::at(Label(0))
+                .nav(Nav::UpExact(1))
+                .next(Label(1))
+                .into(),
+            MatchIR::at(Label(1))
+                .nav(Nav::UpExact(1))
+                .next(Label(2))
+                .into(),
+            MatchIR::terminal(Label(2)).into(),
+        ],
+        def_entries: Default::default(),
+        preamble_entry: Label(0),
+    };
+
+    collapse_up(&mut result);
+
+    assert_eq!(result.instructions.len(), 2);
+
+    let InstructionIR::Match(m) = &result.instructions[0] else {
+        panic!("expected Match");
+    };
+    assert_eq!(m.nav, Nav::UpExact(2));
+}
+
+#[test]
+fn collapse_up_with_effects_no_merge() {
+    // Up(1) with post_effects → Up(1) should NOT merge
+    use crate::bytecode::EffectIR;
+    use plotnik_bytecode::EffectOpcode;
+
+    let mut result = CompileResult {
+        instructions: vec![
+            MatchIR::at(Label(0)).nav(Nav::Up(1)).next(Label(1)).into(),
+            MatchIR::at(Label(1))
+                .nav(Nav::Up(1))
+                .post_effects(vec![EffectIR::simple(EffectOpcode::Null, 0)])
+                .next(Label(2))
+                .into(),
+            MatchIR::terminal(Label(2)).into(),
+        ],
+        def_entries: Default::default(),
+        preamble_entry: Label(0),
+    };
+
+    collapse_up(&mut result);
+
+    // Should stay 3 instructions (effectful Up can't be absorbed)
+    assert_eq!(result.instructions.len(), 3);
+}
+
+#[test]
+fn collapse_up_max_63() {
+    // Up(60) → Up(10) should become Up(63) (capped)
+    let mut result = CompileResult {
+        instructions: vec![
+            MatchIR::at(Label(0)).nav(Nav::Up(60)).next(Label(1)).into(),
+            MatchIR::at(Label(1)).nav(Nav::Up(10)).next(Label(2)).into(),
+            MatchIR::terminal(Label(2)).into(),
+        ],
+        def_entries: Default::default(),
+        preamble_entry: Label(0),
+    };
+
+    collapse_up(&mut result);
+
+    // Capped at 63, remaining Up(7) stays separate
+    assert_eq!(result.instructions.len(), 2);
+
+    let InstructionIR::Match(m) = &result.instructions[0] else {
+        panic!("expected Match");
+    };
+    assert_eq!(m.nav, Nav::Up(63));
+}
+
+#[test]
+fn collapse_up_branching_no_merge() {
+    // Up(1) with multiple successors should NOT merge
+    let mut result = CompileResult {
+        instructions: vec![
+            MatchIR::at(Label(0))
+                .nav(Nav::Up(1))
+                .next_many(vec![Label(1), Label(2)])
+                .into(),
+            MatchIR::at(Label(1)).nav(Nav::Up(1)).next(Label(3)).into(),
+            MatchIR::at(Label(2)).nav(Nav::Up(1)).next(Label(3)).into(),
+            MatchIR::terminal(Label(3)).into(),
+        ],
+        def_entries: Default::default(),
+        preamble_entry: Label(0),
+    };
+
+    collapse_up(&mut result);
+
+    // Branching instruction can't merge, but its successors can be processed
+    // Label(0) has 2 successors, so it stays as Up(1)
+    let InstructionIR::Match(m) = &result.instructions[0] else {
+        panic!("expected Match");
+    };
+    assert_eq!(m.nav, Nav::Up(1));
+}
+
+#[test]
+fn collapse_up_no_up_unchanged() {
+    // Non-Up instructions should pass through unchanged
+    let mut result = CompileResult {
+        instructions: vec![
+            MatchIR::at(Label(0)).nav(Nav::Down).next(Label(1)).into(),
+            MatchIR::at(Label(1)).nav(Nav::Next).next(Label(2)).into(),
+            MatchIR::terminal(Label(2)).into(),
+        ],
+        def_entries: Default::default(),
+        preamble_entry: Label(0),
+    };
+
+    collapse_up(&mut result);
+
+    assert_eq!(result.instructions.len(), 3);
+}
diff --git a/crates/plotnik-compiler/src/compile/compiler.rs b/crates/plotnik-compiler/src/compile/compiler.rs
index 5de4be3..32f1bef 100644
--- a/crates/plotnik-compiler/src/compile/compiler.rs
+++ b/crates/plotnik-compiler/src/compile/compiler.rs
@@ -13,6 +13,7 @@ use crate::parser::Expr;
 use plotnik_bytecode::Nav;
 
 use super::capture::CaptureEffects;
+use super::collapse_up::collapse_up;
 use super::dce::remove_unreachable;
 use super::epsilon_elim::eliminate_epsilons;
 use super::lower::lower;
@@ -87,6 +88,9 @@ impl<'a> Compiler<'a> {
         // Remove unreachable instructions (bypassed epsilons, etc.)
         remove_unreachable(&mut result);
 
+        // Collapse consecutive Up instructions of the same mode
+        collapse_up(&mut result);
+
         // Lower to bytecode-compatible form (cascade overflows)
         lower(&mut result);
 
diff --git a/crates/plotnik-compiler/src/compile/mod.rs b/crates/plotnik-compiler/src/compile/mod.rs
index bc9f48c..43cb941 100644
--- a/crates/plotnik-compiler/src/compile/mod.rs
+++ b/crates/plotnik-compiler/src/compile/mod.rs
@@ -15,6 +15,7 @@
 //! - `sequences`: Sequence and alternation compilation
 
 mod capture;
+mod collapse_up;
 mod compiler;
 mod dce;
 pub(crate) mod epsilon_elim;
@@ -30,6 +31,8 @@ mod verify;
 #[cfg(test)]
 mod capture_tests;
 #[cfg(test)]
+mod collapse_up_tests;
+#[cfg(test)]
 mod compile_tests;
 #[cfg(test)]
 mod lower_tests;
diff --git a/crates/plotnik-compiler/src/emit/snapshots/plotnik_compiler__emit__emit_tests__captures_deeply_nested.snap b/crates/plotnik-compiler/src/emit/snapshots/plotnik_compiler__emit__emit_tests__captures_deeply_nested.snap
index 8b77769..7071366 100644
--- a/crates/plotnik-compiler/src/emit/snapshots/plotnik_compiler__emit__emit_tests__captures_deeply_nested.snap
+++ b/crates/plotnik-compiler/src/emit/snapshots/plotnik_compiler__emit__emit_tests__captures_deeply_nested.snap
@@ -50,7 +50,5 @@ Test:
   08   ▽   (array) [Node Set(M1)]           10
   10   ▽   (array) [Node Set(M2)]           12
   12   ▽   (identifier) [Node Set(M3)]      14
-  14   △   _                                15
-  15   △   _                                16
-  16   △   _                                17
-  17                                        ▶
+  14   △³  _                                15
+  15                                        ▶
diff --git a/crates/plotnik-compiler/src/emit/snapshots/plotnik_compiler__emit__emit_tests__captures_nested_flat.snap b/crates/plotnik-compiler/src/emit/snapshots/plotnik_compiler__emit__emit_tests__captures_nested_flat.snap
index 004db2b..0dec4bc 100644
--- a/crates/plotnik-compiler/src/emit/snapshots/plotnik_compiler__emit__emit_tests__captures_nested_flat.snap
+++ b/crates/plotnik-compiler/src/emit/snapshots/plotnik_compiler__emit__emit_tests__captures_nested_flat.snap
@@ -43,6 +43,5 @@ Test:
   06   !   (array) [Node Set(M0)]           08
   08   ▽   (array) [Node Set(M1)]           10
   10   ▽   (identifier) [Node Set(M2)]      12
-  12   △   _                                13
-  13   △   _                                14
-  14                                        ▶
+  12   △²  _                                13
+  13                                        ▶