feat(codegen): support narrow immutables

crates/ast/src/ast/ty.rs

@@ -231,15 +231,18 @@ impl TypeSize {
     /// Panics if `bits` is not a multiple of 8 or greater than 256.
     #[inline]
     #[track_caller]
-    pub fn new_int_bits(bits: u16) -> Self {
-        Self::try_new_int_bits(bits).unwrap_or_else(|| panic!("invalid integer size: {bits}"))
+    pub const fn new_int_bits(bits: u16) -> Self {
+        match Self::try_new_int_bits(bits) {
+            Some(size) => size,
+            None => panic!("invalid integer size"),
+        }
     }
 
     /// Creates a new `TypeSize` for an integer type from **bits**.
     ///
     /// Returns None if `bits` is not a multiple of 8 or greater than 256.
     #[inline]
-    pub fn try_new_int_bits(bits: u16) -> Option<Self> {
+    pub const fn try_new_int_bits(bits: u16) -> Option<Self> {
         if bits.is_multiple_of(8) { Self::new(bits) } else { None }
     }
 
@@ -265,15 +268,18 @@ impl TypeSize {
     /// Panics if `bytes` is not in the range 1..=32.
     #[inline]
     #[track_caller]
-    pub fn new_fb_bytes(bytes: u8) -> Self {
-        Self::try_new_fb_bytes(bytes).unwrap_or_else(|| panic!("invalid fixed-bytes size: {bytes}"))
+    pub const fn new_fb_bytes(bytes: u8) -> Self {
+        match Self::try_new_fb_bytes(bytes) {
+            Some(size) => size,
+            None => panic!("invalid fixed-bytes size"),
+        }
     }
 
     /// Creates a new `TypeSize` for a fixed-bytes type from **bytes**.
     ///
     /// Returns None if `bytes` is not in the range 1..=32.
     #[inline]
-    pub fn try_new_fb_bytes(bytes: u8) -> Option<Self> {
+    pub const fn try_new_fb_bytes(bytes: u8) -> Option<Self> {
         if bytes == 0 {
             return None;
         }

crates/codegen/src/analysis/validator.rs

@@ -88,10 +88,30 @@ impl<'a> Validator<'a> {
 
     fn validate_function(&mut self, module: &Module, func: &Function) {
         self.validate_function_body(func);
+        self.validate_immutables(module, func);
         self.validate_calls(module, func);
         self.validate_function_phase(module, func);
     }
 
+    fn validate_immutables(&mut self, module: &Module, func: &Function) {
+        for (inst, instruction) in func.instructions.iter_enumerated() {
+            let InstKind::LoadImmutable { id, ty } = instruction.kind else { continue };
+            match module.get_immutable_type(id) {
+                Some(expected) if expected != ty => self.emit(format_args!(
+                    "inst{} loads immutable {} as `{ty}`, expected `{expected}`",
+                    inst.index(),
+                    id.index()
+                )),
+                None => self.emit(format_args!(
+                    "inst{} loads nonexistent immutable {}",
+                    inst.index(),
+                    id.index()
+                )),
+                _ => {}
+            }
+        }
+    }
+
     fn validate_function_body(&mut self, func: &Function) {
         let errors_before = self.error_count;
         let num_values = func.values.len();

crates/codegen/src/backend/evm/assembler/mod.rs

@@ -11,7 +11,7 @@ use crate::{
         ir::{self, assembly},
         op,
     },
-    mir::IMMUTABLE_WORD_SIZE,
+    mir::{ImmutableId, TypeSize},
 };
 use alloy_primitives::U256;
 use solar_config::{EvmVersion, OptimizationMode};
@@ -23,25 +23,32 @@ const EVM_WORD_BYTES: usize = 32;
 mod id_counter;
 pub(in crate::backend::evm) use id_counter::IdCounter;
 
-pub(super) use assembly::{AsmInst, AsmInstKind, PushValueId};
+pub(super) use assembly::{AsmInst, AsmInstKind, ImmutablePushId, PushValueId};
 pub(crate) use assembly::{DeferredConst, Label};
 
 mod local_interner;
 pub(in crate::backend::evm) use local_interner::LocalInterner;
 
 use assembly::Program as AssemblyProgram;
 
-/// A `PUSH32` immutable placeholder emitted into the assembled bytecode.
-///
-/// TODO: Track placeholder byte width here when smaller immutable references
-/// are supported.
+/// An immutable placeholder emitted into the assembled bytecode.
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub(crate) struct ImmutableRef {
-    /// The immutable's byte offset identifier.
-    pub id: u32,
-    /// Byte offset of the `PUSH32` opcode in the assembled bytecode.
-    /// The 32 placeholder bytes start one byte later.
+    /// The immutable identifier.
+    pub id: ImmutableId,
+    /// Byte offset of the `PUSH<N>` opcode in the assembled bytecode.
+    /// The placeholder bytes start one byte later.
     pub code_offset: usize,
+    /// Type size encoded by the placeholder.
+    pub type_size: TypeSize,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+pub(in crate::backend::evm) struct ImmutablePush {
+    // Unlike a deferred constant, this value is unknown until the constructor runs.
+    // Assembly must therefore retain its fixed width and emit a patch relocation.
+    id: ImmutableId,
+    type_size: TypeSize,
 }
 
 /// Result of assembly.
@@ -61,6 +68,7 @@ pub(in crate::backend::evm) struct PreparedAssembly {
     program: AssemblyProgram,
     evm_ir: Option<ir::Module>,
     push_values: LocalInterner<U256, PushValueId>,
+    immutable_pushes: LocalInterner<ImmutablePush, ImmutablePushId>,
     next_label: IdCounter<Label>,
     deferred_values: FxHashMap<DeferredConst, U256>,
 }
@@ -99,6 +107,8 @@ pub(crate) struct Assembler {
     deferred_relocations: Vec<(ir::BlockId, usize, DeferredConst)>,
     /// Interned push immediates too large for inline storage.
     push_values: LocalInterner<U256, PushValueId>,
+    /// Interned immutable placeholders.
+    immutable_pushes: LocalInterner<ImmutablePush, ImmutablePushId>,
     /// Next label ID.
     next_label: IdCounter<Label>,
     /// Next deferred constant ID.
@@ -127,6 +137,7 @@ impl Assembler {
             label_relocations: Vec::new(),
             deferred_relocations: Vec::new(),
             push_values: LocalInterner::new(),
+            immutable_pushes: LocalInterner::new(),
             next_label: IdCounter::new(),
             next_deferred: IdCounter::new(),
             deferred_values: FxHashMap::default(),
@@ -143,6 +154,7 @@ impl Assembler {
         self.label_relocations.clear();
         self.deferred_relocations.clear();
         self.push_values.clear();
+        self.immutable_pushes.clear();
         self.next_label.clear();
         self.next_deferred.clear();
         self.deferred_values.clear();
@@ -187,9 +199,9 @@ impl Assembler {
         self.deferred_values.insert(id, value);
     }
 
-    /// Emits a `PUSH32` zero placeholder for the immutable identified by `id`.
-    pub(crate) fn emit_push_immutable(&mut self, id: u32) {
-        self.push_ir_instruction(ir::Instruction::push_immutable(id));
+    /// Emits a `PUSH<N>` zero placeholder for the immutable identified by `id`.
+    pub(crate) fn emit_push_immutable(&mut self, id: ImmutableId, type_size: TypeSize) {
+        self.push_ir_instruction(ir::Instruction::push_immutable(id, type_size));
     }
 
     /// Defines a label and emits a `JUMPDEST` at the current position.
@@ -243,10 +255,22 @@ impl Assembler {
         AsmInst::push(self.push_values.intern(value))
     }
 
+    pub(in crate::backend::evm) fn immutable_push_inst(
+        &mut self,
+        id: ImmutableId,
+        type_size: TypeSize,
+    ) -> AsmInst {
+        AsmInst::push_immutable(self.immutable_pushes.intern(ImmutablePush { id, type_size }))
+    }
+
     pub(super) fn push_value(&self, index: PushValueId) -> U256 {
         *self.push_values.get(index)
     }
 
+    fn immutable_push(&self, index: ImmutablePushId) -> ImmutablePush {
+        *self.immutable_pushes.get(index)
+    }
+
     fn finish_evm_ir(&mut self) -> Option<(ir::Module, Vec<Option<Label>>)> {
         let mut module = std::mem::replace(&mut self.program, Self::new_ir_module());
         self.current_block = None;
@@ -344,6 +368,7 @@ impl Assembler {
             evm_ir,
             program,
             push_values: std::mem::take(&mut self.push_values),
+            immutable_pushes: std::mem::take(&mut self.immutable_pushes),
             next_label: std::mem::take(&mut self.next_label),
             deferred_values: std::mem::take(&mut self.deferred_values),
         }
@@ -369,6 +394,7 @@ impl Assembler {
         deferred_values: &[(DeferredConst, U256)],
     ) -> AssembledCode {
         self.push_values = prepared.push_values.clone();
+        self.immutable_pushes = prepared.immutable_pushes.clone();
         self.next_label = prepared.next_label.clone();
         self.deferred_values.clone_from(&prepared.deferred_values);
         self.deferred_values.extend(deferred_values.iter().copied());
@@ -470,9 +496,8 @@ impl Assembler {
                 AsmInstKind::PushDeferred(_) => {
                     unreachable!("deferred constants must be resolved before assembly");
                 }
-                AsmInstKind::PushImmutable(_) => {
-                    // PUSH32 opcode plus 32 placeholder bytes.
-                    offset += 33;
+                AsmInstKind::PushImmutable(id) => {
+                    offset += 1 + usize::from(self.immutable_push(id).type_size.bytes());
                 }
                 AsmInstKind::Label(label) => {
                     label_offsets.insert(label, offset);
@@ -526,7 +551,7 @@ impl Assembler {
                     unreachable!("deferred constants must be resolved before assembly");
                 }
                 AsmInstKind::PushImmutable(id) => {
-                    out.emit_push_immutable(id);
+                    out.emit_push_immutable(self.immutable_push(id));
                 }
                 AsmInstKind::Label(_) => {
                     out.emit_op(op::JUMPDEST);
@@ -572,10 +597,15 @@ impl BytecodeAssembler {
         self.bytecode.push(opcode);
     }
 
-    fn emit_push_immutable(&mut self, id: u32) {
-        self.immutable_refs.push(ImmutableRef { id, code_offset: self.bytecode.len() });
-        self.bytecode.push(op::PUSH32);
-        self.bytecode.extend(std::iter::repeat_n(0, IMMUTABLE_WORD_SIZE));
+    fn emit_push_immutable(&mut self, push: ImmutablePush) {
+        self.immutable_refs.push(ImmutableRef {
+            id: push.id,
+            code_offset: self.bytecode.len(),
+            type_size: push.type_size,
+        });
+        let byte_width = push.type_size.bytes();
+        self.bytecode.push(op::push(byte_width));
+        self.bytecode.extend(std::iter::repeat_n(0, usize::from(byte_width)));
     }
 
     fn encoded_push_len(&self, value: U256) -> usize {
@@ -697,6 +727,35 @@ mod tests {
         assert_eq!(*asm.push_values.get(PushValueId::from_usize(0)), large);
     }
 
+    #[test]
+    fn immutable_push_uses_declared_width() {
+        let mut asm = Assembler::new();
+        let narrow = ImmutableId::new(3);
+        let address = ImmutableId::new(4);
+        let narrow_size = TypeSize::new_int_bits(8);
+        let address_size = TypeSize::new_int_bits(160);
+
+        asm.emit_push_immutable(narrow, narrow_size);
+        asm.emit_push_immutable(address, address_size);
+        let result = asm.assemble();
+
+        assert_data_eq!(
+            disassemble(&result.bytecode),
+            str![[r#"
+PUSH1 0x00
+PUSH20 0x0000000000000000000000000000000000000000
+
+"#]]
+        );
+        assert_eq!(
+            result.immutable_refs,
+            [
+                ImmutableRef { id: narrow, code_offset: 0, type_size: narrow_size },
+                ImmutableRef { id: address, code_offset: 2, type_size: address_size },
+            ]
+        );
+    }
+
     #[test]
     fn assembler_can_be_reused_after_assembly() {
         let mut asm = Assembler::new();
@@ -714,6 +773,7 @@ PUSH4 0x80000000
         );
         assert!(asm.program.blocks.is_empty());
         assert_eq!(asm.push_values.len(), 0);
+        assert_eq!(asm.immutable_pushes.len(), 0);
 
         asm.emit_push(U256::from(2));
         let second = asm.assemble();