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

[Krzysztof Parzyszek 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);
----------------
kparzysz wrote:

This starts with an actual directive `D`.  It is then decomposed into either leaf or composite constructs, so presumably there shouldn't be a situation where there is a composite construct followed by something else.  I think we should assert on that (instead of returning an empty list), because it would indicate some kind of an internal inconsistency and this is an API function.

We don't need to return `Output`---I meant it to be usable in cases like
```
for (auto &c : getLeafOrCompositeConstructs(D, Output)) {
  ...
}
```

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