1//===- trie-node.h - XRay Call Stack Data Structure -----------------------===//
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// This file provides a data structure and routines for working with call stacks
10// of instrumented functions.
11//
12//===----------------------------------------------------------------------===//
13
14#ifndef LLVM_TOOLS_LLVM_XRAY_STACK_TRIE_H
15#define LLVM_TOOLS_LLVM_XRAY_STACK_TRIE_H
16
17#include <forward_list>
18#include <numeric>
19
20#include "llvm/ADT/DenseMap.h"
21#include "llvm/ADT/STLExtras.h"
22#include "llvm/ADT/SmallVector.h"
23
24/// A type to represent a trie of invocations. It is useful to construct a
25/// graph of these nodes from reading an XRay trace, such that each function
26/// call can be placed in a larger context.
27///
28/// The template parameter allows users of the template to attach their own
29/// data elements to each node in the invocation graph.
30template <typename AssociatedData> struct TrieNode {
31  /// The function ID.
32  int32_t FuncId;
33
34  /// The caller of this function.
35  TrieNode<AssociatedData> *Parent;
36
37  /// The callees from this function.
38  llvm::SmallVector<TrieNode<AssociatedData> *, 4> Callees;
39
40  /// Additional parameterized data on each node.
41  AssociatedData ExtraData;
42};
43
44/// Merges together two TrieNodes with like function ids, aggregating their
45/// callee lists and durations. The caller must provide storage where new merged
46/// nodes can be allocated in the form of a linked list.
47template <typename T, typename Callable>
48TrieNode<T> *
49mergeTrieNodes(const TrieNode<T> &Left, const TrieNode<T> &Right,
50               /*Non-deduced pointer type for nullptr compatibility*/
51               std::remove_reference_t<TrieNode<T> *> NewParent,
52               std::forward_list<TrieNode<T>> &NodeStore,
53               Callable &&MergeCallable) {
54  llvm::function_ref<T(const T &, const T &)> MergeFn(
55      std::forward<Callable>(MergeCallable));
56  assert(Left.FuncId == Right.FuncId);
57  NodeStore.push_front(TrieNode<T>{
58      Left.FuncId, NewParent, {}, MergeFn(Left.ExtraData, Right.ExtraData)});
59  auto I = NodeStore.begin();
60  auto *Node = &*I;
61
62  // Build a map of callees from the left side.
63  llvm::DenseMap<int32_t, TrieNode<T> *> LeftCalleesByFuncId;
64  for (auto *Callee : Left.Callees) {
65    LeftCalleesByFuncId[Callee->FuncId] = Callee;
66  }
67
68  // Iterate through the right side, either merging with the map values or
69  // directly adding to the Callees vector. The iteration also removes any
70  // merged values from the left side map.
71  // TODO: Unroll into iterative and explicit stack for efficiency.
72  for (auto *Callee : Right.Callees) {
73    auto iter = LeftCalleesByFuncId.find(Callee->FuncId);
74    if (iter != LeftCalleesByFuncId.end()) {
75      Node->Callees.push_back(
76          mergeTrieNodes(*(iter->second), *Callee, Node, NodeStore, MergeFn));
77      LeftCalleesByFuncId.erase(iter);
78    } else {
79      Node->Callees.push_back(Callee);
80    }
81  }
82
83  // Add any callees that weren't found in the right side.
84  for (auto MapPairIter : LeftCalleesByFuncId) {
85    Node->Callees.push_back(MapPairIter.second);
86  }
87
88  return Node;
89}
90
91#endif // LLVM_TOOLS_LLVM_XRAY_STACK_TRIE_H
92