From 58b041e973def728b3a67d91e92150bd5271ab6e Mon Sep 17 00:00:00 2001
From: Tyker <tyker1@outlook.com>
Date: Mon, 29 Apr 2024 22:31:17 +0200
Subject: [PATCH] Add support for Shlx, Shrx, Sarx

---
 src/capstone2llvmir/x86/x86.cpp      | 30 ++++++++++++++++++++++++++++
 src/capstone2llvmir/x86/x86_impl.h   |  1 +
 src/capstone2llvmir/x86/x86_init.cpp |  6 +++---
 3 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/src/capstone2llvmir/x86/x86.cpp b/src/capstone2llvmir/x86/x86.cpp
index 6f03aaede..1a6501479 100644
--- a/src/capstone2llvmir/x86/x86.cpp
+++ b/src/capstone2llvmir/x86/x86.cpp
@@ -3349,6 +3349,36 @@ void Capstone2LlvmIrTranslatorX86_impl::translateShrd(cs_insn* i, cs_x86* xi, ll
 	storeRegister(X86_REG_OF, ofV, bodyIrb);
 }
 
+/**
+ * X86_INS_SHLX, X86_INS_SHRX, X86_INS_SARX
+ */
+void Capstone2LlvmIrTranslatorX86_impl::translateShiftX(cs_insn* i, cs_x86* xi, llvm::IRBuilder<>& irb)
+{
+	EXPECT_IS_TERNARY(i, xi, irb);
+
+	std::tie(op0, op1, op2) = loadOpTernary(xi, irb);
+
+	auto OpKind = [&] {
+		switch (i->id)
+		{
+		case X86_INS_SHR: return llvm::Instruction::BinaryOps::LShr;
+		case X86_INS_SAR: return llvm::Instruction::BinaryOps::AShr;
+		case X86_INS_SHL: return llvm::Instruction::BinaryOps::Shl;
+		}
+	}();
+
+	unsigned bitwidth = op0->getType()->getIntegerBitWidth();
+	auto* is_overflow = irb.CreateICmpUGE(op2, llvm::ConstantInt::get(op1->getType(), bitwidth));
+	auto* mask = llvm::ConstantInt::get(op1->getType(), bitwidth == 64 ? 0x3f : 0x1f);
+	op1 = irb.CreateAnd(op1, mask);
+	auto* shift = irb.CreateBinOp(OpKind, op1, op2);
+	auto* res = irb.CreateSelect(is_overflow, llvm::ConstantInt::get(op2->getType(), 0), shift);
+
+	// X variants of shifts have no flag impact
+
+	storeOp(xi->operands[0], res, irb);
+}
+
 /**
  * X86_INS_RCR
  */
diff --git a/src/capstone2llvmir/x86/x86_impl.h b/src/capstone2llvmir/x86/x86_impl.h
index daff80b9e..73cf56976 100644
--- a/src/capstone2llvmir/x86/x86_impl.h
+++ b/src/capstone2llvmir/x86/x86_impl.h
@@ -346,6 +346,7 @@ class Capstone2LlvmIrTranslatorX86_impl :
 		void translateShiftRight(cs_insn* i, cs_x86* xi, llvm::IRBuilder<>& irb);
 		void translateShld(cs_insn* i, cs_x86* xi, llvm::IRBuilder<>& irb);
 		void translateShrd(cs_insn* i, cs_x86* xi, llvm::IRBuilder<>& irb);
+		void translateShiftX(cs_insn* i, cs_x86* xi, llvm::IRBuilder<>& irb);
 		void translateStc(cs_insn* i, cs_x86* xi, llvm::IRBuilder<>& irb);
 		void translateStd(cs_insn* i, cs_x86* xi, llvm::IRBuilder<>& irb);
 		void translateStoreString(cs_insn* i, cs_x86* xi, llvm::IRBuilder<>& irb);
diff --git a/src/capstone2llvmir/x86/x86_init.cpp b/src/capstone2llvmir/x86/x86_init.cpp
index c5d167c48..dadb730af 100644
--- a/src/capstone2llvmir/x86/x86_init.cpp
+++ b/src/capstone2llvmir/x86/x86_init.cpp
@@ -1237,7 +1237,7 @@ Capstone2LlvmIrTranslatorX86_impl::_i2fm =
 		{X86_INS_SAL, &Capstone2LlvmIrTranslatorX86_impl::translateShiftLeft},
 		{X86_INS_SALC, &Capstone2LlvmIrTranslatorX86_impl::translateSalc},
 		{X86_INS_SAR, &Capstone2LlvmIrTranslatorX86_impl::translateShiftRight},
-		{X86_INS_SARX, nullptr},
+		{X86_INS_SARX, &Capstone2LlvmIrTranslatorX86_impl::translateShiftX},
 		{X86_INS_SBB, &Capstone2LlvmIrTranslatorX86_impl::translateSbb},
 		{X86_INS_SCASB, &Capstone2LlvmIrTranslatorX86_impl::translateScanString},
 		{X86_INS_SCASD, &Capstone2LlvmIrTranslatorX86_impl::translateScanString},
@@ -1270,10 +1270,10 @@ Capstone2LlvmIrTranslatorX86_impl::_i2fm =
 		{X86_INS_SHA256RNDS2, nullptr},
 		{X86_INS_SHL, &Capstone2LlvmIrTranslatorX86_impl::translateShiftLeft},
 		{X86_INS_SHLD, &Capstone2LlvmIrTranslatorX86_impl::translateShld},
-		{X86_INS_SHLX, nullptr},
+		{X86_INS_SHLX, &Capstone2LlvmIrTranslatorX86_impl::translateShiftX},
 		{X86_INS_SHR, &Capstone2LlvmIrTranslatorX86_impl::translateShiftRight},
 		{X86_INS_SHRD, &Capstone2LlvmIrTranslatorX86_impl::translateShrd},
-		{X86_INS_SHRX, nullptr},
+		{X86_INS_SHRX, &Capstone2LlvmIrTranslatorX86_impl::translateShiftX},
 		{X86_INS_SHUFPD, nullptr},
 		{X86_INS_SHUFPS, nullptr},
 		{X86_INS_SIDT, nullptr},