1 //===- SymbolTable.cpp - MLIR Symbol Table Class --------------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "mlir/IR/SymbolTable.h"
10 #include "mlir/IR/Builders.h"
11 #include "mlir/IR/OpImplementation.h"
12 #include "llvm/ADT/SetVector.h"
13 #include "llvm/ADT/SmallPtrSet.h"
14 #include "llvm/ADT/SmallString.h"
15 #include "llvm/ADT/StringSwitch.h"
16 
17 using namespace mlir;
18 
19 /// Return true if the given operation is unknown and may potentially define a
20 /// symbol table.
isPotentiallyUnknownSymbolTable(Operation * op)21 static bool isPotentiallyUnknownSymbolTable(Operation *op) {
22   return op->getNumRegions() == 1 && !op->getDialect();
23 }
24 
25 /// Returns the string name of the given symbol, or None if this is not a
26 /// symbol.
getNameIfSymbol(Operation * symbol)27 static Optional<StringRef> getNameIfSymbol(Operation *symbol) {
28   auto nameAttr =
29       symbol->getAttrOfType<StringAttr>(SymbolTable::getSymbolAttrName());
30   return nameAttr ? nameAttr.getValue() : Optional<StringRef>();
31 }
getNameIfSymbol(Operation * symbol,Identifier symbolAttrNameId)32 static Optional<StringRef> getNameIfSymbol(Operation *symbol,
33                                            Identifier symbolAttrNameId) {
34   auto nameAttr = symbol->getAttrOfType<StringAttr>(symbolAttrNameId);
35   return nameAttr ? nameAttr.getValue() : Optional<StringRef>();
36 }
37 
38 /// Computes the nested symbol reference attribute for the symbol 'symbolName'
39 /// that are usable within the symbol table operations from 'symbol' as far up
40 /// to the given operation 'within', where 'within' is an ancestor of 'symbol'.
41 /// Returns success if all references up to 'within' could be computed.
42 static LogicalResult
collectValidReferencesFor(Operation * symbol,StringRef symbolName,Operation * within,SmallVectorImpl<SymbolRefAttr> & results)43 collectValidReferencesFor(Operation *symbol, StringRef symbolName,
44                           Operation *within,
45                           SmallVectorImpl<SymbolRefAttr> &results) {
46   assert(within->isAncestor(symbol) && "expected 'within' to be an ancestor");
47   MLIRContext *ctx = symbol->getContext();
48 
49   auto leafRef = FlatSymbolRefAttr::get(ctx, symbolName);
50   results.push_back(leafRef);
51 
52   // Early exit for when 'within' is the parent of 'symbol'.
53   Operation *symbolTableOp = symbol->getParentOp();
54   if (within == symbolTableOp)
55     return success();
56 
57   // Collect references until 'symbolTableOp' reaches 'within'.
58   SmallVector<FlatSymbolRefAttr, 1> nestedRefs(1, leafRef);
59   Identifier symbolNameId =
60       Identifier::get(SymbolTable::getSymbolAttrName(), ctx);
61   do {
62     // Each parent of 'symbol' should define a symbol table.
63     if (!symbolTableOp->hasTrait<OpTrait::SymbolTable>())
64       return failure();
65     // Each parent of 'symbol' should also be a symbol.
66     Optional<StringRef> symbolTableName =
67         getNameIfSymbol(symbolTableOp, symbolNameId);
68     if (!symbolTableName)
69       return failure();
70     results.push_back(SymbolRefAttr::get(ctx, *symbolTableName, nestedRefs));
71 
72     symbolTableOp = symbolTableOp->getParentOp();
73     if (symbolTableOp == within)
74       break;
75     nestedRefs.insert(nestedRefs.begin(),
76                       FlatSymbolRefAttr::get(ctx, *symbolTableName));
77   } while (true);
78   return success();
79 }
80 
81 /// Walk all of the operations within the given set of regions, without
82 /// traversing into any nested symbol tables. Stops walking if the result of the
83 /// callback is anything other than `WalkResult::advance`.
84 static Optional<WalkResult>
walkSymbolTable(MutableArrayRef<Region> regions,function_ref<Optional<WalkResult> (Operation *)> callback)85 walkSymbolTable(MutableArrayRef<Region> regions,
86                 function_ref<Optional<WalkResult>(Operation *)> callback) {
87   SmallVector<Region *, 1> worklist(llvm::make_pointer_range(regions));
88   while (!worklist.empty()) {
89     for (Operation &op : worklist.pop_back_val()->getOps()) {
90       Optional<WalkResult> result = callback(&op);
91       if (result != WalkResult::advance())
92         return result;
93 
94       // If this op defines a new symbol table scope, we can't traverse. Any
95       // symbol references nested within 'op' are different semantically.
96       if (!op.hasTrait<OpTrait::SymbolTable>()) {
97         for (Region &region : op.getRegions())
98           worklist.push_back(&region);
99       }
100     }
101   }
102   return WalkResult::advance();
103 }
104 
105 //===----------------------------------------------------------------------===//
106 // SymbolTable
107 //===----------------------------------------------------------------------===//
108 
109 /// Build a symbol table with the symbols within the given operation.
SymbolTable(Operation * symbolTableOp)110 SymbolTable::SymbolTable(Operation *symbolTableOp)
111     : symbolTableOp(symbolTableOp) {
112   assert(symbolTableOp->hasTrait<OpTrait::SymbolTable>() &&
113          "expected operation to have SymbolTable trait");
114   assert(symbolTableOp->getNumRegions() == 1 &&
115          "expected operation to have a single region");
116   assert(llvm::hasSingleElement(symbolTableOp->getRegion(0)) &&
117          "expected operation to have a single block");
118 
119   Identifier symbolNameId = Identifier::get(SymbolTable::getSymbolAttrName(),
120                                             symbolTableOp->getContext());
121   for (auto &op : symbolTableOp->getRegion(0).front()) {
122     Optional<StringRef> name = getNameIfSymbol(&op, symbolNameId);
123     if (!name)
124       continue;
125 
126     auto inserted = symbolTable.insert({*name, &op});
127     (void)inserted;
128     assert(inserted.second &&
129            "expected region to contain uniquely named symbol operations");
130   }
131 }
132 
133 /// Look up a symbol with the specified name, returning null if no such name
134 /// exists. Names never include the @ on them.
lookup(StringRef name) const135 Operation *SymbolTable::lookup(StringRef name) const {
136   return symbolTable.lookup(name);
137 }
138 
139 /// Erase the given symbol from the table.
erase(Operation * symbol)140 void SymbolTable::erase(Operation *symbol) {
141   Optional<StringRef> name = getNameIfSymbol(symbol);
142   assert(name && "expected valid 'name' attribute");
143   assert(symbol->getParentOp() == symbolTableOp &&
144          "expected this operation to be inside of the operation with this "
145          "SymbolTable");
146 
147   auto it = symbolTable.find(*name);
148   if (it != symbolTable.end() && it->second == symbol) {
149     symbolTable.erase(it);
150     symbol->erase();
151   }
152 }
153 
154 // TODO: Consider if this should be renamed to something like insertOrUpdate
155 /// Insert a new symbol into the table and associated operation if not already
156 /// there and rename it as necessary to avoid collisions.
insert(Operation * symbol,Block::iterator insertPt)157 void SymbolTable::insert(Operation *symbol, Block::iterator insertPt) {
158   // The symbol cannot be the child of another op and must be the child of the
159   // symbolTableOp after this.
160   //
161   // TODO: consider if SymbolTable's constructor should behave the same.
162   if (!symbol->getParentOp()) {
163     auto &body = symbolTableOp->getRegion(0).front();
164     if (insertPt == Block::iterator()) {
165       insertPt = Block::iterator(body.end());
166     } else {
167       assert((insertPt == body.end() ||
168               insertPt->getParentOp() == symbolTableOp) &&
169              "expected insertPt to be in the associated module operation");
170     }
171     // Insert before the terminator, if any.
172     if (insertPt == Block::iterator(body.end()) && !body.empty() &&
173         std::prev(body.end())->hasTrait<OpTrait::IsTerminator>())
174       insertPt = std::prev(body.end());
175 
176     body.getOperations().insert(insertPt, symbol);
177   }
178   assert(symbol->getParentOp() == symbolTableOp &&
179          "symbol is already inserted in another op");
180 
181   // Add this symbol to the symbol table, uniquing the name if a conflict is
182   // detected.
183   StringRef name = getSymbolName(symbol);
184   if (symbolTable.insert({name, symbol}).second)
185     return;
186   // If the symbol was already in the table, also return.
187   if (symbolTable.lookup(name) == symbol)
188     return;
189   // If a conflict was detected, then the symbol will not have been added to
190   // the symbol table. Try suffixes until we get to a unique name that works.
191   SmallString<128> nameBuffer(name);
192   unsigned originalLength = nameBuffer.size();
193 
194   // Iteratively try suffixes until we find one that isn't used.
195   do {
196     nameBuffer.resize(originalLength);
197     nameBuffer += '_';
198     nameBuffer += std::to_string(uniquingCounter++);
199   } while (!symbolTable.insert({nameBuffer, symbol}).second);
200   setSymbolName(symbol, nameBuffer);
201 }
202 
203 /// Returns the name of the given symbol operation.
getSymbolName(Operation * symbol)204 StringRef SymbolTable::getSymbolName(Operation *symbol) {
205   Optional<StringRef> name = getNameIfSymbol(symbol);
206   assert(name && "expected valid symbol name");
207   return *name;
208 }
209 /// Sets the name of the given symbol operation.
setSymbolName(Operation * symbol,StringRef name)210 void SymbolTable::setSymbolName(Operation *symbol, StringRef name) {
211   symbol->setAttr(getSymbolAttrName(),
212                   StringAttr::get(symbol->getContext(), name));
213 }
214 
215 /// Returns the visibility of the given symbol operation.
getSymbolVisibility(Operation * symbol)216 SymbolTable::Visibility SymbolTable::getSymbolVisibility(Operation *symbol) {
217   // If the attribute doesn't exist, assume public.
218   StringAttr vis = symbol->getAttrOfType<StringAttr>(getVisibilityAttrName());
219   if (!vis)
220     return Visibility::Public;
221 
222   // Otherwise, switch on the string value.
223   return StringSwitch<Visibility>(vis.getValue())
224       .Case("private", Visibility::Private)
225       .Case("nested", Visibility::Nested)
226       .Case("public", Visibility::Public);
227 }
228 /// Sets the visibility of the given symbol operation.
setSymbolVisibility(Operation * symbol,Visibility vis)229 void SymbolTable::setSymbolVisibility(Operation *symbol, Visibility vis) {
230   MLIRContext *ctx = symbol->getContext();
231 
232   // If the visibility is public, just drop the attribute as this is the
233   // default.
234   if (vis == Visibility::Public) {
235     symbol->removeAttr(Identifier::get(getVisibilityAttrName(), ctx));
236     return;
237   }
238 
239   // Otherwise, update the attribute.
240   assert((vis == Visibility::Private || vis == Visibility::Nested) &&
241          "unknown symbol visibility kind");
242 
243   StringRef visName = vis == Visibility::Private ? "private" : "nested";
244   symbol->setAttr(getVisibilityAttrName(), StringAttr::get(ctx, visName));
245 }
246 
247 /// Returns the nearest symbol table from a given operation `from`. Returns
248 /// nullptr if no valid parent symbol table could be found.
getNearestSymbolTable(Operation * from)249 Operation *SymbolTable::getNearestSymbolTable(Operation *from) {
250   assert(from && "expected valid operation");
251   if (isPotentiallyUnknownSymbolTable(from))
252     return nullptr;
253 
254   while (!from->hasTrait<OpTrait::SymbolTable>()) {
255     from = from->getParentOp();
256 
257     // Check that this is a valid op and isn't an unknown symbol table.
258     if (!from || isPotentiallyUnknownSymbolTable(from))
259       return nullptr;
260   }
261   return from;
262 }
263 
264 /// Walks all symbol table operations nested within, and including, `op`. For
265 /// each symbol table operation, the provided callback is invoked with the op
266 /// and a boolean signifying if the symbols within that symbol table can be
267 /// treated as if all uses are visible. `allSymUsesVisible` identifies whether
268 /// all of the symbol uses of symbols within `op` are visible.
walkSymbolTables(Operation * op,bool allSymUsesVisible,function_ref<void (Operation *,bool)> callback)269 void SymbolTable::walkSymbolTables(
270     Operation *op, bool allSymUsesVisible,
271     function_ref<void(Operation *, bool)> callback) {
272   bool isSymbolTable = op->hasTrait<OpTrait::SymbolTable>();
273   if (isSymbolTable) {
274     SymbolOpInterface symbol = dyn_cast<SymbolOpInterface>(op);
275     allSymUsesVisible |= !symbol || symbol.isPrivate();
276   } else {
277     // Otherwise if 'op' is not a symbol table, any nested symbols are
278     // guaranteed to be hidden.
279     allSymUsesVisible = true;
280   }
281 
282   for (Region &region : op->getRegions())
283     for (Block &block : region)
284       for (Operation &nestedOp : block)
285         walkSymbolTables(&nestedOp, allSymUsesVisible, callback);
286 
287   // If 'op' had the symbol table trait, visit it after any nested symbol
288   // tables.
289   if (isSymbolTable)
290     callback(op, allSymUsesVisible);
291 }
292 
293 /// Returns the operation registered with the given symbol name with the
294 /// regions of 'symbolTableOp'. 'symbolTableOp' is required to be an operation
295 /// with the 'OpTrait::SymbolTable' trait. Returns nullptr if no valid symbol
296 /// was found.
lookupSymbolIn(Operation * symbolTableOp,StringRef symbol)297 Operation *SymbolTable::lookupSymbolIn(Operation *symbolTableOp,
298                                        StringRef symbol) {
299   assert(symbolTableOp->hasTrait<OpTrait::SymbolTable>());
300   Region &region = symbolTableOp->getRegion(0);
301   if (region.empty())
302     return nullptr;
303 
304   // Look for a symbol with the given name.
305   Identifier symbolNameId = Identifier::get(SymbolTable::getSymbolAttrName(),
306                                             symbolTableOp->getContext());
307   for (auto &op : region.front())
308     if (getNameIfSymbol(&op, symbolNameId) == symbol)
309       return &op;
310   return nullptr;
311 }
lookupSymbolIn(Operation * symbolTableOp,SymbolRefAttr symbol)312 Operation *SymbolTable::lookupSymbolIn(Operation *symbolTableOp,
313                                        SymbolRefAttr symbol) {
314   SmallVector<Operation *, 4> resolvedSymbols;
315   if (failed(lookupSymbolIn(symbolTableOp, symbol, resolvedSymbols)))
316     return nullptr;
317   return resolvedSymbols.back();
318 }
319 
320 /// Internal implementation of `lookupSymbolIn` that allows for specialized
321 /// implementations of the lookup function.
lookupSymbolInImpl(Operation * symbolTableOp,SymbolRefAttr symbol,SmallVectorImpl<Operation * > & symbols,function_ref<Operation * (Operation *,StringRef)> lookupSymbolFn)322 static LogicalResult lookupSymbolInImpl(
323     Operation *symbolTableOp, SymbolRefAttr symbol,
324     SmallVectorImpl<Operation *> &symbols,
325     function_ref<Operation *(Operation *, StringRef)> lookupSymbolFn) {
326   assert(symbolTableOp->hasTrait<OpTrait::SymbolTable>());
327 
328   // Lookup the root reference for this symbol.
329   symbolTableOp = lookupSymbolFn(symbolTableOp, symbol.getRootReference());
330   if (!symbolTableOp)
331     return failure();
332   symbols.push_back(symbolTableOp);
333 
334   // If there are no nested references, just return the root symbol directly.
335   ArrayRef<FlatSymbolRefAttr> nestedRefs = symbol.getNestedReferences();
336   if (nestedRefs.empty())
337     return success();
338 
339   // Verify that the root is also a symbol table.
340   if (!symbolTableOp->hasTrait<OpTrait::SymbolTable>())
341     return failure();
342 
343   // Otherwise, lookup each of the nested non-leaf references and ensure that
344   // each corresponds to a valid symbol table.
345   for (FlatSymbolRefAttr ref : nestedRefs.drop_back()) {
346     symbolTableOp = lookupSymbolFn(symbolTableOp, ref.getValue());
347     if (!symbolTableOp || !symbolTableOp->hasTrait<OpTrait::SymbolTable>())
348       return failure();
349     symbols.push_back(symbolTableOp);
350   }
351   symbols.push_back(lookupSymbolFn(symbolTableOp, symbol.getLeafReference()));
352   return success(symbols.back());
353 }
354 
355 LogicalResult
lookupSymbolIn(Operation * symbolTableOp,SymbolRefAttr symbol,SmallVectorImpl<Operation * > & symbols)356 SymbolTable::lookupSymbolIn(Operation *symbolTableOp, SymbolRefAttr symbol,
357                             SmallVectorImpl<Operation *> &symbols) {
358   auto lookupFn = [](Operation *symbolTableOp, StringRef symbol) {
359     return lookupSymbolIn(symbolTableOp, symbol);
360   };
361   return lookupSymbolInImpl(symbolTableOp, symbol, symbols, lookupFn);
362 }
363 
364 /// Returns the operation registered with the given symbol name within the
365 /// closes parent operation with the 'OpTrait::SymbolTable' trait. Returns
366 /// nullptr if no valid symbol was found.
lookupNearestSymbolFrom(Operation * from,StringRef symbol)367 Operation *SymbolTable::lookupNearestSymbolFrom(Operation *from,
368                                                 StringRef symbol) {
369   Operation *symbolTableOp = getNearestSymbolTable(from);
370   return symbolTableOp ? lookupSymbolIn(symbolTableOp, symbol) : nullptr;
371 }
lookupNearestSymbolFrom(Operation * from,SymbolRefAttr symbol)372 Operation *SymbolTable::lookupNearestSymbolFrom(Operation *from,
373                                                 SymbolRefAttr symbol) {
374   Operation *symbolTableOp = getNearestSymbolTable(from);
375   return symbolTableOp ? lookupSymbolIn(symbolTableOp, symbol) : nullptr;
376 }
377 
operator <<(raw_ostream & os,SymbolTable::Visibility visibility)378 raw_ostream &mlir::operator<<(raw_ostream &os,
379                               SymbolTable::Visibility visibility) {
380   switch (visibility) {
381   case SymbolTable::Visibility::Public:
382     return os << "public";
383   case SymbolTable::Visibility::Private:
384     return os << "private";
385   case SymbolTable::Visibility::Nested:
386     return os << "nested";
387   }
388   llvm_unreachable("Unexpected visibility");
389 }
390 
391 //===----------------------------------------------------------------------===//
392 // SymbolTable Trait Types
393 //===----------------------------------------------------------------------===//
394 
verifySymbolTable(Operation * op)395 LogicalResult detail::verifySymbolTable(Operation *op) {
396   if (op->getNumRegions() != 1)
397     return op->emitOpError()
398            << "Operations with a 'SymbolTable' must have exactly one region";
399   if (!llvm::hasSingleElement(op->getRegion(0)))
400     return op->emitOpError()
401            << "Operations with a 'SymbolTable' must have exactly one block";
402 
403   // Check that all symbols are uniquely named within child regions.
404   DenseMap<Attribute, Location> nameToOrigLoc;
405   for (auto &block : op->getRegion(0)) {
406     for (auto &op : block) {
407       // Check for a symbol name attribute.
408       auto nameAttr =
409           op.getAttrOfType<StringAttr>(mlir::SymbolTable::getSymbolAttrName());
410       if (!nameAttr)
411         continue;
412 
413       // Try to insert this symbol into the table.
414       auto it = nameToOrigLoc.try_emplace(nameAttr, op.getLoc());
415       if (!it.second)
416         return op.emitError()
417             .append("redefinition of symbol named '", nameAttr.getValue(), "'")
418             .attachNote(it.first->second)
419             .append("see existing symbol definition here");
420     }
421   }
422 
423   // Verify any nested symbol user operations.
424   SymbolTableCollection symbolTable;
425   auto verifySymbolUserFn = [&](Operation *op) -> Optional<WalkResult> {
426     if (SymbolUserOpInterface user = dyn_cast<SymbolUserOpInterface>(op))
427       return WalkResult(user.verifySymbolUses(symbolTable));
428     return WalkResult::advance();
429   };
430 
431   Optional<WalkResult> result =
432       walkSymbolTable(op->getRegions(), verifySymbolUserFn);
433   return success(result && !result->wasInterrupted());
434 }
435 
verifySymbol(Operation * op)436 LogicalResult detail::verifySymbol(Operation *op) {
437   // Verify the name attribute.
438   if (!op->getAttrOfType<StringAttr>(mlir::SymbolTable::getSymbolAttrName()))
439     return op->emitOpError() << "requires string attribute '"
440                              << mlir::SymbolTable::getSymbolAttrName() << "'";
441 
442   // Verify the visibility attribute.
443   if (Attribute vis = op->getAttr(mlir::SymbolTable::getVisibilityAttrName())) {
444     StringAttr visStrAttr = vis.dyn_cast<StringAttr>();
445     if (!visStrAttr)
446       return op->emitOpError() << "requires visibility attribute '"
447                                << mlir::SymbolTable::getVisibilityAttrName()
448                                << "' to be a string attribute, but got " << vis;
449 
450     if (!llvm::is_contained(ArrayRef<StringRef>{"public", "private", "nested"},
451                             visStrAttr.getValue()))
452       return op->emitOpError()
453              << "visibility expected to be one of [\"public\", \"private\", "
454                 "\"nested\"], but got "
455              << visStrAttr;
456   }
457   return success();
458 }
459 
460 //===----------------------------------------------------------------------===//
461 // Symbol Use Lists
462 //===----------------------------------------------------------------------===//
463 
464 /// Walk all of the symbol references within the given operation, invoking the
465 /// provided callback for each found use. The callbacks takes as arguments: the
466 /// use of the symbol, and the nested access chain to the attribute within the
467 /// operation dictionary. An access chain is a set of indices into nested
468 /// container attributes. For example, a symbol use in an attribute dictionary
469 /// that looks like the following:
470 ///
471 ///    {use = [{other_attr, @symbol}]}
472 ///
473 /// May have the following access chain:
474 ///
475 ///     [0, 0, 1]
476 ///
walkSymbolRefs(Operation * op,function_ref<WalkResult (SymbolTable::SymbolUse,ArrayRef<int>)> callback)477 static WalkResult walkSymbolRefs(
478     Operation *op,
479     function_ref<WalkResult(SymbolTable::SymbolUse, ArrayRef<int>)> callback) {
480   // Check to see if the operation has any attributes.
481   DictionaryAttr attrDict = op->getAttrDictionary();
482   if (attrDict.empty())
483     return WalkResult::advance();
484 
485   // A worklist of a container attribute and the current index into the held
486   // attribute list.
487   SmallVector<Attribute, 1> attrWorklist(1, attrDict);
488   SmallVector<int, 1> curAccessChain(1, /*Value=*/-1);
489 
490   // Process the symbol references within the given nested attribute range.
491   auto processAttrs = [&](int &index, auto attrRange) -> WalkResult {
492     for (Attribute attr : llvm::drop_begin(attrRange, index)) {
493       /// Check for a nested container attribute, these will also need to be
494       /// walked.
495       if (attr.isa<ArrayAttr, DictionaryAttr>()) {
496         attrWorklist.push_back(attr);
497         curAccessChain.push_back(-1);
498         return WalkResult::advance();
499       }
500 
501       // Invoke the provided callback if we find a symbol use and check for a
502       // requested interrupt.
503       if (auto symbolRef = attr.dyn_cast<SymbolRefAttr>())
504         if (callback({op, symbolRef}, curAccessChain).wasInterrupted())
505           return WalkResult::interrupt();
506 
507       // Make sure to keep the index counter in sync.
508       ++index;
509     }
510 
511     // Pop this container attribute from the worklist.
512     attrWorklist.pop_back();
513     curAccessChain.pop_back();
514     return WalkResult::advance();
515   };
516 
517   WalkResult result = WalkResult::advance();
518   do {
519     Attribute attr = attrWorklist.back();
520     int &index = curAccessChain.back();
521     ++index;
522 
523     // Process the given attribute, which is guaranteed to be a container.
524     if (auto dict = attr.dyn_cast<DictionaryAttr>())
525       result = processAttrs(index, make_second_range(dict.getValue()));
526     else
527       result = processAttrs(index, attr.cast<ArrayAttr>().getValue());
528   } while (!attrWorklist.empty() && !result.wasInterrupted());
529   return result;
530 }
531 
532 /// Walk all of the uses, for any symbol, that are nested within the given
533 /// regions, invoking the provided callback for each. This does not traverse
534 /// into any nested symbol tables.
walkSymbolUses(MutableArrayRef<Region> regions,function_ref<WalkResult (SymbolTable::SymbolUse,ArrayRef<int>)> callback)535 static Optional<WalkResult> walkSymbolUses(
536     MutableArrayRef<Region> regions,
537     function_ref<WalkResult(SymbolTable::SymbolUse, ArrayRef<int>)> callback) {
538   return walkSymbolTable(regions, [&](Operation *op) -> Optional<WalkResult> {
539     // Check that this isn't a potentially unknown symbol table.
540     if (isPotentiallyUnknownSymbolTable(op))
541       return llvm::None;
542 
543     return walkSymbolRefs(op, callback);
544   });
545 }
546 /// Walk all of the uses, for any symbol, that are nested within the given
547 /// operation 'from', invoking the provided callback for each. This does not
548 /// traverse into any nested symbol tables.
walkSymbolUses(Operation * from,function_ref<WalkResult (SymbolTable::SymbolUse,ArrayRef<int>)> callback)549 static Optional<WalkResult> walkSymbolUses(
550     Operation *from,
551     function_ref<WalkResult(SymbolTable::SymbolUse, ArrayRef<int>)> callback) {
552   // If this operation has regions, and it, as well as its dialect, isn't
553   // registered then conservatively fail. The operation may define a
554   // symbol table, so we can't opaquely know if we should traverse to find
555   // nested uses.
556   if (isPotentiallyUnknownSymbolTable(from))
557     return llvm::None;
558 
559   // Walk the uses on this operation.
560   if (walkSymbolRefs(from, callback).wasInterrupted())
561     return WalkResult::interrupt();
562 
563   // Only recurse if this operation is not a symbol table. A symbol table
564   // defines a new scope, so we can't walk the attributes from within the symbol
565   // table op.
566   if (!from->hasTrait<OpTrait::SymbolTable>())
567     return walkSymbolUses(from->getRegions(), callback);
568   return WalkResult::advance();
569 }
570 
571 namespace {
572 /// This class represents a single symbol scope. A symbol scope represents the
573 /// set of operations nested within a symbol table that may reference symbols
574 /// within that table. A symbol scope does not contain the symbol table
575 /// operation itself, just its contained operations. A scope ends at leaf
576 /// operations or another symbol table operation.
577 struct SymbolScope {
578   /// Walk the symbol uses within this scope, invoking the given callback.
579   /// This variant is used when the callback type matches that expected by
580   /// 'walkSymbolUses'.
581   template <typename CallbackT,
582             typename std::enable_if_t<!std::is_same<
583                 typename llvm::function_traits<CallbackT>::result_t,
584                 void>::value> * = nullptr>
walk__anon8066b17f0511::SymbolScope585   Optional<WalkResult> walk(CallbackT cback) {
586     if (Region *region = limit.dyn_cast<Region *>())
587       return walkSymbolUses(*region, cback);
588     return walkSymbolUses(limit.get<Operation *>(), cback);
589   }
590   /// This variant is used when the callback type matches a stripped down type:
591   /// void(SymbolTable::SymbolUse use)
592   template <typename CallbackT,
593             typename std::enable_if_t<std::is_same<
594                 typename llvm::function_traits<CallbackT>::result_t,
595                 void>::value> * = nullptr>
walk__anon8066b17f0511::SymbolScope596   Optional<WalkResult> walk(CallbackT cback) {
597     return walk([=](SymbolTable::SymbolUse use, ArrayRef<int>) {
598       return cback(use), WalkResult::advance();
599     });
600   }
601 
602   /// The representation of the symbol within this scope.
603   SymbolRefAttr symbol;
604 
605   /// The IR unit representing this scope.
606   llvm::PointerUnion<Operation *, Region *> limit;
607 };
608 } // end anonymous namespace
609 
610 /// Collect all of the symbol scopes from 'symbol' to (inclusive) 'limit'.
collectSymbolScopes(Operation * symbol,Operation * limit)611 static SmallVector<SymbolScope, 2> collectSymbolScopes(Operation *symbol,
612                                                        Operation *limit) {
613   StringRef symName = SymbolTable::getSymbolName(symbol);
614   assert(!symbol->hasTrait<OpTrait::SymbolTable>() || symbol != limit);
615 
616   // Compute the ancestors of 'limit'.
617   SetVector<Operation *, SmallVector<Operation *, 4>,
618             SmallPtrSet<Operation *, 4>>
619       limitAncestors;
620   Operation *limitAncestor = limit;
621   do {
622     // Check to see if 'symbol' is an ancestor of 'limit'.
623     if (limitAncestor == symbol) {
624       // Check that the nearest symbol table is 'symbol's parent. SymbolRefAttr
625       // doesn't support parent references.
626       if (SymbolTable::getNearestSymbolTable(limit->getParentOp()) ==
627           symbol->getParentOp())
628         return {{SymbolRefAttr::get(symbol->getContext(), symName), limit}};
629       return {};
630     }
631 
632     limitAncestors.insert(limitAncestor);
633   } while ((limitAncestor = limitAncestor->getParentOp()));
634 
635   // Try to find the first ancestor of 'symbol' that is an ancestor of 'limit'.
636   Operation *commonAncestor = symbol->getParentOp();
637   do {
638     if (limitAncestors.count(commonAncestor))
639       break;
640   } while ((commonAncestor = commonAncestor->getParentOp()));
641   assert(commonAncestor && "'limit' and 'symbol' have no common ancestor");
642 
643   // Compute the set of valid nested references for 'symbol' as far up to the
644   // common ancestor as possible.
645   SmallVector<SymbolRefAttr, 2> references;
646   bool collectedAllReferences = succeeded(
647       collectValidReferencesFor(symbol, symName, commonAncestor, references));
648 
649   // Handle the case where the common ancestor is 'limit'.
650   if (commonAncestor == limit) {
651     SmallVector<SymbolScope, 2> scopes;
652 
653     // Walk each of the ancestors of 'symbol', calling the compute function for
654     // each one.
655     Operation *limitIt = symbol->getParentOp();
656     for (size_t i = 0, e = references.size(); i != e;
657          ++i, limitIt = limitIt->getParentOp()) {
658       assert(limitIt->hasTrait<OpTrait::SymbolTable>());
659       scopes.push_back({references[i], &limitIt->getRegion(0)});
660     }
661     return scopes;
662   }
663 
664   // Otherwise, we just need the symbol reference for 'symbol' that will be
665   // used within 'limit'. This is the last reference in the list we computed
666   // above if we were able to collect all references.
667   if (!collectedAllReferences)
668     return {};
669   return {{references.back(), limit}};
670 }
collectSymbolScopes(Operation * symbol,Region * limit)671 static SmallVector<SymbolScope, 2> collectSymbolScopes(Operation *symbol,
672                                                        Region *limit) {
673   auto scopes = collectSymbolScopes(symbol, limit->getParentOp());
674 
675   // If we collected some scopes to walk, make sure to constrain the one for
676   // limit to the specific region requested.
677   if (!scopes.empty())
678     scopes.back().limit = limit;
679   return scopes;
680 }
681 template <typename IRUnit>
collectSymbolScopes(StringRef symbol,IRUnit * limit)682 static SmallVector<SymbolScope, 1> collectSymbolScopes(StringRef symbol,
683                                                        IRUnit *limit) {
684   return {{SymbolRefAttr::get(limit->getContext(), symbol), limit}};
685 }
686 
687 /// Returns true if the given reference 'SubRef' is a sub reference of the
688 /// reference 'ref', i.e. 'ref' is a further qualified reference.
isReferencePrefixOf(SymbolRefAttr subRef,SymbolRefAttr ref)689 static bool isReferencePrefixOf(SymbolRefAttr subRef, SymbolRefAttr ref) {
690   if (ref == subRef)
691     return true;
692 
693   // If the references are not pointer equal, check to see if `subRef` is a
694   // prefix of `ref`.
695   if (ref.isa<FlatSymbolRefAttr>() ||
696       ref.getRootReference() != subRef.getRootReference())
697     return false;
698 
699   auto refLeafs = ref.getNestedReferences();
700   auto subRefLeafs = subRef.getNestedReferences();
701   return subRefLeafs.size() < refLeafs.size() &&
702          subRefLeafs == refLeafs.take_front(subRefLeafs.size());
703 }
704 
705 //===----------------------------------------------------------------------===//
706 // SymbolTable::getSymbolUses
707 
708 /// The implementation of SymbolTable::getSymbolUses below.
709 template <typename FromT>
getSymbolUsesImpl(FromT from)710 static Optional<SymbolTable::UseRange> getSymbolUsesImpl(FromT from) {
711   std::vector<SymbolTable::SymbolUse> uses;
712   auto walkFn = [&](SymbolTable::SymbolUse symbolUse, ArrayRef<int>) {
713     uses.push_back(symbolUse);
714     return WalkResult::advance();
715   };
716   auto result = walkSymbolUses(from, walkFn);
717   return result ? Optional<SymbolTable::UseRange>(std::move(uses)) : llvm::None;
718 }
719 
720 /// Get an iterator range for all of the uses, for any symbol, that are nested
721 /// within the given operation 'from'. This does not traverse into any nested
722 /// symbol tables, and will also only return uses on 'from' if it does not
723 /// also define a symbol table. This is because we treat the region as the
724 /// boundary of the symbol table, and not the op itself. This function returns
725 /// None if there are any unknown operations that may potentially be symbol
726 /// tables.
getSymbolUses(Operation * from)727 auto SymbolTable::getSymbolUses(Operation *from) -> Optional<UseRange> {
728   return getSymbolUsesImpl(from);
729 }
getSymbolUses(Region * from)730 auto SymbolTable::getSymbolUses(Region *from) -> Optional<UseRange> {
731   return getSymbolUsesImpl(MutableArrayRef<Region>(*from));
732 }
733 
734 //===----------------------------------------------------------------------===//
735 // SymbolTable::getSymbolUses
736 
737 /// The implementation of SymbolTable::getSymbolUses below.
738 template <typename SymbolT, typename IRUnitT>
getSymbolUsesImpl(SymbolT symbol,IRUnitT * limit)739 static Optional<SymbolTable::UseRange> getSymbolUsesImpl(SymbolT symbol,
740                                                          IRUnitT *limit) {
741   std::vector<SymbolTable::SymbolUse> uses;
742   for (SymbolScope &scope : collectSymbolScopes(symbol, limit)) {
743     if (!scope.walk([&](SymbolTable::SymbolUse symbolUse) {
744           if (isReferencePrefixOf(scope.symbol, symbolUse.getSymbolRef()))
745             uses.push_back(symbolUse);
746         }))
747       return llvm::None;
748   }
749   return SymbolTable::UseRange(std::move(uses));
750 }
751 
752 /// Get all of the uses of the given symbol that are nested within the given
753 /// operation 'from', invoking the provided callback for each. This does not
754 /// traverse into any nested symbol tables. This function returns None if there
755 /// are any unknown operations that may potentially be symbol tables.
getSymbolUses(StringRef symbol,Operation * from)756 auto SymbolTable::getSymbolUses(StringRef symbol, Operation *from)
757     -> Optional<UseRange> {
758   return getSymbolUsesImpl(symbol, from);
759 }
getSymbolUses(Operation * symbol,Operation * from)760 auto SymbolTable::getSymbolUses(Operation *symbol, Operation *from)
761     -> Optional<UseRange> {
762   return getSymbolUsesImpl(symbol, from);
763 }
getSymbolUses(StringRef symbol,Region * from)764 auto SymbolTable::getSymbolUses(StringRef symbol, Region *from)
765     -> Optional<UseRange> {
766   return getSymbolUsesImpl(symbol, from);
767 }
getSymbolUses(Operation * symbol,Region * from)768 auto SymbolTable::getSymbolUses(Operation *symbol, Region *from)
769     -> Optional<UseRange> {
770   return getSymbolUsesImpl(symbol, from);
771 }
772 
773 //===----------------------------------------------------------------------===//
774 // SymbolTable::symbolKnownUseEmpty
775 
776 /// The implementation of SymbolTable::symbolKnownUseEmpty below.
777 template <typename SymbolT, typename IRUnitT>
symbolKnownUseEmptyImpl(SymbolT symbol,IRUnitT * limit)778 static bool symbolKnownUseEmptyImpl(SymbolT symbol, IRUnitT *limit) {
779   for (SymbolScope &scope : collectSymbolScopes(symbol, limit)) {
780     // Walk all of the symbol uses looking for a reference to 'symbol'.
781     if (scope.walk([&](SymbolTable::SymbolUse symbolUse, ArrayRef<int>) {
782           return isReferencePrefixOf(scope.symbol, symbolUse.getSymbolRef())
783                      ? WalkResult::interrupt()
784                      : WalkResult::advance();
785         }) != WalkResult::advance())
786       return false;
787   }
788   return true;
789 }
790 
791 /// Return if the given symbol is known to have no uses that are nested within
792 /// the given operation 'from'. This does not traverse into any nested symbol
793 /// tables. This function will also return false if there are any unknown
794 /// operations that may potentially be symbol tables.
symbolKnownUseEmpty(StringRef symbol,Operation * from)795 bool SymbolTable::symbolKnownUseEmpty(StringRef symbol, Operation *from) {
796   return symbolKnownUseEmptyImpl(symbol, from);
797 }
symbolKnownUseEmpty(Operation * symbol,Operation * from)798 bool SymbolTable::symbolKnownUseEmpty(Operation *symbol, Operation *from) {
799   return symbolKnownUseEmptyImpl(symbol, from);
800 }
symbolKnownUseEmpty(StringRef symbol,Region * from)801 bool SymbolTable::symbolKnownUseEmpty(StringRef symbol, Region *from) {
802   return symbolKnownUseEmptyImpl(symbol, from);
803 }
symbolKnownUseEmpty(Operation * symbol,Region * from)804 bool SymbolTable::symbolKnownUseEmpty(Operation *symbol, Region *from) {
805   return symbolKnownUseEmptyImpl(symbol, from);
806 }
807 
808 //===----------------------------------------------------------------------===//
809 // SymbolTable::replaceAllSymbolUses
810 
811 /// Rebuild the given attribute container after replacing all references to a
812 /// symbol with the updated attribute in 'accesses'.
rebuildAttrAfterRAUW(Attribute container,ArrayRef<std::pair<SmallVector<int,1>,SymbolRefAttr>> accesses,unsigned depth)813 static Attribute rebuildAttrAfterRAUW(
814     Attribute container,
815     ArrayRef<std::pair<SmallVector<int, 1>, SymbolRefAttr>> accesses,
816     unsigned depth) {
817   // Given a range of Attributes, update the ones referred to by the given
818   // access chains to point to the new symbol attribute.
819   auto updateAttrs = [&](auto &&attrRange) {
820     auto attrBegin = std::begin(attrRange);
821     for (unsigned i = 0, e = accesses.size(); i != e;) {
822       ArrayRef<int> access = accesses[i].first;
823       Attribute &attr = *std::next(attrBegin, access[depth]);
824 
825       // Check to see if this is a leaf access, i.e. a SymbolRef.
826       if (access.size() == depth + 1) {
827         attr = accesses[i].second;
828         ++i;
829         continue;
830       }
831 
832       // Otherwise, this is a container. Collect all of the accesses for this
833       // index and recurse. The recursion here is bounded by the size of the
834       // largest access array.
835       auto nestedAccesses = accesses.drop_front(i).take_while([&](auto &it) {
836         ArrayRef<int> nextAccess = it.first;
837         return nextAccess.size() > depth + 1 &&
838                nextAccess[depth] == access[depth];
839       });
840       attr = rebuildAttrAfterRAUW(attr, nestedAccesses, depth + 1);
841 
842       // Skip over all of the accesses that refer to the nested container.
843       i += nestedAccesses.size();
844     }
845   };
846 
847   if (auto dictAttr = container.dyn_cast<DictionaryAttr>()) {
848     auto newAttrs = llvm::to_vector<4>(dictAttr.getValue());
849     updateAttrs(make_second_range(newAttrs));
850     return DictionaryAttr::get(dictAttr.getContext(), newAttrs);
851   }
852   auto newAttrs = llvm::to_vector<4>(container.cast<ArrayAttr>().getValue());
853   updateAttrs(newAttrs);
854   return ArrayAttr::get(container.getContext(), newAttrs);
855 }
856 
857 /// Generates a new symbol reference attribute with a new leaf reference.
generateNewRefAttr(SymbolRefAttr oldAttr,FlatSymbolRefAttr newLeafAttr)858 static SymbolRefAttr generateNewRefAttr(SymbolRefAttr oldAttr,
859                                         FlatSymbolRefAttr newLeafAttr) {
860   if (oldAttr.isa<FlatSymbolRefAttr>())
861     return newLeafAttr;
862   auto nestedRefs = llvm::to_vector<2>(oldAttr.getNestedReferences());
863   nestedRefs.back() = newLeafAttr;
864   return SymbolRefAttr::get(oldAttr.getContext(), oldAttr.getRootReference(),
865                             nestedRefs);
866 }
867 
868 /// The implementation of SymbolTable::replaceAllSymbolUses below.
869 template <typename SymbolT, typename IRUnitT>
870 static LogicalResult
replaceAllSymbolUsesImpl(SymbolT symbol,StringRef newSymbol,IRUnitT * limit)871 replaceAllSymbolUsesImpl(SymbolT symbol, StringRef newSymbol, IRUnitT *limit) {
872   // A collection of operations along with their new attribute dictionary.
873   std::vector<std::pair<Operation *, DictionaryAttr>> updatedAttrDicts;
874 
875   // The current operation being processed.
876   Operation *curOp = nullptr;
877 
878   // The set of access chains into the attribute dictionary of the current
879   // operation, as well as the replacement attribute to use.
880   SmallVector<std::pair<SmallVector<int, 1>, SymbolRefAttr>, 1> accessChains;
881 
882   // Generate a new attribute dictionary for the current operation by replacing
883   // references to the old symbol.
884   auto generateNewAttrDict = [&] {
885     auto oldDict = curOp->getAttrDictionary();
886     auto newDict = rebuildAttrAfterRAUW(oldDict, accessChains, /*depth=*/0);
887     return newDict.cast<DictionaryAttr>();
888   };
889 
890   // Generate a new attribute to replace the given attribute.
891   MLIRContext *ctx = limit->getContext();
892   FlatSymbolRefAttr newLeafAttr = FlatSymbolRefAttr::get(ctx, newSymbol);
893   for (SymbolScope &scope : collectSymbolScopes(symbol, limit)) {
894     SymbolRefAttr newAttr = generateNewRefAttr(scope.symbol, newLeafAttr);
895     auto walkFn = [&](SymbolTable::SymbolUse symbolUse,
896                       ArrayRef<int> accessChain) {
897       SymbolRefAttr useRef = symbolUse.getSymbolRef();
898       if (!isReferencePrefixOf(scope.symbol, useRef))
899         return WalkResult::advance();
900 
901       // If we have a valid match, check to see if this is a proper
902       // subreference. If it is, then we will need to generate a different new
903       // attribute specifically for this use.
904       SymbolRefAttr replacementRef = newAttr;
905       if (useRef != scope.symbol) {
906         if (scope.symbol.isa<FlatSymbolRefAttr>()) {
907           replacementRef =
908               SymbolRefAttr::get(ctx, newSymbol, useRef.getNestedReferences());
909         } else {
910           auto nestedRefs = llvm::to_vector<4>(useRef.getNestedReferences());
911           nestedRefs[scope.symbol.getNestedReferences().size() - 1] =
912               newLeafAttr;
913           replacementRef =
914               SymbolRefAttr::get(ctx, useRef.getRootReference(), nestedRefs);
915         }
916       }
917 
918       // If there was a previous operation, generate a new attribute dict
919       // for it. This means that we've finished processing the current
920       // operation, so generate a new dictionary for it.
921       if (curOp && symbolUse.getUser() != curOp) {
922         updatedAttrDicts.push_back({curOp, generateNewAttrDict()});
923         accessChains.clear();
924       }
925 
926       // Record this access.
927       curOp = symbolUse.getUser();
928       accessChains.push_back({llvm::to_vector<1>(accessChain), replacementRef});
929       return WalkResult::advance();
930     };
931     if (!scope.walk(walkFn))
932       return failure();
933 
934     // Check to see if we have a dangling op that needs to be processed.
935     if (curOp) {
936       updatedAttrDicts.push_back({curOp, generateNewAttrDict()});
937       curOp = nullptr;
938     }
939   }
940 
941   // Update the attribute dictionaries as necessary.
942   for (auto &it : updatedAttrDicts)
943     it.first->setAttrs(it.second);
944   return success();
945 }
946 
947 /// Attempt to replace all uses of the given symbol 'oldSymbol' with the
948 /// provided symbol 'newSymbol' that are nested within the given operation
949 /// 'from'. This does not traverse into any nested symbol tables. If there are
950 /// any unknown operations that may potentially be symbol tables, no uses are
951 /// replaced and failure is returned.
replaceAllSymbolUses(StringRef oldSymbol,StringRef newSymbol,Operation * from)952 LogicalResult SymbolTable::replaceAllSymbolUses(StringRef oldSymbol,
953                                                 StringRef newSymbol,
954                                                 Operation *from) {
955   return replaceAllSymbolUsesImpl(oldSymbol, newSymbol, from);
956 }
replaceAllSymbolUses(Operation * oldSymbol,StringRef newSymbol,Operation * from)957 LogicalResult SymbolTable::replaceAllSymbolUses(Operation *oldSymbol,
958                                                 StringRef newSymbol,
959                                                 Operation *from) {
960   return replaceAllSymbolUsesImpl(oldSymbol, newSymbol, from);
961 }
replaceAllSymbolUses(StringRef oldSymbol,StringRef newSymbol,Region * from)962 LogicalResult SymbolTable::replaceAllSymbolUses(StringRef oldSymbol,
963                                                 StringRef newSymbol,
964                                                 Region *from) {
965   return replaceAllSymbolUsesImpl(oldSymbol, newSymbol, from);
966 }
replaceAllSymbolUses(Operation * oldSymbol,StringRef newSymbol,Region * from)967 LogicalResult SymbolTable::replaceAllSymbolUses(Operation *oldSymbol,
968                                                 StringRef newSymbol,
969                                                 Region *from) {
970   return replaceAllSymbolUsesImpl(oldSymbol, newSymbol, from);
971 }
972 
973 //===----------------------------------------------------------------------===//
974 // SymbolTableCollection
975 //===----------------------------------------------------------------------===//
976 
lookupSymbolIn(Operation * symbolTableOp,StringRef symbol)977 Operation *SymbolTableCollection::lookupSymbolIn(Operation *symbolTableOp,
978                                                  StringRef symbol) {
979   return getSymbolTable(symbolTableOp).lookup(symbol);
980 }
lookupSymbolIn(Operation * symbolTableOp,SymbolRefAttr name)981 Operation *SymbolTableCollection::lookupSymbolIn(Operation *symbolTableOp,
982                                                  SymbolRefAttr name) {
983   SmallVector<Operation *, 4> symbols;
984   if (failed(lookupSymbolIn(symbolTableOp, name, symbols)))
985     return nullptr;
986   return symbols.back();
987 }
988 /// A variant of 'lookupSymbolIn' that returns all of the symbols referenced by
989 /// a given SymbolRefAttr. Returns failure if any of the nested references could
990 /// not be resolved.
991 LogicalResult
lookupSymbolIn(Operation * symbolTableOp,SymbolRefAttr name,SmallVectorImpl<Operation * > & symbols)992 SymbolTableCollection::lookupSymbolIn(Operation *symbolTableOp,
993                                       SymbolRefAttr name,
994                                       SmallVectorImpl<Operation *> &symbols) {
995   auto lookupFn = [this](Operation *symbolTableOp, StringRef symbol) {
996     return lookupSymbolIn(symbolTableOp, symbol);
997   };
998   return lookupSymbolInImpl(symbolTableOp, name, symbols, lookupFn);
999 }
1000 
1001 /// Returns the operation registered with the given symbol name within the
1002 /// closest parent operation of, or including, 'from' with the
1003 /// 'OpTrait::SymbolTable' trait. Returns nullptr if no valid symbol was
1004 /// found.
lookupNearestSymbolFrom(Operation * from,StringRef symbol)1005 Operation *SymbolTableCollection::lookupNearestSymbolFrom(Operation *from,
1006                                                           StringRef symbol) {
1007   Operation *symbolTableOp = SymbolTable::getNearestSymbolTable(from);
1008   return symbolTableOp ? lookupSymbolIn(symbolTableOp, symbol) : nullptr;
1009 }
1010 Operation *
lookupNearestSymbolFrom(Operation * from,SymbolRefAttr symbol)1011 SymbolTableCollection::lookupNearestSymbolFrom(Operation *from,
1012                                                SymbolRefAttr symbol) {
1013   Operation *symbolTableOp = SymbolTable::getNearestSymbolTable(from);
1014   return symbolTableOp ? lookupSymbolIn(symbolTableOp, symbol) : nullptr;
1015 }
1016 
1017 /// Lookup, or create, a symbol table for an operation.
getSymbolTable(Operation * op)1018 SymbolTable &SymbolTableCollection::getSymbolTable(Operation *op) {
1019   auto it = symbolTables.try_emplace(op, nullptr);
1020   if (it.second)
1021     it.first->second = std::make_unique<SymbolTable>(op);
1022   return *it.first->second;
1023 }
1024 
1025 //===----------------------------------------------------------------------===//
1026 // SymbolUserMap
1027 //===----------------------------------------------------------------------===//
1028 
SymbolUserMap(SymbolTableCollection & symbolTable,Operation * symbolTableOp)1029 SymbolUserMap::SymbolUserMap(SymbolTableCollection &symbolTable,
1030                              Operation *symbolTableOp)
1031     : symbolTable(symbolTable) {
1032   // Walk each of the symbol tables looking for discardable callgraph nodes.
1033   SmallVector<Operation *> symbols;
1034   auto walkFn = [&](Operation *symbolTableOp, bool allUsesVisible) {
1035     for (Operation &nestedOp : symbolTableOp->getRegion(0).getOps()) {
1036       auto symbolUses = SymbolTable::getSymbolUses(&nestedOp);
1037       assert(symbolUses && "expected uses to be valid");
1038 
1039       for (const SymbolTable::SymbolUse &use : *symbolUses) {
1040         symbols.clear();
1041         (void)symbolTable.lookupSymbolIn(symbolTableOp, use.getSymbolRef(),
1042                                          symbols);
1043         for (Operation *symbolOp : symbols)
1044           symbolToUsers[symbolOp].insert(use.getUser());
1045       }
1046     }
1047   };
1048   // We just set `allSymUsesVisible` to false here because it isn't necessary
1049   // for building the user map.
1050   SymbolTable::walkSymbolTables(symbolTableOp, /*allSymUsesVisible=*/false,
1051                                 walkFn);
1052 }
1053 
replaceAllUsesWith(Operation * symbol,StringRef newSymbolName)1054 void SymbolUserMap::replaceAllUsesWith(Operation *symbol,
1055                                        StringRef newSymbolName) {
1056   auto it = symbolToUsers.find(symbol);
1057   if (it == symbolToUsers.end())
1058     return;
1059   SetVector<Operation *> &users = it->second;
1060 
1061   // Replace the uses within the users of `symbol`.
1062   for (Operation *user : users)
1063     (void)SymbolTable::replaceAllSymbolUses(symbol, newSymbolName, user);
1064 
1065   // Move the current users of `symbol` to the new symbol if it is in the
1066   // symbol table.
1067   Operation *newSymbol =
1068       symbolTable.lookupSymbolIn(symbol->getParentOp(), newSymbolName);
1069   if (newSymbol != symbol) {
1070     // Transfer over the users to the new symbol.
1071     auto newIt = symbolToUsers.find(newSymbol);
1072     if (newIt == symbolToUsers.end())
1073       symbolToUsers.try_emplace(newSymbol, std::move(users));
1074     else
1075       newIt->second.set_union(users);
1076     symbolToUsers.erase(symbol);
1077   }
1078 }
1079 
1080 //===----------------------------------------------------------------------===//
1081 // Visibility parsing implementation.
1082 //===----------------------------------------------------------------------===//
1083 
parseOptionalVisibilityKeyword(OpAsmParser & parser,NamedAttrList & attrs)1084 ParseResult impl::parseOptionalVisibilityKeyword(OpAsmParser &parser,
1085                                                  NamedAttrList &attrs) {
1086   StringRef visibility;
1087   if (parser.parseOptionalKeyword(&visibility, {"public", "private", "nested"}))
1088     return failure();
1089 
1090   StringAttr visibilityAttr = parser.getBuilder().getStringAttr(visibility);
1091   attrs.push_back(parser.getBuilder().getNamedAttr(
1092       SymbolTable::getVisibilityAttrName(), visibilityAttr));
1093   return success();
1094 }
1095 
1096 //===----------------------------------------------------------------------===//
1097 // Symbol Interfaces
1098 //===----------------------------------------------------------------------===//
1099 
1100 /// Include the generated symbol interfaces.
1101 #include "mlir/IR/SymbolInterfaces.cpp.inc"
1102