CensoredUsername / dynasm-rs

A dynasm-like tool for rust.
https://censoredusername.github.io/dynasm-rs/language/index.html
Mozilla Public License 2.0
705 stars 52 forks source link

Dynamic Memory Operands #88

Open paul1999 opened 4 months ago

paul1999 commented 4 months ago

Greetings!

I am currently porting a simple JIT, used in a compiler construction class, from C++ to Rust. In C++, I used the asmjit project. In asmjit, memory operands and registers can be used as values, i.e. I can assign a memory operands to a variable and later use it in an assembler call

auto foo = x86::qword_ptr(x86::rbp, -off_s);
as.mov(foo, 42);

If I want to use a value stored in memory at the moment, I have to repeat

let a = Foo::from_rbp_offset(8);
let b = Foo::from_rbp_offset(16);
dynasm!(self.as
    ; mov rax, QWORD [rbp - a.offset]
    ; add rax, QWORD [rbp - b.offset]);

I.e. I have to repeat the memory reference each time.

It would be really handy to have a feature similar to the asmjit one in this crate, i.e. being able to do something like

enum Location {
    Register(reg),
    RegisterStaticOffset(reg, offset),
    RegisterDynamicOffset(reg, reg),
}

impl MemoryReference for Foo {
    fn location(&self) -> Location;
}

fn load_to_rax(mut self, loc: Foo) {
    dynasm!(self.as
        ; mov rax, QWORD [loc]);
} 

Because then you could encapsulate where certain values are stored (relative to rbp, relative to rsp, etc.) in the Trait Implementation.

I do not know too much about Rust macros, so I apologize if something like this is not feasible to implement. On the other hand, I'd be willing to tinker with this myself, I would just need a pointer (or reference) in the right direction.

Best regards!

CensoredUsername commented 4 months ago

Hi!

I'm thinking about if this is reasonably implementable in dynasm-rs. Dynasm-rs does most of its work at compile time. Basically anything but non-constant immediates, dynamic register support and relocations are resolved at compile time already.

This includes things like instruction variant selection and memory reference variant selection. I think instruction variant selection generally just depends on the presence of a memory reference, not its contents, so pushing that logic to compile time might be possible. We'd also need syntax to escape this behavior and emit the correct code to handle this. something like:

dynasm!(assembler; mov rax, QWORD [dyn foo])

Would be fine. We can't just use [foo] as that would be interpreted as just dereferencing the address foo.

That should then be interpreted by dynasm-rs to insert a piece of code that emits the correct byte(s), or in the case of aarch64, or them with the rest of the instruction word.

On the API design, the different architectures have very different requirements. x64/x86 are variable width, even in memory references, so we cannot just insert something like assembler.push(memoryref.code()), it'd prefer a method on the memoryref that takes the assembler as argument. Meanwhile aarch64 essentially ORs all fields together so it'd explicitly want something like assembler.push(rest_of_operand | memoryref.code()), and taking the assembler as argument wouldn't work as it isn't emitting separate bytes from the rest of the instruction. x64 might just need separate handling and APIs there. Alternatively, we'd just need to stick to always emitting a SIB byte for x64. But that still leaves the need to conditionally emit the displacement as well.

Overall, I think that it is a possible feature to implement, but it definitely isn't the easiest. I hope that helps you a bit.

paul1999 commented 4 months ago

Thanks for the reply!

I think instruction variant selection generally just depends on the presence of a memory reference, not its contents, so pushing that logic to compile time might be possible.

Looking at some common instructions, this generally seems to be the case: the opcode seems to depend on operand order (i.e. is the memory operand source or target) and sometimes on operand size (e.g. mov r8, r/m8 uses 0x8A, while mov r64, r/m64 uses 0x8B).

However, both operand order and operand size should be known at compile time.

That should then be interpreted by dynasm-rs to insert a piece of code that emits the correct byte(s), or in the case of aarch64, or them with the rest of the instruction word.

On the API design, the different architectures have very different requirements. x64/x86 are variable width, even in memory references, so we cannot just insert something like assembler.push(memoryref.code()), it'd prefer a method on the memoryref that takes the assembler as argument. Meanwhile aarch64 essentially ORs all fields together so it'd explicitly want something like assembler.push(rest_of_operand | memoryref.code()), and taking the assembler as argument wouldn't work as it isn't emitting separate bytes from the rest of the instruction. x64 might just need separate handling and APIs there. Alternatively, we'd just need to stick to always emitting a SIB byte for x64. But that still leaves the need to conditionally emit the displacement as well.

Based on that, I've sketched out the following:

// Runtime library
mod dynasmrt::x64 {
    enum Scale {
        One,
        Two,
        Four,
        Eight,
    }

    enum Displacement {
        Byte(i8),
        Word(i16),
        DWord(i32),
    }

    #[derive(Clone, Copy)]
    enum MemoryReference<Base: Register, Index: Register> {
        Displacement(Displacement),
        Base(Base),
        BaseDisplacement(Base, Displacement),
        BaseIndex(Base, Index, Scale),
        BaseIndexDisplacement(Base, Index, Scale, Displacement),
    }

    impl<Base: Register, Index: Register> MemoryReference<Base, Index> {
        fn push_rex_prefix(&self, assembler: &mut Assembler) {
            // TODO How to determine whether REX prefix is needed? Method on Register?
        }

        fn push_operands(&self, assembler: &mut Assembler) {
            // ...
        }
    }

    trait MemoryReferencable {
        fn reference() -> MemoryReference;
    }
}

// Syntax
dynasm!(assembler
    ; mov rax, [dyn foo]
    ; mov QWORD [dyn bar], rax);

// Generated code for x64

(foo as MemoryReferencable).reference().push_rex_prefix(assembler);
assembler.push_u8(0x8B);
(foo as MemoryReferencable).reference().push_operands(assembler);

(bar as MemoryReferencable).reference().push_rex_prefix(assembler);
assembler.push_u8(0x89);
(bar as MemoryReferencable).reference().push_operands(prefix);

However, I think you'd need something to know whether to actually emit a REX prefix. I guess a method a la

trait Register {
    fn needs_rex(&self) -> bool
}

would work?

CensoredUsername commented 4 months ago

Damn, I hadn't even thought about that register bits also might need to end up in the REX prefix. That makes it even messier.

So this'd require possibly dynamic rewriting of the REX / VEX / XOP prefixes, possible changes to the modrm byte, and conditional emitting of a SIB byte and a displacement in x64 mode.

That's quite a few changes, and that's only for one of the architecture. I'm not sure if I'd recommend putting in all that work to be honest. Right now you can already get quite far using dynamic registers. It's definitely possible, but I'm wondering if at that point you'd be better off just lifting all of the logic to runtime to simplify interactions, as this would add even more special cases to the compile time and run-time interaction.