1 // Copyright 2020 The Marl Authors.
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //     https://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 #include "marl/dag.h"
16 
17 #include "marl_test.h"
18 
19 using namespace testing;
20 
21 namespace {
22 
23 struct Data {
24   std::mutex mutex;
25   std::vector<std::string> order;
26 
push__anonfd15334b0111::Data27   void push(std::string&& s) {
28     std::unique_lock<std::mutex> lock(mutex);
29     order.emplace_back(std::move(s));
30   }
31 };
32 
33 template <typename T>
slice(const std::vector<T> & in,size_t from,size_t to)34 std::vector<T> slice(const std::vector<T>& in, size_t from, size_t to) {
35   return {in.begin() + from, in.begin() + to};
36 }
37 
38 }  // namespace
39 
40 //  [A] --> [B] --> [C]                                                        |
TEST_P(WithBoundScheduler,DAGChainNoArg)41 TEST_P(WithBoundScheduler, DAGChainNoArg) {
42   marl::DAG<>::Builder builder;
43 
44   Data data;
45   builder.root()
46       .then([&] { data.push("A"); })
47       .then([&] { data.push("B"); })
48       .then([&] { data.push("C"); });
49 
50   auto dag = builder.build();
51   dag->run();
52 
53   ASSERT_THAT(data.order, ElementsAre("A", "B", "C"));
54 }
55 
56 //  [A] --> [B] --> [C]                                                        |
TEST_P(WithBoundScheduler,DAGChain)57 TEST_P(WithBoundScheduler, DAGChain) {
58   marl::DAG<Data&>::Builder builder;
59 
60   builder.root()
61       .then([](Data& data) { data.push("A"); })
62       .then([](Data& data) { data.push("B"); })
63       .then([](Data& data) { data.push("C"); });
64 
65   auto dag = builder.build();
66 
67   Data data;
68   dag->run(data);
69 
70   ASSERT_THAT(data.order, ElementsAre("A", "B", "C"));
71 }
72 
73 //  [A] --> [B] --> [C]                                                        |
TEST_P(WithBoundScheduler,DAGRunRepeat)74 TEST_P(WithBoundScheduler, DAGRunRepeat) {
75   marl::DAG<Data&>::Builder builder;
76 
77   builder.root()
78       .then([](Data& data) { data.push("A"); })
79       .then([](Data& data) { data.push("B"); })
80       .then([](Data& data) { data.push("C"); });
81 
82   auto dag = builder.build();
83 
84   Data dataA, dataB;
85   dag->run(dataA);
86   dag->run(dataB);
87   dag->run(dataA);
88 
89   ASSERT_THAT(dataA.order, ElementsAre("A", "B", "C", "A", "B", "C"));
90   ASSERT_THAT(dataB.order, ElementsAre("A", "B", "C"));
91 }
92 
93 //           /--> [A]                                                          |
94 //  [root] --|--> [B]                                                          |
95 //           \--> [C]                                                          |
TEST_P(WithBoundScheduler,DAGFanOutFromRoot)96 TEST_P(WithBoundScheduler, DAGFanOutFromRoot) {
97   marl::DAG<Data&>::Builder builder;
98 
99   auto root = builder.root();
100   root.then([](Data& data) { data.push("A"); });
101   root.then([](Data& data) { data.push("B"); });
102   root.then([](Data& data) { data.push("C"); });
103 
104   auto dag = builder.build();
105 
106   Data data;
107   dag->run(data);
108 
109   ASSERT_THAT(data.order, UnorderedElementsAre("A", "B", "C"));
110 }
111 
112 //                /--> [A]                                                     |
113 // [root] -->[N]--|--> [B]                                                     |
114 //                \--> [C]                                                     |
TEST_P(WithBoundScheduler,DAGFanOutFromNonRoot)115 TEST_P(WithBoundScheduler, DAGFanOutFromNonRoot) {
116   marl::DAG<Data&>::Builder builder;
117 
118   auto root = builder.root();
119   auto node = root.then([](Data& data) { data.push("N"); });
120   node.then([](Data& data) { data.push("A"); });
121   node.then([](Data& data) { data.push("B"); });
122   node.then([](Data& data) { data.push("C"); });
123 
124   auto dag = builder.build();
125 
126   Data data;
127   dag->run(data);
128 
129   ASSERT_THAT(data.order, UnorderedElementsAre("N", "A", "B", "C"));
130   ASSERT_EQ(data.order[0], "N");
131   ASSERT_THAT(slice(data.order, 1, 4), UnorderedElementsAre("A", "B", "C"));
132 }
133 
134 //          /--> [A0] --\        /--> [C0] --\        /--> [E0] --\            |
135 // [root] --|--> [A1] --|-->[B]--|--> [C1] --|-->[D]--|--> [E1] --|-->[F]      |
136 //                               \--> [C2] --/        |--> [E2] --|            |
137 //                                                    \--> [E3] --/            |
TEST_P(WithBoundScheduler,DAGFanOutFanIn)138 TEST_P(WithBoundScheduler, DAGFanOutFanIn) {
139   marl::DAG<Data&>::Builder builder;
140 
141   auto root = builder.root();
142   auto a0 = root.then([](Data& data) { data.push("A0"); });
143   auto a1 = root.then([](Data& data) { data.push("A1"); });
144 
145   auto b = builder.node([](Data& data) { data.push("B"); }, {a0, a1});
146 
147   auto c0 = b.then([](Data& data) { data.push("C0"); });
148   auto c1 = b.then([](Data& data) { data.push("C1"); });
149   auto c2 = b.then([](Data& data) { data.push("C2"); });
150 
151   auto d = builder.node([](Data& data) { data.push("D"); }, {c0, c1, c2});
152 
153   auto e0 = d.then([](Data& data) { data.push("E0"); });
154   auto e1 = d.then([](Data& data) { data.push("E1"); });
155   auto e2 = d.then([](Data& data) { data.push("E2"); });
156   auto e3 = d.then([](Data& data) { data.push("E3"); });
157 
158   builder.node([](Data& data) { data.push("F"); }, {e0, e1, e2, e3});
159 
160   auto dag = builder.build();
161 
162   Data data;
163   dag->run(data);
164 
165   ASSERT_THAT(data.order,
166               UnorderedElementsAre("A0", "A1", "B", "C0", "C1", "C2", "D", "E0",
167                                    "E1", "E2", "E3", "F"));
168   ASSERT_THAT(slice(data.order, 0, 2), UnorderedElementsAre("A0", "A1"));
169   ASSERT_THAT(data.order[2], "B");
170   ASSERT_THAT(slice(data.order, 3, 6), UnorderedElementsAre("C0", "C1", "C2"));
171   ASSERT_THAT(data.order[6], "D");
172   ASSERT_THAT(slice(data.order, 7, 11),
173               UnorderedElementsAre("E0", "E1", "E2", "E3"));
174   ASSERT_THAT(data.order[11], "F");
175 }
176