[flang-commits] [flang] [llvm] [mlir] [Frontend][OpenMP] Implement getLeafOrCompositeConstructs (PR #89104)

[Sergio Afonso via flang-commits]

================
@@ -34,6 +71,40 @@ ArrayRef<Directive> getLeafConstructs(Directive D) {
   return ArrayRef(&Row[2], static_cast<int>(Row[1]));
 }
 
+ArrayRef<Directive> getLeafConstructsOrSelf(Directive D) {
+  if (auto Leafs = getLeafConstructs(D); !Leafs.empty())
+    return Leafs;
+  auto Idx = static_cast<size_t>(D);
+  assert(Idx < Directive_enumSize && "Invalid directive");
+  const auto *Row = LeafConstructTable[LeafConstructTableOrdering[Idx]];
+  // The first entry in the row is the directive itself.
+  return ArrayRef(&Row[0], &Row[0] + 1);
+}
+
+ArrayRef<Directive>
+getLeafOrCompositeConstructs(Directive D, SmallVectorImpl<Directive> &Output) {
+  using ArrayTy = ArrayRef<Directive>;
+  using IteratorTy = ArrayTy::iterator;
+  ArrayRef<Directive> Leafs = getLeafConstructsOrSelf(D);
+
+  IteratorTy Iter = Leafs.begin();
+  do {
+    auto Range = getFirstCompositeRange(llvm::make_range(Iter, Leafs.end()));
+    // All directives before the range are leaf constructs.
+    for (; Iter != Range.begin(); ++Iter)
+      Output.push_back(*Iter);
+    if (!Range.empty()) {
+      Directive Comp =
+          getCompoundConstruct(ArrayTy(Range.begin(), Range.end()));
+      assert(Comp != OMPD_unknown);
+      Output.push_back(Comp);
----------------
skatrak wrote:

I think we should check here that `Range.end() == Leafs.end()` and return an empty list in that case (or some other way to signal an invalid input), because otherwise we'd be allowing something like `[distribute, simd, parallel]` to be passed in and then return [distribute-simd, parallel]. Or is this also intended to be very lenient and callers are responsible for only dealing with valid composite constructs? Since it's checked that there's actually a valid composite construct, I think this other check should be done as well.

Do we need to return the `Output` argument?

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