Add jump target alignment check for JAL and JALR instructions

- Implement IsTargetMisaligned to verify 4-byte alignment when the compressed extension is disabled.
- Trigger InstructionAddressMisaligned trap on invalid jump targets in JAL and JALR.
- Add unit tests for misaligned jumps in both RV32 and RV64 environments.
NB: This existed for one variant of jump on RV32 -- this completes the rest.
PiperOrigin-RevId: 862918956
Change-Id: I10ac399ff6b1a78c0a001e7be8b26ee5f6763fb2
diff --git a/riscv/riscv_i_instructions.cc b/riscv/riscv_i_instructions.cc
index feaf111..baf0afe 100644
--- a/riscv/riscv_i_instructions.cc
+++ b/riscv/riscv_i_instructions.cc
@@ -40,6 +40,28 @@
               /*epc*/ inst->address(), inst);
 }
 
+namespace {
+// Helper function to check target alignment for jumps.
+// Returns true if trap was taken, false otherwise.
+template <typename UIntReg>
+bool IsTargetMisaligned(UIntReg target, const Instruction* instruction) {
+  auto* state = static_cast<RiscVState*>(instruction->state());
+  auto res =
+      state->csr_set()->GetCsr(static_cast<uint64_t>(RiscVCsrEnum::kMIsa));
+  bool has_c_extension =
+      res.ok() && (res.value() != nullptr) &&
+      ((res.value()->GetUint64() &
+        static_cast<uint64_t>(IsaExtensions::kCompressed)) != 0);
+  if (!has_c_extension && ((target & 0x3) != 0)) {
+    state->Trap(/*is_interrupt*/ false, /*trap_value*/ target,
+                *ExceptionCode::kInstructionAddressMisaligned,
+                instruction->address(), instruction);
+    return true;
+  }
+  return false;
+}
+}  // namespace
+
 // The following instruction semantic functions implement basic alu operations.
 // They are used for both register-register and register-immediate versions of
 // the corresponding instructions.
@@ -126,6 +148,9 @@
 void RiscVIJal(const Instruction* instruction) {
   UIntReg offset = generic::GetInstructionSource<UIntReg>(instruction, 0);
   UIntReg target = (offset + instruction->address()) & ~0x1;
+
+  if (IsTargetMisaligned(target, instruction)) return;
+
   UIntReg return_address = instruction->address() + instruction->size();
   auto* db = instruction->Destination(0)->AllocateDataBuffer();
   db->SetSubmit<UIntReg>(0, target);
@@ -141,24 +166,13 @@
   UIntReg reg_base = generic::GetInstructionSource<UIntReg>(instruction, 0);
   UIntReg offset = generic::GetInstructionSource<UIntReg>(instruction, 1);
   UIntReg target = (offset + reg_base) & ~0x1;
-  // Check target for proper alignment.
-  auto* state = static_cast<RiscVState*>(instruction->state());
-  auto res =
-      state->csr_set()->GetCsr(static_cast<uint64_t>(RiscVCsrEnum::kMIsa));
-  bool has_c_extension = true;
-  if (res.ok() || res.value() != nullptr) {
-    has_c_extension = (res.value()->GetUint64() &
-                       static_cast<uint64_t>(IsaExtensions::kCompressed)) != 0;
-  }
-  if (!has_c_extension && ((target & 0x3) != 0)) {
-    state->Trap(/*is_interrupt*/ false, instruction->address(),
-                *ExceptionCode::kInstructionAddressMisaligned,
-                instruction->address(), instruction);
-    return;
-  }
+
+  if (IsTargetMisaligned(target, instruction)) return;
+
   UIntReg return_address = instruction->address() + instruction->size();
   auto* db = instruction->Destination(0)->AllocateDataBuffer();
   db->SetSubmit<UIntReg>(0, target);
+  auto* state = static_cast<RiscVState*>(instruction->state());
   state->set_branch(true);
   auto* reg = static_cast<generic::RegisterDestinationOperand<UIntReg>*>(
                   instruction->Destination(1))
@@ -293,6 +307,9 @@
   UIntReg offset = generic::GetInstructionSource<UIntReg>(instruction, 0);
   UIntReg target = offset + instruction->address();
   target &= (std::numeric_limits<UIntReg>::max() << 1);
+
+  if (IsTargetMisaligned(target, instruction)) return;
+
   UIntReg return_address = instruction->address() + instruction->size();
   auto* db = instruction->Destination(0)->AllocateDataBuffer();
   db->SetSubmit<UIntReg>(0, target);
@@ -309,24 +326,13 @@
   UIntReg offset = generic::GetInstructionSource<UIntReg>(instruction, 1);
   UIntReg target = offset + reg_base;
   target &= (std::numeric_limits<UIntReg>::max() << 1);
-  // Check target for proper alignment.
-  auto* state = static_cast<RiscVState*>(instruction->state());
-  auto res =
-      state->csr_set()->GetCsr(static_cast<uint64_t>(RiscVCsrEnum::kMIsa));
-  bool has_c_extension = true;
-  if (res.ok() || res.value() != nullptr) {
-    has_c_extension = (res.value()->GetUint64() &
-                       static_cast<uint64_t>(IsaExtensions::kCompressed)) != 0;
-  }
-  if (!has_c_extension && ((target & 0x3) != 0)) {
-    state->Trap(/*is_interrupt*/ false, instruction->address(),
-                *ExceptionCode::kInstructionAddressMisaligned,
-                instruction->address(), instruction);
-    return;
-  }
+
+  if (IsTargetMisaligned(target, instruction)) return;
+
   UIntReg return_address = instruction->address() + instruction->size();
   auto* db = instruction->Destination(0)->AllocateDataBuffer();
   db->SetSubmit<UIntReg>(0, target);
+  auto* state = static_cast<RiscVState*>(instruction->state());
   state->set_branch(true);
   auto* reg = static_cast<generic::RegisterDestinationOperand<UIntReg>*>(
                   instruction->Destination(1))
diff --git a/riscv/test/BUILD b/riscv/test/BUILD
index fc60503..a959815 100644
--- a/riscv/test/BUILD
+++ b/riscv/test/BUILD
@@ -990,6 +990,20 @@
     ],
 )
 
+cc_test(
+    name = "riscv64_i_instructions_test",
+    srcs = ["riscv64_i_instructions_test.cc"],
+    deps = [
+        "//riscv:riscv_g",
+        "//riscv:riscv_state",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_googletest//:gtest_main",
+        "@com_google_mpact-sim//mpact/sim/generic:core",
+        "@com_google_mpact-sim//mpact/sim/generic:instruction",
+        "@com_google_mpact-sim//mpact/sim/util/memory",
+    ],
+)
+
 config_setting(
     name = "arm_cpu",
     values = {"cpu": "arm"},
diff --git a/riscv/test/riscv64_i_instructions_test.cc b/riscv/test/riscv64_i_instructions_test.cc
new file mode 100644
index 0000000..ad4bf9a
--- /dev/null
+++ b/riscv/test/riscv64_i_instructions_test.cc
@@ -0,0 +1,207 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <cstdint>
+#include <string>
+#include <tuple>
+#include <vector>
+
+#include "absl/strings/string_view.h"
+#include "googlemock/include/gmock/gmock.h"
+#include "mpact/sim/generic/data_buffer.h"
+#include "mpact/sim/generic/immediate_operand.h"
+#include "mpact/sim/generic/instruction.h"
+#include "mpact/sim/util/memory/flat_demand_memory.h"
+#include "riscv/riscv_csr.h"
+#include "riscv/riscv_i_instructions.h"
+#include "riscv/riscv_register.h"
+#include "riscv/riscv_state.h"
+
+// This file contains tests for individual RiscV64I instructions.
+
+namespace {
+
+using ::mpact::sim::generic::ImmediateOperand;
+using ::mpact::sim::generic::Instruction;
+using ::mpact::sim::riscv::RiscVState;
+using ::mpact::sim::riscv::RiscVXlen;
+using ::mpact::sim::riscv::RV64Register;
+using ::mpact::sim::util::FlatDemandMemory;
+
+constexpr char kX1[] = "x1";
+constexpr char kX3[] = "x3";
+
+constexpr uint32_t kInstAddress = 0x2468;
+constexpr uint32_t kOffset = 0x246;
+constexpr uint32_t kMemAddress = 0x1000;
+
+class RV64IInstructionTest : public testing::Test {
+ public:
+  RV64IInstructionTest() {
+    memory_ = new FlatDemandMemory();
+    state_ = new RiscVState("test", RiscVXlen::RV64, memory_);
+    instruction_ = new Instruction(kInstAddress, state_);
+    instruction_->set_size(4);
+    state_->set_on_trap([this](bool is_interrupt, uint64_t trap_value,
+                               uint64_t exception_code, uint64_t epc,
+                               const Instruction* inst) {
+      trap_taken_ = true;
+      trap_is_interrupt_ = is_interrupt;
+      trap_value_ = trap_value;
+      trap_code_ = exception_code;
+      trap_epc_ = epc;
+      return true;
+    });
+  }
+  ~RV64IInstructionTest() override {
+    delete memory_;
+    delete state_;
+    delete instruction_;
+  }
+
+  // Appends the source and destination operands for the register names
+  // given in the two vectors.
+  void AppendRegisterOperands(Instruction* inst,
+                              const std::vector<std::string>& sources,
+                              const std::vector<std::string>& destinations) {
+    for (auto& reg_name : sources) {
+      auto* reg = state_->GetRegister<RV64Register>(reg_name).first;
+      inst->AppendSource(reg->CreateSourceOperand());
+    }
+    for (auto& reg_name : destinations) {
+      auto* reg = state_->GetRegister<RV64Register>(reg_name).first;
+      inst->AppendDestination(reg->CreateDestinationOperand(0));
+    }
+  }
+
+  void AppendRegisterOperands(const std::vector<std::string>& sources,
+                              const std::vector<std::string>& destinations) {
+    AppendRegisterOperands(instruction_, sources, destinations);
+  }
+
+  // Appends immediate source operands with the given values.
+  template <typename T>
+  void AppendImmediateOperands(const std::vector<T>& values) {
+    for (auto value : values) {
+      auto* src = new ImmediateOperand<T>(value);
+      instruction_->AppendSource(src);
+    }
+  }
+
+  // Takes a vector of tuples of register names and values. Fetches each
+  // named register and sets it to the corresponding value.
+  template <typename T>
+  void SetRegisterValues(const std::vector<std::tuple<std::string, T>> values) {
+    for (auto& [reg_name, value] : values) {
+      auto* reg = state_->GetRegister<RV64Register>(reg_name).first;
+      auto* db = state_->db_factory()->Allocate<RV64Register::ValueType>(1);
+      db->Set<T>(0, value);
+      reg->SetDataBuffer(db);
+      db->DecRef();
+    }
+  }
+
+  // Initializes the semantic function of the instruction object.
+  void SetSemanticFunction(Instruction::SemanticFunction fcn) {
+    instruction_->set_semantic_function(fcn);
+  }
+
+  // Returns the value of the named register.
+  template <typename T>
+  T GetRegisterValue(absl::string_view reg_name) {
+    auto* reg = state_->GetRegister<RV64Register>(reg_name).first;
+    return reg->data_buffer()->Get<T>(0);
+  }
+
+  FlatDemandMemory* memory_;
+  RiscVState* state_;
+  Instruction* instruction_;
+  bool trap_taken_ = false;
+  uint64_t trap_epc_ = 0;
+  uint64_t trap_value_ = 0;
+  uint64_t trap_code_ = 0;
+  bool trap_is_interrupt_ = false;
+};
+
+TEST_F(RV64IInstructionTest, RV64IJal) {
+  AppendRegisterOperands({}, {RiscVState::kPcName, kX3});
+  AppendImmediateOperands<uint64_t>({kOffset});
+  SetSemanticFunction(&::mpact::sim::riscv::RV64::RiscVIJal);
+
+  instruction_->Execute(nullptr);
+  EXPECT_EQ(GetRegisterValue<uint64_t>(RiscVState::kPcName),
+            instruction_->address() + kOffset);
+  EXPECT_EQ(GetRegisterValue<uint64_t>(kX3), instruction_->address() + 4);
+}
+
+TEST_F(RV64IInstructionTest, RV64IJalMisaligned) {
+  AppendRegisterOperands({}, {RiscVState::kPcName, kX3});
+  AppendImmediateOperands<uint64_t>({2});
+  SetSemanticFunction(&::mpact::sim::riscv::RV64::RiscVIJal);
+  // Get misa and clear C bit
+  auto misa = state_->csr_set()
+                  ->GetCsr(static_cast<uint64_t>(
+                      ::mpact::sim::riscv::RiscVCsrEnum::kMIsa))
+                  .value();
+  misa->Set(
+      misa->GetUint64() &
+      ~static_cast<uint64_t>(::mpact::sim::riscv::IsaExtensions::kCompressed));
+  trap_taken_ = false;
+  instruction_->Execute(nullptr);
+  EXPECT_TRUE(trap_taken_);
+  EXPECT_FALSE(trap_is_interrupt_);
+  EXPECT_EQ(
+      trap_code_,
+      static_cast<uint64_t>(
+          ::mpact::sim::riscv::ExceptionCode::kInstructionAddressMisaligned));
+  EXPECT_EQ(trap_epc_, instruction_->address());
+}
+
+TEST_F(RV64IInstructionTest, RV64IJalr) {
+  AppendRegisterOperands({kX1}, {RiscVState::kPcName, kX3});
+  AppendImmediateOperands<uint64_t>({kOffset});
+  SetSemanticFunction(&::mpact::sim::riscv::RV64::RiscVIJalr);
+
+  SetRegisterValues<uint64_t>({{kX1, kMemAddress}});
+  instruction_->Execute(nullptr);
+  EXPECT_EQ(GetRegisterValue<uint64_t>(RiscVState::kPcName),
+            kOffset + kMemAddress);
+  EXPECT_EQ(GetRegisterValue<uint64_t>(kX3), instruction_->address() + 4);
+}
+
+TEST_F(RV64IInstructionTest, RV64IJalrMisaligned) {
+  AppendRegisterOperands({kX1}, {RiscVState::kPcName, kX3});
+  AppendImmediateOperands<uint64_t>({2});
+  SetSemanticFunction(&::mpact::sim::riscv::RV64::RiscVIJalr);
+  SetRegisterValues<uint64_t>({{kX1, 0}});
+  // Get misa and clear C bit
+  auto misa = state_->csr_set()
+                  ->GetCsr(static_cast<uint64_t>(
+                      ::mpact::sim::riscv::RiscVCsrEnum::kMIsa))
+                  .value();
+  misa->Set(
+      misa->GetUint64() &
+      ~static_cast<uint64_t>(::mpact::sim::riscv::IsaExtensions::kCompressed));
+  trap_taken_ = false;
+  instruction_->Execute(nullptr);
+  EXPECT_TRUE(trap_taken_);
+  EXPECT_FALSE(trap_is_interrupt_);
+  EXPECT_EQ(
+      trap_code_,
+      static_cast<uint64_t>(
+          ::mpact::sim::riscv::ExceptionCode::kInstructionAddressMisaligned));
+  EXPECT_EQ(trap_epc_, instruction_->address());
+}
+
+}  // namespace
\ No newline at end of file
diff --git a/riscv/test/riscv_i_instructions_test.cc b/riscv/test/riscv_i_instructions_test.cc
index be6f98e..63e960d 100644
--- a/riscv/test/riscv_i_instructions_test.cc
+++ b/riscv/test/riscv_i_instructions_test.cc
@@ -25,6 +25,7 @@
 #include "mpact/sim/generic/immediate_operand.h"
 #include "mpact/sim/generic/instruction.h"
 #include "mpact/sim/util/memory/flat_demand_memory.h"
+#include "riscv/riscv_csr.h"
 #include "riscv/riscv_register.h"
 #include "riscv/riscv_state.h"
 
@@ -62,6 +63,16 @@
     state_ = new RiscVState("test", RiscVXlen::RV32, memory_);
     instruction_ = new Instruction(kInstAddress, state_);
     instruction_->set_size(4);
+    state_->set_on_trap([this](bool is_interrupt, uint64_t trap_value,
+                               uint64_t exception_code, uint64_t epc,
+                               const Instruction* inst) {
+      trap_taken_ = true;
+      trap_is_interrupt_ = is_interrupt;
+      trap_value_ = trap_value;
+      trap_code_ = exception_code;
+      trap_epc_ = epc;
+      return true;
+    });
   }
   ~RV32IInstructionTest() override {
     delete memory_;
@@ -126,6 +137,11 @@
   FlatDemandMemory* memory_;
   RiscVState* state_;
   Instruction* instruction_;
+  bool trap_taken_ = false;
+  uint64_t trap_epc_ = 0;
+  uint64_t trap_value_ = 0;
+  uint64_t trap_code_ = 0;
+  bool trap_is_interrupt_ = false;
 };
 
 // Almost all the tests below follow the same pattern. There are two phases.
@@ -273,6 +289,29 @@
   EXPECT_EQ(GetRegisterValue<uint32_t>(kX3), instruction_->address() + 4);
 }
 
+TEST_F(RV32IInstructionTest, RV32IJalMisaligned) {
+  AppendRegisterOperands({}, {RiscVState::kPcName, kX3});
+  AppendImmediateOperands<uint32_t>({2});
+  SetSemanticFunction(&::mpact::sim::riscv::RV32::RiscVIJal);
+  // Get misa and clear C bit
+  auto misa = state_->csr_set()
+                  ->GetCsr(static_cast<uint64_t>(
+                      ::mpact::sim::riscv::RiscVCsrEnum::kMIsa))
+                  .value();
+  misa->Set(
+      misa->GetUint64() &
+      ~static_cast<uint64_t>(::mpact::sim::riscv::IsaExtensions::kCompressed));
+  trap_taken_ = false;
+  instruction_->Execute(nullptr);
+  EXPECT_TRUE(trap_taken_);
+  EXPECT_FALSE(trap_is_interrupt_);
+  EXPECT_EQ(
+      trap_code_,
+      static_cast<uint64_t>(
+          ::mpact::sim::riscv::ExceptionCode::kInstructionAddressMisaligned));
+  EXPECT_EQ(trap_epc_, instruction_->address());
+}
+
 TEST_F(RV32IInstructionTest, RV32IJalr) {
   AppendRegisterOperands({kX1}, {RiscVState::kPcName, kX3});
   AppendImmediateOperands<uint32_t>({kOffset});
@@ -285,6 +324,30 @@
   EXPECT_EQ(GetRegisterValue<uint32_t>(kX3), instruction_->address() + 4);
 }
 
+TEST_F(RV32IInstructionTest, RV32IJalrMisaligned) {
+  AppendRegisterOperands({kX1}, {RiscVState::kPcName, kX3});
+  AppendImmediateOperands<uint32_t>({2});
+  SetSemanticFunction(&::mpact::sim::riscv::RV32::RiscVIJalr);
+  SetRegisterValues<uint32_t>({{kX1, 0}});
+  // Get misa and clear C bit
+  auto misa = state_->csr_set()
+                  ->GetCsr(static_cast<uint64_t>(
+                      ::mpact::sim::riscv::RiscVCsrEnum::kMIsa))
+                  .value();
+  misa->Set(
+      misa->GetUint64() &
+      ~static_cast<uint64_t>(::mpact::sim::riscv::IsaExtensions::kCompressed));
+  trap_taken_ = false;
+  instruction_->Execute(nullptr);
+  EXPECT_TRUE(trap_taken_);
+  EXPECT_FALSE(trap_is_interrupt_);
+  EXPECT_EQ(
+      trap_code_,
+      static_cast<uint64_t>(
+          ::mpact::sim::riscv::ExceptionCode::kInstructionAddressMisaligned));
+  EXPECT_EQ(trap_epc_, instruction_->address());
+}
+
 TEST_F(RV32IInstructionTest, RV32IBeq) {
   AppendRegisterOperands({kX1, kX2}, {RiscVState::kPcName});
   AppendImmediateOperands<int32_t>({kOffset});