[flang-commits] [flang] [llvm] [mlir] [Frontend][OpenMP] Refactor getLeafConstructs, add getCompoundConstruct (PR #87247)

[Krzysztof Parzyszek via flang-commits]

================
@@ -8,12 +8,80 @@
 
 #include "llvm/Frontend/OpenMP/OMP.h"
 
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/Support/ErrorHandling.h"
 
+#include <algorithm>
+#include <iterator>
+#include <type_traits>
+
 using namespace llvm;
-using namespace omp;
+using namespace llvm::omp;
 
 #define GEN_DIRECTIVES_IMPL
 #include "llvm/Frontend/OpenMP/OMP.inc"
+
+namespace llvm::omp {
+ArrayRef<Directive> getLeafConstructs(Directive D) {
+  auto Idx = static_cast<std::size_t>(D);
+  if (Idx >= Directive_enumSize)
+    return {};
+  const auto *Row = LeafConstructTable[LeafConstructTableOrdering[Idx]];
+  return ArrayRef(&Row[2], &Row[2] + static_cast<int>(Row[1]));
+}
+
+Directive getCompoundConstruct(ArrayRef<Directive> Parts) {
+  if (Parts.empty())
+    return OMPD_unknown;
+
+  // Parts don't have to be leafs, so expand them into leafs first.
+  // Store the expanded leafs in the same format as rows in the leaf
+  // table (generated by tablegen).
+  SmallVector<Directive> RawLeafs(2);
+  for (Directive P : Parts) {
+    ArrayRef<Directive> Ls = getLeafConstructs(P);
+    if (!Ls.empty())
+      RawLeafs.append(Ls.begin(), Ls.end());
+    else
+      RawLeafs.push_back(P);
+  }
+
+  // RawLeafs will be used as key in the binary search. The search doesn't
+  // guarantee that the exact same entry will be found (since RawLeafs may
+  // not correspond to any compound directive). Because of that, we will
+  // need to compare the search result with the given set of leafs.
+  // Also, if there is only one leaf in the list, it corresponds to itself,
+  // no search is necessary.
+  auto GivenLeafs{ArrayRef<Directive>(RawLeafs).drop_front(2)};
+  if (GivenLeafs.size() == 1)
+    return GivenLeafs.front();
+  RawLeafs[1] = static_cast<Directive>(GivenLeafs.size());
+
+  auto Iter = std::lower_bound(
+      LeafConstructTable, LeafConstructTableEndDirective,
+      static_cast<std::decay_t<decltype(*LeafConstructTable)>>(RawLeafs.data()),
+      [](const auto *RowA, const auto *RowB) {
----------------
kparzysz wrote:

Done.  I kept the autos below because they trivially follow the types of RowA and RowB.

https://github.com/llvm/llvm-project/pull/87247