feat: use PolkaVM disassembler (#6)

Integrate the PolkaVM disassembler to fix `--asm` output

Co-authored-by: Cyrill Leutwiler <bigcyrill@hotmail.com>

crates/solidity/src/build/contract.rs

@@ -8,6 +8,9 @@ use std::path::Path;
 use serde::Deserialize;
 use serde::Serialize;
 
+use polkavm_common::program::ProgramBlob;
+use polkavm_disassembler::{Disassembler, DisassemblyFormat};
+
 use crate::solc::combined_json::contract::Contract as CombinedJsonContract;
 use crate::solc::standard_json::output::contract::Contract as StandardJsonOutputContract;
 
@@ -53,9 +56,14 @@ impl Contract {
         overwrite: bool,
     ) -> anyhow::Result<()> {
         let file_name = Self::short_path(self.path.as_str());
+        let bytescode = self.build.bytecode;
 
         if output_assembly {
-            let file_name = format!("{}.{}", file_name, revive_common::EXTENSION_POLKAVM_ASSEMBLY);
+            let file_name = format!(
+                "{}.{}",
+                file_name,
+                revive_common::EXTENSION_POLKAVM_ASSEMBLY
+            );
             let mut file_path = path.to_owned();
             file_path.push(file_name);
 
@@ -64,11 +72,39 @@ impl Contract {
                     "Refusing to overwrite an existing file {file_path:?} (use --overwrite to force)."
                 );
             } else {
+                let program_blob = ProgramBlob::parse(bytescode.as_slice()).map_err(|error| {
+                    anyhow::anyhow!(format!("Failed to parse program blob: {}", error))
+                })?;
+
+                let mut disassembler = Disassembler::new(&program_blob, DisassemblyFormat::Guest)
+                    .map_err(|error| {
+                    anyhow::anyhow!(format!(
+                        "Failed to create disassembler for contract '{:?}'\n\nDue to:\n{}",
+                        &file_path, error
+                    ))
+                })?;
+
+                let mut disassembled_code = Vec::new();
+                disassembler
+                    .disassemble_into(&mut disassembled_code)
+                    .map_err(|error| {
+                        anyhow::anyhow!(format!(
+                            "Failed to disassemble contract '{:?}'\n\nDue to:\n{}\n\nGas details:{:?}\n",
+                            &file_path, error, disassembler.display_gas()
+                        ))
+                    })?;
+
+                let assembly_text = String::from_utf8(disassembled_code)
+                    .map_err(|error| anyhow::anyhow!(format!(
+                        "Failed to convert disassembled code to string for contract '{:?}'\n\nDue to:\n{}",
+                        &file_path, error
+                    )))?;
+
                 File::create(&file_path)
                     .map_err(|error| {
                         anyhow::anyhow!("File {:?} creating error: {}", file_path, error)
                     })?
-                    .write_all(self.build.assembly_text.as_bytes())
+                    .write_all(assembly_text.as_bytes())
                     .map_err(|error| {
                         anyhow::anyhow!("File {:?} writing error: {}", file_path, error)
                     })?;
@@ -89,7 +125,7 @@ impl Contract {
                     .map_err(|error| {
                         anyhow::anyhow!("File {:?} creating error: {}", file_path, error)
                     })?
-                    .write_all(self.build.bytecode.as_slice())
+                    .write_all(bytescode.as_slice())
                     .map_err(|error| {
                         anyhow::anyhow!("File {:?} writing error: {}", file_path, error)
                     })?;

crates/solidity/src/evmla/ethereal_ir/entry_link.rs

@@ -23,10 +23,7 @@ impl<D> revive_llvm_context::PolkaVMWriteLLVM<D> for EntryLink
 where
     D: revive_llvm_context::PolkaVMDependency + Clone,
 {
-    fn into_llvm(
-        self,
-        context: &mut revive_llvm_context::PolkaVMContext<D>,
-    ) -> anyhow::Result<()> {
+    fn into_llvm(self, context: &mut revive_llvm_context::PolkaVMContext<D>) -> anyhow::Result<()> {
         let target = context
             .get_function(EtherealIR::DEFAULT_ENTRY_FUNCTION_NAME)
             .expect("Always exists")

crates/solidity/src/evmla/ethereal_ir/function/block/element/mod.rs

@@ -1168,8 +1168,7 @@ where
                     .map(Some)
             }
             InstructionName::NUMBER => {
-                revive_llvm_context::polkavm_evm_contract_context::block_number(context)
-                    .map(Some)
+                revive_llvm_context::polkavm_evm_contract_context::block_number(context).map(Some)
             }
             InstructionName::BLOCKHASH => {
                 let arguments = self.pop_arguments_llvm(context);
@@ -1244,9 +1243,7 @@ where
                             .to_llvm()
                             .into_pointer_value();
                         context.build_store(
-                            revive_llvm_context::PolkaVMPointer::new_stack_field(
-                                context, pointer,
-                            ),
+                            revive_llvm_context::PolkaVMPointer::new_stack_field(context, pointer),
                             value,
                         )?;
                     }
@@ -1297,9 +1294,7 @@ where
                         assert_eq!(arguments.len(), 1);
                         context.build_store(pointer, arguments.remove(0))?;
                     }
-                    revive_llvm_context::PolkaVMFunctionReturn::Compound {
-                        pointer, ..
-                    } => {
+                    revive_llvm_context::PolkaVMFunctionReturn::Compound { pointer, .. } => {
                         for (index, argument) in arguments.into_iter().enumerate() {
                             let element_pointer = context.build_gep(
                                 pointer,

crates/solidity/src/evmla/ethereal_ir/function/block/mod.rs

@@ -117,10 +117,7 @@ impl<D> revive_llvm_context::PolkaVMWriteLLVM<D> for Block
 where
     D: revive_llvm_context::PolkaVMDependency + Clone,
 {
-    fn into_llvm(
-        self,
-        context: &mut revive_llvm_context::PolkaVMContext<D>,
-    ) -> anyhow::Result<()> {
+    fn into_llvm(self, context: &mut revive_llvm_context::PolkaVMContext<D>) -> anyhow::Result<()> {
         context.set_code_type(self.key.code_type);
 
         for element in self.elements.into_iter() {

crates/solidity/src/evmla/ethereal_ir/function/mod.rs

@@ -252,12 +252,10 @@ impl Function {
                             destination.to_owned() - num::BigUint::from(1u64 << 32),
                         )
                     }
-                    Element::Tag(destination) => {
-                        revive_llvm_context::PolkaVMFunctionBlockKey::new(
-                            code_type,
-                            destination.to_owned(),
-                        )
-                    }
+                    Element::Tag(destination) => revive_llvm_context::PolkaVMFunctionBlockKey::new(
+                        code_type,
+                        destination.to_owned(),
+                    ),
                     Element::ReturnAddress(output_size) => {
                         block_element.instruction =
                             Instruction::recursive_return(1 + output_size, instruction);
@@ -316,12 +314,10 @@ impl Function {
                             destination.to_owned() - num::BigUint::from(1u64 << 32),
                         )
                     }
-                    Element::Tag(destination) => {
-                        revive_llvm_context::PolkaVMFunctionBlockKey::new(
-                            code_type,
-                            destination.to_owned(),
-                        )
-                    }
+                    Element::Tag(destination) => revive_llvm_context::PolkaVMFunctionBlockKey::new(
+                        code_type,
+                        destination.to_owned(),
+                    ),
                     element => {
                         return Err(anyhow::anyhow!(
                             "The {} instruction expected a tag or return address, found {}",
@@ -346,8 +342,7 @@ impl Function {
                 ..
             } => {
                 let tag: num::BigUint = tag.parse().expect("Always valid");
-                let block_key =
-                    revive_llvm_context::PolkaVMFunctionBlockKey::new(code_type, tag);
+                let block_key = revive_llvm_context::PolkaVMFunctionBlockKey::new(code_type, tag);
 
                 queue_element.predecessor = Some((queue_element.block_key.clone(), instance));
                 queue_element.block_key = block_key.clone();
@@ -1014,21 +1009,16 @@ impl Function {
         block_stack: &mut Stack,
         block_element: &mut BlockElement,
         version: &semver::Version,
-    ) -> anyhow::Result<(
-        revive_llvm_context::PolkaVMFunctionBlockKey,
-        Vec<Element>,
-    )> {
+    ) -> anyhow::Result<(revive_llvm_context::PolkaVMFunctionBlockKey, Vec<Element>)> {
         let return_address_offset = block_stack.elements.len() - 2 - recursive_function.input_size;
         let input_arguments_offset = return_address_offset + 1;
         let callee_tag_offset = input_arguments_offset + recursive_function.input_size;
 
         let return_address = match block_stack.elements[return_address_offset] {
-            Element::Tag(ref return_address) => {
-                revive_llvm_context::PolkaVMFunctionBlockKey::new(
-                    block_key.code_type,
-                    return_address.to_owned(),
-                )
-            }
+            Element::Tag(ref return_address) => revive_llvm_context::PolkaVMFunctionBlockKey::new(
+                block_key.code_type,
+                return_address.to_owned(),
+            ),
             ref element => anyhow::bail!("Expected the function return address, found {}", element),
         };
         let mut stack = Stack::with_capacity(1 + recursive_function.input_size);
@@ -1177,17 +1167,16 @@ where
             output_size,
             Some(inkwell::module::Linkage::Private),
         )?;
-        function.borrow_mut().set_evmla_data(
-            revive_llvm_context::PolkaVMFunctionEVMLAData::new(self.stack_size),
-        );
+        function
+            .borrow_mut()
+            .set_evmla_data(revive_llvm_context::PolkaVMFunctionEVMLAData::new(
+                self.stack_size,
+            ));
 
         Ok(())
     }
 
-    fn into_llvm(
-        self,
-        context: &mut revive_llvm_context::PolkaVMContext<D>,
-    ) -> anyhow::Result<()> {
+    fn into_llvm(self, context: &mut revive_llvm_context::PolkaVMContext<D>) -> anyhow::Result<()> {
         context.set_current_function(self.name.as_str())?;
 
         for (key, blocks) in self.blocks.iter() {

crates/solidity/src/evmla/ethereal_ir/mod.rs

@@ -81,10 +81,7 @@ impl EtherealIR {
                 &instructions[offset..],
             )?;
             blocks.insert(
-                revive_llvm_context::PolkaVMFunctionBlockKey::new(
-                    code_type,
-                    block.key.tag.clone(),
-                ),
+                revive_llvm_context::PolkaVMFunctionBlockKey::new(code_type, block.key.tag.clone()),
                 block,
             );
             offset += size;
@@ -111,10 +108,7 @@ where
         Ok(())
     }
 
-    fn into_llvm(
-        self,
-        context: &mut revive_llvm_context::PolkaVMContext<D>,
-    ) -> anyhow::Result<()> {
+    fn into_llvm(self, context: &mut revive_llvm_context::PolkaVMContext<D>) -> anyhow::Result<()> {
         context.evmla_mut().stack = vec![];
 
         self.entry_function.into_llvm(context)?;

crates/solidity/src/project/contract/ir/evmla.rs

@@ -41,10 +41,7 @@ where
         self.assembly.declare(context)
     }
 
-    fn into_llvm(
-        self,
-        context: &mut revive_llvm_context::PolkaVMContext<D>,
-    ) -> anyhow::Result<()> {
+    fn into_llvm(self, context: &mut revive_llvm_context::PolkaVMContext<D>) -> anyhow::Result<()> {
         self.assembly.into_llvm(context)
     }
 }

crates/solidity/src/project/contract/ir/mod.rs

@@ -83,10 +83,7 @@ where
         }
     }
 
-    fn into_llvm(
-        self,
-        context: &mut revive_llvm_context::PolkaVMContext<D>,
-    ) -> anyhow::Result<()> {
+    fn into_llvm(self, context: &mut revive_llvm_context::PolkaVMContext<D>) -> anyhow::Result<()> {
         match self {
             Self::Yul(inner) => inner.into_llvm(context),
             Self::EVMLA(inner) => inner.into_llvm(context),

crates/solidity/src/project/contract/ir/yul.rs

@@ -42,10 +42,7 @@ where
         self.object.declare(context)
     }
 
-    fn into_llvm(
-        self,
-        context: &mut revive_llvm_context::PolkaVMContext<D>,
-    ) -> anyhow::Result<()> {
+    fn into_llvm(self, context: &mut revive_llvm_context::PolkaVMContext<D>) -> anyhow::Result<()> {
         self.object.into_llvm(context)
     }
 }

crates/solidity/src/project/contract/mod.rs

@@ -148,8 +148,7 @@ impl Contract {
                 context.set_yul_data(yul_data);
             }
             IR::EVMLA(_) => {
-                let evmla_data =
-                    revive_llvm_context::PolkaVMContextEVMLAData::new(version.default);
+                let evmla_data = revive_llvm_context::PolkaVMContextEVMLAData::new(version.default);
                 context.set_evmla_data(evmla_data);
             }
             IR::LLVMIR(_) => {}
@@ -201,10 +200,7 @@ where
         self.ir.declare(context)
     }
 
-    fn into_llvm(
-        self,
-        context: &mut revive_llvm_context::PolkaVMContext<D>,
-    ) -> anyhow::Result<()> {
+    fn into_llvm(self, context: &mut revive_llvm_context::PolkaVMContext<D>) -> anyhow::Result<()> {
         self.ir.into_llvm(context)
     }
 }

crates/solidity/src/tests/cli-tests/tests/asm.test.ts

@@ -14,19 +14,23 @@ describe("Run with --asm by default", () => {
   });
 
   it("--asm output is presented", () => {
-      expect(result.output).toMatch(/(__entry:)/i);
+    const expectedPatterns = [/(deploy)/i, /(call)/i, /(seal_return)/i];
+
+    for (const pattern of expectedPatterns) {
+      expect(result.output).toMatch(pattern);
+    }
   });
 
-
   it("solc exit code == zksolc exit code", () => {
-      const command = `solc ${paths.pathToBasicSolContract} --asm`;
-      const solcResult = executeCommand(command);
-      expect(solcResult.exitCode).toBe(result.exitCode);
+    const command = `solc ${paths.pathToBasicSolContract} --asm`;
+    const solcResult = executeCommand(command);
+    expect(solcResult.exitCode).toBe(result.exitCode);
   });
 
   it("run invalid: zksolc --asm", () => {
     expect(resultInvalid.output).toMatch(/(No input sources specified|Compilation aborted)/i);
   });
+  
   it("Invalid command exit code = 1", () => {
     expect(resultInvalid.exitCode).toBe(1);
   });
@@ -36,4 +40,4 @@ describe("Run with --asm by default", () => {
     const solcResult = executeCommand(command);
     expect(solcResult.exitCode).toBe(resultInvalid.exitCode);
   });
-});
+});

crates/solidity/src/tests/cli-tests/tsconfig.json

@@ -3,6 +3,5 @@
     "target": "ES6",
     "module": "CommonJS",
     "outDir": "./dist",
-    "rootDir": "./src"
   }
 }

crates/solidity/src/yul/parser/statement/block.rs

@@ -132,10 +132,7 @@ impl<D> revive_llvm_context::PolkaVMWriteLLVM<D> for Block
 where
     D: revive_llvm_context::PolkaVMDependency + Clone,
 {
-    fn into_llvm(
-        self,
-        context: &mut revive_llvm_context::PolkaVMContext<D>,
-    ) -> anyhow::Result<()> {
+    fn into_llvm(self, context: &mut revive_llvm_context::PolkaVMContext<D>) -> anyhow::Result<()> {
         let current_function = context.current_function().borrow().name().to_owned();
         let current_block = context.basic_block();
 

crates/solidity/src/yul/parser/statement/code.rs

@@ -59,10 +59,7 @@ impl<D> revive_llvm_context::PolkaVMWriteLLVM<D> for Code
 where
     D: revive_llvm_context::PolkaVMDependency + Clone,
 {
-    fn into_llvm(
-        self,
-        context: &mut revive_llvm_context::PolkaVMContext<D>,
-    ) -> anyhow::Result<()> {
+    fn into_llvm(self, context: &mut revive_llvm_context::PolkaVMContext<D>) -> anyhow::Result<()> {
         self.block.into_llvm(context)?;
 
         Ok(())

crates/solidity/src/yul/parser/statement/expression/function_call/mod.rs

@@ -911,9 +911,7 @@ impl FunctionCall {
                 Ok(Some(context.integer_const(256, 0).as_basic_value_enum()))
             }
 
-            Name::CallValue => {
-                revive_llvm_context::polkavm_evm_ether_gas::value(context).map(Some)
-            }
+            Name::CallValue => revive_llvm_context::polkavm_evm_ether_gas::value(context).map(Some),
             Name::Gas => revive_llvm_context::polkavm_evm_ether_gas::gas(context).map(Some),
             Name::Balance => {
                 let arguments = self.pop_arguments_llvm::<D, 1>(context)?;
@@ -940,8 +938,7 @@ impl FunctionCall {
                     .map(Some)
             }
             Name::Number => {
-                revive_llvm_context::polkavm_evm_contract_context::block_number(context)
-                    .map(Some)
+                revive_llvm_context::polkavm_evm_contract_context::block_number(context).map(Some)
             }
             Name::BlockHash => {
                 let arguments = self.pop_arguments_llvm::<D, 1>(context)?;

crates/solidity/src/yul/parser/statement/expression/mod.rs

@@ -137,9 +137,7 @@ impl Expression {
 
                 match constant {
                     Some(constant) => Ok(Some(
-                        revive_llvm_context::PolkaVMArgument::new_with_constant(
-                            value, constant,
-                        ),
+                        revive_llvm_context::PolkaVMArgument::new_with_constant(value, constant),
                     )),
                     None => Ok(Some(value.into())),
                 }

crates/solidity/src/yul/parser/statement/for_loop.rs

@@ -64,10 +64,7 @@ impl<D> revive_llvm_context::PolkaVMWriteLLVM<D> for ForLoop
 where
     D: revive_llvm_context::PolkaVMDependency + Clone,
 {
-    fn into_llvm(
-        self,
-        context: &mut revive_llvm_context::PolkaVMContext<D>,
-    ) -> anyhow::Result<()> {
+    fn into_llvm(self, context: &mut revive_llvm_context::PolkaVMContext<D>) -> anyhow::Result<()> {
         self.initializer.into_llvm(context)?;
 
         let condition_block = context.append_basic_block("for_condition");

crates/solidity/src/yul/parser/statement/function_definition.rs

@@ -109,9 +109,10 @@ impl FunctionDefinition {
 
             arguments.remove(0);
         }
-        if identifier.inner.contains(
-            revive_llvm_context::PolkaVMFunction::ZKSYNC_NEAR_CALL_ABI_EXCEPTION_HANDLER,
-        ) && !arguments.is_empty()
+        if identifier
+            .inner
+            .contains(revive_llvm_context::PolkaVMFunction::ZKSYNC_NEAR_CALL_ABI_EXCEPTION_HANDLER)
+            && !arguments.is_empty()
         {
             return Err(ParserError::InvalidNumberOfArguments {
                 location,

crates/solidity/src/yul/parser/statement/if_conditional.rs

@@ -52,10 +52,7 @@ impl<D> revive_llvm_context::PolkaVMWriteLLVM<D> for IfConditional
 where
     D: revive_llvm_context::PolkaVMDependency + Clone,
 {
-    fn into_llvm(
-        self,
-        context: &mut revive_llvm_context::PolkaVMContext<D>,
-    ) -> anyhow::Result<()> {
+    fn into_llvm(self, context: &mut revive_llvm_context::PolkaVMContext<D>) -> anyhow::Result<()> {
         let condition = self
             .condition
             .into_llvm(context)?

crates/solidity/src/yul/parser/statement/object.rs

@@ -219,16 +219,11 @@ where
         Ok(())
     }
 
-    fn into_llvm(
-        self,
-        context: &mut revive_llvm_context::PolkaVMContext<D>,
-    ) -> anyhow::Result<()> {
+    fn into_llvm(self, context: &mut revive_llvm_context::PolkaVMContext<D>) -> anyhow::Result<()> {
         if self.identifier.ends_with("_deployed") {
-            revive_llvm_context::PolkaVMRuntimeCodeFunction::new(self.code)
-                .into_llvm(context)?;
+            revive_llvm_context::PolkaVMRuntimeCodeFunction::new(self.code).into_llvm(context)?;
         } else {
-            revive_llvm_context::PolkaVMDeployCodeFunction::new(self.code)
-                .into_llvm(context)?;
+            revive_llvm_context::PolkaVMDeployCodeFunction::new(self.code).into_llvm(context)?;
         }
 
         match self.inner_object {

crates/solidity/src/yul/parser/statement/switch/mod.rs

@@ -122,10 +122,7 @@ impl<D> revive_llvm_context::PolkaVMWriteLLVM<D> for Switch
 where
     D: revive_llvm_context::PolkaVMDependency + Clone,
 {
-    fn into_llvm(
-        self,
-        context: &mut revive_llvm_context::PolkaVMContext<D>,
-    ) -> anyhow::Result<()> {
+    fn into_llvm(self, context: &mut revive_llvm_context::PolkaVMContext<D>) -> anyhow::Result<()> {
         let scrutinee = self.expression.into_llvm(context)?;
 
         if self.cases.is_empty() {

crates/solidity/src/zksolc/arguments.rs

@@ -212,7 +212,9 @@ impl Arguments {
 
         if self.yul || self.llvm_ir || self.zkasm {
             if self.base_path.is_some() {
-                anyhow::bail!("`base-path` is not used in Yul, LLVM IR and PolkaVM assembly modes.");
+                anyhow::bail!(
+                    "`base-path` is not used in Yul, LLVM IR and PolkaVM assembly modes."
+                );
             }
             if !self.include_paths.is_empty() {
                 anyhow::bail!(

crates/solidity/src/zksolc/main.rs

@@ -6,6 +6,9 @@ use std::str::FromStr;
 
 use self::arguments::Arguments;
 
+use polkavm_common::program::ProgramBlob;
+use polkavm_disassembler::{Disassembler, DisassemblyFormat};
+
 /// The rayon worker stack size.
 const RAYON_WORKER_STACK_SIZE: usize = 16 * 1024 * 1024;
 
@@ -190,18 +193,41 @@ fn main_inner() -> anyhow::Result<()> {
         );
     } else if arguments.output_assembly || arguments.output_binary {
         for (path, contract) in build.contracts.into_iter() {
+            let bytescode = contract.build.bytecode;
+
             if arguments.output_assembly {
-                println!(
-                    "Contract `{}` assembly:\n\n{}",
-                    path, contract.build.assembly_text
-                );
+                let program_blob = ProgramBlob::parse(bytescode.as_slice()).map_err(|error| {
+                    anyhow::anyhow!(format!("Failed to parse program blob: {}", error))
+                })?;
+
+                let mut disassembler = Disassembler::new(&program_blob, DisassemblyFormat::Guest)
+                    .map_err(|error| {
+                    anyhow::anyhow!(format!(
+                        "Failed to create disassembler for contract '{}'\n\nDue to:\n{}",
+                        path, error
+                    ))
+                })?;
+
+                let mut disassembled_code = Vec::new();
+                disassembler
+                    .disassemble_into(&mut disassembled_code)
+                    .map_err(|error| {
+                        anyhow::anyhow!(format!(
+                            "Failed to disassemble contract '{}'\n\nDue to:\n{}\n\nGas details:{:?}\n",
+                            path, error, disassembler.display_gas()
+                        ))
+                    })?;
+
+                let assembly_text = String::from_utf8(disassembled_code)
+                    .map_err(|error| anyhow::anyhow!(format!(
+                        "Failed to convert disassembled code to string for contract '{}'\n\nDue to:\n{}",
+                        path, error
+                    )))?;
+
+                println!("Contract `{}` assembly:\n\n{}", path, assembly_text);
             }
             if arguments.output_binary {
-                println!(
-                    "Contract `{}` bytecode: 0x{}",
-                    path,
-                    hex::encode(contract.build.bytecode)
-                );
+                println!("Contract `{}` bytecode: 0x{}", path, hex::encode(bytescode));
             }
         }
     } else {