Feat/cli label indices

huff_cli/src/huffc.rs

@@ -18,8 +18,8 @@ use huff_tests::{
 use huff_utils::{
     file_provider::FileSystemFileProvider,
     prelude::{
-        export_interfaces, gen_sol_interfaces, str_to_bytes32, unpack_files, AstSpan, CodegenError,
-        CodegenErrorKind, CompilerError, FileSource, Literal, OutputLocation, Span,
+        export_interfaces, gen_sol_interfaces, str_to_bytes32, unpack_files, AstSpan, BytecodeRes,
+        CodegenError, CodegenErrorKind, CompilerError, FileSource, Literal, OutputLocation, Span,
     },
 };
 use isatty::stdout_isatty;
@@ -82,6 +82,10 @@ struct Huff {
     #[clap(short = 'v', long = "verbose")]
     verbose: bool,
 
+    /// Prints out the jump label PC indices for the specified contract.
+    #[clap(short = 'x', long = "label-indices")]
+    label_indices: bool,
+
     /// Override / set constants for the compilation environment.
     #[clap(short = 'c', long = "constants", multiple_values = true)]
     constants: Option<Vec<String>>,
@@ -204,6 +208,60 @@ fn main() {
         file_provider: Arc::new(FileSystemFileProvider {}),
     };
 
+    if cli.label_indices {
+        match compiler.grab_contracts() {
+            Ok(contracts) => {
+                if contracts.len() > 1 {
+                    eprintln!(
+                        "{}",
+                        Paint::red("Multiple contracts found. Please specify a single contract and try again.")
+                    );
+                    std::process::exit(1);
+                }
+
+                if let Some(contract) = contracts.first() {
+                    let m_macro = contract.find_macro_by_name("MAIN").unwrap();
+
+                    // Recurse through the macro and generate bytecode
+                    let bytecode_res: BytecodeRes = Codegen::macro_to_bytecode(
+                        m_macro.clone(),
+                        contract,
+                        &mut vec![m_macro],
+                        0,
+                        &mut Vec::default(),
+                        false,
+                        None,
+                    )
+                    .unwrap();
+
+                    // Format the label indices nicely
+                    let label_indices = bytecode_res
+                        .label_indices
+                        .iter()
+                        .map(|(label, index)| format!("{}: {:#04x}", label, index))
+                        .collect::<Vec<_>>()
+                        .join("\n  ");
+
+                    println!("\n------ JUMP LABEL INDICES ------");
+                    println!("  {}", label_indices);
+                    println!("--------------------------------\n")
+                } else {
+                    eprintln!(
+                        "{}",
+                        Paint::red("No contract found. Please specify a contract and try again.")
+                    );
+                    std::process::exit(1);
+                }
+            }
+            Err(e) => {
+                tracing::error!(target: "cli", "PARSER ERRORED!");
+                eprintln!("{}", Paint::red(e));
+                std::process::exit(1);
+            }
+        }
+        return
+    }
+
     if let Some(TestCommands::Test { format, match_ }) = cli.test {
         match compiler.grab_contracts() {
             Ok(contracts) => {