diff --git a/src/main.zig b/src/main.zig
index b62e3e4..84a2557 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -1,6 +1,8 @@
 const std = @import("std");
+const log = std.log;
 const page_size = std.mem.page_size;
 const assert = std.debug.assert;
+const expect = std.testing.expect;
 const expectEqual = std.testing.expectEqual;
 
 const Section = struct {
@@ -56,11 +58,34 @@ const Rex = enum(u8) {
     const upper_4_bits = 0b100_0000;
 };
 
+const GPRegister = enum(u4) {
+    a = 0,
+    c = 1,
+    d = 2,
+    b = 3,
+    sp = 4,
+    bp = 5,
+    si = 6,
+    di = 7,
+    r8 = 8,
+    r9 = 9,
+    r10 = 10,
+    r11 = 11,
+    r12 = 12,
+    r13 = 13,
+    r14 = 14,
+    r15 = 15,
+};
+
+const prefix_lock = 0xf0;
+const prefix_repne_nz = 0xf2;
+const prefix_rep = 0xf3;
 const prefix_rex_w = [1]u8{@intFromEnum(Rex.w)};
 const prefix_16_bit_operand = [1]u8{0x66};
 
 const ret = [1]u8{0xc3};
-const movabs_to_register_a = [1]u8{0xb8};
+const mov_a_imm = [1]u8{0xb8};
+const mov_reg_imm8: u8 = 0xb0;
 
 inline fn intToArrayOfBytes(integer: anytype) [@sizeOf(@TypeOf(integer))]u8 {
     comptime {
@@ -70,16 +95,31 @@ inline fn intToArrayOfBytes(integer: anytype) [@sizeOf(@TypeOf(integer))]u8 {
     return @as([@sizeOf(@TypeOf(integer))]u8, @bitCast(integer));
 }
 
-inline fn movU16ToA(integer: u16) [4]u8 {
-    return prefix_16_bit_operand ++ movabs_to_register_a ++ intToArrayOfBytes(integer);
+fn movToAInstructionByteCount(comptime bit_count: u16) type {
+    return [
+        switch (bit_count) {
+            8 => 2,
+            16 => 4,
+            32 => 5,
+            64 => 10,
+            else => @compileError("Not supported"),
+        }
+    ]u8;
 }
 
-inline fn movU32ToA(integer: u32) [5]u8 {
-    return movabs_to_register_a ++ intToArrayOfBytes(integer);
-}
-
-inline fn movU64ToA(integer: u64) [10]u8 {
-    return prefix_rex_w ++ movabs_to_register_a ++ intToArrayOfBytes(integer);
+inline fn movAImm(comptime signedness: std.builtin.Signedness, comptime bit_count: u16, integer: @Type(.{
+    .Int = .{
+        .signedness = signedness,
+        .bits = bit_count,
+    },
+})) movToAInstructionByteCount(bit_count) {
+    return switch (@TypeOf(integer)) {
+        u8, i8 => .{mov_reg_imm8 | @intFromEnum(GPRegister.a)},
+        u16, i16 => prefix_16_bit_operand ++ mov_a_imm,
+        u32, i32 => mov_a_imm,
+        u64, i64 => prefix_rex_w ++ mov_a_imm,
+        else => @compileError("Unsupported"),
+    } ++ intToArrayOfBytes(integer);
 }
 
 test "ret void" {
@@ -90,35 +130,104 @@ test "ret void" {
     function_pointer();
 }
 
-test "ret unsigned integer 16-bit" {
+const integer_types_to_test = [_]type{ u8, u16, u32, u64, i8, i16, i32, i64 };
+
+fn testRetInteger(comptime T: type) !void {
+    comptime {
+        assert(@typeInfo(T) == .Int);
+        assert((T == u8 or T == u16 or T == u32 or T == u64) or (T == i8 or T == i16 or T == i32 or T == i64));
+    }
+
     var image = try Image.create();
-    const expected_number = 0xffff;
-    image.appendCode(&movU16ToA(expected_number));
+    const signedness = @typeInfo(T).Int.signedness;
+    const expected_number = getMaxInteger(T);
+
+    image.appendCode(&movAImm(signedness, @bitSizeOf(T), expected_number));
     image.appendCode(&ret);
 
-    const function_pointer = image.getEntryPoint(fn () callconv(.C) u16);
+    const function_pointer = image.getEntryPoint(fn () callconv(.C) T);
     const result = function_pointer();
-    try expectEqual(result, expected_number);
+    try expect(result == expected_number);
 }
 
-test "ret unsigned integer 32-bit" {
-    var image = try Image.create();
-    const expected_number = 0xffff_ffff;
-    image.appendCode(&movU32ToA(expected_number));
-    image.appendCode(&ret);
+fn getMaxInteger(comptime T: type) T {
+    comptime {
+        assert(@typeInfo(T) == .Int);
+    }
 
-    const function_pointer = image.getEntryPoint(fn () callconv(.C) u32);
-    const result = function_pointer();
-    try expectEqual(result, expected_number);
+    return switch (@typeInfo(T).Int.signedness) {
+        .unsigned => std.math.maxInt(T),
+        .signed => std.math.minInt(T),
+    };
 }
 
-test "ret unsigned integer 64-bit" {
-    var image = try Image.create();
-    const expected_number = 0xffff_ffff_ffff_ffff;
-    image.appendCode(&movU64ToA(expected_number));
-    image.appendCode(&ret);
-
-    const function_pointer = image.getEntryPoint(fn () callconv(.C) u64);
-    const result = function_pointer();
-    try expectEqual(result, expected_number);
+test "ret integer" {
+    inline for (integer_types_to_test) |Int| {
+        try testRetInteger(Int);
+    }
+}
+
+const mov_rm_r = 0x89;
+
+test "ret integer argument" {
+    inline for (integer_types_to_test) |Int| {
+        var image = try Image.create();
+        const number = getMaxInteger(Int);
+        const mov_a_di = switch (Int) {
+            u8, i8 => .{ 0x40, 0x88, 0xf8 },
+            u16, i16 => prefix_16_bit_operand ++ .{ mov_rm_r, 0xf8 },
+            u32, i32 => .{ mov_rm_r, 0xf8 },
+            u64, i64 => prefix_rex_w ++ .{ mov_rm_r, 0xf8 },
+            else => @compileError("Not supported"),
+        };
+
+        image.appendCode(&mov_a_di);
+        image.appendCode(&ret);
+
+        const functionPointer = image.getEntryPoint(fn (Int) callconv(.C) Int);
+        const result = functionPointer(number);
+        try expectEqual(number, result);
+    }
+}
+var r = std.rand.Pcg.init(0xffffffffffffffff);
+
+fn getRandomNumberRange(comptime T: type, min: T, max: T) T {
+    const random = r.random();
+    return switch (@typeInfo(T).Int.signedness) {
+        .signed => random.intRangeAtMost(T, min, max),
+        .unsigned => random.uintAtMost(T, max),
+    };
+}
+
+const sub_rm_r = 0x29;
+
+test "ret sub arguments" {
+    inline for (integer_types_to_test) |Int| {
+        var image = try Image.create();
+        const a = getRandomNumberRange(Int, std.math.minInt(Int) / 2, std.math.maxInt(Int) / 2);
+        const b = getRandomNumberRange(Int, std.math.minInt(Int) / 2, a);
+
+        const mov_a_di = switch (Int) {
+            u8, i8 => .{ 0x40, 0x88, 0xf8 },
+            u16, i16 => prefix_16_bit_operand ++ .{ mov_rm_r, 0xf8 },
+            u32, i32 => .{ mov_rm_r, 0xf8 },
+            u64, i64 => prefix_rex_w ++ .{ mov_rm_r, 0xf8 },
+            else => @compileError("Not supported"),
+        };
+        image.appendCode(&mov_a_di);
+
+        const sub_a_si = switch (Int) {
+            u8, i8 => .{ 0x40, 0x28, 0xf0 },
+            u16, i16 => prefix_16_bit_operand ++ .{ sub_rm_r, 0xf0 },
+            u32, i32 => .{ sub_rm_r, 0xf0 },
+            u64, i64 => prefix_rex_w ++ .{ sub_rm_r, 0xf0 },
+            else => @compileError("Not supported"),
+        };
+        image.appendCode(&sub_a_si);
+        image.appendCode(&ret);
+
+        const functionPointer = image.getEntryPoint(fn (Int, Int) callconv(.C) Int);
+        const result = functionPointer(a, b);
+        try expectEqual(a - b, result);
+    }
 }