PNaCl: Extend ExpandMulWithOverflow pass to handle uadd.with.overflow too

lib/Transforms/NaCl/ExpandArithWithOverflow.cpp

-//===- ExpandMulWithOverflow.cpp - Expand out usage of umul.with.overflow--===//
+//===- ExpandArithWithOverflow.cpp - Expand out uses of *.with.overflow----===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // The llvm.*.with.overflow.*() intrinsics are awkward for PNaCl
 // support because they return structs, and we want to omit struct
 // types from IR in PNaCl's stable ABI.
 //
-// However, llvm.umul.with.overflow.*() is used by Clang to implement
-// an overflow check for C++'s new[] operator.  This pass expands out
-// these uses so that PNaCl does not have to support
-// umul.with.overflow as part of PNaCl's stable ABI.
+// However, llvm.{umul,uadd}.with.overflow.*() are used by Clang to
+// implement an overflow check for C++'s new[] operator.  This pass
+// expands out these uses so that PNaCl does not have to support
+// *.with.overflow as part of PNaCl's stable ABI.
 //
-// This pass only handles multiplication by a constant, which is the
-// only case of umul.with.overflow that is currently generated by
+// This pass only handles adding/multiplying by a constant, which is
+// the only use of *.with.overflow that is currently generated by
 // Clang (unless '-ftrapv' is passed to Clang).
 //
 // X * Const overflows iff X > UINT_MAX / Const, where UINT_MAX is the
-// maximum value for the integer type being multiplied.
+// maximum value for the integer type being used.
+//
+// Similarly, X + Const overflows iff X > UINT_MAX - Const.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/NaCl.h"
 
 using namespace llvm;
 
 namespace {
   // This is a ModulePass so that the pass can easily iterate over all
   // uses of the intrinsics.
-  class ExpandMulWithOverflow : public ModulePass {
+  class ExpandArithWithOverflow : public ModulePass {
   public:
     static char ID; // Pass identification, replacement for typeid
-    ExpandMulWithOverflow() : ModulePass(ID) {
-      initializeExpandMulWithOverflowPass(*PassRegistry::getPassRegistry());
+    ExpandArithWithOverflow() : ModulePass(ID) {
+      initializeExpandArithWithOverflowPass(*PassRegistry::getPassRegistry());
     }
 
     virtual bool runOnModule(Module &M);
   };
 }
 
-char ExpandMulWithOverflow::ID = 0;
-INITIALIZE_PASS(ExpandMulWithOverflow, "expand-mul-with-overflow",
-                "Expand out uses of llvm.umul.with.overflow intrinsics",
+char ExpandArithWithOverflow::ID = 0;
+INITIALIZE_PASS(ExpandArithWithOverflow, "expand-arith-with-overflow",
+                "Expand out some uses of *.with.overflow intrinsics",
                 false, false)
 
 static uint64_t UintTypeMax(unsigned Bits) {
   // Avoid doing 1 << 64 because that is undefined on a uint64_t.
   if (Bits == 64)
     return ~(uint64_t) 0;
   return (((uint64_t) 1) << Bits) - 1;
 }
 
-static bool ExpandForIntSize(Module *M, unsigned Bits) {
+static bool ExpandOpForIntSize(Module *M, unsigned Bits, bool Mul) {
   IntegerType *IntTy = IntegerType::get(M->getContext(), Bits);
   SmallVector<Type *, 1> Types;
   Types.push_back(IntTy);
-  std::string Name = Intrinsic::getName(Intrinsic::umul_with_overflow, Types);
+  Intrinsic::ID ID = (Mul ? Intrinsic::umul_with_overflow
+                          : Intrinsic::uadd_with_overflow);
+  std::string Name = Intrinsic::getName(ID, Types);
   Function *Intrinsic = M->getFunction(Name);
   if (!Intrinsic)
     return false;
   for (Value::use_iterator CallIter = Intrinsic->use_begin(),
          E = Intrinsic->use_end(); CallIter != E; ) {
     CallInst *Call = dyn_cast<CallInst>(*CallIter++);
     if (!Call) {
-      report_fatal_error("ExpandMulWithOverflow: Taking the address of a "
-                         "umul.with.overflow intrinsic is not allowed");
+      report_fatal_error("ExpandArithWithOverflow: Taking the address of a "
+                         "*.with.overflow intrinsic is not allowed");
     }
     Value *VariableArg;
     ConstantInt *ConstantArg;
     if (ConstantInt *C = dyn_cast<ConstantInt>(Call->getArgOperand(0))) {
       VariableArg = Call->getArgOperand(1);
       ConstantArg = C;
     } else if (ConstantInt *C = dyn_cast<ConstantInt>(Call->getArgOperand(1))) {
       VariableArg = Call->getArgOperand(0);
       ConstantArg = C;
     } else {
       errs() << "Use: " << *Call << "\n";
-      report_fatal_error("ExpandMulWithOverflow: At least one argument of "
-                         "umul.with.overflow must be a constant");
+      report_fatal_error("ExpandArithWithOverflow: At least one argument of "
+                         "*.with.overflow must be a constant");
     }
 
-    Value *Mul = BinaryOperator::Create(
-        Instruction::Mul, VariableArg, ConstantArg,
-        Call->getName() + ".mul", Call);
+    Value *ArithResult = BinaryOperator::Create(
+        (Mul ? Instruction::Mul : Instruction::Add), VariableArg, ConstantArg,
+        Call->getName() + ".arith", Call);
 
-    uint64_t ArgMax = UintTypeMax(Bits) / ConstantArg->getZExtValue();
+    uint64_t ArgMax;
+    if (Mul) {
+      ArgMax = UintTypeMax(Bits) / ConstantArg->getZExtValue();
+    } else {
+      ArgMax = UintTypeMax(Bits) - ConstantArg->getZExtValue();
+    }
     Value *Overflow = new ICmpInst(
         Call, CmpInst::ICMP_UGT, VariableArg, ConstantInt::get(IntTy, ArgMax),
         Call->getName() + ".overflow");
 
     for (Value::use_iterator FieldIter = Call->use_begin(),
            E = Call->use_end(); FieldIter != E; ) {
       User *U = *FieldIter++;
       ExtractValueInst *Field = dyn_cast<ExtractValueInst>(U);
       if (!Field) {
         errs() << "Use: " << *U << "\n";
         report_fatal_error(
-            "ExpandMulWithOverflow: Use is not an extractvalue");
+            "ExpandArithWithOverflow: Use is not an extractvalue");
       }
       if (Field->getNumIndices() != 1) {
-        report_fatal_error("ExpandMulWithOverflow: Unexpected indices");
+        report_fatal_error("ExpandArithWithOverflow: Unexpected indices");
       }
       unsigned Index = Field->getIndices()[0];
       if (Index == 0) {
-        Field->replaceAllUsesWith(Mul);
+        Field->replaceAllUsesWith(ArithResult);
       } else if (Index == 1) {
         Field->replaceAllUsesWith(Overflow);
       } else {
-        report_fatal_error("ExpandMulWithOverflow: Unexpected index");
+        report_fatal_error("ExpandArithWithOverflow: Unexpected index");
       }
       Field->eraseFromParent();
     }
     Call->eraseFromParent();
   }
   Intrinsic->eraseFromParent();
   return true;
 }
 
-bool ExpandMulWithOverflow::runOnModule(Module &M) {
+static bool ExpandForIntSize(Module *M, unsigned Bits) {
+  bool Modified = false;
+  Modified |= ExpandOpForIntSize(M, Bits, true); // Expand umul
+  Modified |= ExpandOpForIntSize(M, Bits, false); // Expand uadd
+  return Modified;
+}
+
+bool ExpandArithWithOverflow::runOnModule(Module &M) {
   bool Modified = false;
   Modified |= ExpandForIntSize(&M, 64);
   Modified |= ExpandForIntSize(&M, 32);
   Modified |= ExpandForIntSize(&M, 16);
   Modified |= ExpandForIntSize(&M, 8);
   return Modified;
 }
 
-ModulePass *llvm::createExpandMulWithOverflowPass() {
-  return new ExpandMulWithOverflow();
+ModulePass *llvm::createExpandArithWithOverflowPass() {
+  return new ExpandArithWithOverflow();
 }