[mlir][sparse] Cleaning up OOB implementation details for `VarSet`

[aartbik]

Looks good, but double check the clang-format presubmit error

[wrengr]

git-clang-format

mlir/lib/Dialect/SparseTensor/IR/Detail/Var.h

 
 //===----------------------------------------------------------------------===//
 // Forward-decl so that we can declare methods of `Ranks` and `VarSet`.
 class DimLvlExpr;
 
 //===----------------------------------------------------------------------===//
 class Ranks final {
   // Not using `VarKindArray` since `EnumeratedArray` doesn't support constexpr.
-  // TODO(wrengr): to what extent do we actually care about constexpr here?
   unsigned impl[3];
 
   static constexpr unsigned to_index(VarKind vk) {
     assert(isWF(vk) && "unknown VarKind");
     return static_cast<unsigned>(to_underlying(vk));
   }
 
 public:
public:
   constexpr bool isValid(Var var) const {
     return var.getNum() < getRank(var.getKind());
   }
   bool isValid(DimLvlExpr expr) const;
 };
 static_assert(IsZeroCostAbstraction<Ranks>);
 
 //===----------------------------------------------------------------------===//
+/// Efficient representation of a set of `Var`.
+///
+/// NOTE: For the `contains`/`occursIn` methods: if variables occurring in
+/// the method parameter are OOB for the `VarSet`, then these methods will
+/// always return false.  However, for the `add` methods: OOB parameters
+/// cause undefined behavior.  Currently the `add` methods will raise an
+/// assertion error; though we may change that behavior in the future
+/// (e.g., to resize the underlying bitvectors).
 class VarSet final {
   // If we're willing to give up the possibility of resizing the
   // individual bitvectors, then we could flatten this into a single
   // bitvector (akin to how `mlir::presburger::PresburgerSpace` does it);
   // however, doing so would greatly complicate the implementation of the
   // `occursIn(VarSet)` method.
   VarKindArray<llvm::SmallBitVector> impl;
 
 public:
   explicit VarSet(Ranks const &ranks);
 
-  // TODO(wrengr): can we come up with a single name that works for all three of
-  // these?
   bool contains(Var var) const;
   bool occursIn(VarSet const &vars) const;
   bool occursIn(DimLvlExpr expr) const;
 
   void add(Var var);
-  // TODO(wrengr): void add(VarSet const& vars);
+  void add(VarSet const &vars);
   void add(DimLvlExpr expr);
 };
 
 //===----------------------------------------------------------------------===//
 /// A record of metadata for/about a variable, used by `VarEnv`.
 /// The principal goal of this record is to enable `VarEnv` to be used for
 /// incremental parsing; in particular, `VarInfo` allows the `Var::Num` to
 /// remain unknown, since each record is instead identified by `VarInfo::ID`.
class VarEnv final {
   /// Map from `VarInfo::ID` to shared storage for the actual `VarInfo` objects.
   SmallVector<VarInfo> vars;
   /// Map from variable names to their `VarInfo::ID`.
   llvm::StringMap<VarInfo::ID> ids;
 
   VarInfo::ID nextID() const { return static_cast<VarInfo::ID>(vars.size()); }
 
 public:
-  // NOTE TO Wren: initializer needed!
   VarEnv() : nextNum(0) {}
 
   /// Gets the underlying storage for the `VarInfo` identified by
   /// the `VarInfo::ID`.
   ///
   /// NOTE: The returned reference can become dangling if the `VarEnv`
   /// object is mutated during the lifetime of the pointer.  Therefore,
   /// client code should not store the reference nor otherwise allow it

mlir/lib/Dialect/SparseTensor/IR/Detail/Var.cpp

 //===----------------------------------------------------------------------===//
 // `VarSet` implementation.
 //===----------------------------------------------------------------------===//
 
 static constexpr const VarKind everyVarKind[] = {
     VarKind::Dimension, VarKind::Symbol, VarKind::Level};
 
 VarSet::VarSet(Ranks const &ranks) {
-  // FIXME(wrengr): will this DWIM, or do we need to worry about
-  // `reserve` causing resizing/dangling issues?
   for (const auto vk : everyVarKind)
     impl[vk].reserve(ranks.getRank(vk));
 }
 
 bool VarSet::contains(Var var) const {
-  // FIXME(wrengr): this implementation will raise assertion failure on OOB;
-  // but perhaps we'd rather have this return false on OOB?  That's
-  // necessary for consistency with the `anyCommon` implementation of
-  // `occursIn(VarSet)`.
+  // NOTE: We make sure to return false on OOB, for consistency with
+  // the `anyCommon` implementation of `VarSet::occursIn(VarSet)`.
+  // However beware that, as always with silencing OOB, this can hide
+  // bugs in client code.
   const llvm::SmallBitVector &bits = impl[var.getKind()];
-  // NOTE TO Wren: did this to avoid OOB but perhaps it is result of bug
-  if (var.getNum() >= bits.size())
-    return false;
-  return bits[var.getNum()];
+  const auto num = var.getNum();
+  // FIXME(wrengr): If we `assert(num < bits.size())` then
+  // "roundtrip_encoding.mlir" will fail.  So we need to figure out
+  // where exactly the OOB `var` is coming from, to determine whether
+  // that's a logic bug or not.
+  return num < bits.size() && bits[num];
 }
 
 bool VarSet::occursIn(VarSet const &other) const {
   for (const auto vk : everyVarKind)
     if (impl[vk].anyCommon(other.impl[vk]))
       return true;
   return false;
 }
case AffineExprKind::CeilDiv: {
     (void)op;
     return occursIn(lhs) || occursIn(rhs);
   }
   }
   llvm_unreachable("unknown AffineExprKind");
 }
 
 void VarSet::add(Var var) {
-  // FIXME(wrengr): this implementation will raise assertion failure on OOB;
-  // but perhaps we'd rather have this be a noop on OOB?  or to grow
-  // the underlying bitvectors on OOB?
+  // NOTE: `SmallBitVactor::operator[]` will raise assertion errors for OOB.
   impl[var.getKind()][var.getNum()] = true;
 }
 
-// TODO(wrengr): void VarSet::add(VarSet const& other);
+void VarSet::add(VarSet const &other) {
+  // NOTE: `SmallBitVector::operator&=` will implicitly resize
+  // the bitvector (unlike `BitVector::operator&=`), so we add an
+  // assertion against OOB for consistency with the implementation
+  // of `VarSet::add(Var)`.
+  for (const auto vk : everyVarKind) {
+    assert(impl[vk].size() >= other.impl[vk].size());
+    impl[vk] &= other.impl[vk];
+  }
+}
 
 void VarSet::add(DimLvlExpr expr) {
   if (!expr)
     return;
   switch (expr.getAffineKind()) {
   case AffineExprKind::Constant:
     return;
   case AffineExprKind::SymbolId: