From 2fccad278882729b72758407ee608b0ede9a0f25 Mon Sep 17 00:00:00 2001
From: Antonio Frighetto <10052132+antoniofrighetto@users.noreply.github.com>
Date: Fri, 29 Mar 2024 09:49:56 +0100
Subject: [PATCH] Add support for nowrap flags to `trunc` (#1024)

---
 ir/instr.cpp                       | 31 ++++++++++++++++++++++++++++--
 ir/instr.h                         |  5 ++---
 llvm_util/llvm2alive.cpp           |  8 +++++++-
 tests/alive-tv/nsw.trunc.srctgt.ll |  9 +++++++++
 tests/alive-tv/nuw.trunc.srctgt.ll |  9 +++++++++
 5 files changed, 56 insertions(+), 6 deletions(-)
 create mode 100644 tests/alive-tv/nsw.trunc.srctgt.ll
 create mode 100644 tests/alive-tv/nuw.trunc.srctgt.ll

diff --git a/ir/instr.cpp b/ir/instr.cpp
index 6a3af0fea..796d5e7ef 100644
--- a/ir/instr.cpp
+++ b/ir/instr.cpp
@@ -1507,6 +1507,21 @@ unique_ptr<Instr> TestOp::dup(Function &f, const string &suffix) const {
   return make_unique<TestOp>(getType(), getName() + suffix, *lhs, *rhs, op);
 }
 
+ConversionOp::ConversionOp(Type &type, std::string &&name, Value &val, Op op,
+                           unsigned flags)
+    : Instr(type, std::move(name)), val(&val), op(op), flags(flags) {
+  switch (op) {
+  case ZExt:
+    assert((flags & NNEG) == flags);
+    break;
+  case Trunc:
+    assert((flags & (NSW | NUW)) == flags);
+    break;
+  default:
+    assert(flags == 0);
+    break;
+  }
+}
 
 vector<Value*> ConversionOp::operands() const {
   return { val };
@@ -1538,6 +1553,10 @@ void ConversionOp::print(ostream &os) const {
   os << getName() << " = " << str;
   if (flags & NNEG)
     os << "nneg ";
+  if (flags & NSW)
+    os << "nsw ";
+  if (flags & NUW)
+    os << "nuw ";
   os << *val << print_type(getType(), " to ", "");
 }
 
@@ -1558,8 +1577,16 @@ StateValue ConversionOp::toSMT(State &s) const {
     };
     break;
   case Trunc:
-    fn = [](auto &&val, auto &to_type) -> StateValue {
-      return {val.trunc(to_type.bits()), true};
+    fn = [this](auto &&val, auto &to_type) -> StateValue {
+      AndExpr non_poison;
+      unsigned orig_bits = val.bits();
+      unsigned trunc_bits = to_type.bits();
+      expr val_truncated = val.trunc(trunc_bits);
+      if (flags & NUW)
+        non_poison.add(val_truncated.zext(orig_bits - trunc_bits) == val);
+      if (flags & NSW)
+        non_poison.add(val_truncated.sext(orig_bits - trunc_bits) == val);
+      return {std::move(val_truncated), non_poison()};
     };
     break;
   case BitCast:
diff --git a/ir/instr.h b/ir/instr.h
index 15465fbd4..3927bb1ed 100644
--- a/ir/instr.h
+++ b/ir/instr.h
@@ -259,7 +259,7 @@ class TestOp final : public Instr {
 class ConversionOp final : public Instr {
 public:
   enum Op { SExt, ZExt, Trunc, BitCast, Ptr2Int, Int2Ptr };
-  enum Flags { None = 0, NNEG = 1 << 0 };
+  enum Flags { None = 0, NNEG = 1 << 0, NSW = 1 << 1, NUW = 1 << 2 };
 
 private:
   Value *val;
@@ -268,8 +268,7 @@ class ConversionOp final : public Instr {
 
 public:
   ConversionOp(Type &type, std::string &&name, Value &val, Op op,
-               unsigned flags = None)
-    : Instr(type, std::move(name)), val(&val), op(op), flags(flags) {}
+               unsigned flags = None);
 
   Op getOp() const { return op; }
   Value& getValue() const { return *val; }
diff --git a/llvm_util/llvm2alive.cpp b/llvm_util/llvm2alive.cpp
index 201f6e377..6ac92afc9 100644
--- a/llvm_util/llvm2alive.cpp
+++ b/llvm_util/llvm2alive.cpp
@@ -12,8 +12,9 @@
 #include "llvm/IR/GetElementPtrTypeIterator.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/InlineAsm.h"
-#include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/InstVisitor.h"
+#include "llvm/IR/InstrTypes.h"
+#include "llvm/IR/Instructions.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/Support/ModRef.h"
 #include <sstream>
@@ -308,6 +309,11 @@ class llvm2alive_ : public llvm::InstVisitor<llvm2alive_, unique_ptr<Instr>> {
         if (const auto *NNI = dyn_cast<llvm::PossiblyNonNegInst>(&i)) {
           if (NNI->hasNonNeg())
             flags |= ConversionOp::NNEG;
+        } else if (const auto *TI = dyn_cast<llvm::TruncInst>(&i)) {
+          if (TI->hasNoUnsignedWrap())
+            flags |= ConversionOp::NUW;
+          if (TI->hasNoSignedWrap())
+            flags |= ConversionOp::NSW;
         }
         RETURN_IDENTIFIER(
           make_unique<ConversionOp>(*ty, value_name(i), *val, op, flags));
diff --git a/tests/alive-tv/nsw.trunc.srctgt.ll b/tests/alive-tv/nsw.trunc.srctgt.ll
new file mode 100644
index 000000000..84b10fb5e
--- /dev/null
+++ b/tests/alive-tv/nsw.trunc.srctgt.ll
@@ -0,0 +1,9 @@
+define i8 @src(i8 %x) {
+  %trunc = trunc nsw i8 %x to i1
+  %sext = sext i1 %trunc to i8
+  ret i8 %sext
+}
+
+define i8 @tgt(i8 %x) {
+  ret i8 %x
+}
diff --git a/tests/alive-tv/nuw.trunc.srctgt.ll b/tests/alive-tv/nuw.trunc.srctgt.ll
new file mode 100644
index 000000000..1687d3528
--- /dev/null
+++ b/tests/alive-tv/nuw.trunc.srctgt.ll
@@ -0,0 +1,9 @@
+define i8 @src(i8 %x) {
+  %trunc = trunc nuw i8 %x to i1
+  %zext = zext i1 %trunc to i8
+  ret i8 %zext
+}
+
+define i8 @tgt(i8 %x) {
+  ret i8 %x
+}