IR support for switch instructions

forc-pkg/Cargo.toml

@@ -37,6 +37,7 @@ serde_with.workspace = true
 sway-core.workspace = true
 sway-error.workspace = true
 sway-features.workspace = true
+sway-ir.workspace = true
 sway-types.workspace = true
 sway-utils.workspace = true
 tar.workspace = true

forc-pkg/src/pkg.rs

@@ -1940,6 +1940,82 @@ pub fn compile(
     Ok(compiled_package)
 }
 
+/// Compiles the given Sway-IR script.
+pub fn compile_ir(
+    main_file: &Path,
+    engines: &Engines,
+    experimental: ExperimentalFeatures,
+    source_map: &mut SourceMap,
+) -> Result<BuiltPackageBytecode> {
+    // Read main_file into a string
+    let source = fs::read_to_string(main_file)
+        .with_context(|| format!("Failed to read Sway-IR file: {}", main_file.display()))?;
+
+    let sway_ir = sway_ir::parser::parse(
+        &source,
+        engines.se(),
+        experimental,
+        sway_ir::context::Backtrace::default(),
+    )?;
+
+    let handler = Handler::default();
+    let asm = sway_core::asm_generation::from_ir::compile_ir_context_to_finalized_asm(
+        &handler, &sway_ir, None,
+    );
+
+    let fail = |handler: Handler| {
+        let (errors, warnings, infos) = handler.consume();
+        print_on_failure(engines.se(), false, &infos, &warnings, &errors, false);
+        bail!("Failed to compile {}", main_file.display());
+    };
+    let asm = match asm {
+        Err(_) => return fail(handler),
+        Ok(asm) => asm,
+    };
+
+    let panic_occurrences = sway_core::PanicOccurrences::default();
+    let panicking_call_occurrences = sway_core::PanickingCallOccurrences::default();
+    let mut compiled_asm = sway_core::CompiledAsm {
+        finalized_asm: asm,
+        panic_occurrences,
+        panicking_call_occurrences,
+    };
+
+    let entries = compiled_asm
+        .finalized_asm
+        .entries
+        .iter()
+        .map(|finalized_entry| PkgEntry::from_finalized_entry(finalized_entry, engines))
+        .collect::<anyhow::Result<_>>()?;
+
+    let sway_build_config = sway_core::BuildConfig::dummy_for_asm_generation();
+
+    let bc_res = sway_core::asm_to_bytecode(
+        &handler,
+        &mut compiled_asm,
+        source_map,
+        engines.se(),
+        &sway_build_config,
+    );
+    let errored = handler.has_errors();
+
+    let compiled = match bc_res {
+        Ok(compiled) if !errored => compiled,
+        _ => return fail(handler),
+    };
+
+    let (_, _warnings, _infos) = handler.consume();
+    // TODO Print infos and warnings?
+    // TODO: Set configurables offset metadata if needed.
+
+    let bytecode = BuiltPackageBytecode {
+        bytes: compiled.bytecode,
+        entries,
+    };
+
+    Ok(bytecode)
+}
+
 /// Reports assembly information for a compiled package to an external `dyno` process through `stdout`.
 fn report_assembly_information(
     compiled_asm: &sway_core::CompiledAsm,
@@ -1982,6 +2058,8 @@ fn report_assembly_information(
                 }
             }
 
+            sway_core::asm_generation::Datum::WordArray(words) => (words.len() * 8) as u64,
+
             sway_core::asm_generation::Datum::Collection(items) => {
                 items.iter().map(calculate_entry_size).sum()
             }

sway-core/src/asm_generation/evm/evm_asm_builder.rs

@@ -319,6 +319,11 @@ impl<'ir, 'eng> EvmAsmBuilder<'ir, 'eng> {
                 } => {
                     self.compile_conditional_branch(handler, cond_value, true_block, false_block)?
                 }
+                InstOp::Switch {
+                    discriminant,
+                    cases,
+                    default,
+                } => self.compile_switch(handler, instr_val, discriminant, cases, default)?,
                 InstOp::ContractCall {
                     params,
                     coins,
@@ -444,6 +449,17 @@ impl<'ir, 'eng> EvmAsmBuilder<'ir, 'eng> {
         todo!();
     }
 
+    fn compile_switch(
+        &mut self,
+        handler: &Handler,
+        instr_val: &Value,
+        discriminant: &Value,
+        cases: &[(u64, BranchToWithArgs)],
+        default: &Option<BranchToWithArgs>,
+    ) -> Result<(), ErrorEmitted> {
+        todo!();
+    }
+
     fn compile_branch_to_phi_value(&mut self, to_block: &BranchToWithArgs) {
         todo!();
     }

sway-core/src/asm_generation/finalized_asm.rs

@@ -323,6 +323,19 @@ fn to_bytecode_mut(
                     }
                     println!("\"");
                 }
+                Datum::WordArray(ws) => {
+                    print!(".words as hex (");
+
+                    let mut first = true;
+                    for w in ws {
+                        if !first {
+                            print!(", ");
+                        }
+                        first = false;
+                        print!("{:02X?}", w.to_be_bytes());
+                    }
+                    println!("), len i{}", ws.len());
+                }
                 Datum::Slice(bs) => {
                     print!(".slice as hex ({bs:02X?}), len i{}, as ascii \"", bs.len());
 

sway-core/src/asm_generation/fuel/allocated_abstract_instruction_set.rs

@@ -285,56 +285,126 @@ impl AllocatedAbstractInstructionSet {
                     owning_span,
                     comment,
                 }),
-                Either::Right(org_op) => match org_op {
-                    ControlFlowOp::Jump { to, type_ } => {
-                        let target_offset = label_offsets.get(&to).unwrap().offs;
-                        let ops = if matches!(type_, JumpType::Call) {
-                            compile_call(
-                                data_section,
-                                curr_offset,
-                                target_offset,
-                                far_jump_sizes.get(&op_idx).copied(),
-                                comment,
+                Either::Right(org_op) => {
+                    match org_op {
+                        ControlFlowOp::Jump { to, type_ } => {
+                            let target_offset = label_offsets.get(&to).unwrap().offs;
+                            let ops = if matches!(type_, JumpType::Call) {
+                                compile_call(
+                                    data_section,
+                                    curr_offset,
+                                    target_offset,
+                                    far_jump_sizes.get(&op_idx).copied(),
+                                    comment,
+                                    owning_span,
+                                )
+                            } else {
+                                compile_jump(
+                                    data_section,
+                                    curr_offset,
+                                    target_offset,
+                                    match type_ {
+                                        JumpType::NotZero(cond) => Some(cond),
+                                        _ => None,
+                                    },
+                                    far_jump_sizes.contains_key(&op_idx),
+                                    comment,
+                                    owning_span,
+                                )
+                            };
+                            debug_assert_eq!(ops.len() as u64, op_size);
+                            realized_ops.extend(ops);
+                        }
+                        ControlFlowOp::Switch {
+                            discriminant,
+                            cases,
+                        } => {
+                            let target_offsets = cases
+                            .iter()
+                            .map(|label| {
+                                let label_offset = label_offsets
+                                    .get(label)
+                                    .map(|b| b.offs)
+                                    .expect("Switch case label not found in label_offsets");
+                                label_offset.checked_sub(curr_offset).expect(
+                                    "Switch case target offset is before switch instruction",
+                                ).checked_sub(op_size).expect("Switch case target offset right at switch instruction")
+                            })
+                            .collect::<Vec<u64>>();
+                            // Insert a data section entry for the switch targets.
+                            let data_id = data_section.insert_data_value(Entry::new_word_array(
+                                target_offsets,
+                                EntryName::NonConfigurable,
+                                None,
+                            ));
+                            realized_ops.push(RealizedOp {
+                                opcode: AllocatedInstruction::AddrDataId(
+                                    AllocatedRegister::Constant(ConstantRegister::Scratch),
+                                    data_id,
+                                ),
+                                owning_span: owning_span.clone(),
+                                comment: "load switch target table address".into(),
+                            });
+                            // Multiply discriminant by 8 (since each address is 8 bytes) and add to the base address.
+                            realized_ops.push(RealizedOp {
+                                opcode: AllocatedInstruction::SLLI(
+                                    discriminant.clone(),
+                                    discriminant.clone(),
+                                    VirtualImmediate12::new(3),
+                                ),
+                                owning_span: owning_span.clone(),
+                                comment: "multiply discriminant by 8".into(),
+                            });
+                            realized_ops.push(RealizedOp {
+                                opcode: AllocatedInstruction::ADD(
+                                    AllocatedRegister::Constant(ConstantRegister::Scratch),
+                                    discriminant,
+                                    AllocatedRegister::Constant(ConstantRegister::Scratch),
+                                ),
+                                owning_span: owning_span.clone(),
+                                comment: "add discriminant to switch target table address".into(),
+                            });
+                            realized_ops.push(RealizedOp {
+                                opcode: AllocatedInstruction::LW(
+                                    AllocatedRegister::Constant(ConstantRegister::Scratch),
+                                    AllocatedRegister::Constant(ConstantRegister::Scratch),
+                                    VirtualImmediate12::new(0),
+                                ),
+                                owning_span: owning_span.clone(),
+                                comment: "load switch target address".into(),
+                            });
+                            // Finally, jump to the loaded address.
+                            realized_ops.push(RealizedOp {
+                                opcode: AllocatedInstruction::JMPF(
+                                    AllocatedRegister::Constant(ConstantRegister::Scratch),
+                                    VirtualImmediate18::new(0),
+                                ),
                                 owning_span,
-                            )
-                        } else {
-                            compile_jump(
-                                data_section,
-                                curr_offset,
-                                target_offset,
-                                match type_ {
-                                    JumpType::NotZero(cond) => Some(cond),
-                                    _ => None,
-                                },
-                                far_jump_sizes.contains_key(&op_idx),
                                 comment,
-                                owning_span,
-                            )
-                        };
-                        debug_assert_eq!(ops.len() as u64, op_size);
-                        realized_ops.extend(ops);
-                    }
-                    ControlFlowOp::DataSectionOffsetPlaceholder => {
-                        realized_ops.push(RealizedOp {
-                            opcode: AllocatedInstruction::DataSectionOffsetPlaceholder,
-                            owning_span: None,
-                            comment: String::new(),
-                        });
-                    }
-                    ControlFlowOp::ConfigurablesOffsetPlaceholder => {
-                        realized_ops.push(RealizedOp {
-                            opcode: AllocatedInstruction::ConfigurablesOffsetPlaceholder,
-                            owning_span: None,
-                            comment: String::new(),
-                        });
-                    }
-                    ControlFlowOp::Comment => continue,
-                    ControlFlowOp::Label(..) => continue,
+                            });
+                        }
+                        ControlFlowOp::DataSectionOffsetPlaceholder => {
+                            realized_ops.push(RealizedOp {
+                                opcode: AllocatedInstruction::DataSectionOffsetPlaceholder,
+                                owning_span: None,
+                                comment: String::new(),
+                            });
+                        }
+                        ControlFlowOp::ConfigurablesOffsetPlaceholder => {
+                            realized_ops.push(RealizedOp {
+                                opcode: AllocatedInstruction::ConfigurablesOffsetPlaceholder,
+                                owning_span: None,
+                                comment: String::new(),
+                            });
+                        }
+                        ControlFlowOp::Comment => continue,
+                        ControlFlowOp::Label(..) => continue,
 
-                    ControlFlowOp::PushAll(_) | ControlFlowOp::PopAll(_) => {
-                        unreachable!("still don't belong in organisational ops")
+                        ControlFlowOp::PushAll(_) | ControlFlowOp::PopAll(_) => {
+                            unreachable!("still don't belong in organisational ops")
+                        }
                     }
-                },
+                }
             };
             curr_offset += op_size;
         }
@@ -397,6 +467,10 @@ impl AllocatedAbstractInstructionSet {
                 JumpType::Call => 3,
             },
 
+            // A switch expands to AddrDataId (2 opcodes) + scale discriminant by word size (1 opcode)
+            // + add discriminant (1 opcode) + load (1 opcode) + jump (1 opcode) = 6 opcodes
+            Either::Right(Switch { .. }) => 6,
+
             Either::Right(Comment) => 0,
 
             Either::Right(DataSectionOffsetPlaceholder) => {
@@ -454,6 +528,10 @@ impl AllocatedAbstractInstructionSet {
             // Far jumps must be handled separately, as they require two instructions.
             Either::Right(Jump { .. }) => 1,
 
+            // A switch expands to AddrDataId (2 opcodes) + scale discriminant by word size (1 opcode)
+            // + add discriminant (1 opcode) + load (1 opcode) + jump (1 opcode) = 6 opcodes
+            Either::Right(Switch { .. }) => 6,
+
             Either::Right(Comment) => 0,
 
             Either::Right(DataSectionOffsetPlaceholder) => {
@@ -492,7 +570,10 @@ impl AllocatedAbstractInstructionSet {
 
         for (op_idx, op) in self.ops.iter().enumerate() {
             // If we're seeing a control flow op then it's the end of the block.
-            if let Either::Right(ControlFlowOp::Label(_) | ControlFlowOp::Jump { .. }) = op.opcode {
+            if let Either::Right(
+                ControlFlowOp::Label(_) | ControlFlowOp::Jump { .. } | ControlFlowOp::Switch { .. },
+            ) = op.opcode
+            {
                 if let Some((lab, _idx, offs)) = cur_basic_block {
                     // Insert the previous basic block.
                     labelled_blocks.insert(lab, BasicBlock { offs });
@@ -575,7 +656,10 @@ impl AllocatedAbstractInstructionSet {
 
         for (op_idx, op) in self.ops.iter().enumerate() {
             // If we're seeing a control flow op then it's the end of the block.
-            if let Either::Right(ControlFlowOp::Label(_) | ControlFlowOp::Jump { .. }) = op.opcode {
+            if let Either::Right(
+                ControlFlowOp::Label(_) | ControlFlowOp::Jump { .. } | ControlFlowOp::Switch { .. },
+            ) = op.opcode
+            {
                 if let Some((lab, _idx, offs)) = cur_basic_block {
                     // Insert the previous basic block.
                     labelled_blocks.insert(lab, BasicBlock { offs });