1 //===- Visitors.cpp - MLIR Visitor Utilities ------------------------------===//
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/Visitors.h"
10 #include "mlir/IR/Operation.h"
11 
12 using namespace mlir;
13 
14 /// Walk all of the regions/blocks/operations nested under and including the
15 /// given operation.
16 void detail::walk(Operation *op, function_ref<void(Region *)> callback) {
17   for (auto &region : op->getRegions()) {
18     callback(&region);
19     for (auto &block : region) {
20       for (auto &nestedOp : block)
21         walk(&nestedOp, callback);
22     }
23   }
24 }
25 
26 void detail::walk(Operation *op, function_ref<void(Block *)> callback) {
27   for (auto &region : op->getRegions()) {
28     for (auto &block : region) {
29       callback(&block);
30       for (auto &nestedOp : block)
31         walk(&nestedOp, callback);
32     }
33   }
34 }
35 
36 void detail::walk(Operation *op, function_ref<void(Operation *op)> callback) {
37   // TODO: This walk should be iterative over the operations.
38   for (auto &region : op->getRegions()) {
39     for (auto &block : region) {
40       // Early increment here in the case where the operation is erased.
41       for (auto &nestedOp : llvm::make_early_inc_range(block))
42         walk(&nestedOp, callback);
43     }
44   }
45   callback(op);
46 }
47 
48 /// Walk all of the regions/blocks/operations nested under and including the
49 /// given operation. These functions walk operations until an interrupt result
50 /// is returned by the callback.
51 WalkResult detail::walk(Operation *op,
52                         function_ref<WalkResult(Region *op)> callback) {
53   for (auto &region : op->getRegions()) {
54     if (callback(&region).wasInterrupted())
55       return WalkResult::interrupt();
56     for (auto &block : region) {
57       for (auto &nestedOp : block)
58         walk(&nestedOp, callback);
59     }
60   }
61   return WalkResult::advance();
62 }
63 
64 WalkResult detail::walk(Operation *op,
65                         function_ref<WalkResult(Block *op)> callback) {
66   for (auto &region : op->getRegions()) {
67     for (auto &block : region) {
68       if (callback(&block).wasInterrupted())
69         return WalkResult::interrupt();
70       for (auto &nestedOp : block)
71         walk(&nestedOp, callback);
72     }
73   }
74   return WalkResult::advance();
75 }
76 
77 WalkResult detail::walk(Operation *op,
78                         function_ref<WalkResult(Operation *op)> callback) {
79   // TODO: This walk should be iterative over the operations.
80   for (auto &region : op->getRegions()) {
81     for (auto &block : region) {
82       // Early increment here in the case where the operation is erased.
83       for (auto &nestedOp : llvm::make_early_inc_range(block)) {
84         if (walk(&nestedOp, callback).wasInterrupted())
85           return WalkResult::interrupt();
86       }
87     }
88   }
89   return callback(op);
90 }
91