huff_neo_utils/
bytecode.rs

1//! Bytecode Traits
2//!
3//! Abstract translating state into bytecode.
4
5use crate::ast::span::AstSpan;
6use crate::scope::ScopeId;
7use crate::{
8    evm_version::EVMVersion,
9    prelude::{Statement, TableDefinition},
10};
11use std::{
12    collections::BTreeMap,
13    fmt::{self, Display},
14};
15
16/// Push opcode for jump instructions
17#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
18pub enum PushOpcode {
19    /// PUSH1 opcode (0x60) - pushes 1 byte onto stack
20    Push1,
21    /// PUSH2 opcode (0x61) - pushes 2 bytes onto stack
22    Push2,
23    /// PUSH3 opcode (0x62) - pushes 3 bytes onto stack; needed for offsets beyond 64 KiB, which
24    /// initcode can reach from Amsterdam (EIP-7954)
25    Push3,
26}
27
28impl PushOpcode {
29    /// Convert to hex string representation
30    pub fn to_hex(&self) -> &'static str {
31        match self {
32            PushOpcode::Push1 => "60",
33            PushOpcode::Push2 => "61",
34            PushOpcode::Push3 => "62",
35        }
36    }
37
38    /// Get the number of bytes this push opcode accepts (not including the opcode itself)
39    pub fn byte_size(&self) -> usize {
40        match self {
41            PushOpcode::Push1 => 1,
42            PushOpcode::Push2 => 2,
43            PushOpcode::Push3 => 3,
44        }
45    }
46
47    /// Check if a value can be represented by this push opcode
48    pub fn can_represent(&self, value: usize) -> bool {
49        value < 1 << (8 * self.byte_size())
50    }
51
52    /// Get the optimal push opcode for a given value
53    pub fn optimal_for(value: usize) -> Self {
54        if value <= 0xFF {
55            PushOpcode::Push1
56        } else if value <= 0xFFFF {
57            PushOpcode::Push2
58        } else {
59            PushOpcode::Push3
60        }
61    }
62
63    /// Formats a value as the zero-padded push data of this opcode
64    pub fn format_value(&self, value: usize) -> String {
65        format!("{value:0width$x}", width = self.byte_size() * 2)
66    }
67
68    /// Parse from hex string (for tests and compatibility)
69    pub fn from_hex(s: &str) -> Option<Self> {
70        match s {
71            "60" => Some(PushOpcode::Push1),
72            "61" => Some(PushOpcode::Push2),
73            "62" => Some(PushOpcode::Push3),
74            _ => None,
75        }
76    }
77}
78
79/// Data for a jump placeholder that needs label resolution
80#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
81pub struct JumpPlaceholderData {
82    /// The label name this jump references
83    pub label: String,
84    /// The push opcode to use for the jump destination
85    pub push_opcode: PushOpcode,
86    /// The opcode suffix (e.g., "565b" for Jump+Jumpdest, or "" for simple jumps)
87    pub suffix: String,
88    /// The resolved target offset (None if not yet resolved)
89    pub target_offset: Option<usize>,
90}
91
92impl JumpPlaceholderData {
93    /// Create a new unresolved jump placeholder
94    pub fn new(label: String, push_opcode: PushOpcode, suffix: String) -> Self {
95        Self { label, push_opcode, suffix, target_offset: None }
96    }
97
98    /// Create a jump placeholder with a resolved target
99    pub fn with_target(label: String, push_opcode: PushOpcode, suffix: String, target_offset: usize) -> Self {
100        Self { label, push_opcode, suffix, target_offset: Some(target_offset) }
101    }
102
103    /// Resolve this jump to a target offset
104    pub fn resolve(&mut self, target_offset: usize) {
105        self.target_offset = Some(target_offset);
106    }
107
108    /// Generate the hex string representation
109    /// If resolved, returns the actual bytecode. If unresolved, returns placeholder with "xxxx"
110    pub fn to_hex(&self) -> String {
111        let value_hex = if let Some(offset) = self.target_offset {
112            // Resolved: format as actual hex value
113            self.push_opcode.format_value(offset)
114        } else {
115            // Unresolved: use placeholder
116            "x".repeat(self.push_opcode.byte_size() * 2)
117        };
118
119        format!("{}{}{}", self.push_opcode.to_hex(), value_hex, self.suffix)
120    }
121}
122
123/// Data for a circular codesize placeholder
124#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
125pub struct CircularCodesizePlaceholderData {
126    /// The macro name that contains this circular reference
127    pub macro_name: String,
128    /// The push opcode to use (starts as PUSH1 and only grows until the size fits)
129    pub push_opcode: PushOpcode,
130    /// The resolved macro size (None if not yet resolved)
131    pub resolved_value: Option<usize>,
132    /// Code range `[start, end)` of the macro invocation whose size is pushed, in the same offsets
133    /// as the surrounding segments. Set once that invocation is fully generated; without it the
134    /// size of the whole code is pushed.
135    pub measured_range: Option<(usize, usize)>,
136    /// Size of the named macro from the previous compile pass, used when the code being generated
137    /// does not contain the named macro, so its size cannot be measured here
138    pub size_hint: Option<usize>,
139}
140
141impl CircularCodesizePlaceholderData {
142    /// Create a new unresolved circular codesize placeholder
143    /// Starts with PUSH1 optimistically, may grow during resolution
144    pub fn new(macro_name: String, size_hint: Option<usize>) -> Self {
145        Self { macro_name, push_opcode: PushOpcode::Push1, resolved_value: None, measured_range: None, size_hint }
146    }
147
148    /// The size this placeholder pushes: the measured invocation, else the size hint, else the
149    /// whole code it is part of
150    pub fn size(&self, measured: impl FnOnce((usize, usize)) -> usize, total_length: usize) -> usize {
151        self.measured_range.map(measured).or(self.size_hint).unwrap_or(total_length)
152    }
153
154    /// Resolve this codesize to a specific size
155    ///
156    /// Keeps the current width if the size fits, since the surrounding code is already laid out for
157    /// it. Returns the number of bytes the push data grew by.
158    pub fn resolve(&mut self, size: usize) -> usize {
159        let old_size = self.push_opcode.byte_size();
160        if !self.push_opcode.can_represent(size) {
161            self.push_opcode = PushOpcode::optimal_for(size);
162        }
163        self.resolved_value = Some(size);
164        self.push_opcode.byte_size() - old_size
165    }
166
167    /// Generate the hex string representation
168    /// If resolved, returns the actual bytecode. If unresolved, returns placeholder with "cccc"
169    pub fn to_hex(&self) -> String {
170        if let Some(size) = self.resolved_value {
171            // Resolved: format as actual push instruction
172            format!("{}{}", self.push_opcode.to_hex(), self.push_opcode.format_value(size))
173        } else {
174            // Unresolved: use "cccc" placeholder (2 bytes = PUSH1 size)
175            "cccc".to_string()
176        }
177    }
178}
179
180/// The bytecode of `__argsstart()` until the layout is known: a PUSH2 whose value the compiler fills in.
181/// Its value is not hex, so it cannot be mistaken for code or for a jump placeholder.
182pub const ARGS_START_PLACEHOLDER: &str = "61yyyy";
183
184/// Data for an __ASSERT_PC placeholder
185/// This placeholder generates no bytecode but validates the PC during bytecode generation
186#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
187pub struct AssertPcPlaceholderData {
188    /// The expected program counter value
189    pub expected_offset: usize,
190    /// Source location for error reporting
191    pub span: AstSpan,
192}
193
194impl AssertPcPlaceholderData {
195    /// Create a new assert PC placeholder
196    pub fn new(expected_offset: usize, span: AstSpan) -> Self {
197        Self { expected_offset, span }
198    }
199}
200
201/// Bytecode with different types for placeholders
202#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
203pub enum Bytes {
204    /// Raw bytecode - ready to use
205    Raw(String),
206    /// Jump placeholder - "xxxx" needs label resolution
207    JumpPlaceholder(JumpPlaceholderData),
208    /// Circular codesize placeholder - "cccc" for self-referential __codesize
209    CircularCodesizePlaceholder(CircularCodesizePlaceholderData),
210    /// `__argsstart()` placeholder - a PUSH2 filled with the offset of the constructor arguments,
211    /// with the span of the call for errors outside the constructor
212    ArgsStartPlaceholder(AstSpan),
213    /// Assert PC placeholder - validates program counter after jump relaxation
214    AssertPcPlaceholder(AssertPcPlaceholderData),
215}
216
217impl Bytes {
218    /// Get the underlying hex string (generates it for placeholders)
219    pub fn as_str(&self) -> String {
220        match self {
221            Bytes::Raw(s) => s.clone(),
222            Bytes::JumpPlaceholder(data) => data.to_hex(),
223            Bytes::CircularCodesizePlaceholder(data) => data.to_hex(),
224            Bytes::ArgsStartPlaceholder(_) => ARGS_START_PLACEHOLDER.to_string(),
225            Bytes::AssertPcPlaceholder(_) => String::new(), // Generates no bytecode
226        }
227    }
228
229    /// Convert to hex representation (same as as_str but more explicit name)
230    pub fn to_hex(&self) -> String {
231        self.as_str()
232    }
233
234    /// Get the length in bytes (not hex characters)
235    /// Each pair of hex characters represents one byte
236    pub fn len(&self) -> usize {
237        match self {
238            Bytes::Raw(s) => s.len() / 2,
239            Bytes::JumpPlaceholder(data) => {
240                // 1 byte for opcode + N bytes for value + suffix bytes
241                1 + data.push_opcode.byte_size() + (data.suffix.len() / 2)
242            }
243            Bytes::CircularCodesizePlaceholder(data) => {
244                // 1 byte for opcode + N bytes for value
245                1 + data.push_opcode.byte_size()
246            }
247            Bytes::ArgsStartPlaceholder(_) => ARGS_START_PLACEHOLDER.len() / 2,
248            Bytes::AssertPcPlaceholder(_) => 0, // Generates no bytecode
249        }
250    }
251
252    /// Check if the bytecode is empty
253    pub fn is_empty(&self) -> bool {
254        self.len() == 0
255    }
256}
257
258impl Display for Bytes {
259    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
260        write!(f, "{}", self.as_str())
261    }
262}
263
264/// A segment of bytecode with its offset
265#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
266pub struct BytecodeSegment {
267    /// The bytecode offset where this segment starts
268    pub offset: usize,
269    /// The hex-encoded bytecode string
270    pub bytes: Bytes,
271}
272
273impl BytecodeSegment {
274    /// Create a new bytecode segment
275    pub fn new(offset: usize, bytes: Bytes) -> Self {
276        Self { offset, bytes }
277    }
278}
279
280/// Collection of bytecode segments with helper methods
281#[derive(Debug, Default, Clone, PartialEq, Eq, PartialOrd, Ord)]
282pub struct BytecodeSegments {
283    segments: Vec<BytecodeSegment>,
284}
285
286impl BytecodeSegments {
287    /// Create a new empty collection
288    pub fn new() -> Self {
289        Self { segments: Vec::new() }
290    }
291
292    /// Push a bytecode segment
293    pub fn push(&mut self, segment: BytecodeSegment) {
294        self.segments.push(segment);
295    }
296
297    /// Push a bytecode segment with an offset (without scope information)
298    pub fn push_with_offset(&mut self, offset: usize, bytes: Bytes) {
299        self.segments.push(BytecodeSegment::new(offset, bytes));
300    }
301
302    /// Extend with another collection
303    pub fn extend(&mut self, other: BytecodeSegments) {
304        self.segments.extend(other.segments);
305    }
306
307    /// Get an iterator over the segments
308    pub fn iter(&self) -> impl Iterator<Item = &BytecodeSegment> {
309        self.segments.iter()
310    }
311
312    /// Get a mutable iterator over the segments
313    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut BytecodeSegment> {
314        self.segments.iter_mut()
315    }
316
317    /// Get the number of segments
318    pub fn len(&self) -> usize {
319        self.segments.len()
320    }
321
322    /// Check if the collection is empty
323    pub fn is_empty(&self) -> bool {
324        self.segments.is_empty()
325    }
326
327    /// Calculate the cumulative byte offsets for each segment
328    /// Returns a vector where each element is the byte offset of the corresponding segment
329    pub fn calculate_offsets(&self) -> Vec<usize> {
330        let mut offsets = Vec::with_capacity(self.segments.len());
331        let mut current_offset = 0;
332
333        for segment in &self.segments {
334            offsets.push(current_offset);
335            current_offset += segment.bytes.len();
336        }
337
338        offsets
339    }
340
341    /// Calculate the total byte size of all segments
342    pub fn total_byte_size(&self) -> usize {
343        self.segments.iter().map(|seg| seg.bytes.len()).sum()
344    }
345
346    /// Find all jump placeholders with their byte offsets
347    /// Returns a vector of (offset, jump_data) tuples
348    pub fn find_jump_placeholders(&self) -> Vec<(usize, &JumpPlaceholderData)> {
349        let mut jumps = Vec::new();
350        let mut offset = 0;
351
352        for segment in &self.segments {
353            if let Bytes::JumpPlaceholder(ref data) = segment.bytes {
354                jumps.push((offset, data));
355            }
356            offset += segment.bytes.len();
357        }
358
359        jumps
360    }
361
362    /// Resolve all jump placeholders using the provided jump table and label indices
363    /// Returns a vector of unmatched jumps that couldn't be resolved
364    pub fn resolve_jumps(&mut self, jump_table: &JumpTable, label_indices: &LabelIndices) -> Result<Vec<Jump>, LabelError> {
365        let mut unmatched_jumps = Vec::new();
366
367        for segment in &mut self.segments {
368            let code_index = segment.offset;
369
370            // Check if a jump table exists at this segment's offset
371            if let Some(jt) = jump_table.get(&code_index) {
372                // Try to resolve the jump placeholder
373                if let Bytes::JumpPlaceholder(ref mut data) = segment.bytes {
374                    // Get the first jump from the table (there should only be one per segment)
375                    if let Some(jump) = jt.first() {
376                        // Verify the label names match (debugging/consistency check)
377                        // This should never happen - it indicates a compiler bug
378                        assert_eq!(
379                            data.label, jump.label,
380                            "Label mismatch at offset {:#x}: placeholder has \"{}\", jump table has \"{}\"",
381                            code_index, data.label, jump.label
382                        );
383
384                        // Look up the label's target offset
385                        match label_indices.get(jump.label.as_str(), &jump.scope_id) {
386                            Ok(Some(target_offset)) => {
387                                // Verify the offset fits in the chosen push opcode
388                                if !data.push_opcode.can_represent(target_offset) {
389                                    return Err(LabelError::JumpTargetTooLarge {
390                                        label: jump.label.clone(),
391                                        target: target_offset,
392                                        opcode: format!("{:?}", data.push_opcode),
393                                    });
394                                }
395                                // Resolve the jump
396                                data.resolve(target_offset);
397                            }
398                            Ok(None) => {
399                                // Unmatched jump - add to list
400                                unmatched_jumps.push(Jump {
401                                    label: jump.label.clone(),
402                                    bytecode_index: code_index,
403                                    span: jump.span.clone(),
404                                    scope_id: jump.scope_id.clone(),
405                                });
406                            }
407                            Err(err_msg) => {
408                                return Err(err_msg);
409                            }
410                        }
411                    }
412                }
413            }
414        }
415
416        Ok(unmatched_jumps)
417    }
418
419    /// Marks `[start, end)` as the code measured by the circular codesize placeholders of
420    /// `macro_name` that do not measure an invocation yet.
421    ///
422    /// Called when an invocation of `macro_name` is fully generated, so a placeholder measures the
423    /// innermost invocation of the macro it names.
424    pub fn measure_circular_codesize(&mut self, macro_name: &str, start: usize, end: usize) {
425        for segment in &mut self.segments {
426            if let Bytes::CircularCodesizePlaceholder(ref mut data) = segment.bytes
427                && data.macro_name == macro_name
428                && data.measured_range.is_none()
429            {
430                data.measured_range = Some((start, end));
431            }
432        }
433    }
434
435    /// Fills all circular codesize placeholders with the current size of the code they measure.
436    ///
437    /// Returns the number of bytes the placeholders grew by.
438    pub fn resolve_circular_codesize(&mut self) -> usize {
439        let total_length = self.total_byte_size();
440        let mut growth = 0;
441        for segment in &mut self.segments {
442            if let Bytes::CircularCodesizePlaceholder(ref mut data) = segment.bytes {
443                let size = data.size(|(start, end)| end - start, total_length);
444                growth += data.resolve(size);
445            }
446        }
447        growth
448    }
449}
450
451impl IntoIterator for BytecodeSegments {
452    type Item = BytecodeSegment;
453    type IntoIter = std::vec::IntoIter<BytecodeSegment>;
454
455    fn into_iter(self) -> Self::IntoIter {
456        self.segments.into_iter()
457    }
458}
459
460impl<'a> IntoIterator for &'a BytecodeSegments {
461    type Item = &'a BytecodeSegment;
462    type IntoIter = std::slice::Iter<'a, BytecodeSegment>;
463
464    fn into_iter(self) -> Self::IntoIter {
465        self.segments.iter()
466    }
467}
468
469/// Intermediate Bytecode Representation
470#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
471pub struct IRBytes<'a> {
472    /// The type of IRByte
473    pub ty: IRByteType,
474    /// The Span of the IRBytes
475    pub span: &'a AstSpan,
476}
477
478/// IRBytes Type
479#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
480pub enum IRByteType {
481    /// Bytes
482    Bytes(Bytes),
483    /// Macro Statement to be expanded
484    Statement(Box<Statement>),
485    /// A Constant to be referenced
486    Constant(String),
487    /// An Arg Call needs to use the calling macro context.
488    /// Macro name and argument name
489    ArgCall(String, String),
490}
491
492/// Full Intermediate Bytecode Representation
493#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
494pub struct IRBytecode<'a>(pub Vec<IRBytes<'a>>);
495
496/// ToIRBytecode
497///
498/// Converts a stateful object to intermediate bytecode
499pub trait ToIRBytecode<E> {
500    /// Translates `self` to intermediate bytecode representation
501    fn to_irbytecode(&self, evm_version: &EVMVersion) -> Result<IRBytecode<'_>, E>;
502}
503
504/// Full Bytecode
505#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
506pub struct Bytecode(pub String);
507
508/// ToBytecode
509///
510/// Converts a stateful object to bytecode
511pub trait ToBytecode<'a, E> {
512    /// Translates `self` to a bytecode string
513    fn to_bytecode(&self) -> Result<Bytecode, E>;
514}
515
516impl From<Vec<Bytes>> for Bytecode {
517    fn from(b: Vec<Bytes>) -> Self {
518        Bytecode(b.iter().fold("".to_string(), |acc, b| format!("{acc}{}", b.as_str())))
519    }
520}
521
522/// Result type for [huff_codegen](../../huff_codegen)'s
523/// [`recurse_bytecode`](../../huff_codegen/src/lib.rs#recurse_bytecode)
524#[derive(Debug, Default, Clone, PartialEq, Eq, PartialOrd, Ord)]
525pub struct BytecodeRes {
526    /// Resulting bytes
527    pub bytes: BytecodeSegments,
528    /// Source spans for each bytecode segment (parallel to bytes vec)
529    pub spans: Vec<Option<(usize, usize)>>,
530    /// Jump Indices
531    pub label_indices: LabelIndices,
532    /// Unmatched Jumps
533    pub unmatched_jumps: Jumps,
534    /// Table Instances
535    pub table_instances: Jumps,
536    /// Utilized Tables
537    pub utilized_tables: Vec<TableDefinition>,
538    /// Embedded Tables (table name -> first embedding offset)
539    pub embedded_tables: BTreeMap<String, usize>,
540}
541
542impl Display for BytecodeRes {
543    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
544        write!(
545            f,
546            r#"BytecodeRes(
547            bytes: [{}],
548            label_indices: {:?},
549            unmatched_jumps: {:?}
550            table_instances: {:?}
551        )"#,
552            self.bytes.iter().fold("".to_string(), |acc, b| format!("{acc}{}", b.bytes.as_str())),
553            self.label_indices,
554            self.unmatched_jumps,
555            self.table_instances
556        )
557    }
558}
559
560/// A Jump
561#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
562pub struct Jump {
563    /// Jump's Label
564    pub label: String,
565    /// Index of jump within bytecode
566    pub bytecode_index: usize,
567    /// The Jump Span
568    pub span: AstSpan,
569    /// The scope where the jump was made
570    pub scope_id: ScopeId,
571}
572
573impl Default for Jump {
574    fn default() -> Self {
575        Self { label: String::new(), bytecode_index: 0, span: AstSpan::default(), scope_id: ScopeId::new(vec![], 0) }
576    }
577}
578
579/// Type for a vec of `Jump`s
580pub type Jumps = Vec<Jump>;
581
582/// Scoped label information
583#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
584pub struct ScopedLabel {
585    /// The bytecode offset of the label
586    pub offset: usize,
587    /// The scope where the label was defined
588    pub scope_id: ScopeId,
589}
590
591/// Error type for label operations
592#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
593pub enum LabelError {
594    /// A label was defined multiple times in the same scope
595    #[error("Duplicate label '{0}' in the same scope")]
596    DuplicateLabelInScope(String),
597    /// A label is defined in multiple sibling scopes, creating ambiguity
598    #[error("Duplicate label '{0}' across sibling scopes")]
599    DuplicateLabelAcrossSiblings(String),
600    /// A jump target is too large for the specified opcode
601    #[error("Jump target {target:#x} too large for {opcode} in label \"{label}\"")]
602    JumpTargetTooLarge {
603        /// The label name
604        label: String,
605        /// The target offset that was too large
606        target: usize,
607        /// The opcode that couldn't represent the target
608        opcode: String,
609    },
610}
611
612/// Scoped label indices that track labels with their scope information
613#[derive(Debug, Default, Clone, PartialEq, Eq, PartialOrd, Ord)]
614pub struct ScopedLabelIndices {
615    /// Map of label name to all its definitions with scope info
616    labels: BTreeMap<String, Vec<ScopedLabel>>,
617}
618
619impl ScopedLabelIndices {
620    /// Create a new empty scoped label indices
621    pub fn new() -> Self {
622        Self { labels: BTreeMap::new() }
623    }
624
625    /// Insert a label with scope information
626    /// Returns an error if the label already exists in the same scope
627    pub fn insert(&mut self, name: String, offset: usize, scope_id: &ScopeId) -> Result<(), LabelError> {
628        let entry = self.labels.entry(name.clone()).or_default();
629
630        // Check if label already exists at the same scope
631        for existing in entry.iter() {
632            if existing.scope_id == *scope_id {
633                return Err(LabelError::DuplicateLabelInScope(name));
634            }
635        }
636
637        entry.push(ScopedLabel { offset, scope_id: scope_id.clone() });
638        Ok(())
639    }
640
641    /// Get the label offset for a given name, considering scope
642    /// Returns the label from the deepest matching scope
643    /// Returns an error if duplicate labels are found in sibling scopes when cross-referencing
644    pub fn get(&self, name: &str, current_scope: &ScopeId) -> Result<Option<usize>, LabelError> {
645        let current_scope_depth = current_scope.depth;
646        let current_scope_path = &current_scope.path;
647
648        self.labels.get(name).map_or(Ok(None), |labels| {
649            // First, try to find a label in the current scope or parent scopes
650            let in_scope = labels
651                .iter()
652                .filter(|label| {
653                    // Label is accessible if it's at the same depth or shallower
654                    // and the scope path matches up to the label's depth
655                    label.scope_id.depth <= current_scope_depth
656                        && label.scope_id.path.len() <= current_scope_path.len()
657                        && label.scope_id.path.iter().zip(current_scope_path.iter()).all(|(a, b)| a == b)
658                })
659                .max_by_key(|label| label.scope_id.depth)
660                .map(|label| label.offset);
661
662            if in_scope.is_some() {
663                return Ok(in_scope);
664            }
665
666            // If not found in current/parent scopes, search in child scopes
667            // This allows accessing labels defined in macros invoked from the current scope
668            let child_labels: Vec<&ScopedLabel> = labels
669                .iter()
670                .filter(|label| {
671                    // Check if this label is in a child scope (starts with our scope path)
672                    label.scope_id.path.len() > current_scope_path.len()
673                        && current_scope_path.iter().zip(label.scope_id.path.iter()).all(|(a, b)| a == b)
674                })
675                .collect();
676
677            if !child_labels.is_empty() {
678                // Check for duplicates in child scopes
679                if child_labels.len() > 1 {
680                    // Multiple children define this label - use the first one (depth-first order)
681                    // This matches the order of macro expansion
682                }
683                return Ok(child_labels.first().map(|label| label.offset));
684            }
685
686            // If not found in current/parent/child scopes, search in sibling scopes
687            // This allows accessing labels defined in sibling macros invoked from the same parent
688            // We need to check siblings at all levels up the tree
689            for depth in (0..current_scope_path.len()).rev() {
690                // Get the scope at this depth
691                let search_scope = &current_scope_path[..depth + 1];
692
693                // Find all labels in cousin/sibling scopes that share this common ancestor
694                let cousin_labels: Vec<&ScopedLabel> = labels
695                    .iter()
696                    .filter(|label| {
697                        // Must have at least the depth we're checking
698                        if label.scope_id.path.len() <= depth {
699                            return false;
700                        }
701
702                        // Must share the common ancestor path
703                        if label.scope_id.path[..depth] != search_scope[..depth] {
704                            return false;
705                        }
706
707                        // If we're at the immediate parent level, exclude our own path
708                        if depth == current_scope_path.len() - 1 {
709                            // Same depth siblings
710                            label.scope_id.path.len() == current_scope_path.len() && label.scope_id.path != *current_scope_path
711                        } else {
712                            // Different branch from this ancestor - must diverge after depth
713                            if label.scope_id.path.len() > depth {
714                                // The path diverges at depth+1 (different child of common ancestor)
715                                label.scope_id.path.get(depth) != current_scope_path.get(depth)
716                            } else {
717                                false
718                            }
719                        }
720                    })
721                    .collect();
722
723                if !cousin_labels.is_empty() {
724                    // Check for duplicates across cousins only when we're doing cross-reference
725                    if cousin_labels.len() > 1 {
726                        // Multiple cousins define this label - check if they're in different branches
727                        let unique_scopes: std::collections::HashSet<_> = cousin_labels.iter().map(|label| &label.scope_id.path).collect();
728
729                        if unique_scopes.len() > 1 {
730                            // Labels are in different cousin scopes - ambiguous reference
731                            return Err(LabelError::DuplicateLabelAcrossSiblings(name.to_string()));
732                        }
733                    }
734
735                    return Ok(cousin_labels.first().map(|label| label.offset));
736                }
737            }
738
739            Ok(None)
740        })
741    }
742
743    /// Get any label offset for a given name, regardless of scope
744    /// Used for jump tables which may reference labels in any scope
745    /// Returns the first matching label if there are multiple
746    pub fn get_any(&self, name: &str) -> Option<usize> {
747        self.labels.get(name).and_then(|labels| labels.first().map(|label| label.offset))
748    }
749
750    /// Extend this indices with another set of indices
751    pub fn extend(&mut self, other: ScopedLabelIndices) {
752        for (name, new_labels) in other.labels {
753            self.labels.entry(name).or_default().extend(new_labels);
754        }
755    }
756
757    /// Check if the indices are empty
758    pub fn is_empty(&self) -> bool {
759        self.labels.is_empty()
760    }
761
762    /// Iterate over all labels
763    pub fn iter(&self) -> impl Iterator<Item = (&String, &Vec<ScopedLabel>)> {
764        self.labels.iter()
765    }
766
767    /// Update label offsets using a mapping from old offsets to new offsets
768    /// This is used after jump relaxation to update label positions
769    pub fn update_offsets(&mut self, offset_mapping: &BTreeMap<usize, usize>) {
770        for labels in self.labels.values_mut() {
771            for label in labels.iter_mut() {
772                if let Some(&new_offset) = offset_mapping.get(&label.offset) {
773                    label.offset = new_offset;
774                }
775            }
776        }
777    }
778}
779
780/// Type to map `Jump` labels to their bytecode indices (for backward compatibility)
781pub type LabelIndices = ScopedLabelIndices;
782
783/// Type for a map of bytecode indexes to `Jumps`. Represents a Jump Table.
784pub type JumpTable = BTreeMap<usize, Jumps>;
huff_neo_utils/bytecode.rs

huff_neo_utils/
bytecode.rs
