534ffef80978c310bcbd4ffa6b75645d653a5245
[platform/core/ml/nnfw.git] / runtime / onert / core / src / ir / GraphIterator.h
1 /*
2  * Copyright (c) 2020 Samsung Electronics Co., Ltd. All Rights Reserved
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16
17 #ifndef __ONERT_IR_GRAPH_ITERATOR_H__
18 #define __ONERT_IR_GRAPH_ITERATOR_H__
19
20 #include <type_traits>
21
22 #include "ir/Index.h"
23
24 namespace onert
25 {
26 namespace ir
27 {
28
29 class Graph;
30 class Operation;
31 class LoweredGraph;
32 class OpSequence;
33
34 template <bool is_const> class Iterator
35 {
36 public:
37   using GraphRef = typename std::conditional<is_const, const Graph &, Graph &>::type;
38   using IndexRef = const OperationIndex &;
39   using NodeRef = typename std::conditional<is_const, const Operation &, Operation &>::type;
40   using IterFn = std::function<void(IndexRef, NodeRef)>;
41
42 public:
43   virtual ~Iterator() = default;
44   virtual void iterate(GraphRef graph, const IterFn &fn) const = 0;
45 };
46
47 template <bool is_const = false> class DefaultIterator final : public Iterator<is_const>
48 {
49 public:
50   using GraphRef = typename Iterator<is_const>::GraphRef;
51   using IndexRef = typename Iterator<is_const>::IndexRef;
52   using NodeRef = typename Iterator<is_const>::NodeRef;
53   using IterFn = typename Iterator<is_const>::IterFn;
54
55 public:
56   void iterate(GraphRef graph, const IterFn &fn) const;
57 };
58 using DefaultConstIterator = DefaultIterator<true>;
59
60 template <bool is_const = false> class PostDfsIterator final : public Iterator<is_const>
61 {
62 public:
63   using GraphRef = typename Iterator<is_const>::GraphRef;
64   using IndexRef = typename Iterator<is_const>::IndexRef;
65   using NodeRef = typename Iterator<is_const>::NodeRef;
66   using IterFn = typename Iterator<is_const>::IterFn;
67   using LoweredGraphRef =
68       typename std::conditional<is_const, const LoweredGraph &, LoweredGraph &>::type;
69   using OpSequenceRef = typename std::conditional<is_const, const OpSequence &, OpSequence &>::type;
70   using OpSeqIndexRef = const OpSequenceIndex &;
71   using OpSeqIterFn = std::function<void(OpSeqIndexRef, OpSequenceRef)>;
72
73 public:
74   void iterate(GraphRef graph, const IterFn &fn) const;
75   void iterateOpSeqs(LoweredGraphRef lowered_graph, const OpSeqIterFn &f) const;
76 };
77 using PostDfsConstIterator = PostDfsIterator<true>;
78
79 } // namespace ir
80 } // namespace onert
81
82 #endif // __ONERT_IR_GRAPH_ITERATOR_H__