From 22eee2acf7262b3cd721a095a8673800ecb5da5e Mon Sep 17 00:00:00 2001
From: Paul Cercueil <paul@crapouillou.net>
Date: Mon, 25 Apr 2022 21:06:57 +0100
Subject: [PATCH] git subrepo pull --force deps/lightrec

subrepo:
  subdir:   "deps/lightrec"
  merged:   "de06670b"
upstream:
  origin:   "https://github.com/pcercuei/lightrec.git"
  branch:   "master"
  commit:   "de06670b"
git-subrepo:
  version:  "0.4.3"
  origin:   "https://github.com/ingydotnet/git-subrepo.git"
  commit:   "2f68596"
---
 deps/lightrec/.gitrepo            |   4 +-
 deps/lightrec/disassembler.c      |  48 +++++-
 deps/lightrec/disassembler.h      |  20 ++-
 deps/lightrec/emitter.c           | 274 ++++++++++++++++++++++++------
 deps/lightrec/lightning-wrapper.h |   2 +
 deps/lightrec/lightrec-private.h  |   1 +
 deps/lightrec/lightrec.c          |  41 +++--
 deps/lightrec/optimizer.c         |  77 ++++++---
 8 files changed, 364 insertions(+), 103 deletions(-)

diff --git a/deps/lightrec/.gitrepo b/deps/lightrec/.gitrepo
index 5adaf6ba..ae2fc77c 100644
--- a/deps/lightrec/.gitrepo
+++ b/deps/lightrec/.gitrepo
@@ -6,7 +6,7 @@
 [subrepo]
 	remote = https://github.com/pcercuei/lightrec.git
 	branch = master
-	commit = b9f1b040bb1333508caee3795375e13b8cf8f684
-	parent = 37d9bf8315be570a350cd44876ae14f9b0eff20b
+	commit = de06670b257004d98d30e8585a4e6530e77d3acd
+	parent = e24732050e902bd5402b2b7da7c391d2ca8fa799
 	method = merge
 	cmdver = 0.4.3
diff --git a/deps/lightrec/disassembler.c b/deps/lightrec/disassembler.c
index c357a308..43ac6772 100644
--- a/deps/lightrec/disassembler.c
+++ b/deps/lightrec/disassembler.c
@@ -104,10 +104,17 @@ static const char *opcode_flags[] = {
 };
 
 static const char *opcode_io_flags[] = {
-	"memory I/O",
-	"hardware I/O",
 	"self-modifying code",
 	"no invalidation",
+	"no mask",
+};
+
+static const char *opcode_io_modes[] = {
+	"Memory access",
+	"I/O access",
+	"RAM access",
+	"BIOS access",
+	"Scratchpad access",
 };
 
 static const char *opcode_branch_flags[] = {
@@ -122,10 +129,11 @@ static const char *opcode_multdiv_flags[] = {
 };
 
 static int print_flags(char *buf, size_t len, u16 flags,
-		       const char **array, size_t array_size)
+		       const char **array, size_t array_size,
+		       bool is_io)
 {
-	const char *flag_name;
-	unsigned int i;
+	const char *flag_name, *io_mode_name;
+	unsigned int i, io_mode;
 	size_t count = 0, bytes;
 	bool first = true;
 
@@ -149,6 +157,23 @@ static int print_flags(char *buf, size_t len, u16 flags,
 		count += bytes;
 	}
 
+	if (is_io) {
+		io_mode = LIGHTREC_FLAGS_GET_IO_MODE(flags);
+		if (io_mode > 0) {
+			io_mode_name = opcode_io_modes[io_mode - 1];
+
+			if (first)
+				bytes = snprintf(buf, len, "(%s", io_mode_name);
+			else
+				bytes = snprintf(buf, len, ", %s", io_mode_name);
+
+			first = false;
+			buf += bytes;
+			len -= bytes;
+			count += bytes;
+		}
+	}
+
 	if (!first)
 		count += snprintf(buf, len, ")");
 	else
@@ -257,7 +282,8 @@ static int print_op_cp(union code c, char *buf, size_t len, unsigned int cp)
 }
 
 static int print_op(union code c, u32 pc, char *buf, size_t len,
-		    const char ***flags_ptr, size_t *nb_flags)
+		    const char ***flags_ptr, size_t *nb_flags,
+		    bool *is_io)
 {
 	if (c.opcode == 0)
 		return snprintf(buf, len, "nop     ");
@@ -324,6 +350,7 @@ static int print_op(union code c, u32 pc, char *buf, size_t len,
 	case OP_SWR:
 		*flags_ptr = opcode_io_flags;
 		*nb_flags = ARRAY_SIZE(opcode_io_flags);
+		*is_io = true;
 		return snprintf(buf, len, "%s%s,%hd(%s)",
 				std_opcodes[c.i.op],
 				lightrec_reg_name(c.i.rt),
@@ -363,6 +390,7 @@ void lightrec_print_disassembly(const struct block *block, const u32 *code)
 	char buf[256], buf2[256], buf3[256];
 	unsigned int i;
 	u32 pc, branch_pc;
+	bool is_io;
 
 	for (i = 0; i < block->nb_ops; i++) {
 		op = &block->opcode_list[i];
@@ -370,19 +398,21 @@ void lightrec_print_disassembly(const struct block *block, const u32 *code)
 		pc = block->pc + (i << 2);
 
 		count = print_op((union code)code[i], pc, buf, sizeof(buf),
-				 &flags_ptr, &nb_flags);
+				 &flags_ptr, &nb_flags, &is_io);
 
 		flags_ptr = NULL;
 		nb_flags = 0;
+		is_io = false;
 		count2 = print_op(op->c, branch_pc, buf2, sizeof(buf2),
-				  &flags_ptr, &nb_flags);
+				  &flags_ptr, &nb_flags, &is_io);
 
 		if (code[i] == op->c.opcode) {
 			*buf2 = '\0';
 			count2 = 0;
 		}
 
-		print_flags(buf3, sizeof(buf3), op->flags, flags_ptr, nb_flags);
+		print_flags(buf3, sizeof(buf3), op->flags, flags_ptr, nb_flags,
+			    is_io);
 
 		printf("0x%08x (0x%x)\t%s%*c%s%*c%s\n", pc, i << 2,
 		       buf, 30 - (int)count, ' ', buf2, 30 - (int)count2, ' ', buf3);
diff --git a/deps/lightrec/disassembler.h b/deps/lightrec/disassembler.h
index ae2af7ea..9abc28ef 100644
--- a/deps/lightrec/disassembler.h
+++ b/deps/lightrec/disassembler.h
@@ -23,10 +23,22 @@
 #define LIGHTREC_SYNC		BIT(4)
 
 /* Flags for load/store opcodes */
-#define LIGHTREC_DIRECT_IO	BIT(5)
-#define LIGHTREC_HW_IO		BIT(6)
-#define LIGHTREC_SMC		BIT(7)
-#define LIGHTREC_NO_INVALIDATE	BIT(8)
+#define LIGHTREC_SMC		BIT(5)
+#define LIGHTREC_NO_INVALIDATE	BIT(6)
+#define LIGHTREC_NO_MASK	BIT(7)
+
+/* I/O mode for load/store opcodes */
+#define LIGHTREC_IO_MODE_LSB	8
+#define LIGHTREC_IO_MODE(x)	((x) << LIGHTREC_IO_MODE_LSB)
+#define LIGHTREC_IO_UNKNOWN	0x0
+#define LIGHTREC_IO_DIRECT	0x1
+#define LIGHTREC_IO_HW		0x2
+#define LIGHTREC_IO_RAM		0x3
+#define LIGHTREC_IO_BIOS	0x4
+#define LIGHTREC_IO_SCRATCH	0x5
+#define LIGHTREC_IO_MASK	LIGHTREC_IO_MODE(0x7)
+#define LIGHTREC_FLAGS_GET_IO_MODE(x) \
+	(((x) & LIGHTREC_IO_MASK) >> LIGHTREC_IO_MODE_LSB)
 
 /* Flags for branches */
 #define LIGHTREC_EMULATE_BRANCH	BIT(5)
diff --git a/deps/lightrec/emitter.c b/deps/lightrec/emitter.c
index 578af874..fa74cc09 100644
--- a/deps/lightrec/emitter.c
+++ b/deps/lightrec/emitter.c
@@ -76,7 +76,7 @@ static void lightrec_emit_end_of_block(struct lightrec_cstate *state,
 	}
 
 	if (offset + !!(op->flags & LIGHTREC_NO_DS) < block->nb_ops - 1)
-		state->branches[state->nb_branches++] = jit_jmpi();
+		state->branches[state->nb_branches++] = jit_b();
 }
 
 void lightrec_emit_eob(struct lightrec_cstate *state, const struct block *block,
@@ -95,30 +95,17 @@ void lightrec_emit_eob(struct lightrec_cstate *state, const struct block *block,
 	jit_movi(JIT_V0, block->pc + (offset << 2));
 	jit_subi(LIGHTREC_REG_CYCLE, LIGHTREC_REG_CYCLE, cycles);
 
-	state->branches[state->nb_branches++] = jit_jmpi();
+	state->branches[state->nb_branches++] = jit_b();
 }
 
 static u8 get_jr_jalr_reg(struct lightrec_cstate *state, const struct block *block, u16 offset)
 {
 	struct regcache *reg_cache = state->reg_cache;
 	jit_state_t *_jit = block->_jit;
-	const struct opcode *op = &block->opcode_list[offset],
-			    *next = &block->opcode_list[offset + 1];
-	u8 rs = lightrec_request_reg_in(reg_cache, _jit, op->r.rs, JIT_V0);
-
-	/* If the source register is already mapped to JIT_R0 or JIT_R1, and the
-	 * delay slot is a I/O operation, unload the register, since JIT_R0 and
-	 * JIT_R1 are explicitely used by the I/O opcode generators. */
-	if ((rs == JIT_R0 || rs == JIT_R1) &&
-	    !(op->flags & LIGHTREC_NO_DS) &&
-	    opcode_is_io(next->c) &&
-	    !(next->flags & (LIGHTREC_NO_INVALIDATE | LIGHTREC_DIRECT_IO))) {
-		lightrec_unload_reg(reg_cache, _jit, rs);
-		lightrec_free_reg(reg_cache, rs);
-
-		rs = lightrec_request_reg_in(reg_cache, _jit, op->r.rs, JIT_V0);
-	}
+	const struct opcode *op = &block->opcode_list[offset];
+	u8 rs;
 
+	rs = lightrec_request_reg_in(reg_cache, _jit, op->r.rs, JIT_V0);
 	lightrec_lock_reg(reg_cache, _jit, rs);
 
 	return rs;
@@ -225,7 +212,7 @@ static void rec_b(struct lightrec_cstate *state, const struct block *block, u16
 
 		branch->target = target_offset;
 		if (is_forward)
-			branch->branch = jit_jmpi();
+			branch->branch = jit_b();
 		else
 			branch->branch = jit_bgti(LIGHTREC_REG_CYCLE, 0);
 	}
@@ -867,10 +854,8 @@ static void rec_alu_div(struct lightrec_cstate *state,
 	}
 
 	if (!no_check) {
-		lightrec_regcache_mark_live(reg_cache, _jit);
-
 		/* Jump above the div-by-zero handler */
-		to_end = jit_jmpi();
+		to_end = jit_b();
 
 		jit_patch(branch);
 
@@ -985,22 +970,26 @@ static void call_to_c_wrapper(struct lightrec_cstate *state, const struct block
 {
 	struct regcache *reg_cache = state->reg_cache;
 	jit_state_t *_jit = block->_jit;
-	u8 tmp, tmp3;
+	u8 tmp, tmp2;
 
-	if (with_arg)
-		tmp3 = lightrec_alloc_reg(reg_cache, _jit, JIT_R1);
 	tmp = lightrec_alloc_reg_temp(reg_cache, _jit);
-
 	jit_ldxi(tmp, LIGHTREC_REG_STATE,
 		 offsetof(struct lightrec_state, wrappers_eps[wrapper]));
-	if (with_arg)
-		jit_movi(tmp3, arg);
 
+	if (with_arg) {
+		tmp2 = lightrec_alloc_reg_temp(reg_cache, _jit);
+		jit_movi(tmp2, arg);
+
+		jit_stxi_i(offsetof(struct lightrec_state, c_wrapper_arg),
+			   LIGHTREC_REG_STATE, tmp2);
+
+		lightrec_free_reg(reg_cache, tmp2);
+	}
+
+	lightrec_regcache_mark_live(reg_cache, _jit);
 	jit_callr(tmp);
 
 	lightrec_free_reg(reg_cache, tmp);
-	if (with_arg)
-		lightrec_free_reg(reg_cache, tmp3);
 	lightrec_regcache_mark_live(reg_cache, _jit);
 }
 
@@ -1012,7 +1001,7 @@ static void rec_io(struct lightrec_cstate *state,
 	jit_state_t *_jit = block->_jit;
 	union code c = block->opcode_list[offset].c;
 	u16 flags = block->opcode_list[offset].flags;
-	bool is_tagged = flags & (LIGHTREC_HW_IO | LIGHTREC_DIRECT_IO);
+	bool is_tagged = LIGHTREC_FLAGS_GET_IO_MODE(flags);
 	u32 lut_entry;
 
 	jit_note(__FILE__, __LINE__);
@@ -1033,6 +1022,111 @@ static void rec_io(struct lightrec_cstate *state,
 	}
 }
 
+static void rec_store_memory(struct lightrec_cstate *cstate,
+			     const struct block *block,
+			     u16 offset, jit_code_t code,
+			     uintptr_t addr_offset, u32 addr_mask,
+			     bool invalidate)
+{
+	struct regcache *reg_cache = cstate->reg_cache;
+	struct opcode *op = &block->opcode_list[offset];
+	jit_state_t *_jit = block->_jit;
+	union code c = op->c;
+	u8 rs, rt, tmp, tmp2, tmp3, addr_reg, addr_reg2;
+	s16 imm = (s16)c.i.imm;
+	s32 simm = (s32)imm << (__WORDSIZE / 32 - 1);
+	s32 lut_offt = offsetof(struct lightrec_state, code_lut);
+	bool no_mask = op->flags & LIGHTREC_NO_MASK;
+	bool add_imm = c.i.imm && invalidate && simm + lut_offt != (s16)(simm + lut_offt);
+	bool need_tmp = !no_mask || addr_offset || add_imm;
+	bool need_tmp2 = addr_offset || invalidate;
+
+	rt = lightrec_alloc_reg_in(reg_cache, _jit, c.i.rt, 0);
+	rs = lightrec_alloc_reg_in(reg_cache, _jit, c.i.rs, 0);
+	if (need_tmp)
+		tmp = lightrec_alloc_reg_temp(reg_cache, _jit);
+
+	addr_reg = rs;
+
+	if (add_imm) {
+		jit_addi(tmp, addr_reg, (s16)c.i.imm);
+		addr_reg = tmp;
+		imm = 0;
+	} else if (simm) {
+		lut_offt += simm;
+	}
+
+	if (!no_mask) {
+		jit_andi(tmp, addr_reg, addr_mask);
+		addr_reg = tmp;
+	}
+
+	if (need_tmp2)
+		tmp2 = lightrec_alloc_reg_temp(reg_cache, _jit);
+
+	if (addr_offset) {
+		jit_addi(tmp2, addr_reg, addr_offset);
+		addr_reg2 = tmp2;
+	} else {
+		addr_reg2 = addr_reg;
+	}
+
+	jit_new_node_www(code, imm, addr_reg2, rt);
+	lightrec_free_reg(reg_cache, rt);
+
+	if (invalidate) {
+		tmp3 = lightrec_alloc_reg_in(reg_cache, _jit, 0, 0);
+
+		if (c.i.op != OP_SW) {
+			jit_andi(tmp2, addr_reg, ~3);
+			addr_reg = tmp2;
+		}
+
+		if (__WORDSIZE == 64) {
+			jit_lshi(tmp2, addr_reg, 1);
+			addr_reg = tmp2;
+		}
+
+		if (__WORDSIZE == 64 || addr_reg != rs || c.i.rs != 0) {
+			jit_addr(tmp2, addr_reg, LIGHTREC_REG_STATE);
+			addr_reg = tmp2;
+		}
+
+		jit_stxi(lut_offt, addr_reg, tmp3);
+
+		lightrec_free_reg(reg_cache, tmp3);
+	}
+
+	if (need_tmp2)
+		lightrec_free_reg(reg_cache, tmp2);
+	if (need_tmp)
+		lightrec_free_reg(reg_cache, tmp);
+	lightrec_free_reg(reg_cache, rs);
+}
+
+static void rec_store_ram(struct lightrec_cstate *cstate,
+			  const struct block *block,
+			  u16 offset, jit_code_t code,
+			  bool invalidate)
+{
+	_jit_note(block->_jit, __FILE__, __LINE__);
+
+	return rec_store_memory(cstate, block, offset, code,
+				cstate->state->offset_ram,
+				RAM_SIZE - 1, invalidate);
+}
+
+static void rec_store_scratch(struct lightrec_cstate *cstate,
+			      const struct block *block,
+			      u16 offset, jit_code_t code)
+{
+	_jit_note(block->_jit, __FILE__, __LINE__);
+
+	return rec_store_memory(cstate, block, offset, code,
+				cstate->state->offset_scratch,
+				0x1fffffff, false);
+}
+
 static void rec_store_direct_no_invalidate(struct lightrec_cstate *cstate,
 					   const struct block *block,
 					   u16 offset, jit_code_t code)
@@ -1070,11 +1164,9 @@ static void rec_store_direct_no_invalidate(struct lightrec_cstate *cstate,
 	if (state->offset_ram != state->offset_scratch) {
 		to_not_ram = jit_bmsi(tmp, BIT(28));
 
-		lightrec_regcache_mark_live(reg_cache, _jit);
-
 		jit_movi(tmp2, state->offset_ram);
 
-		to_end = jit_jmpi();
+		to_end = jit_b();
 		jit_patch(to_not_ram);
 
 		jit_movi(tmp2, state->offset_scratch);
@@ -1125,8 +1217,6 @@ static void rec_store_direct(struct lightrec_cstate *cstate, const struct block
 
 	to_not_ram = jit_bgti(tmp2, ram_size);
 
-	lightrec_regcache_mark_live(reg_cache, _jit);
-
 	/* Compute the offset to the code LUT */
 	jit_andi(tmp, tmp2, (RAM_SIZE - 1) & ~3);
 	if (__WORDSIZE == 64)
@@ -1139,7 +1229,7 @@ static void rec_store_direct(struct lightrec_cstate *cstate, const struct block
 	if (state->offset_ram != state->offset_scratch) {
 		jit_movi(tmp, state->offset_ram);
 
-		to_end = jit_jmpi();
+		to_end = jit_b();
 	}
 
 	jit_patch(to_not_ram);
@@ -1167,16 +1257,25 @@ static void rec_store(struct lightrec_cstate *state,
 		      const struct block *block, u16 offset, jit_code_t code)
 {
 	u16 flags = block->opcode_list[offset].flags;
+	bool no_invalidate = (flags & LIGHTREC_NO_INVALIDATE) ||
+		state->state->invalidate_from_dma_only;
 
-	if (flags & LIGHTREC_NO_INVALIDATE) {
-		rec_store_direct_no_invalidate(state, block, offset, code);
-	} else if (flags & LIGHTREC_DIRECT_IO) {
-		if (state->state->invalidate_from_dma_only)
+	switch (LIGHTREC_FLAGS_GET_IO_MODE(flags)) {
+	case LIGHTREC_IO_RAM:
+		rec_store_ram(state, block, offset, code, !no_invalidate);
+		break;
+	case LIGHTREC_IO_SCRATCH:
+		rec_store_scratch(state, block, offset, code);
+		break;
+	case LIGHTREC_IO_DIRECT:
+		if (no_invalidate)
 			rec_store_direct_no_invalidate(state, block, offset, code);
 		else
 			rec_store_direct(state, block, offset, code);
-	} else {
+		break;
+	default:
 		rec_io(state, block, offset, true, false);
+		break;
 	}
 }
 
@@ -1223,6 +1322,74 @@ static void rec_SWC2(struct lightrec_cstate *state,
 	rec_io(state, block, offset, false, false);
 }
 
+static void rec_load_memory(struct lightrec_cstate *cstate,
+			    const struct block *block,
+			    u16 offset, jit_code_t code, bool is_unsigned,
+			    uintptr_t addr_offset, u32 addr_mask)
+{
+	struct regcache *reg_cache = cstate->reg_cache;
+	struct opcode *op = &block->opcode_list[offset];
+	jit_state_t *_jit = block->_jit;
+	u8 rs, rt, addr_reg, flags = REG_EXT;
+	union code c = op->c;
+
+	if (!c.i.rt)
+		return;
+
+	if (is_unsigned)
+		flags |= REG_ZEXT;
+
+	rs = lightrec_alloc_reg_in(reg_cache, _jit, c.i.rs, 0);
+	rt = lightrec_alloc_reg_out(reg_cache, _jit, c.i.rt, flags);
+
+	if (!(op->flags & LIGHTREC_NO_MASK)) {
+		jit_andi(rt, rs, addr_mask);
+		addr_reg = rt;
+	} else {
+		addr_reg = rs;
+	}
+
+	if (addr_offset) {
+		jit_addi(rt, addr_reg, addr_offset);
+		addr_reg = rt;
+	}
+
+	jit_new_node_www(code, rt, addr_reg, (s16)c.i.imm);
+
+	lightrec_free_reg(reg_cache, rs);
+	lightrec_free_reg(reg_cache, rt);
+}
+
+static void rec_load_ram(struct lightrec_cstate *cstate,
+			 const struct block *block,
+			 u16 offset, jit_code_t code, bool is_unsigned)
+{
+	_jit_note(block->_jit, __FILE__, __LINE__);
+
+	rec_load_memory(cstate, block, offset, code, is_unsigned,
+			cstate->state->offset_ram, RAM_SIZE - 1);
+}
+
+static void rec_load_bios(struct lightrec_cstate *cstate,
+			  const struct block *block,
+			  u16 offset, jit_code_t code, bool is_unsigned)
+{
+	_jit_note(block->_jit, __FILE__, __LINE__);
+
+	rec_load_memory(cstate, block, offset, code, is_unsigned,
+			cstate->state->offset_bios, 0x1fffffff);
+}
+
+static void rec_load_scratch(struct lightrec_cstate *cstate,
+			     const struct block *block,
+			     u16 offset, jit_code_t code, bool is_unsigned)
+{
+	_jit_note(block->_jit, __FILE__, __LINE__);
+
+	rec_load_memory(cstate, block, offset, code, is_unsigned,
+			cstate->state->offset_scratch, 0x1fffffff);
+}
+
 static void rec_load_direct(struct lightrec_cstate *cstate, const struct block *block,
 			    u16 offset, jit_code_t code, bool is_unsigned)
 {
@@ -1276,15 +1443,13 @@ static void rec_load_direct(struct lightrec_cstate *cstate, const struct block *
 	} else {
 		to_not_ram = jit_bmsi(addr_reg, BIT(28));
 
-		lightrec_regcache_mark_live(reg_cache, _jit);
-
 		/* Convert to KUNSEG and avoid RAM mirrors */
 		jit_andi(rt, addr_reg, RAM_SIZE - 1);
 
 		if (state->offset_ram)
 			jit_movi(tmp, state->offset_ram);
 
-		to_end = jit_jmpi();
+		to_end = jit_b();
 
 		jit_patch(to_not_ram);
 
@@ -1297,7 +1462,7 @@ static void rec_load_direct(struct lightrec_cstate *cstate, const struct block *
 		jit_movi(tmp, state->offset_bios);
 
 		if (state->offset_bios != state->offset_scratch) {
-			to_end2 = jit_jmpi();
+			to_end2 = jit_b();
 
 			jit_patch(to_not_bios);
 
@@ -1328,10 +1493,23 @@ static void rec_load(struct lightrec_cstate *state, const struct block *block,
 {
 	u16 flags = block->opcode_list[offset].flags;
 
-	if (flags & LIGHTREC_DIRECT_IO)
+	switch (LIGHTREC_FLAGS_GET_IO_MODE(flags)) {
+	case LIGHTREC_IO_RAM:
+		rec_load_ram(state, block, offset, code, is_unsigned);
+		break;
+	case LIGHTREC_IO_BIOS:
+		rec_load_bios(state, block, offset, code, is_unsigned);
+		break;
+	case LIGHTREC_IO_SCRATCH:
+		rec_load_scratch(state, block, offset, code, is_unsigned);
+		break;
+	case LIGHTREC_IO_DIRECT:
 		rec_load_direct(state, block, offset, code, is_unsigned);
-	else
+		break;
+	default:
 		rec_io(state, block, offset, false, true);
+		break;
+	}
 }
 
 static void rec_LB(struct lightrec_cstate *state, const struct block *block, u16 offset)
diff --git a/deps/lightrec/lightning-wrapper.h b/deps/lightrec/lightning-wrapper.h
index 7eeb15f0..b0e8bf3b 100644
--- a/deps/lightrec/lightning-wrapper.h
+++ b/deps/lightrec/lightning-wrapper.h
@@ -19,4 +19,6 @@
 
 #endif
 
+#define jit_b()			jit_beqr(0, 0)
+
 #endif /* __LIGHTNING_WRAPPER_H__ */
diff --git a/deps/lightrec/lightrec-private.h b/deps/lightrec/lightrec-private.h
index 4b797a1d..75940b3f 100644
--- a/deps/lightrec/lightrec-private.h
+++ b/deps/lightrec/lightrec-private.h
@@ -126,6 +126,7 @@ struct lightrec_state {
 	u32 target_cycle;
 	u32 exit_flags;
 	u32 old_cycle_counter;
+	u32 c_wrapper_arg;
 	struct block *dispatcher, *c_wrapper_block;
 	void *c_wrappers[C_WRAPPERS_COUNT];
 	void *wrappers_eps[C_WRAPPERS_COUNT];
diff --git a/deps/lightrec/lightrec.c b/deps/lightrec/lightrec.c
index d4ab419f..9889272a 100644
--- a/deps/lightrec/lightrec.c
+++ b/deps/lightrec/lightrec.c
@@ -250,13 +250,13 @@ u32 lightrec_rw(struct lightrec_state *state, union code op,
 	}
 
 	if (unlikely(map->ops)) {
-		if (flags)
-			*flags |= LIGHTREC_HW_IO;
+		if (flags && !LIGHTREC_FLAGS_GET_IO_MODE(*flags))
+			*flags |= LIGHTREC_IO_MODE(LIGHTREC_IO_HW);
 
 		ops = map->ops;
 	} else {
-		if (flags)
-			*flags |= LIGHTREC_DIRECT_IO;
+		if (flags && !LIGHTREC_FLAGS_GET_IO_MODE(*flags))
+			*flags |= LIGHTREC_IO_MODE(LIGHTREC_IO_DIRECT);
 
 		ops = &lightrec_default_ops;
 	}
@@ -323,16 +323,17 @@ static void lightrec_rw_helper(struct lightrec_state *state,
 	}
 }
 
-static void lightrec_rw_cb(struct lightrec_state *state, union code op)
+static void lightrec_rw_cb(struct lightrec_state *state)
 {
-	lightrec_rw_helper(state, op, NULL, NULL);
+	lightrec_rw_helper(state, (union code)state->c_wrapper_arg, NULL, NULL);
 }
 
-static void lightrec_rw_generic_cb(struct lightrec_state *state, u32 arg)
+static void lightrec_rw_generic_cb(struct lightrec_state *state)
 {
 	struct block *block;
 	struct opcode *op;
 	bool was_tagged;
+	u32 arg = state->c_wrapper_arg;
 	u16 offset = (u16)arg;
 
 	block = lightrec_find_block_from_lut(state->block_cache,
@@ -344,7 +345,7 @@ static void lightrec_rw_generic_cb(struct lightrec_state *state, u32 arg)
 	}
 
 	op = &block->opcode_list[offset];
-	was_tagged = op->flags & (LIGHTREC_HW_IO | LIGHTREC_DIRECT_IO);
+	was_tagged = LIGHTREC_FLAGS_GET_IO_MODE(op->flags);
 
 	lightrec_rw_helper(state, op->c, &op->flags, block);
 
@@ -529,8 +530,10 @@ void lightrec_mtc(struct lightrec_state *state, union code op, u32 data)
 		lightrec_mtc2(state, op.r.rd, data);
 }
 
-static void lightrec_mtc_cb(struct lightrec_state *state, union code op)
+static void lightrec_mtc_cb(struct lightrec_state *state)
 {
+	union code op = (union code) state->c_wrapper_arg;
+
 	lightrec_mtc(state, op, state->regs.gpr[op.r.rt]);
 }
 
@@ -558,12 +561,17 @@ void lightrec_cp(struct lightrec_state *state, union code op)
 	(*state->ops.cop2_op)(state, op.opcode);
 }
 
-static void lightrec_syscall_cb(struct lightrec_state *state, union code op)
+static void lightrec_cp_cb(struct lightrec_state *state)
+{
+	lightrec_cp(state, (union code) state->c_wrapper_arg);
+}
+
+static void lightrec_syscall_cb(struct lightrec_state *state)
 {
 	lightrec_set_exit_flags(state, LIGHTREC_EXIT_SYSCALL);
 }
 
-static void lightrec_break_cb(struct lightrec_state *state, union code op)
+static void lightrec_break_cb(struct lightrec_state *state)
 {
 	lightrec_set_exit_flags(state, LIGHTREC_EXIT_BREAK);
 }
@@ -669,12 +677,11 @@ static void * get_next_block_func(struct lightrec_state *state, u32 pc)
 }
 
 static s32 c_function_wrapper(struct lightrec_state *state, s32 cycles_delta,
-			      void (*f)(struct lightrec_state *, u32 d),
-			      u32 d)
+			      void (*f)(struct lightrec_state *))
 {
 	state->current_cycle = state->target_cycle - cycles_delta;
 
-	(*f)(state, d);
+	(*f)(state);
 
 	return state->target_cycle - state->current_cycle;
 }
@@ -753,7 +760,6 @@ static struct block * generate_wrapper(struct lightrec_state *state)
 	jit_pushargr(LIGHTREC_REG_STATE);
 	jit_pushargr(LIGHTREC_REG_CYCLE);
 	jit_pushargr(JIT_R0);
-	jit_pushargr(JIT_R1);
 	jit_finishi(c_function_wrapper);
 	jit_retval_i(LIGHTREC_REG_CYCLE);
 
@@ -1130,8 +1136,7 @@ static bool lightrec_block_is_fully_tagged(const struct block *block)
 		case OP_SWR:
 		case OP_LWC2:
 		case OP_SWC2:
-			if (!(op->flags & (LIGHTREC_DIRECT_IO |
-					   LIGHTREC_HW_IO)))
+			if (!LIGHTREC_FLAGS_GET_IO_MODE(op->flags))
 				return false;
 		default: /* fall-through */
 			continue;
@@ -1522,7 +1527,7 @@ struct lightrec_state * lightrec_init(char *argv0,
 	state->c_wrappers[C_WRAPPER_RW] = lightrec_rw_cb;
 	state->c_wrappers[C_WRAPPER_RW_GENERIC] = lightrec_rw_generic_cb;
 	state->c_wrappers[C_WRAPPER_MTC] = lightrec_mtc_cb;
-	state->c_wrappers[C_WRAPPER_CP] = lightrec_cp;
+	state->c_wrappers[C_WRAPPER_CP] = lightrec_cp_cb;
 	state->c_wrappers[C_WRAPPER_SYSCALL] = lightrec_syscall_cb;
 	state->c_wrappers[C_WRAPPER_BREAK] = lightrec_break_cb;
 
diff --git a/deps/lightrec/optimizer.c b/deps/lightrec/optimizer.c
index f719192b..7304abca 100644
--- a/deps/lightrec/optimizer.c
+++ b/deps/lightrec/optimizer.c
@@ -277,7 +277,8 @@ static bool reg_is_dead(const struct opcode *list, unsigned int offset, u8 reg)
 			return true;
 
 		if (has_delay_slot(list[i].c)) {
-			if (list[i].flags & LIGHTREC_NO_DS)
+			if (list[i].flags & LIGHTREC_NO_DS ||
+			    opcode_reads_register(list[i + 1].c, reg))
 				return false;
 
 			return opcode_writes_register(list[i + 1].c, reg);
@@ -459,16 +460,18 @@ bool load_in_delay_slot(union code op)
 	return false;
 }
 
-static u32 lightrec_propagate_consts(const struct opcode *op, u32 known, u32 *v)
+static u32 lightrec_propagate_consts(const struct opcode *op,
+				     const struct opcode *prev,
+				     u32 known, u32 *v)
 {
-	union code c = op->c;
+	union code c = prev->c;
 
 	/* Register $zero is always, well, zero */
 	known |= BIT(0);
 	v[0] = 0;
 
 	if (op->flags & LIGHTREC_SYNC)
-		return 0;
+		return BIT(0);
 
 	switch (c.i.op) {
 	case OP_SPECIAL:
@@ -817,15 +820,19 @@ static void lightrec_optimize_sll_sra(struct opcode *list, unsigned int offset)
 static int lightrec_transform_ops(struct lightrec_state *state, struct block *block)
 {
 	struct opcode *list = block->opcode_list;
-	struct opcode *op;
+	struct opcode *prev, *op = NULL;
 	u32 known = BIT(0);
 	u32 values[32] = { 0 };
 	unsigned int i;
 	int reader;
 
 	for (i = 0; i < block->nb_ops; i++) {
+		prev = op;
 		op = &list[i];
 
+		if (prev)
+			known = lightrec_propagate_consts(op, prev, known, values);
+
 		/* Transform all opcodes detected as useless to real NOPs
 		 * (0x0: SLL r0, r0, #0) */
 		if (op->opcode != 0 && is_nop(op->c)) {
@@ -933,8 +940,6 @@ static int lightrec_transform_ops(struct lightrec_state *state, struct block *bl
 		default: /* fall-through */
 			break;
 		}
-
-		known = lightrec_propagate_consts(op, known, values);
 	}
 
 	return 0;
@@ -1229,15 +1234,20 @@ static int lightrec_early_unload(struct lightrec_state *state, struct block *blo
 static int lightrec_flag_io(struct lightrec_state *state, struct block *block)
 {
 	const struct lightrec_mem_map *map;
-	struct opcode *list;
+	struct opcode *prev2, *prev = NULL, *list = NULL;
 	u32 known = BIT(0);
 	u32 values[32] = { 0 };
 	unsigned int i;
 	u32 val;
 
 	for (i = 0; i < block->nb_ops; i++) {
+		prev2 = prev;
+		prev = list;
 		list = &block->opcode_list[i];
 
+		if (prev)
+			known = lightrec_propagate_consts(list, prev, known, values);
+
 		switch (list->i.op) {
 		case OP_SB:
 		case OP_SH:
@@ -1279,24 +1289,47 @@ static int lightrec_flag_io(struct lightrec_state *state, struct block *block)
 		case OP_LWR:
 		case OP_LWC2:
 			if (OPT_FLAG_IO && (known & BIT(list->i.rs))) {
-				val = kunseg(values[list->i.rs] + (s16) list->i.imm);
-				map = lightrec_get_map(state, NULL, val);
+				if (prev && prev->i.op == OP_LUI &&
+				    !(prev2 && has_delay_slot(prev2->c)) &&
+				    prev->i.rt == list->i.rs &&
+				    list->i.rt == list->i.rs &&
+				    prev->i.imm & 0x8000) {
+					pr_debug("Convert LUI at offset 0x%x to kuseg\n",
+						 i - 1 << 2);
+
+					val = kunseg(prev->i.imm << 16);
+					prev->i.imm = val >> 16;
+					values[list->i.rs] = val;
+				}
+
+				val = values[list->i.rs] + (s16) list->i.imm;
+				map = lightrec_get_map(state, NULL, kunseg(val));
 
 				if (!map || map->ops ||
 				    map == &state->maps[PSX_MAP_PARALLEL_PORT]) {
-					pr_debug("Flagging opcode %u as accessing I/O registers\n",
+					pr_debug("Flagging opcode %u as I/O access\n",
 						 i);
-					list->flags |= LIGHTREC_HW_IO;
-				} else {
-					pr_debug("Flaging opcode %u as direct memory access\n", i);
-					list->flags |= LIGHTREC_DIRECT_IO;
+					list->flags |= LIGHTREC_IO_MODE(LIGHTREC_IO_HW);
+					break;
+				}
+
+				if (val - map->pc < map->length)
+					list->flags |= LIGHTREC_NO_MASK;
+
+				if (map == &state->maps[PSX_MAP_KERNEL_USER_RAM]) {
+					pr_debug("Flaging opcode %u as RAM access\n", i);
+					list->flags |= LIGHTREC_IO_MODE(LIGHTREC_IO_RAM);
+				} else if (map == &state->maps[PSX_MAP_BIOS]) {
+					pr_debug("Flaging opcode %u as BIOS access\n", i);
+					list->flags |= LIGHTREC_IO_MODE(LIGHTREC_IO_BIOS);
+				} else if (map == &state->maps[PSX_MAP_SCRATCH_PAD]) {
+					pr_debug("Flaging opcode %u as scratchpad access\n", i);
+					list->flags |= LIGHTREC_IO_MODE(LIGHTREC_IO_SCRATCH);
 				}
 			}
 		default: /* fall-through */
 			break;
 		}
-
-		known = lightrec_propagate_consts(list, known, values);
 	}
 
 	return 0;
@@ -1483,15 +1516,19 @@ static bool lightrec_always_skip_div_check(void)
 
 static int lightrec_flag_mults_divs(struct lightrec_state *state, struct block *block)
 {
-	struct opcode *list;
+	struct opcode *prev, *list = NULL;
 	u8 reg_hi, reg_lo;
 	unsigned int i;
 	u32 known = BIT(0);
 	u32 values[32] = { 0 };
 
 	for (i = 0; i < block->nb_ops - 1; i++) {
+		prev = list;
 		list = &block->opcode_list[i];
 
+		if (prev)
+			known = lightrec_propagate_consts(list, prev, known, values);
+
 		if (list->i.op != OP_SPECIAL)
 			continue;
 
@@ -1507,14 +1544,12 @@ static int lightrec_flag_mults_divs(struct lightrec_state *state, struct block *
 		case OP_SPECIAL_MULTU:
 			break;
 		default:
-			known = lightrec_propagate_consts(list, known, values);
 			continue;
 		}
 
 		/* Don't support opcodes in delay slots */
 		if ((i && has_delay_slot(block->opcode_list[i - 1].c)) ||
 		    (list->flags & LIGHTREC_NO_DS)) {
-			known = lightrec_propagate_consts(list, known, values);
 			continue;
 		}
 
@@ -1558,8 +1593,6 @@ static int lightrec_flag_mults_divs(struct lightrec_state *state, struct block *
 		} else {
 			list->r.imm = 0;
 		}
-
-		known = lightrec_propagate_consts(list, known, values);
 	}
 
 	return 0;
-- 
2.39.2