[ARM] Allow zext in ARMCodeGenPrepare

Treat zext instructions as roots, like we do for truncs.

Differential Revision: https://reviews.llvm.org/D50759

llvm-svn: 339868
diff --git a/llvm/lib/Target/ARM/ARMCodeGenPrepare.cpp b/llvm/lib/Target/ARM/ARMCodeGenPrepare.cpp
index 164b255..296a22f 100644
--- a/llvm/lib/Target/ARM/ARMCodeGenPrepare.cpp
+++ b/llvm/lib/Target/ARM/ARMCodeGenPrepare.cpp
@@ -181,6 +181,8 @@
     return UsesNarrowValue(Return->getReturnValue());
   if (auto *Trunc = dyn_cast<TruncInst>(V))
     return UsesNarrowValue(Trunc->getOperand(0));
+  if (auto *ZExt = dyn_cast<ZExtInst>(V))
+    return UsesNarrowValue(ZExt->getOperand(0));
   if (auto *ICmp = dyn_cast<ICmpInst>(V))
     return ICmp->isSigned();
 
@@ -422,7 +424,8 @@
     if (!isa<Instruction>(V) || !isa<IntegerType>(V->getType()))
       return nullptr;
 
-    if ((!Promoted.count(V) && !NewInsts.count(V)) || !TruncTysMap.count(V))
+    if ((!Promoted.count(V) && !NewInsts.count(V)) || !TruncTysMap.count(V) ||
+        Leaves.count(V))
       return nullptr;
 
     Type *TruncTy = TruncTysMap[V];
@@ -463,7 +466,7 @@
       }
     }
   }
-  LLVM_DEBUG(dbgs() << "ARM CGP: Mutation complete.\n");
+  LLVM_DEBUG(dbgs() << "ARM CGP: Mutation complete:\n");
 }
 
 /// We accept most instructions, as well as Arguments and ConstantInsts. We
@@ -492,10 +495,12 @@
       isa<LoadInst>(V))
     return isSupportedType(V);
 
-  // Currently, Trunc is the only cast we support.
   if (auto *Trunc = dyn_cast<TruncInst>(V))
     return isSupportedType(Trunc->getOperand(0));
 
+  if (auto *ZExt = dyn_cast<ZExtInst>(V))
+    return isSupportedType(ZExt->getOperand(0));
+
   // Special cases for calls as we need to check for zeroext
   // TODO We should accept calls even if they don't have zeroext, as they can
   // still be roots.