1 /*
2  * Copyright (C) 2024 The Android Open Source Project
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 #include "src/trace_processor/perfetto_sql/intrinsics/functions/graph_traversal.h"
18 
19 #include <algorithm>
20 #include <cstddef>
21 #include <cstdint>
22 #include <memory>
23 #include <optional>
24 #include <utility>
25 #include <vector>
26 
27 #include "perfetto/base/logging.h"
28 #include "perfetto/base/status.h"
29 #include "perfetto/ext/base/circular_queue.h"
30 #include "perfetto/public/compiler.h"
31 #include "src/trace_processor/containers/string_pool.h"
32 #include "src/trace_processor/perfetto_sql/engine/perfetto_sql_engine.h"
33 #include "src/trace_processor/perfetto_sql/intrinsics/functions/tables_py.h"
34 #include "src/trace_processor/perfetto_sql/intrinsics/types/array.h"
35 #include "src/trace_processor/perfetto_sql/intrinsics/types/node.h"
36 #include "src/trace_processor/sqlite/bindings/sqlite_aggregate_function.h"
37 #include "src/trace_processor/sqlite/bindings/sqlite_result.h"
38 #include "src/trace_processor/sqlite/bindings/sqlite_value.h"
39 #include "src/trace_processor/util/status_macros.h"
40 
41 namespace perfetto::trace_processor {
42 namespace tables {
43 TreeTable::~TreeTable() = default;
44 }  // namespace tables
45 
46 namespace {
47 
48 struct State {
49   uint32_t id;
50   std::optional<uint32_t> parent_id;
51 };
52 
53 // An SQL aggregate-function which performs a DFS from a given start node in a
54 // graph and returns all the nodes which are reachable from the start node.
55 //
56 // Note: this function is not intended to be used directly from SQL: instead
57 // macros exist in the standard library, wrapping it and making it
58 // user-friendly.
59 struct Dfs : public SqliteAggregateFunction<Dfs> {
60   static constexpr char kName[] = "__intrinsic_dfs";
61   static constexpr int kArgCount = 2;
62   using UserDataContext = StringPool;
63 
Stepperfetto::trace_processor::__anon37a616900111::Dfs64   static void Step(sqlite3_context* ctx, int argc, sqlite3_value** argv) {
65     PERFETTO_DCHECK(argc == kArgCount);
66 
67     auto* graph = sqlite::value::Pointer<perfetto_sql::Graph>(argv[0], "GRAPH");
68     auto table = std::make_unique<tables::TreeTable>(GetUserData(ctx));
69     if (!graph) {
70       return sqlite::result::UniquePointer(ctx, std::move(table), "TABLE");
71     }
72     PERFETTO_DCHECK(!graph->empty());
73 
74     // If the array is empty, be forgiving and return an empty array. We could
75     // return an error here but in 99% of cases, the caller will simply want
76     // an empty table instead.
77     auto* start_ids =
78         sqlite::value::Pointer<perfetto_sql::IntArray>(argv[1], "ARRAY<LONG>");
79     if (!start_ids) {
80       return sqlite::result::UniquePointer(ctx, std::move(table), "TABLE");
81     }
82     PERFETTO_DCHECK(!start_ids->empty());
83 
84     std::vector<bool> visited(graph->size());
85     std::vector<State> stack;
86     for (int64_t x : *start_ids) {
87       stack.emplace_back(State{static_cast<uint32_t>(x), std::nullopt});
88     }
89     while (!stack.empty()) {
90       State state = stack.back();
91       stack.pop_back();
92 
93       auto& node = (*graph)[state.id];
94       if (visited[state.id]) {
95         continue;
96       }
97       table->Insert({state.id, state.parent_id});
98       visited[state.id] = true;
99 
100       const auto& children = node.outgoing_edges;
101       for (auto it = children.rbegin(); it != children.rend(); ++it) {
102         stack.emplace_back(State{*it, state.id});
103       }
104     }
105     return sqlite::result::UniquePointer(ctx, std::move(table), "TABLE");
106   }
107 };
108 
109 // An SQL aggregate-function which performs a BFS from a given start node in a
110 // graph and returns all the nodes which are reachable from the start node.
111 //
112 // Note: this function is not intended to be used directly from SQL: instead
113 // macros exist in the standard library, wrapping it and making it
114 // user-friendly.
115 struct Bfs : public SqliteAggregateFunction<Bfs> {
116   static constexpr char kName[] = "__intrinsic_bfs";
117   static constexpr int kArgCount = 2;
118   using UserDataContext = StringPool;
119 
Stepperfetto::trace_processor::__anon37a616900111::Bfs120   static void Step(sqlite3_context* ctx, int argc, sqlite3_value** argv) {
121     PERFETTO_DCHECK(argc == kArgCount);
122 
123     auto* graph = sqlite::value::Pointer<perfetto_sql::Graph>(argv[0], "GRAPH");
124     auto table = std::make_unique<tables::TreeTable>(GetUserData(ctx));
125     if (!graph) {
126       return sqlite::result::UniquePointer(ctx, std::move(table), "TABLE");
127     }
128     PERFETTO_DCHECK(!graph->empty());
129 
130     // If the array is empty, be forgiving and return an empty array. We could
131     // return an error here but in 99% of cases, the caller will simply want
132     // an empty table instead.
133     auto* start_ids =
134         sqlite::value::Pointer<perfetto_sql::IntArray>(argv[1], "ARRAY<LONG>");
135     if (!start_ids) {
136       return sqlite::result::UniquePointer(ctx, std::move(table), "TABLE");
137     }
138     PERFETTO_DCHECK(!start_ids->empty());
139 
140     std::vector<bool> visited(graph->size());
141     base::CircularQueue<State> queue;
142     for (int64_t raw_id : *start_ids) {
143       auto id = static_cast<uint32_t>(raw_id);
144       if (id >= graph->size() || visited[id]) {
145         continue;
146       }
147       visited[id] = true;
148       queue.emplace_back(State{id, std::nullopt});
149     }
150     while (!queue.empty()) {
151       State state = queue.front();
152       queue.pop_front();
153       table->Insert({state.id, state.parent_id});
154 
155       auto& node = (*graph)[state.id];
156       for (uint32_t n : node.outgoing_edges) {
157         if (visited[n]) {
158           continue;
159         }
160         visited[n] = true;
161         queue.emplace_back(State{n, state.id});
162       }
163     }
164     return sqlite::result::UniquePointer(ctx, std::move(table), "TABLE");
165   }
166 };
167 
168 }  // namespace
169 
RegisterGraphTraversalFunctions(PerfettoSqlEngine & engine,StringPool & pool)170 base::Status RegisterGraphTraversalFunctions(PerfettoSqlEngine& engine,
171                                              StringPool& pool) {
172   RETURN_IF_ERROR(engine.RegisterSqliteFunction<Dfs>(&pool));
173   return engine.RegisterSqliteFunction<Bfs>(&pool);
174 }
175 
176 }  // namespace perfetto::trace_processor
177