From ed343c20799575371898b9999cd552b558130f58 Mon Sep 17 00:00:00 2001
From: Yurii Rashkovskii <yrashk@gmail.com>
Date: Fri, 25 Feb 2022 10:44:54 -0800
Subject: [PATCH] Problem: re-assembling wasm code to get label offsets

Solution: use new version's `assemble_options` to retrieve it

See https://github.com/icedland/iced/discussions/271#discussioncomment-2239492
---
 Cargo.lock                 |  4 ++--
 parawasm/src/x86_64/mod.rs | 28 ++++++++++++++++++++++------
 parawasm_tests/Cargo.toml  |  2 +-
 3 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 5a61343..139ff6e 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -380,9 +380,9 @@ dependencies = [
 
 [[package]]
 name = "iced-x86"
-version = "1.16.0"
+version = "1.17.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0909b90ac9fb561b496133348f34719bd63aaeac03d8fafaeef13752721a7ad9"
+checksum = "158f5204401d08f91d19176112146d75e99b3cf745092e268fa7be33e09adcec"
 dependencies = [
  "hashbrown 0.12.0",
  "lazy_static",
diff --git a/parawasm/src/x86_64/mod.rs b/parawasm/src/x86_64/mod.rs
index 6d8ef2b..40d798d 100644
--- a/parawasm/src/x86_64/mod.rs
+++ b/parawasm/src/x86_64/mod.rs
@@ -10,7 +10,7 @@ use core::ops::{Deref, DerefMut};
 use iced_x86::code_asm::{
     dword_ptr, qword_ptr, r8, r9, rax, rbp, rcx, rdi, rdx, rsi, rsp, AsmRegister64, CodeAssembler,
 };
-use iced_x86::IcedError;
+use iced_x86::{BlockEncoderOptions, IcedError};
 use wasmparser_nostd::*;
 
 mod instructions;
@@ -207,7 +207,16 @@ impl Compiler for X86_64Compiler {
                         Payload::ImportSection(is) => {
                             for i in is {
                                 let import = i?;
-                                let offset = assembler.assemble(0)?.len();
+                                let mut current_label = assembler.create_label();
+                                assembler.set_label(&mut current_label)?;
+                                assembler.zero_bytes()?;
+                                let offset = assembler
+                                    .assemble_options(
+                                        0,
+                                        BlockEncoderOptions::RETURN_NEW_INSTRUCTION_OFFSETS,
+                                    )?
+                                    .label_ip(&current_label)?
+                                    as usize;
                                 let reference = (
                                     import.module.to_owned(),
                                     import.field.map(str::to_owned),
@@ -365,15 +374,22 @@ impl Compiler for X86_64Compiler {
         // Bind labels
         for (idx, instruction) in assembler.take_instructions().into_iter().enumerate() {
             if let Some((_, label)) = label_indices.iter_mut().find(|(i, _)| *i == idx) {
+                assembler.set_label(label)?;
+                assembler.zero_bytes()?;
                 // If this is a label pointing to a function, record function body entry point
                 if let Some((_label, index)) =
                     function_bodies.iter().find(|(label_, _)| label_ == label)
                 {
-                    module
-                        .function_bodies
-                        .insert(*index, assembler.assemble(0)?.len());
+                    module.function_bodies.insert(
+                        *index,
+                        assembler
+                            .assemble_options(
+                                0,
+                                BlockEncoderOptions::RETURN_NEW_INSTRUCTION_OFFSETS,
+                            )?
+                            .label_ip(label)? as usize,
+                    );
                 }
-                assembler.set_label(label)?;
             }
             assembler.add_instruction(instruction)?;
         }
diff --git a/parawasm_tests/Cargo.toml b/parawasm_tests/Cargo.toml
index f942d93..3c029dc 100644
--- a/parawasm_tests/Cargo.toml
+++ b/parawasm_tests/Cargo.toml
@@ -14,6 +14,6 @@ wat = "1.0.41"
 byteorder = { version = "1.4.3", default-features = false }
 
 [dependencies.iced-x86]
-version = "1.16.0"
+version = "1.17.0"
 default-features = false
 features = ["no_std", "encoder", "decoder", "intel", "code_asm"]