From b22e86146a314a2aff0b4d6d089e8c63a845ab92 Mon Sep 17 00:00:00 2001
From: Manasij Mukherjee <manasij7479@gmail.com>
Date: Sat, 16 Nov 2019 01:57:07 -0700
Subject: [PATCH] Restrict input bits for analysis implementations in the
 presence of PC

---
 include/souper/Infer/Pruning.h |  1 +
 lib/Infer/Pruning.cpp          | 54 ++++++++++++++++++++++++++++------
 2 files changed, 46 insertions(+), 9 deletions(-)

diff --git a/include/souper/Infer/Pruning.h b/include/souper/Infer/Pruning.h
index 8d709b8f1..d121b1688 100644
--- a/include/souper/Infer/Pruning.h
+++ b/include/souper/Infer/Pruning.h
@@ -62,6 +62,7 @@ class PruningManager {
   std::vector<ValueCache> generateInputSets(std::vector<Inst *> &Inputs);
   // For the LHS contained in @SC, check if the given input in @Cache is valid.
   bool isInputValid(ValueCache &Cache);
+  std::unordered_map<Inst *, llvm::APInt> computeInputRestrictions();
   Inst *Ante;
 };
 
diff --git a/lib/Infer/Pruning.cpp b/lib/Infer/Pruning.cpp
index fe3bd202d..db99990d6 100644
--- a/lib/Infer/Pruning.cpp
+++ b/lib/Infer/Pruning.cpp
@@ -119,9 +119,35 @@ std::pair<llvm::APInt, llvm::APInt> knownBitsNarrowing
   return {KnownNotZero, KnownNotOne};
 }
 
+std::unordered_map<Inst *, llvm::APInt>
+PruningManager::computeInputRestrictions() {
+  std::unordered_map<Inst *, std::pair<llvm::APInt, llvm::APInt>> Seen;
+  // ^ SeenZero, SeenOne
+
+  for (auto &&V : InputVars) {
+    Seen[V].first = llvm::APInt(V->Width, 0);
+    Seen[V].second = llvm::APInt(V->Width, 0);
+  }
+
+  for (auto &&VC : InputVals) {
+    for (auto &&V : InputVars) {
+      if (VC[V].hasValue()) {
+        auto Val = VC[V].getValue();
+        Seen[V].first |= ~Val;
+        Seen[V].second |= Val;
+      }
+    }
+  }
+  std::unordered_map<Inst *, llvm::APInt> Result;
+  for (auto &&V : InputVars) {
+    Result[V] = ~(Seen[V].first & Seen[V].second);
+  }
+  return Result;
+}
+
 // TODO : Comment out debug stmts and conditions before benchmarking
 bool PruningManager::isInfeasible(souper::Inst *RHS,
-                                 unsigned StatsLevel) {
+                                  unsigned StatsLevel) {
   bool HasHole = !isConcrete(RHS, false, true);
   bool RHSIsConcrete = isConcrete(RHS);
 
@@ -133,15 +159,25 @@ bool PruningManager::isInfeasible(souper::Inst *RHS,
   getConstants(RHS, Constants);
 
   if (!Constants.empty()) {
-    auto RestrictedBits = RestrictedBitsAnalysis().findRestrictedBits(RHS);
-    if ((~RestrictedBits & (LHSKnownBitsNoSpec.Zero | LHSKnownBitsNoSpec.One)) != 0) {
-//     if (RestrictedBits == 0 && (LHSKB.Zero != 0 || LHSKB.One != 0)) {
-      if (StatsLevel > 2) {
-        llvm::errs() << "  pruned using restricted bits analysis.\n";
-        llvm::errs() << "  LHSKB : " << KnownBitsAnalysis::knownBitsString(LHSKnownBitsNoSpec) << "\n";
-        llvm::errs() << "  RB    : " << RestrictedBits.toString(2, false) << "\n";
+    auto RBA = RestrictedBitsAnalysis();
+    if (!SC.PCs.empty()) {
+      // `Unrestricts` bits for which PC validated inputs
+      // have both zero and one values. Default is all `restricted`
+      // Instead of treating inputs as completely unrestricted,
+      // this imposes restrictions on inputs if there are no valid
+      // inputs to demonstrate otherwise.
+      RBA.RBCache = computeInputRestrictions();
+    }
+    if (SC.PCs.empty() || !InputVals.empty()) {
+      auto RestrictedBits = RBA.findRestrictedBits(RHS);
+      if ((~RestrictedBits & (LHSKnownBitsNoSpec.Zero | LHSKnownBitsNoSpec.One)) != 0) {
+        if (StatsLevel > 2) {
+          llvm::errs() << "  pruned using restricted bits analysis.\n";
+          llvm::errs() << "  LHSKB : " << KnownBitsAnalysis::knownBitsString(LHSKnownBitsNoSpec) << "\n";
+          llvm::errs() << "  RB    : " << RestrictedBits.toString(2, false) << "\n";
+        }
+        return true;
       }
-      return true;
     }
 
 //     auto LHSCR = ConstantRangeAnalysis().findConstantRange(SC.LHS, BlankCI, false);