redpiler: allow nodes to map to multiple blocks

Cargo.toml

@@ -93,3 +93,4 @@ indexmap = { version = "2.13", features = ["serde"] }
 line-index = "0.1"
 anstream = "1.0"
 owo-colors = "4.3"
+smallvec = { version = "1.15", features = ["union", "serde"] }

crates/redpiler/Cargo.toml

@@ -17,3 +17,4 @@ itertools = { workspace = true }
 rustc-hash = { workspace = true }
 enum_dispatch = { workspace = true }
 indexmap = { workspace = true }
+smallvec = { workspace = true }

crates/redpiler/src/backend/direct/compile.rs

@@ -8,6 +8,7 @@ use mchprs_world::TickEntry;
 use petgraph::visit::EdgeRef;
 use petgraph::Direction;
 use rustc_hash::FxHashMap;
+use smallvec::SmallVec;
 use std::sync::Arc;
 use tracing::trace;
 
@@ -27,7 +28,7 @@ fn compile_node(
     node_idx: NodeIdx,
     nodes_len: usize,
     nodes_map: &FxHashMap<NodeIdx, usize>,
-    noteblock_info: &mut Vec<(BlockPos, Instrument, u8)>,
+    noteblock_info: &mut Vec<(SmallVec<[BlockPos; 1]>, Instrument, u8)>,
     forward_links: &mut ForwardLinks,
     stats: &mut FinalGraphStats,
 ) -> Node {
@@ -124,7 +125,11 @@ fn compile_node(
         CNodeType::Constant => NodeType::Constant,
         CNodeType::NoteBlock { instrument, note } => {
             let noteblock_id = noteblock_info.len().try_into().unwrap();
-            noteblock_info.push((node.block.unwrap().0, *instrument, *note));
+            noteblock_info.push((
+                node.block.iter().copied().map(|(pos, _)| pos).collect(),
+                *instrument,
+                *note,
+            ));
             NodeType::NoteBlock { noteblock_id }
         }
     };
@@ -178,13 +183,19 @@ pub fn compile(
 
     backend.blocks = graph
         .node_weights()
-        .map(|node| node.block.map(|(pos, id)| (pos, Block::from_id(id))))
+        .map(|node| {
+            node.block
+                .iter()
+                .copied()
+                .map(|(pos, id)| (pos, Block::from_id(id)))
+                .collect()
+        })
         .collect();
     backend.nodes = Nodes::new(nodes);
 
     // Create a mapping from block pos to backend NodeId
     for i in 0..backend.blocks.len() {
-        if let Some((pos, _)) = backend.blocks[i] {
+        for (pos, _) in backend.blocks[i].iter().copied() {
             backend.pos_map.insert(pos, backend.nodes.get(i));
         }
     }

crates/redpiler/src/backend/direct/mod.rs

@@ -17,6 +17,8 @@ use mchprs_redstone::{bool_to_ss, noteblock};
 use mchprs_world::{TickEntry, TickPriority, World};
 use node::{Node, NodeId, NodeType, Nodes};
 use rustc_hash::FxHashMap;
+use smallvec::SmallVec;
+use std::fmt::Write;
 use std::sync::Arc;
 use std::{fmt, mem};
 use tracing::{debug, warn};
@@ -46,7 +48,7 @@ impl TickScheduler {
     const NUM_PRIORITIES: usize = 4;
     const NUM_QUEUES: usize = 16;
 
-    fn reset<W: World>(&mut self, world: &mut W, blocks: &[Option<(BlockPos, Block)>]) {
+    fn reset<W: World>(&mut self, world: &mut W, blocks: &[impl AsRef<[(BlockPos, Block)]>]) {
         for (idx, queues) in self.queues_deque.iter().enumerate() {
             let delay = if self.pos >= idx {
                 idx + Self::NUM_QUEUES
@@ -55,11 +57,14 @@ impl TickScheduler {
             } - self.pos;
             for (entries, priority) in queues.0.iter().zip(Self::priorities()) {
                 for node in entries {
-                    let Some((pos, _)) = blocks[node.index()] else {
+                    let node_blocks = blocks[node.index()].as_ref();
+                    if node_blocks.is_empty() {
                         warn!("Cannot schedule tick for node {:?} because block information is missing", node);
                         continue;
                     };
-                    world.schedule_tick(pos, delay as u32, priority);
+                    for (pos, _) in node_blocks.iter().copied() {
+                        world.schedule_tick(pos, delay as u32, priority);
+                    }
                 }
             }
         }
@@ -112,11 +117,11 @@ enum Event {
 pub struct DirectBackend {
     nodes: Nodes,
     forward_links: ForwardLinks,
-    blocks: Vec<Option<(BlockPos, Block)>>,
+    blocks: Vec<SmallVec<[(BlockPos, Block); 1]>>,
     pos_map: FxHashMap<BlockPos, NodeId>,
     scheduler: TickScheduler,
     events: Vec<Event>,
-    noteblock_info: Vec<(BlockPos, Instrument, u8)>,
+    noteblock_info: Vec<(SmallVec<[BlockPos; 1]>, Instrument, u8)>,
 }
 
 impl DirectBackend {
@@ -183,18 +188,17 @@ impl JITBackend for DirectBackend {
         let nodes = std::mem::take(&mut self.nodes);
 
         for (i, node) in nodes.into_inner().iter().enumerate() {
-            let Some((pos, block)) = self.blocks[i] else {
-                continue;
-            };
-            if matches!(node.ty, NodeType::Comparator { .. }) {
-                let block_entity = BlockEntity::Comparator {
-                    output_strength: node.output_power,
-                };
-                world.set_block_entity(pos, block_entity);
-            }
+            for (pos, block) in self.blocks[i].iter().copied() {
+                if matches!(node.ty, NodeType::Comparator { .. }) {
+                    let block_entity = BlockEntity::Comparator {
+                        output_strength: node.output_power,
+                    };
+                    world.set_block_entity(pos, block_entity);
+                }
 
-            if io_only && !node.is_io {
-                world.set_block(pos, block);
+                if io_only && !node.is_io {
+                    world.set_block(pos, block);
+                }
             }
         }
 
@@ -247,16 +251,19 @@ impl JITBackend for DirectBackend {
         for event in self.events.drain(..) {
             match event {
                 Event::NoteBlockPlay { noteblock_id } => {
-                    let (pos, instrument, note) = self.noteblock_info[noteblock_id as usize];
-                    noteblock::play_note(world, pos, instrument, note);
+                    let (positions, instrument, note) = &self.noteblock_info[noteblock_id as usize];
+                    for pos in positions.iter().copied() {
+                        noteblock::play_note(world, pos, *instrument, *note);
+                    }
                 }
             }
         }
         for (i, node) in self.nodes.inner_mut().iter_mut().enumerate() {
-            let Some((pos, block)) = &mut self.blocks[i] else {
+            if !node.changed || (io_only && !node.is_io) {
                 continue;
-            };
-            if node.changed && (!io_only || node.is_io) {
+            }
+            node.changed = false;
+            for (pos, block) in &mut self.blocks[i] {
                 if let Some(powered) = block_powered_mut(block) {
                     *powered = node.powered
                 }
@@ -268,7 +275,6 @@ impl JITBackend for DirectBackend {
                 }
                 world.set_block(*pos, *block);
             }
-            node.changed = false;
         }
     }
 
@@ -378,8 +384,15 @@ impl fmt::Display for DirectBackend {
                 NodeType::Constant => format!("Constant({})", node.output_power),
                 NodeType::NoteBlock { .. } => "NoteBlock".to_string(),
             };
-            let pos = if let Some((pos, _)) = self.blocks[id] {
-                format!("{}, {}, {}", pos.x, pos.y, pos.z)
+            let pos = if !self.blocks[id].is_empty() {
+                let mut string = String::new();
+                for (idx, (pos, _)) in self.blocks[id].iter().enumerate() {
+                    if idx != 0 {
+                        write!(&mut string, "; ")?;
+                    }
+                    write!(&mut string, "{}, {}, {}", pos.x, pos.y, pos.z)?;
+                }
+                string
             } else {
                 "No Pos".to_string()
             };

crates/redpiler/src/compile_graph.rs

@@ -1,6 +1,7 @@
 use mchprs_blocks::blocks::{ComparatorMode, Instrument};
 use mchprs_blocks::BlockPos;
 use petgraph::stable_graph::{NodeIndex, StableGraph};
+use smallvec::SmallVec;
 
 pub type NodeIdx = NodeIndex;
 
@@ -91,7 +92,7 @@ pub struct Annotations {}
 #[derive(Debug)]
 pub struct CompileNode {
     pub ty: NodeType,
-    pub block: Option<(BlockPos, u32)>,
+    pub block: SmallVec<[(BlockPos, u32); 1]>,
     pub name: Option<String>,
     pub state: NodeState,
 

crates/redpiler/src/passes/frontend/identify_nodes.rs

@@ -18,6 +18,7 @@ use mchprs_redstone::{self, comparator, noteblock, wire};
 use mchprs_world::{for_each_block_optimized, World};
 use rustc_hash::{FxHashMap, FxHashSet};
 use serde_json::Value;
+use smallvec::smallvec;
 use tracing::warn;
 
 pub struct IdentifyNodes;
@@ -95,7 +96,7 @@ fn for_pos<W: World>(
 
     let node_idx = graph.add_node(CompileNode {
         ty,
-        block: Some((pos, id)),
+        block: smallvec![(pos, id)],
         name: None,
         state,
 

crates/redpiler/src/passes/frontend/input_search.rs

@@ -78,8 +78,8 @@ impl<'a, W: World> InputSearchState<'a, W> {
     fn new(world: &'a W, graph: &'a mut CompileGraph) -> InputSearchState<'a, W> {
         let mut pos_map = FxHashMap::default();
         for id in graph.node_indices() {
-            if let Some((pos, _)) = graph[id].block {
-                pos_map.insert(pos, id);
+            for (pos, _) in &graph[id].block {
+                pos_map.insert(*pos, id);
             }
         }
 
@@ -353,7 +353,7 @@ impl<'a, W: World> InputSearchState<'a, W> {
                 continue;
             }
             let node = &self.graph[idx];
-            if let Some(block) = node.block {
+            if let Some(block) = node.block.first().copied() {
                 self.search_node(idx, block);
             }
         }

crates/redpiler/src/passes/misc/export_graph.rs

@@ -67,16 +67,21 @@ fn convert_node(
             CNodeType::Constant => NodeType::Constant,
             CNodeType::NoteBlock { .. } => NodeType::NoteBlock,
         },
-        block: node.block.map(|(pos, id)| {
-            (
-                BlockPos {
-                    x: pos.x,
-                    y: pos.y,
-                    z: pos.z,
-                },
-                id,
-            )
-        }),
+        block: node
+            .block
+            .iter()
+            .copied()
+            .map(|(pos, id)| {
+                (
+                    BlockPos {
+                        x: pos.x,
+                        y: pos.y,
+                        z: pos.z,
+                    },
+                    id,
+                )
+            })
+            .collect(),
         state: NodeState {
             output_strength: node.state.output_strength,
             powered: node.state.powered,

crates/redpiler/src/passes/opt/coalesce.rs

@@ -104,5 +104,7 @@ fn coalesce(graph: &mut CompileGraph, node: NodeIdx, into: NodeIdx) {
         let weight = graph.remove_edge(edge_idx).unwrap();
         graph.add_edge(into, dest, weight);
     }
-    graph.remove_node(node);
+    if let Some(mut node) = graph.remove_node(node) {
+        graph[into].block.append(&mut node.block);
+    }
 }

crates/redpiler/src/passes/opt/constant_coalesce.rs

@@ -53,7 +53,7 @@ impl<W: World> Pass<W> for ConstantCoalesce {
                     Entry::Vacant(entry) => {
                         let constant_idx = graph.add_node(CompileNode {
                             ty: NodeType::Constant,
-                            block: None,
+                            block: Default::default(),
                             name: None,
                             state: NodeState::ss(ss),
                             is_input: false,

crates/redpiler/src/ril.rs

@@ -192,8 +192,14 @@ fn dump_node(f: &mut impl fmt::Write, ctx: &FmtContext<'_>) -> fmt::Result {
         }
     }?;
 
-    if let Some((pos, _)) = node.block {
-        write!(f, "  # Loc: {}", pos)?;
+    if node.block.len() > 0 {
+        write!(f, "  # Loc: ")?;
+        for (idx, (pos, _)) in node.block.iter().copied().enumerate() {
+            if idx != 0 {
+                write!(f, ", ")?;
+            }
+            write!(f, "{}", pos)?;
+        }
     }
 
     Ok(())
@@ -629,7 +635,7 @@ impl RILModule {
                 ty: component.node_ty.clone(),
                 state: component.node_state.clone(),
                 name: Some(component.name.clone()),
-                block: None,
+                block: Default::default(),
                 is_input: component.node_ty.is_normally_input(),
                 is_output: component.node_ty.is_normally_output(),
                 annotations: Default::default(),

crates/redpiler_graph/Cargo.toml

@@ -15,3 +15,4 @@ repository.workspace = true
 [dependencies]
 serde = { workspace = true, features = ["derive"] }
 bincode = { workspace = true }
+smallvec = { workspace = true }

crates/redpiler_graph/src/lib.rs

@@ -1,5 +1,6 @@
 use bincode::{BincodeRead, Result};
 use serde::{Deserialize, Serialize};
+use smallvec::SmallVec;
 
 pub type NodeId = usize;
 
@@ -55,7 +56,7 @@ pub struct NodeState {
 pub struct Node {
     pub ty: NodeType,
     /// Position and protocol id for block
-    pub block: Option<(BlockPos, u32)>,
+    pub block: SmallVec<[(BlockPos, u32); 1]>,
     pub state: NodeState,
 
     pub facing_diode: bool,