[Attributor] Make IntegerState more flexible

To make IntegerState more flexible but also less error prone we split it
up into (1) incrementing, (2) decrementing, and (3) bit-tracking states.
This adds functionality compared to before and disallows misuse, e.g.,
"incrementing" updates on a bit-tracking state.

Part of the change is a single operator in the base class which
simplifies helper functions that deal with states.

There are certain functional changes but all of which should actually be
corrections.

2 files changed, 192 insertions, 98 deletions

diff --git a/llvm/include/llvm/Transforms/IPO/Attributor.h b/llvm/include/llvm/Transforms/IPO/Attributor.h
index 3dbe0fcd76e..2451ae1a2d2 100644
--- a/llvm/include/llvm/Transforms/IPO/Attributor.h
+++ b/llvm/include/llvm/Transforms/IPO/Attributor.h
@@ -72,7 +72,8 @@
 // - Define a class (transitively) inheriting from AbstractAttribute and one
 //   (which could be the same) that (transitively) inherits from AbstractState.
 //   For the latter, consider the already available BooleanState and
-//   IntegerState if they fit your needs, e.g., you require only a bit-encoding.
+//   {Inc,Dec,Bit}IntegerState if they fit your needs, e.g., you require only a
+//   number tracking or bit-encoding.
 // - Implement all pure methods. Also use overloading if the attribute is not
 //   conforming with the "default" behavior: A (set of) LLVM-IR attribute(s) for
 //   an argument, call site argument, function return value, or function. See
@@ -1027,9 +1028,10 @@ private:
 ///
 /// All methods need to be implemented by the subclass. For the common use case,
 /// a single boolean state or a bit-encoded state, the BooleanState and
-/// IntegerState classes are already provided. An abstract attribute can inherit
-/// from them to get the abstract state interface and additional methods to
-/// directly modify the state based if needed. See the class comments for help.
+/// {Inc,Dec,Bit}IntegerState classes are already provided. An abstract
+/// attribute can inherit from them to get the abstract state interface and
+/// additional methods to directly modify the state based if needed. See the
+/// class comments for help.
 struct AbstractState {
   virtual ~AbstractState() {}
 
@@ -1068,15 +1070,15 @@ struct AbstractState {
 /// force/inidicate a fixpoint. If an optimistic one is indicated, the known
 /// state will catch up with the assumed one, for a pessimistic fixpoint it is
 /// the other way around.
-struct IntegerState : public AbstractState {
-  /// Underlying integer type, we assume 32 bits to be enough.
-  using base_t = uint32_t;
+template <typename base_ty, base_ty BestState, base_ty WorstState>
+struct IntegerStateBase : public AbstractState {
+  using base_t = base_ty;
 
-  /// Initialize the (best) state.
-  IntegerState(base_t BestState = ~0) : Assumed(BestState) {}
+  /// Return the best possible representable state.
+  static constexpr base_t getBestState() { return BestState; }
 
   /// Return the worst possible representable state.
-  static constexpr base_t getWorstState() { return 0; }
+  static constexpr base_t getWorstState() { return WorstState; }
 
   /// See AbstractState::isValidState()
   /// NOTE: For now we simply pretend that the worst possible state is invalid.
@@ -1103,6 +1105,58 @@ struct IntegerState : public AbstractState {
   /// Return the assumed state encoding.
   base_t getAssumed() const { return Assumed; }
 
+  /// Equality for IntegerStateBase.
+  bool
+  operator==(const IntegerStateBase<base_t, BestState, WorstState> &R) const {
+    return this->getAssumed() == R.getAssumed() &&
+           this->getKnown() == R.getKnown();
+  }
+
+  /// Inequality for IntegerStateBase.
+  bool
+  operator!=(const IntegerStateBase<base_t, BestState, WorstState> &R) const {
+    return !(*this == R);
+  }
+
+  /// "Clamp" this state with \p R. The result is subtype dependent but it is
+  /// intended that only information assumed in both states will be assumed in
+  /// this one afterwards.
+  void operator^=(const IntegerStateBase<base_t, BestState, WorstState> &R) {
+    handleNewAssumedValue(R.getAssumed());
+  }
+
+  void operator|=(const IntegerStateBase<base_t, BestState, WorstState> &R) {
+    joinOR(R.getAssumed(), R.getKnown());
+  }
+
+  void operator&=(const IntegerStateBase<base_t, BestState, WorstState> &R) {
+    joinAND(R.getAssumed(), R.getKnown());
+  }
+
+protected:
+  /// Handle a new assumed value \p Value. Subtype dependent.
+  virtual void handleNewAssumedValue(base_t Value) = 0;
+
+  /// Handle a new known value \p Value. Subtype dependent.
+  virtual void handleNewKnownValue(base_t Value) = 0;
+
+  /// Handle a  value \p Value. Subtype dependent.
+  virtual void joinOR(base_t AssumedValue, base_t KnownValue) = 0;
+
+  /// Handle a new assumed value \p Value. Subtype dependent.
+  virtual void joinAND(base_t AssumedValue, base_t KnownValue) = 0;
+
+  /// The known state encoding in an integer of type base_t.
+  base_t Known = getWorstState();
+
+  /// The assumed state encoding in an integer of type base_t.
+  base_t Assumed = getBestState();
+};
+
+/// Specialization of the integer state for a bit-wise encoding.
+struct BitIntegerState : public IntegerStateBase<uint32_t, ~0u, 0> {
+  using base_t = IntegerStateBase::base_t;
+
   /// Return true if the bits set in \p BitsEncoding are "known bits".
   bool isKnown(base_t BitsEncoding) const {
     return (Known & BitsEncoding) == BitsEncoding;
@@ -1114,7 +1168,7 @@ struct IntegerState : public AbstractState {
   }
 
   /// Add the bits in \p BitsEncoding to the "known bits".
-  IntegerState &addKnownBits(base_t Bits) {
+  BitIntegerState &addKnownBits(base_t Bits) {
     // Make sure we never miss any "known bits".
     Assumed |= Bits;
     Known |= Bits;
@@ -1122,92 +1176,145 @@ struct IntegerState : public AbstractState {
   }
 
   /// Remove the bits in \p BitsEncoding from the "assumed bits" if not known.
-  IntegerState &removeAssumedBits(base_t BitsEncoding) {
-    // Make sure we never loose any "known bits".
-    Assumed = (Assumed & ~BitsEncoding) | Known;
-    return *this;
+  BitIntegerState &removeAssumedBits(base_t BitsEncoding) {
+    return intersectAssumedBits(~BitsEncoding);
   }
 
   /// Remove the bits in \p BitsEncoding from the "known bits".
-  IntegerState &removeKnownBits(base_t BitsEncoding) {
+  BitIntegerState &removeKnownBits(base_t BitsEncoding) {
     Known = (Known & ~BitsEncoding);
     return *this;
   }
 
   /// Keep only "assumed bits" also set in \p BitsEncoding but all known ones.
-  IntegerState &intersectAssumedBits(base_t BitsEncoding) {
+  BitIntegerState &intersectAssumedBits(base_t BitsEncoding) {
     // Make sure we never loose any "known bits".
     Assumed = (Assumed & BitsEncoding) | Known;
     return *this;
   }
 
+private:
+  void handleNewAssumedValue(base_t Value) override {
+    intersectAssumedBits(Value);
+  }
+  void handleNewKnownValue(base_t Value) override { addKnownBits(Value); }
+  void joinOR(base_t AssumedValue, base_t KnownValue) override {
+    Known |= KnownValue;
+    Assumed |= AssumedValue;
+  }
+  void joinAND(base_t AssumedValue, base_t KnownValue) override {
+    Known &= KnownValue;
+    Assumed &= AssumedValue;
+  }
+};
+
+/// Specialization of the integer state for an increasing value, hence ~0u is
+/// the best state and 0 the worst.
+struct IncIntegerState : public IntegerStateBase<uint32_t, ~0u, 0> {
+  using base_t = IntegerStateBase::base_t;
+
   /// Take minimum of assumed and \p Value.
-  IntegerState &takeAssumedMinimum(base_t Value) {
+  IncIntegerState &takeAssumedMinimum(base_t Value) {
     // Make sure we never loose "known value".
     Assumed = std::max(std::min(Assumed, Value), Known);
     return *this;
   }
 
   /// Take maximum of known and \p Value.
-  IntegerState &takeKnownMaximum(base_t Value) {
+  IncIntegerState &takeKnownMaximum(base_t Value) {
     // Make sure we never loose "known value".
     Assumed = std::max(Value, Assumed);
     Known = std::max(Value, Known);
     return *this;
   }
 
-  /// Equality for IntegerState.
-  bool operator==(const IntegerState &R) const {
-    return this->getAssumed() == R.getAssumed() &&
-           this->getKnown() == R.getKnown();
+private:
+  void handleNewAssumedValue(base_t Value) override {
+    takeAssumedMinimum(Value);
+  }
+  void handleNewKnownValue(base_t Value) override { takeKnownMaximum(Value); }
+  void joinOR(base_t AssumedValue, base_t KnownValue) override {
+    Known = std::max(Known, KnownValue);
+    Assumed = std::max(Assumed, AssumedValue);
   }
+  void joinAND(base_t AssumedValue, base_t KnownValue) override {
+    Known = std::min(Known, KnownValue);
+    Assumed = std::min(Assumed, AssumedValue);
+  }
+};
 
-  /// Inequality for IntegerState.
-  bool operator!=(const IntegerState &R) const { return !(*this == R); }
+/// Specialization of the integer state for a decreasing value, hence 0 is the
+/// best state and ~0u the worst.
+template <typename base_ty = uint32_t>
+struct DecIntegerState : public IntegerStateBase<base_ty, 0, ~base_ty(0)> {
+  using base_t = base_ty;
 
-  /// "Clamp" this state with \p R. The result is the minimum of the assumed
-  /// information but not less than what was known before.
-  ///
-  /// TODO: Consider replacing the operator with a call or using it only when
-  ///       we can also take the maximum of the known information, thus when
-  ///       \p R is not dependent on additional assumed state.
-  IntegerState operator^=(const IntegerState &R) {
-    takeAssumedMinimum(R.Assumed);
+  /// Take maximum of assumed and \p Value.
+  DecIntegerState &takeAssumedMaximum(base_t Value) {
+    // Make sure we never loose "known value".
+    this->Assumed = std::min(std::max(this->Assumed, Value), this->Known);
     return *this;
   }
 
-  /// "Clamp" this state with \p R. The result is the maximum of the known
-  /// information but not more than what was assumed before.
-  IntegerState operator+=(const IntegerState &R) {
-    takeKnownMaximum(R.Known);
+  /// Take minimum of known and \p Value.
+  DecIntegerState &takeKnownMinimum(base_t Value) {
+    // Make sure we never loose "known value".
+    this->Assumed = std::min(Value, this->Assumed);
+    this->Known = std::min(Value, this->Known);
     return *this;
   }
 
-  /// Make this the minimum, known and assumed, of this state and \p R.
-  IntegerState operator&=(const IntegerState &R) {
-    Known = std::min(Known, R.Known);
-    Assumed = std::min(Assumed, R.Assumed);
-    return *this;
+private:
+  void handleNewAssumedValue(base_t Value) override {
+    takeAssumedMaximum(Value);
+  }
+  void handleNewKnownValue(base_t Value) override { takeKnownMinimum(Value); }
+  void joinOR(base_t AssumedValue, base_t KnownValue) override {
+    this->Assumed = std::min(this->Assumed, KnownValue);
+    this->Assumed = std::min(this->Assumed, AssumedValue);
   }
+  void joinAND(base_t AssumedValue, base_t KnownValue) override {
+    this->Assumed = std::max(this->Assumed, KnownValue);
+    this->Assumed = std::max(this->Assumed, AssumedValue);
+  }
+};
 
-  /// Make this the maximum, known and assumed, of this state and \p R.
-  IntegerState operator|=(const IntegerState &R) {
-    Known = std::max(Known, R.Known);
-    Assumed = std::max(Assumed, R.Assumed);
-    return *this;
+/// Simple wrapper for a single bit (boolean) state.
+struct BooleanState : public IntegerStateBase<bool, 1, 0> {
+  using base_t = IntegerStateBase::base_t;
+
+  /// Set the assumed value to \p Value but never below the known one.
+  void setAssumed(bool Value) { Assumed &= (Known | Value); }
+
+  /// Set the known and asssumed value to \p Value.
+  void setKnown(bool Value) {
+    Known |= Value;
+    Assumed |= Value;
   }
 
-private:
-  /// The known state encoding in an integer of type base_t.
-  base_t Known = getWorstState();
+  /// Return true if the state is assumed to hold.
+  bool isAssumed() const { return getAssumed(); }
 
-  /// The assumed state encoding in an integer of type base_t.
-  base_t Assumed;
-};
+  /// Return true if the state is known to hold.
+  bool isKnown() const { return getKnown(); }
 
-/// Simple wrapper for a single bit (boolean) state.
-struct BooleanState : public IntegerState {
-  BooleanState() : IntegerState(1){};
+private:
+  void handleNewAssumedValue(base_t Value) override {
+    if (!Value)
+      Assumed = Known;
+  }
+  void handleNewKnownValue(base_t Value) override {
+    if (Value)
+      Known = (Assumed = Value);
+  }
+  void joinOR(base_t AssumedValue, base_t KnownValue) override {
+    Known |= KnownValue;
+    Assumed |= AssumedValue;
+  }
+  void joinAND(base_t AssumedValue, base_t KnownValue) override {
+    Known &= KnownValue;
+    Assumed &= AssumedValue;
+  }
 };
 
 /// Helper struct necessary as the modular build fails if the virtual method
@@ -1404,7 +1511,10 @@ raw_ostream &operator<<(raw_ostream &OS, ChangeStatus S);
 raw_ostream &operator<<(raw_ostream &OS, IRPosition::Kind);
 raw_ostream &operator<<(raw_ostream &OS, const IRPosition &);
 raw_ostream &operator<<(raw_ostream &OS, const AbstractState &State);
-raw_ostream &operator<<(raw_ostream &OS, const IntegerState &S);
+template <typename base_ty, base_ty BestState, base_ty WorstState>
+raw_ostream &
+operator<<(raw_ostream &OS,
+           const IntegerStateBase<base_ty, BestState, WorstState> &State);
 ///}
 
 struct AttributorPass : public PassInfoMixin<AttributorPass> {
@@ -1656,7 +1766,7 @@ struct AAIsDead : public StateWrapper<BooleanState, AbstractAttribute>,
 struct DerefState : AbstractState {
 
   /// State representing for dereferenceable bytes.
-  IntegerState DerefBytesState;
+  IncIntegerState DerefBytesState;
 
   /// State representing that whether the value is globaly dereferenceable.
   BooleanState GlobalState;
@@ -1700,31 +1810,24 @@ struct DerefState : AbstractState {
            this->GlobalState == R.GlobalState;
   }
 
-  /// Inequality for IntegerState.
+  /// Inequality for DerefState.
   bool operator!=(const DerefState &R) { return !(*this == R); }
 
-  /// See IntegerState::operator^=
+  /// See IntegerStateBase::operator^=
   DerefState operator^=(const DerefState &R) {
     DerefBytesState ^= R.DerefBytesState;
     GlobalState ^= R.GlobalState;
     return *this;
   }
 
-  /// See IntegerState::operator+=
-  DerefState operator+=(const DerefState &R) {
-    DerefBytesState += R.DerefBytesState;
-    GlobalState += R.GlobalState;
-    return *this;
-  }
-
-  /// See IntegerState::operator&=
+  /// See IntegerStateBase::operator&=
   DerefState operator&=(const DerefState &R) {
     DerefBytesState &= R.DerefBytesState;
     GlobalState &= R.GlobalState;
     return *this;
   }
 
-  /// See IntegerState::operator|=
+  /// See IntegerStateBase::operator|=
   DerefState operator|=(const DerefState &R) {
     DerefBytesState |= R.DerefBytesState;
     GlobalState |= R.GlobalState;
@@ -1780,7 +1883,7 @@ struct AADereferenceable
 /// An abstract interface for all align attributes.
 struct AAAlign
     : public IRAttribute<Attribute::Alignment,
-                         StateWrapper<IntegerState, AbstractAttribute>> {
+                         StateWrapper<IncIntegerState, AbstractAttribute>> {
   AAAlign(const IRPosition &IRP) : IRAttribute(IRP) {}
 
   /// Return assumed alignment.
@@ -1799,7 +1902,7 @@ struct AAAlign
 /// An abstract interface for all nocapture attributes.
 struct AANoCapture
     : public IRAttribute<Attribute::NoCapture,
-                         StateWrapper<IntegerState, AbstractAttribute>> {
+                         StateWrapper<BitIntegerState, AbstractAttribute>> {
   AANoCapture(const IRPosition &IRP) : IRAttribute(IRP) {}
 
   /// State encoding bits. A set bit in the state means the property holds.
@@ -1898,7 +2001,7 @@ struct AAHeapToStack : public StateWrapper<BooleanState, AbstractAttribute>,
 /// An abstract interface for all memory related attributes.
 struct AAMemoryBehavior
     : public IRAttribute<Attribute::ReadNone,
-                         StateWrapper<IntegerState, AbstractAttribute>> {
+                         StateWrapper<BitIntegerState, AbstractAttribute>> {
   AAMemoryBehavior(const IRPosition &IRP) : IRAttribute(IRP) {}
 
   /// State encoding bits. A set bit in the state means the property holds.
diff --git a/llvm/lib/Transforms/IPO/Attributor.cpp b/llvm/lib/Transforms/IPO/Attributor.cpp
index 95f47345d8f..a811471d37e 100644
--- a/llvm/lib/Transforms/IPO/Attributor.cpp
+++ b/llvm/lib/Transforms/IPO/Attributor.cpp
@@ -517,30 +517,17 @@ void IRPosition::verify() {
 }
 
 namespace {
-/// Helper functions to clamp a state \p S of type \p StateType with the
+/// Helper function to clamp a state \p S of type \p StateType with the
 /// information in \p R and indicate/return if \p S did change (as-in update is
 /// required to be run again).
-///
-///{
 template <typename StateType>
-ChangeStatus clampStateAndIndicateChange(StateType &S, const StateType &R);
-
-template <>
-ChangeStatus clampStateAndIndicateChange<IntegerState>(IntegerState &S,
-                                                       const IntegerState &R) {
+ChangeStatus clampStateAndIndicateChange(StateType &S, const StateType &R) {
   auto Assumed = S.getAssumed();
   S ^= R;
   return Assumed == S.getAssumed() ? ChangeStatus::UNCHANGED
                                    : ChangeStatus::CHANGED;
 }
 
-template <>
-ChangeStatus clampStateAndIndicateChange<BooleanState>(BooleanState &S,
-                                                       const BooleanState &R) {
-  return clampStateAndIndicateChange<IntegerState>(S, R);
-}
-///}
-
 /// Clamp the information known for all returned values of a function
 /// (identified by \p QueryingAA) into \p S.
 template <typename AAType, typename StateType = typename AAType::StateType>
@@ -1609,7 +1596,7 @@ struct AANonNullImpl : AANonNull {
     bool TrackUse = false;
     getKnownNonNullAndDerefBytesForUse(A, *this, getAssociatedValue(), U, I,
                                        IsNonNull, TrackUse);
-    takeKnownMaximum(IsNonNull);
+    setKnown(IsNonNull);
     return TrackUse;
   }
 
@@ -1661,7 +1648,7 @@ struct AANonNullFloating
 
     const DataLayout &DL = A.getDataLayout();
 
-    auto VisitValueCB = [&](Value &V, AAAlign::StateType &T,
+    auto VisitValueCB = [&](Value &V, AANonNull::StateType &T,
                             bool Stripped) -> bool {
       const auto &AA = A.getAAFor<AANonNull>(*this, IRPosition::value(V));
       if (!Stripped && this == &AA) {
@@ -2463,10 +2450,10 @@ struct AAIsDeadCallSite final : AAIsDeadImpl {
 template <>
 ChangeStatus clampStateAndIndicateChange<DerefState>(DerefState &S,
                                                      const DerefState &R) {
-  ChangeStatus CS0 = clampStateAndIndicateChange<IntegerState>(
+  ChangeStatus CS0 = clampStateAndIndicateChange<IncIntegerState>(
       S.DerefBytesState, R.DerefBytesState);
   ChangeStatus CS1 =
-      clampStateAndIndicateChange<IntegerState>(S.GlobalState, R.GlobalState);
+      clampStateAndIndicateChange<BooleanState>(S.GlobalState, R.GlobalState);
   return CS0 | CS1;
 }
 
@@ -2967,7 +2954,7 @@ struct AANoCaptureImpl : public AANoCapture {
   /// state in memory and through "returning/throwing", respectively.
   static void determineFunctionCaptureCapabilities(const IRPosition &IRP,
                                                    const Function &F,
-                                                   IntegerState &State) {
+                                                   BitIntegerState &State) {
     // TODO: Once we have memory behavior attributes we should use them here.
 
     // If we know we cannot communicate or write to memory, we do not care about
@@ -3036,7 +3023,7 @@ struct AACaptureUseTracker final : public CaptureTracker {
   /// the search is stopped with \p CapturedInMemory and \p CapturedInInteger
   /// conservatively set to true.
   AACaptureUseTracker(Attributor &A, AANoCapture &NoCaptureAA,
-                      const AAIsDead &IsDeadAA, IntegerState &State,
+                      const AAIsDead &IsDeadAA, BitIntegerState &State,
                       SmallVectorImpl<const Value *> &PotentialCopies,
                       unsigned &RemainingUsesToExplore)
       : A(A), NoCaptureAA(NoCaptureAA), IsDeadAA(IsDeadAA), State(State),
@@ -3155,7 +3142,7 @@ private:
   const AAIsDead &IsDeadAA;
 
   /// The state currently updated.
-  IntegerState &State;
+  BitIntegerState &State;
 
   /// Set of potential copies of the tracked value.
   SmallVectorImpl<const Value *> &PotentialCopies;
@@ -3238,7 +3225,7 @@ ChangeStatus AANoCaptureImpl::updateImpl(Attributor &A) {
   while (T.isAssumed(NO_CAPTURE_MAYBE_RETURNED) && Idx < PotentialCopies.size())
     Tracker.valueMayBeCaptured(PotentialCopies[Idx++]);
 
-  AAAlign::StateType &S = getState();
+  AANoCapture::StateType &S = getState();
   auto Assumed = S.getAssumed();
   S.intersectAssumedBits(T.getAssumed());
   return Assumed == S.getAssumed() ? ChangeStatus::UNCHANGED
@@ -3787,7 +3774,7 @@ struct AAMemoryBehaviorImpl : public AAMemoryBehavior {
 
   /// Return the memory behavior information encoded in the IR for \p IRP.
   static void getKnownStateFromValue(const IRPosition &IRP,
-                                     IntegerState &State) {
+                                     BitIntegerState &State) {
     SmallVector<Attribute, 2> Attrs;
     IRP.getAttrs(AttrKinds, Attrs);
     for (const Attribute &Attr : Attrs) {
@@ -4036,7 +4023,8 @@ struct AAMemoryBehaviorCallSite final : AAMemoryBehaviorImpl {
     const IRPosition &FnPos = IRPosition::function(*F);
     auto &FnAA = A.getAAFor<AAMemoryBehavior>(*this, FnPos);
     return clampStateAndIndicateChange(
-        getState(), static_cast<const AAAlign::StateType &>(FnAA.getState()));
+        getState(),
+        static_cast<const AAMemoryBehavior::StateType &>(FnAA.getState()));
   }
 
   /// See AbstractAttribute::trackStatistics()
@@ -4903,7 +4891,10 @@ raw_ostream &llvm::operator<<(raw_ostream &OS, const IRPosition &Pos) {
             << Pos.getAnchorValue().getName() << "@" << Pos.getArgNo() << "]}";
 }
 
-raw_ostream &llvm::operator<<(raw_ostream &OS, const IntegerState &S) {
+template <typename base_ty, base_ty BestState, base_ty WorstState>
+raw_ostream &llvm::
+operator<<(raw_ostream &OS,
+           const IntegerStateBase<base_ty, BestState, WorstState> &S) {
   return OS << "(" << S.getKnown() << "-" << S.getAssumed() << ")"
             << static_cast<const AbstractState &>(S);
 }