SuffixTree.h 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350
  1. //===- llvm/ADT/SuffixTree.h - Tree for substrings --------------*- C++ -*-===//
  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 defines the Suffix Tree class and Suffix Tree Node struct.
  10. //
  11. //===----------------------------------------------------------------------===//
  12. #ifndef LLVM_SUPPORT_SUFFIXTREE_H
  13. #define LLVM_SUPPORT_SUFFIXTREE_H
  14. #include "llvm/ADT/ArrayRef.h"
  15. #include "llvm/ADT/DenseMap.h"
  16. #include "llvm/Support/Allocator.h"
  17. #include <vector>
  18. namespace llvm {
  19. /// Represents an undefined index in the suffix tree.
  20. const unsigned EmptyIdx = -1;
  21. /// A node in a suffix tree which represents a substring or suffix.
  22. ///
  23. /// Each node has either no children or at least two children, with the root
  24. /// being a exception in the empty tree.
  25. ///
  26. /// Children are represented as a map between unsigned integers and nodes. If
  27. /// a node N has a child M on unsigned integer k, then the mapping represented
  28. /// by N is a proper prefix of the mapping represented by M. Note that this,
  29. /// although similar to a trie is somewhat different: each node stores a full
  30. /// substring of the full mapping rather than a single character state.
  31. ///
  32. /// Each internal node contains a pointer to the internal node representing
  33. /// the same string, but with the first character chopped off. This is stored
  34. /// in \p Link. Each leaf node stores the start index of its respective
  35. /// suffix in \p SuffixIdx.
  36. struct SuffixTreeNode {
  37. /// The children of this node.
  38. ///
  39. /// A child existing on an unsigned integer implies that from the mapping
  40. /// represented by the current node, there is a way to reach another
  41. /// mapping by tacking that character on the end of the current string.
  42. llvm::DenseMap<unsigned, SuffixTreeNode *> Children;
  43. /// The start index of this node's substring in the main string.
  44. unsigned StartIdx = EmptyIdx;
  45. /// The end index of this node's substring in the main string.
  46. ///
  47. /// Every leaf node must have its \p EndIdx incremented at the end of every
  48. /// step in the construction algorithm. To avoid having to update O(N)
  49. /// nodes individually at the end of every step, the end index is stored
  50. /// as a pointer.
  51. unsigned *EndIdx = nullptr;
  52. /// For leaves, the start index of the suffix represented by this node.
  53. ///
  54. /// For all other nodes, this is ignored.
  55. unsigned SuffixIdx = EmptyIdx;
  56. /// For internal nodes, a pointer to the internal node representing
  57. /// the same sequence with the first character chopped off.
  58. ///
  59. /// This acts as a shortcut in Ukkonen's algorithm. One of the things that
  60. /// Ukkonen's algorithm does to achieve linear-time construction is
  61. /// keep track of which node the next insert should be at. This makes each
  62. /// insert O(1), and there are a total of O(N) inserts. The suffix link
  63. /// helps with inserting children of internal nodes.
  64. ///
  65. /// Say we add a child to an internal node with associated mapping S. The
  66. /// next insertion must be at the node representing S - its first character.
  67. /// This is given by the way that we iteratively build the tree in Ukkonen's
  68. /// algorithm. The main idea is to look at the suffixes of each prefix in the
  69. /// string, starting with the longest suffix of the prefix, and ending with
  70. /// the shortest. Therefore, if we keep pointers between such nodes, we can
  71. /// move to the next insertion point in O(1) time. If we don't, then we'd
  72. /// have to query from the root, which takes O(N) time. This would make the
  73. /// construction algorithm O(N^2) rather than O(N).
  74. SuffixTreeNode *Link = nullptr;
  75. /// The length of the string formed by concatenating the edge labels from the
  76. /// root to this node.
  77. unsigned ConcatLen = 0;
  78. /// Returns true if this node is a leaf.
  79. bool isLeaf() const { return SuffixIdx != EmptyIdx; }
  80. /// Returns true if this node is the root of its owning \p SuffixTree.
  81. bool isRoot() const { return StartIdx == EmptyIdx; }
  82. /// Return the number of elements in the substring associated with this node.
  83. size_t size() const {
  84. // Is it the root? If so, it's the empty string so return 0.
  85. if (isRoot())
  86. return 0;
  87. assert(*EndIdx != EmptyIdx && "EndIdx is undefined!");
  88. // Size = the number of elements in the string.
  89. // For example, [0 1 2 3] has length 4, not 3. 3-0 = 3, so we have 3-0+1.
  90. return *EndIdx - StartIdx + 1;
  91. }
  92. SuffixTreeNode(unsigned StartIdx, unsigned *EndIdx, SuffixTreeNode *Link)
  93. : StartIdx(StartIdx), EndIdx(EndIdx), Link(Link) {}
  94. SuffixTreeNode() {}
  95. };
  96. /// A data structure for fast substring queries.
  97. ///
  98. /// Suffix trees represent the suffixes of their input strings in their leaves.
  99. /// A suffix tree is a type of compressed trie structure where each node
  100. /// represents an entire substring rather than a single character. Each leaf
  101. /// of the tree is a suffix.
  102. ///
  103. /// A suffix tree can be seen as a type of state machine where each state is a
  104. /// substring of the full string. The tree is structured so that, for a string
  105. /// of length N, there are exactly N leaves in the tree. This structure allows
  106. /// us to quickly find repeated substrings of the input string.
  107. ///
  108. /// In this implementation, a "string" is a vector of unsigned integers.
  109. /// These integers may result from hashing some data type. A suffix tree can
  110. /// contain 1 or many strings, which can then be queried as one large string.
  111. ///
  112. /// The suffix tree is implemented using Ukkonen's algorithm for linear-time
  113. /// suffix tree construction. Ukkonen's algorithm is explained in more detail
  114. /// in the paper by Esko Ukkonen "On-line construction of suffix trees. The
  115. /// paper is available at
  116. ///
  117. /// https://www.cs.helsinki.fi/u/ukkonen/SuffixT1withFigs.pdf
  118. class SuffixTree {
  119. public:
  120. /// Each element is an integer representing an instruction in the module.
  121. llvm::ArrayRef<unsigned> Str;
  122. /// A repeated substring in the tree.
  123. struct RepeatedSubstring {
  124. /// The length of the string.
  125. unsigned Length;
  126. /// The start indices of each occurrence.
  127. std::vector<unsigned> StartIndices;
  128. };
  129. private:
  130. /// Maintains each node in the tree.
  131. llvm::SpecificBumpPtrAllocator<SuffixTreeNode> NodeAllocator;
  132. /// The root of the suffix tree.
  133. ///
  134. /// The root represents the empty string. It is maintained by the
  135. /// \p NodeAllocator like every other node in the tree.
  136. SuffixTreeNode *Root = nullptr;
  137. /// Maintains the end indices of the internal nodes in the tree.
  138. ///
  139. /// Each internal node is guaranteed to never have its end index change
  140. /// during the construction algorithm; however, leaves must be updated at
  141. /// every step. Therefore, we need to store leaf end indices by reference
  142. /// to avoid updating O(N) leaves at every step of construction. Thus,
  143. /// every internal node must be allocated its own end index.
  144. llvm::BumpPtrAllocator InternalEndIdxAllocator;
  145. /// The end index of each leaf in the tree.
  146. unsigned LeafEndIdx = -1;
  147. /// Helper struct which keeps track of the next insertion point in
  148. /// Ukkonen's algorithm.
  149. struct ActiveState {
  150. /// The next node to insert at.
  151. SuffixTreeNode *Node = nullptr;
  152. /// The index of the first character in the substring currently being added.
  153. unsigned Idx = EmptyIdx;
  154. /// The length of the substring we have to add at the current step.
  155. unsigned Len = 0;
  156. };
  157. /// The point the next insertion will take place at in the
  158. /// construction algorithm.
  159. ActiveState Active;
  160. /// Allocate a leaf node and add it to the tree.
  161. ///
  162. /// \param Parent The parent of this node.
  163. /// \param StartIdx The start index of this node's associated string.
  164. /// \param Edge The label on the edge leaving \p Parent to this node.
  165. ///
  166. /// \returns A pointer to the allocated leaf node.
  167. SuffixTreeNode *insertLeaf(SuffixTreeNode &Parent, unsigned StartIdx,
  168. unsigned Edge);
  169. /// Allocate an internal node and add it to the tree.
  170. ///
  171. /// \param Parent The parent of this node. Only null when allocating the root.
  172. /// \param StartIdx The start index of this node's associated string.
  173. /// \param EndIdx The end index of this node's associated string.
  174. /// \param Edge The label on the edge leaving \p Parent to this node.
  175. ///
  176. /// \returns A pointer to the allocated internal node.
  177. SuffixTreeNode *insertInternalNode(SuffixTreeNode *Parent, unsigned StartIdx,
  178. unsigned EndIdx, unsigned Edge);
  179. /// Set the suffix indices of the leaves to the start indices of their
  180. /// respective suffixes.
  181. void setSuffixIndices();
  182. /// Construct the suffix tree for the prefix of the input ending at
  183. /// \p EndIdx.
  184. ///
  185. /// Used to construct the full suffix tree iteratively. At the end of each
  186. /// step, the constructed suffix tree is either a valid suffix tree, or a
  187. /// suffix tree with implicit suffixes. At the end of the final step, the
  188. /// suffix tree is a valid tree.
  189. ///
  190. /// \param EndIdx The end index of the current prefix in the main string.
  191. /// \param SuffixesToAdd The number of suffixes that must be added
  192. /// to complete the suffix tree at the current phase.
  193. ///
  194. /// \returns The number of suffixes that have not been added at the end of
  195. /// this step.
  196. unsigned extend(unsigned EndIdx, unsigned SuffixesToAdd);
  197. public:
  198. /// Construct a suffix tree from a sequence of unsigned integers.
  199. ///
  200. /// \param Str The string to construct the suffix tree for.
  201. SuffixTree(const std::vector<unsigned> &Str);
  202. /// Iterator for finding all repeated substrings in the suffix tree.
  203. struct RepeatedSubstringIterator {
  204. private:
  205. /// The current node we're visiting.
  206. SuffixTreeNode *N = nullptr;
  207. /// The repeated substring associated with this node.
  208. RepeatedSubstring RS;
  209. /// The nodes left to visit.
  210. std::vector<SuffixTreeNode *> ToVisit;
  211. /// The minimum length of a repeated substring to find.
  212. /// Since we're outlining, we want at least two instructions in the range.
  213. /// FIXME: This may not be true for targets like X86 which support many
  214. /// instruction lengths.
  215. const unsigned MinLength = 2;
  216. /// Move the iterator to the next repeated substring.
  217. void advance() {
  218. // Clear the current state. If we're at the end of the range, then this
  219. // is the state we want to be in.
  220. RS = RepeatedSubstring();
  221. N = nullptr;
  222. // Each leaf node represents a repeat of a string.
  223. std::vector<SuffixTreeNode *> LeafChildren;
  224. // Continue visiting nodes until we find one which repeats more than once.
  225. while (!ToVisit.empty()) {
  226. SuffixTreeNode *Curr = ToVisit.back();
  227. ToVisit.pop_back();
  228. LeafChildren.clear();
  229. // Keep track of the length of the string associated with the node. If
  230. // it's too short, we'll quit.
  231. unsigned Length = Curr->ConcatLen;
  232. // Iterate over each child, saving internal nodes for visiting, and
  233. // leaf nodes in LeafChildren. Internal nodes represent individual
  234. // strings, which may repeat.
  235. for (auto &ChildPair : Curr->Children) {
  236. // Save all of this node's children for processing.
  237. if (!ChildPair.second->isLeaf())
  238. ToVisit.push_back(ChildPair.second);
  239. // It's not an internal node, so it must be a leaf. If we have a
  240. // long enough string, then save the leaf children.
  241. else if (Length >= MinLength)
  242. LeafChildren.push_back(ChildPair.second);
  243. }
  244. // The root never represents a repeated substring. If we're looking at
  245. // that, then skip it.
  246. if (Curr->isRoot())
  247. continue;
  248. // Do we have any repeated substrings?
  249. if (LeafChildren.size() >= 2) {
  250. // Yes. Update the state to reflect this, and then bail out.
  251. N = Curr;
  252. RS.Length = Length;
  253. for (SuffixTreeNode *Leaf : LeafChildren)
  254. RS.StartIndices.push_back(Leaf->SuffixIdx);
  255. break;
  256. }
  257. }
  258. // At this point, either NewRS is an empty RepeatedSubstring, or it was
  259. // set in the above loop. Similarly, N is either nullptr, or the node
  260. // associated with NewRS.
  261. }
  262. public:
  263. /// Return the current repeated substring.
  264. RepeatedSubstring &operator*() { return RS; }
  265. RepeatedSubstringIterator &operator++() {
  266. advance();
  267. return *this;
  268. }
  269. RepeatedSubstringIterator operator++(int I) {
  270. RepeatedSubstringIterator It(*this);
  271. advance();
  272. return It;
  273. }
  274. bool operator==(const RepeatedSubstringIterator &Other) const {
  275. return N == Other.N;
  276. }
  277. bool operator!=(const RepeatedSubstringIterator &Other) const {
  278. return !(*this == Other);
  279. }
  280. RepeatedSubstringIterator(SuffixTreeNode *N) : N(N) {
  281. // Do we have a non-null node?
  282. if (N) {
  283. // Yes. At the first step, we need to visit all of N's children.
  284. // Note: This means that we visit N last.
  285. ToVisit.push_back(N);
  286. advance();
  287. }
  288. }
  289. };
  290. typedef RepeatedSubstringIterator iterator;
  291. iterator begin() { return iterator(Root); }
  292. iterator end() { return iterator(nullptr); }
  293. };
  294. } // namespace llvm
  295. #endif // LLVM_SUPPORT_SUFFIXTREE_H