[flang-commits] [flang] [flang][Evaluate] Add `ExtractSubstring` to evaluate/tools.h (PR #81265)

[Krzysztof Parzyszek via flang-commits]

https://github.com/kparzysz created https://github.com/llvm/llvm-project/pull/81265

There are utility functions to extract a DataRef, optionally traversing into a substring, but there is no function to get the subscript itself. This patch adds one.

>From 351e27690238980100075f7a5654c1471f15651f Mon Sep 17 00:00:00 2001
From: Krzysztof Parzyszek <Krzysztof.Parzyszek at amd.com>
Date: Tue, 6 Feb 2024 17:06:14 -0600
Subject: [PATCH] [flang][Evaluate] Add `ExtractSubstring` to evaluate/tools.h

There are utility functions to extract a DataRef, optionally
traversing into a substring, but there is no function to get
the subscript itself. This patch adds one.
---
 flang/include/flang/Evaluate/tools.h | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/flang/include/flang/Evaluate/tools.h b/flang/include/flang/Evaluate/tools.h
index d257da1a709642..e9999974944e88 100644
--- a/flang/include/flang/Evaluate/tools.h
+++ b/flang/include/flang/Evaluate/tools.h
@@ -430,6 +430,29 @@ template <typename A> std::optional<CoarrayRef> ExtractCoarrayRef(const A &x) {
   }
 }
 
+struct ExtractSubstringHelper {
+  template <typename T> static std::optional<Substring> visit(T &&) {
+    return std::nullopt;
+  }
+
+  static std::optional<Substring> visit(const Substring &e) { return e; }
+
+  template <typename T>
+  static std::optional<Substring> visit(const Designator<T> &e) {
+    return std::visit([](auto &&s) { return visit(s); }, e.u);
+  }
+
+  template <typename T>
+  static std::optional<Substring> visit(const Expr<T> &e) {
+    return std::visit([](auto &&s) { return visit(s); }, e.u);
+  }
+};
+
+template <typename A>
+std::optional<Substring> ExtractSubstring(const A &x) {
+  return ExtractSubstringHelper::visit(x);
+}
+
 // If an expression is simply a whole symbol data designator,
 // extract and return that symbol, else null.
 template <typename A> const Symbol *UnwrapWholeSymbolDataRef(const A &x) {