[MLIR] Make SymbolTableCollection methods virtual (llvm#141760)

The `LockedSymbolTable` class not only encapsulate a `SymbolTableCollection`, but also extends it. However, the methods of `SymbolTableCollection` are not marked as `virtual`, and therefore methods receiving a `SymbolTableCollection` would always call the base methods even if the object was a subclass. The proposed changes consist in marking the base methods as `virtual`.

Author: mscuttari

mlir/include/mlir/IR/SymbolTable.h

@@ -282,10 +282,14 @@ raw_ostream &operator<<(raw_ostream &os, SymbolTable::Visibility visibility);
 /// unnecessary tables.
 class SymbolTableCollection {
 public:
+  virtual ~SymbolTableCollection() = default;
+
   /// Look up a symbol with the specified name within the specified symbol table
   /// operation, returning null if no such name exists.
-  Operation *lookupSymbolIn(Operation *symbolTableOp, StringAttr symbol);
-  Operation *lookupSymbolIn(Operation *symbolTableOp, SymbolRefAttr name);
+  virtual Operation *lookupSymbolIn(Operation *symbolTableOp,
+                                    StringAttr symbol);
+  virtual Operation *lookupSymbolIn(Operation *symbolTableOp,
+                                    SymbolRefAttr name);
   template <typename T, typename NameT>
   T lookupSymbolIn(Operation *symbolTableOp, NameT &&name) {
     return dyn_cast_or_null<T>(
@@ -295,15 +299,18 @@ class SymbolTableCollection {
   /// by a given SymbolRefAttr when resolved within the provided symbol table
   /// operation. Returns failure if any of the nested references could not be
   /// resolved.
-  LogicalResult lookupSymbolIn(Operation *symbolTableOp, SymbolRefAttr name,
-                               SmallVectorImpl<Operation *> &symbols);
+  virtual LogicalResult lookupSymbolIn(Operation *symbolTableOp,
+                                       SymbolRefAttr name,
+                                       SmallVectorImpl<Operation *> &symbols);
 
   /// Returns the operation registered with the given symbol name within the
   /// closest parent operation of, or including, 'from' with the
   /// 'OpTrait::SymbolTable' trait. Returns nullptr if no valid symbol was
   /// found.
-  Operation *lookupNearestSymbolFrom(Operation *from, StringAttr symbol);
-  Operation *lookupNearestSymbolFrom(Operation *from, SymbolRefAttr symbol);
+  virtual Operation *lookupNearestSymbolFrom(Operation *from,
+                                             StringAttr symbol);
+  virtual Operation *lookupNearestSymbolFrom(Operation *from,
+                                             SymbolRefAttr symbol);
   template <typename T>
   T lookupNearestSymbolFrom(Operation *from, StringAttr symbol) {
     return dyn_cast_or_null<T>(lookupNearestSymbolFrom(from, symbol));
@@ -314,14 +321,14 @@ class SymbolTableCollection {
   }
 
   /// Lookup, or create, a symbol table for an operation.
-  SymbolTable &getSymbolTable(Operation *op);
+  virtual SymbolTable &getSymbolTable(Operation *op);
 
   /// Invalidate the cached symbol table for an operation.
   /// This is important when doing IR modifications that erase and also create
   /// operations having the 'OpTrait::SymbolTable' trait. If a symbol table of
   /// an erased operation is not invalidated, a new operation sharing the same
   /// address would be associated with outdated, and wrong, information.
-  void invalidateSymbolTable(Operation *op);
+  virtual void invalidateSymbolTable(Operation *op);
 
 private:
   friend class LockedSymbolTableCollection;
@@ -348,13 +355,15 @@ class LockedSymbolTableCollection : public SymbolTableCollection {
 
   /// Look up a symbol with the specified name within the specified symbol table
   /// operation, returning null if no such name exists.
-  Operation *lookupSymbolIn(Operation *symbolTableOp, StringAttr symbol);
+  Operation *lookupSymbolIn(Operation *symbolTableOp,
+                            StringAttr symbol) override;
   /// Look up a symbol with the specified name within the specified symbol table
   /// operation, returning null if no such name exists.
   Operation *lookupSymbolIn(Operation *symbolTableOp, FlatSymbolRefAttr symbol);
   /// Look up a potentially nested symbol within the specified symbol table
   /// operation, returning null if no such symbol exists.
-  Operation *lookupSymbolIn(Operation *symbolTableOp, SymbolRefAttr name);
+  Operation *lookupSymbolIn(Operation *symbolTableOp,
+                            SymbolRefAttr name) override;
 
   /// Lookup a symbol of a particular kind within the specified symbol table,
   /// returning null if the symbol was not found.
@@ -369,14 +378,14 @@ class LockedSymbolTableCollection : public SymbolTableCollection {
   /// operation. Returns failure if any of the nested references could not be
   /// resolved.
   LogicalResult lookupSymbolIn(Operation *symbolTableOp, SymbolRefAttr name,
-                               SmallVectorImpl<Operation *> &symbols);
+                               SmallVectorImpl<Operation *> &symbols) override;
 
 private:
   /// Get the symbol table for the symbol table operation, constructing if it
   /// does not exist. This function provides thread safety over `collection`
   /// by locking when performing the lookup and when inserting
   /// lazily-constructed symbol tables.
-  SymbolTable &getSymbolTable(Operation *symbolTableOp);
+  SymbolTable &getSymbolTable(Operation *symbolTableOp) override;
 
   /// The symbol tables to manage.
   SymbolTableCollection &collection;