1 //===- SubElementInterfaces.cpp - Attr and Type SubElement Interfaces -----===//
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/SubElementInterfaces.h"
10 
11 #include "llvm/ADT/DenseSet.h"
12 
13 using namespace mlir;
14 
15 template <typename InterfaceT>
16 static void walkSubElementsImpl(InterfaceT interface,
17                                 function_ref<void(Attribute)> walkAttrsFn,
18                                 function_ref<void(Type)> walkTypesFn,
19                                 DenseSet<Attribute> &visitedAttrs,
20                                 DenseSet<Type> &visitedTypes) {
21   interface.walkImmediateSubElements(
22       [&](Attribute attr) {
23         // Guard against potentially null inputs. This removes the need for the
24         // derived attribute/type to do it.
25         if (!attr)
26           return;
27 
28         // Avoid infinite recursion when visiting sub attributes later, if this
29         // is a mutable attribute.
30         if (LLVM_UNLIKELY(attr.hasTrait<AttributeTrait::IsMutable>())) {
31           if (!visitedAttrs.insert(attr).second)
32             return;
33         }
34 
35         // Walk any sub elements first.
36         if (auto interface = attr.dyn_cast<SubElementAttrInterface>())
37           walkSubElementsImpl(interface, walkAttrsFn, walkTypesFn, visitedAttrs,
38                               visitedTypes);
39 
40         // Walk this attribute.
41         walkAttrsFn(attr);
42       },
43       [&](Type type) {
44         // Guard against potentially null inputs. This removes the need for the
45         // derived attribute/type to do it.
46         if (!type)
47           return;
48 
49         // Avoid infinite recursion when visiting sub types later, if this
50         // is a mutable type.
51         if (LLVM_UNLIKELY(type.hasTrait<TypeTrait::IsMutable>())) {
52           if (!visitedTypes.insert(type).second)
53             return;
54         }
55 
56         // Walk any sub elements first.
57         if (auto interface = type.dyn_cast<SubElementTypeInterface>())
58           walkSubElementsImpl(interface, walkAttrsFn, walkTypesFn, visitedAttrs,
59                               visitedTypes);
60 
61         // Walk this type.
62         walkTypesFn(type);
63       });
64 }
65 
66 void SubElementAttrInterface::walkSubElements(
67     function_ref<void(Attribute)> walkAttrsFn,
68     function_ref<void(Type)> walkTypesFn) {
69   assert(walkAttrsFn && walkTypesFn && "expected valid walk functions");
70   DenseSet<Attribute> visitedAttrs;
71   DenseSet<Type> visitedTypes;
72   walkSubElementsImpl(*this, walkAttrsFn, walkTypesFn, visitedAttrs,
73                       visitedTypes);
74 }
75 
76 void SubElementTypeInterface::walkSubElements(
77     function_ref<void(Attribute)> walkAttrsFn,
78     function_ref<void(Type)> walkTypesFn) {
79   assert(walkAttrsFn && walkTypesFn && "expected valid walk functions");
80   DenseSet<Attribute> visitedAttrs;
81   DenseSet<Type> visitedTypes;
82   walkSubElementsImpl(*this, walkAttrsFn, walkTypesFn, visitedAttrs,
83                       visitedTypes);
84 }
85 
86 //===----------------------------------------------------------------------===//
87 // SubElementInterface Tablegen definitions
88 //===----------------------------------------------------------------------===//
89 
90 #include "mlir/IR/SubElementAttrInterfaces.cpp.inc"
91 #include "mlir/IR/SubElementTypeInterfaces.cpp.inc"
92