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
19using namespace testing;
20
21namespace {
22
23struct Data {
24 std::mutex mutex;
25 std::vector<std::string> order;
26
27 void push(std::string&& s) {
28 std::unique_lock<std::mutex> lock(mutex);
29 order.emplace_back(std::move(s));
30 }
31};
32
33template <typename T>
34std::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] |
41TEST_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] |
57TEST_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] |
74TEST_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] |
96TEST_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] |
115TEST_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] --/ |
138TEST_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