1 //===-- runtime/findloc.cpp -----------------------------------------------===//
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 // Implements FINDLOC for all required operand types and shapes and result
10 // integer kinds.
11 
12 #include "reduction-templates.h"
13 #include "flang/Runtime/character.h"
14 #include "flang/Runtime/reduction.h"
15 #include <cinttypes>
16 #include <complex>
17 
18 namespace Fortran::runtime {
19 
20 template <TypeCategory CAT1, int KIND1, TypeCategory CAT2, int KIND2>
21 struct Equality {
22   using Type1 = CppTypeFor<CAT1, KIND1>;
23   using Type2 = CppTypeFor<CAT2, KIND2>;
24   bool operator()(const Descriptor &array, const SubscriptValue at[],
25       const Descriptor &target) const {
26     return *array.Element<Type1>(at) == *target.OffsetElement<Type2>();
27   }
28 };
29 
30 template <int KIND1, int KIND2>
31 struct Equality<TypeCategory::Complex, KIND1, TypeCategory::Complex, KIND2> {
32   using Type1 = CppTypeFor<TypeCategory::Complex, KIND1>;
33   using Type2 = CppTypeFor<TypeCategory::Complex, KIND2>;
34   bool operator()(const Descriptor &array, const SubscriptValue at[],
35       const Descriptor &target) const {
36     const Type1 &xz{*array.Element<Type1>(at)};
37     const Type2 &tz{*target.OffsetElement<Type2>()};
38     return xz.real() == tz.real() && xz.imag() == tz.imag();
39   }
40 };
41 
42 template <int KIND1, TypeCategory CAT2, int KIND2>
43 struct Equality<TypeCategory::Complex, KIND1, CAT2, KIND2> {
44   using Type1 = CppTypeFor<TypeCategory::Complex, KIND1>;
45   using Type2 = CppTypeFor<CAT2, KIND2>;
46   bool operator()(const Descriptor &array, const SubscriptValue at[],
47       const Descriptor &target) const {
48     const Type1 &z{*array.Element<Type1>(at)};
49     return z.imag() == 0 && z.real() == *target.OffsetElement<Type2>();
50   }
51 };
52 
53 template <TypeCategory CAT1, int KIND1, int KIND2>
54 struct Equality<CAT1, KIND1, TypeCategory::Complex, KIND2> {
55   using Type1 = CppTypeFor<CAT1, KIND1>;
56   using Type2 = CppTypeFor<TypeCategory::Complex, KIND2>;
57   bool operator()(const Descriptor &array, const SubscriptValue at[],
58       const Descriptor &target) const {
59     const Type2 &z{*target.OffsetElement<Type2>()};
60     return *array.Element<Type1>(at) == z.real() && z.imag() == 0;
61   }
62 };
63 
64 template <int KIND> struct CharacterEquality {
65   using Type = CppTypeFor<TypeCategory::Character, KIND>;
66   bool operator()(const Descriptor &array, const SubscriptValue at[],
67       const Descriptor &target) const {
68     return CharacterScalarCompare<Type>(array.Element<Type>(at),
69                target.OffsetElement<Type>(),
70                array.ElementBytes() / static_cast<unsigned>(KIND),
71                target.ElementBytes() / static_cast<unsigned>(KIND)) == 0;
72   }
73 };
74 
75 struct LogicalEquivalence {
76   bool operator()(const Descriptor &array, const SubscriptValue at[],
77       const Descriptor &target) const {
78     return IsLogicalElementTrue(array, at) ==
79         IsLogicalElementTrue(target, at /*ignored*/);
80   }
81 };
82 
83 template <typename EQUALITY> class LocationAccumulator {
84 public:
85   LocationAccumulator(
86       const Descriptor &array, const Descriptor &target, bool back)
87       : array_{array}, target_{target}, back_{back} {
88     Reinitialize();
89   }
90   void Reinitialize() {
91     // per standard: result indices are all zero if no data
92     for (int j{0}; j < rank_; ++j) {
93       location_[j] = 0;
94     }
95   }
96   template <typename A> void GetResult(A *p, int zeroBasedDim = -1) {
97     if (zeroBasedDim >= 0) {
98       *p = location_[zeroBasedDim] -
99           array_.GetDimension(zeroBasedDim).LowerBound() + 1;
100     } else {
101       for (int j{0}; j < rank_; ++j) {
102         p[j] = location_[j] - array_.GetDimension(j).LowerBound() + 1;
103       }
104     }
105   }
106   template <typename IGNORED> bool AccumulateAt(const SubscriptValue at[]) {
107     if (equality_(array_, at, target_)) {
108       for (int j{0}; j < rank_; ++j) {
109         location_[j] = at[j];
110       }
111       return back_;
112     } else {
113       return true;
114     }
115   }
116 
117 private:
118   const Descriptor &array_;
119   const Descriptor &target_;
120   const bool back_{false};
121   const int rank_{array_.rank()};
122   SubscriptValue location_[maxRank];
123   const EQUALITY equality_{};
124 };
125 
126 template <TypeCategory XCAT, int XKIND, TypeCategory TARGET_CAT>
127 struct TotalNumericFindlocHelper {
128   template <int TARGET_KIND> struct Functor {
129     void operator()(Descriptor &result, const Descriptor &x,
130         const Descriptor &target, int kind, int dim, const Descriptor *mask,
131         bool back, Terminator &terminator) const {
132       using Eq = Equality<XCAT, XKIND, TARGET_CAT, TARGET_KIND>;
133       using Accumulator = LocationAccumulator<Eq>;
134       Accumulator accumulator{x, target, back};
135       DoTotalReduction<void>(x, dim, mask, accumulator, "FINDLOC", terminator);
136       ApplyIntegerKind<LocationResultHelper<Accumulator>::template Functor,
137           void>(kind, terminator, accumulator, result);
138     }
139   };
140 };
141 
142 template <TypeCategory CAT,
143     template <TypeCategory XCAT, int XKIND, TypeCategory TARGET_CAT>
144     class HELPER>
145 struct NumericFindlocHelper {
146   template <int KIND> struct Functor {
147     void operator()(TypeCategory targetCat, int targetKind, Descriptor &result,
148         const Descriptor &x, const Descriptor &target, int kind, int dim,
149         const Descriptor *mask, bool back, Terminator &terminator) const {
150       switch (targetCat) {
151       case TypeCategory::Integer:
152         ApplyIntegerKind<
153             HELPER<CAT, KIND, TypeCategory::Integer>::template Functor, void>(
154             targetKind, terminator, result, x, target, kind, dim, mask, back,
155             terminator);
156         break;
157       case TypeCategory::Real:
158         ApplyFloatingPointKind<
159             HELPER<CAT, KIND, TypeCategory::Real>::template Functor, void>(
160             targetKind, terminator, result, x, target, kind, dim, mask, back,
161             terminator);
162         break;
163       case TypeCategory::Complex:
164         ApplyFloatingPointKind<
165             HELPER<CAT, KIND, TypeCategory::Complex>::template Functor, void>(
166             targetKind, terminator, result, x, target, kind, dim, mask, back,
167             terminator);
168         break;
169       default:
170         terminator.Crash(
171             "FINDLOC: bad target category %d for array category %d",
172             static_cast<int>(targetCat), static_cast<int>(CAT));
173       }
174     }
175   };
176 };
177 
178 template <int KIND> struct CharacterFindlocHelper {
179   void operator()(Descriptor &result, const Descriptor &x,
180       const Descriptor &target, int kind, const Descriptor *mask, bool back,
181       Terminator &terminator) {
182     using Accumulator = LocationAccumulator<CharacterEquality<KIND>>;
183     Accumulator accumulator{x, target, back};
184     DoTotalReduction<void>(x, 0, mask, accumulator, "FINDLOC", terminator);
185     ApplyIntegerKind<LocationResultHelper<Accumulator>::template Functor, void>(
186         kind, terminator, accumulator, result);
187   }
188 };
189 
190 static void LogicalFindlocHelper(Descriptor &result, const Descriptor &x,
191     const Descriptor &target, int kind, const Descriptor *mask, bool back,
192     Terminator &terminator) {
193   using Accumulator = LocationAccumulator<LogicalEquivalence>;
194   Accumulator accumulator{x, target, back};
195   DoTotalReduction<void>(x, 0, mask, accumulator, "FINDLOC", terminator);
196   ApplyIntegerKind<LocationResultHelper<Accumulator>::template Functor, void>(
197       kind, terminator, accumulator, result);
198 }
199 
200 extern "C" {
201 void RTNAME(Findloc)(Descriptor &result, const Descriptor &x,
202     const Descriptor &target, int kind, const char *source, int line,
203     const Descriptor *mask, bool back) {
204   int rank{x.rank()};
205   SubscriptValue extent[1]{rank};
206   result.Establish(TypeCategory::Integer, kind, nullptr, 1, extent,
207       CFI_attribute_allocatable);
208   result.GetDimension(0).SetBounds(1, extent[0]);
209   Terminator terminator{source, line};
210   if (int stat{result.Allocate()}) {
211     terminator.Crash(
212         "FINDLOC: could not allocate memory for result; STAT=%d", stat);
213   }
214   CheckIntegerKind(terminator, kind, "FINDLOC");
215   auto xType{x.type().GetCategoryAndKind()};
216   auto targetType{target.type().GetCategoryAndKind()};
217   RUNTIME_CHECK(terminator, xType.has_value() && targetType.has_value());
218   switch (xType->first) {
219   case TypeCategory::Integer:
220     ApplyIntegerKind<NumericFindlocHelper<TypeCategory::Integer,
221                          TotalNumericFindlocHelper>::template Functor,
222         void>(xType->second, terminator, targetType->first, targetType->second,
223         result, x, target, kind, 0, mask, back, terminator);
224     break;
225   case TypeCategory::Real:
226     ApplyFloatingPointKind<NumericFindlocHelper<TypeCategory::Real,
227                                TotalNumericFindlocHelper>::template Functor,
228         void>(xType->second, terminator, targetType->first, targetType->second,
229         result, x, target, kind, 0, mask, back, terminator);
230     break;
231   case TypeCategory::Complex:
232     ApplyFloatingPointKind<NumericFindlocHelper<TypeCategory::Complex,
233                                TotalNumericFindlocHelper>::template Functor,
234         void>(xType->second, terminator, targetType->first, targetType->second,
235         result, x, target, kind, 0, mask, back, terminator);
236     break;
237   case TypeCategory::Character:
238     RUNTIME_CHECK(terminator,
239         targetType->first == TypeCategory::Character &&
240             targetType->second == xType->second);
241     ApplyCharacterKind<CharacterFindlocHelper, void>(xType->second, terminator,
242         result, x, target, kind, mask, back, terminator);
243     break;
244   case TypeCategory::Logical:
245     RUNTIME_CHECK(terminator, targetType->first == TypeCategory::Logical);
246     LogicalFindlocHelper(result, x, target, kind, mask, back, terminator);
247     break;
248   default:
249     terminator.Crash(
250         "FINDLOC: bad data type code (%d) for array", x.type().raw());
251   }
252 }
253 } // extern "C"
254 
255 // FINDLOC with DIM=
256 
257 template <TypeCategory XCAT, int XKIND, TypeCategory TARGET_CAT>
258 struct PartialNumericFindlocHelper {
259   template <int TARGET_KIND> struct Functor {
260     void operator()(Descriptor &result, const Descriptor &x,
261         const Descriptor &target, int kind, int dim, const Descriptor *mask,
262         bool back, Terminator &terminator) const {
263       using Eq = Equality<XCAT, XKIND, TARGET_CAT, TARGET_KIND>;
264       using Accumulator = LocationAccumulator<Eq>;
265       Accumulator accumulator{x, target, back};
266       ApplyIntegerKind<PartialLocationHelper<Accumulator>::template Functor,
267           void>(kind, terminator, result, x, dim, mask, terminator, "FINDLOC",
268           accumulator);
269     }
270   };
271 };
272 
273 template <int KIND> struct PartialCharacterFindlocHelper {
274   void operator()(Descriptor &result, const Descriptor &x,
275       const Descriptor &target, int kind, int dim, const Descriptor *mask,
276       bool back, Terminator &terminator) {
277     using Accumulator = LocationAccumulator<CharacterEquality<KIND>>;
278     Accumulator accumulator{x, target, back};
279     ApplyIntegerKind<PartialLocationHelper<Accumulator>::template Functor,
280         void>(kind, terminator, result, x, dim, mask, terminator, "FINDLOC",
281         accumulator);
282   }
283 };
284 
285 static void PartialLogicalFindlocHelper(Descriptor &result, const Descriptor &x,
286     const Descriptor &target, int kind, int dim, const Descriptor *mask,
287     bool back, Terminator &terminator) {
288   using Accumulator = LocationAccumulator<LogicalEquivalence>;
289   Accumulator accumulator{x, target, back};
290   ApplyIntegerKind<PartialLocationHelper<Accumulator>::template Functor, void>(
291       kind, terminator, result, x, dim, mask, terminator, "FINDLOC",
292       accumulator);
293 }
294 
295 extern "C" {
296 void RTNAME(FindlocDim)(Descriptor &result, const Descriptor &x,
297     const Descriptor &target, int kind, int dim, const char *source, int line,
298     const Descriptor *mask, bool back) {
299   Terminator terminator{source, line};
300   CheckIntegerKind(terminator, kind, "FINDLOC");
301   auto xType{x.type().GetCategoryAndKind()};
302   auto targetType{target.type().GetCategoryAndKind()};
303   RUNTIME_CHECK(terminator, xType.has_value() && targetType.has_value());
304   switch (xType->first) {
305   case TypeCategory::Integer:
306     ApplyIntegerKind<NumericFindlocHelper<TypeCategory::Integer,
307                          PartialNumericFindlocHelper>::template Functor,
308         void>(xType->second, terminator, targetType->first, targetType->second,
309         result, x, target, kind, dim, mask, back, terminator);
310     break;
311   case TypeCategory::Real:
312     ApplyFloatingPointKind<NumericFindlocHelper<TypeCategory::Real,
313                                PartialNumericFindlocHelper>::template Functor,
314         void>(xType->second, terminator, targetType->first, targetType->second,
315         result, x, target, kind, dim, mask, back, terminator);
316     break;
317   case TypeCategory::Complex:
318     ApplyFloatingPointKind<NumericFindlocHelper<TypeCategory::Complex,
319                                PartialNumericFindlocHelper>::template Functor,
320         void>(xType->second, terminator, targetType->first, targetType->second,
321         result, x, target, kind, dim, mask, back, terminator);
322     break;
323   case TypeCategory::Character:
324     RUNTIME_CHECK(terminator,
325         targetType->first == TypeCategory::Character &&
326             targetType->second == xType->second);
327     ApplyCharacterKind<PartialCharacterFindlocHelper, void>(xType->second,
328         terminator, result, x, target, kind, dim, mask, back, terminator);
329     break;
330   case TypeCategory::Logical:
331     RUNTIME_CHECK(terminator, targetType->first == TypeCategory::Logical);
332     PartialLogicalFindlocHelper(
333         result, x, target, kind, dim, mask, back, terminator);
334     break;
335   default:
336     terminator.Crash(
337         "FINDLOC: bad data type code (%d) for array", x.type().raw());
338   }
339 }
340 } // extern "C"
341 } // namespace Fortran::runtime
342