flang/runtime/findloc.cpp

   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   RT_API_ATTRS bool operator()(const Descriptor &array,
  25       const SubscriptValue at[], 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   RT_API_ATTRS bool operator()(const Descriptor &array,
  35       const SubscriptValue at[], 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   RT_API_ATTRS bool operator()(const Descriptor &array,
  47       const SubscriptValue at[], 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   RT_API_ATTRS bool operator()(const Descriptor &array,
  58       const SubscriptValue at[], 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   RT_API_ATTRS bool operator()(const Descriptor &array,
  67       const SubscriptValue at[], 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   RT_API_ATTRS bool operator()(const Descriptor &array,
  77       const SubscriptValue at[], 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   RT_API_ATTRS LocationAccumulator(
  86       const Descriptor &array, const Descriptor &target, bool back)
  87       : array_{array}, target_{target}, back_{back} {}
  88   RT_API_ATTRS void Reinitialize() { gotAnything_ = false; }
  89   template <typename A>
  90   RT_API_ATTRS void GetResult(A *p, int zeroBasedDim = -1) {
  91     if (zeroBasedDim >= 0) {
  92       *p = gotAnything_ ? location_[zeroBasedDim] -
  93               array_.GetDimension(zeroBasedDim).LowerBound() + 1
  94                         : 0;
  95     } else if (gotAnything_) {
  96       for (int j{0}; j < rank_; ++j) {
  97         p[j] = location_[j] - array_.GetDimension(j).LowerBound() + 1;
  98       }
  99     } else {
 100       // no unmasked hits? result is all zeroes
 101       for (int j{0}; j < rank_; ++j) {
 102         p[j] = 0;
 103       }
 104     }
 105   }
 106   template <typename IGNORED>
 107   RT_API_ATTRS bool AccumulateAt(const SubscriptValue at[]) {
 108     if (equality_(array_, at, target_)) {
 109       gotAnything_ = true;
 110       for (int j{0}; j < rank_; ++j) {
 111         location_[j] = at[j];
 112       }
 113       return back_;
 114     } else {
 115       return true;
 116     }
 117   }
 118
 119 private:
 120   const Descriptor &array_;
 121   const Descriptor &target_;
 122   const bool back_{false};
 123   const int rank_{array_.rank()};
 124   bool gotAnything_{false};
 125   SubscriptValue location_[maxRank];
 126   const EQUALITY equality_{};
 127 };
 128
 129 template <TypeCategory XCAT, int XKIND, TypeCategory TARGET_CAT>
 130 struct TotalNumericFindlocHelper {
 131   template <int TARGET_KIND> struct Functor {
 132     RT_API_ATTRS void operator()(Descriptor &result, const Descriptor &x,
 133         const Descriptor &target, int kind, int dim, const Descriptor *mask,
 134         bool back, Terminator &terminator) const {
 135       using Eq = Equality<XCAT, XKIND, TARGET_CAT, TARGET_KIND>;
 136       using Accumulator = LocationAccumulator<Eq>;
 137       Accumulator accumulator{x, target, back};
 138       DoTotalReduction<void>(x, dim, mask, accumulator, "FINDLOC", terminator);
 139       ApplyIntegerKind<LocationResultHelper<Accumulator>::template Functor,
 140           void>(kind, terminator, accumulator, result);
 141     }
 142   };
 143 };
 144
 145 template <TypeCategory CAT,
 146     template <TypeCategory XCAT, int XKIND, TypeCategory TARGET_CAT>
 147     class HELPER>
 148 struct NumericFindlocHelper {
 149   template <int KIND> struct Functor {
 150     RT_API_ATTRS void operator()(TypeCategory targetCat, int targetKind,
 151         Descriptor &result, const Descriptor &x, const Descriptor &target,
 152         int kind, int dim, const Descriptor *mask, bool back,
 153         Terminator &terminator) const {
 154       switch (targetCat) {
 155       case TypeCategory::Integer:
 156         ApplyIntegerKind<
 157             HELPER<CAT, KIND, TypeCategory::Integer>::template Functor, void>(
 158             targetKind, terminator, result, x, target, kind, dim, mask, back,
 159             terminator);
 160         break;
 161       case TypeCategory::Real:
 162         ApplyFloatingPointKind<
 163             HELPER<CAT, KIND, TypeCategory::Real>::template Functor, void>(
 164             targetKind, terminator, result, x, target, kind, dim, mask, back,
 165             terminator);
 166         break;
 167       case TypeCategory::Complex:
 168         ApplyFloatingPointKind<
 169             HELPER<CAT, KIND, TypeCategory::Complex>::template Functor, void>(
 170             targetKind, terminator, result, x, target, kind, dim, mask, back,
 171             terminator);
 172         break;
 173       default:
 174         terminator.Crash(
 175             "FINDLOC: bad target category %d for array category %d",
 176             static_cast<int>(targetCat), static_cast<int>(CAT));
 177       }
 178     }
 179   };
 180 };
 181
 182 template <int KIND> struct CharacterFindlocHelper {
 183   RT_API_ATTRS void operator()(Descriptor &result, const Descriptor &x,
 184       const Descriptor &target, int kind, const Descriptor *mask, bool back,
 185       Terminator &terminator) {
 186     using Accumulator = LocationAccumulator<CharacterEquality<KIND>>;
 187     Accumulator accumulator{x, target, back};
 188     DoTotalReduction<void>(x, 0, mask, accumulator, "FINDLOC", terminator);
 189     ApplyIntegerKind<LocationResultHelper<Accumulator>::template Functor, void>(
 190         kind, terminator, accumulator, result);
 191   }
 192 };
 193
 194 static RT_API_ATTRS void LogicalFindlocHelper(Descriptor &result,
 195     const Descriptor &x, const Descriptor &target, int kind,
 196     const Descriptor *mask, bool back, Terminator &terminator) {
 197   using Accumulator = LocationAccumulator<LogicalEquivalence>;
 198   Accumulator accumulator{x, target, back};
 199   DoTotalReduction<void>(x, 0, mask, accumulator, "FINDLOC", terminator);
 200   ApplyIntegerKind<LocationResultHelper<Accumulator>::template Functor, void>(
 201       kind, terminator, accumulator, result);
 202 }
 203
 204 extern "C" {
 205 RT_EXT_API_GROUP_BEGIN
 206
 207 void RTDEF(Findloc)(Descriptor &result, const Descriptor &x,
 208     const Descriptor &target, int kind, const char *source, int line,
 209     const Descriptor *mask, bool back) {
 210   int rank{x.rank()};
 211   SubscriptValue extent[1]{rank};
 212   result.Establish(TypeCategory::Integer, kind, nullptr, 1, extent,
 213       CFI_attribute_allocatable);
 214   result.GetDimension(0).SetBounds(1, extent[0]);
 215   Terminator terminator{source, line};
 216   if (int stat{result.Allocate()}) {
 217     terminator.Crash(
 218         "FINDLOC: could not allocate memory for result; STAT=%d", stat);
 219   }
 220   CheckIntegerKind(terminator, kind, "FINDLOC");
 221   auto xType{x.type().GetCategoryAndKind()};
 222   auto targetType{target.type().GetCategoryAndKind()};
 223   RUNTIME_CHECK(terminator, xType.has_value() && targetType.has_value());
 224   switch (xType->first) {
 225   case TypeCategory::Integer:
 226     ApplyIntegerKind<NumericFindlocHelper<TypeCategory::Integer,
 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::Real:
 232     ApplyFloatingPointKind<NumericFindlocHelper<TypeCategory::Real,
 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::Complex:
 238     ApplyFloatingPointKind<NumericFindlocHelper<TypeCategory::Complex,
 239                                TotalNumericFindlocHelper>::template Functor,
 240         void>(xType->second, terminator, targetType->first, targetType->second,
 241         result, x, target, kind, 0, mask, back, terminator);
 242     break;
 243   case TypeCategory::Character:
 244     RUNTIME_CHECK(terminator,
 245         targetType->first == TypeCategory::Character &&
 246             targetType->second == xType->second);
 247     ApplyCharacterKind<CharacterFindlocHelper, void>(xType->second, terminator,
 248         result, x, target, kind, mask, back, terminator);
 249     break;
 250   case TypeCategory::Logical:
 251     RUNTIME_CHECK(terminator, targetType->first == TypeCategory::Logical);
 252     LogicalFindlocHelper(result, x, target, kind, mask, back, terminator);
 253     break;
 254   default:
 255     terminator.Crash(
 256         "FINDLOC: bad data type code (%d) for array", x.type().raw());
 257   }
 258 }
 259
 260 RT_EXT_API_GROUP_END
 261 } // extern "C"
 262
 263 // FINDLOC with DIM=
 264
 265 template <TypeCategory XCAT, int XKIND, TypeCategory TARGET_CAT>
 266 struct PartialNumericFindlocHelper {
 267   template <int TARGET_KIND> struct Functor {
 268     RT_API_ATTRS void operator()(Descriptor &result, const Descriptor &x,
 269         const Descriptor &target, int kind, int dim, const Descriptor *mask,
 270         bool back, Terminator &terminator) const {
 271       using Eq = Equality<XCAT, XKIND, TARGET_CAT, TARGET_KIND>;
 272       using Accumulator = LocationAccumulator<Eq>;
 273       Accumulator accumulator{x, target, back};
 274       ApplyIntegerKind<PartialLocationHelper<Accumulator>::template Functor,
 275           void>(kind, terminator, result, x, dim, mask, terminator, "FINDLOC",
 276           accumulator);
 277     }
 278   };
 279 };
 280
 281 template <int KIND> struct PartialCharacterFindlocHelper {
 282   RT_API_ATTRS void operator()(Descriptor &result, const Descriptor &x,
 283       const Descriptor &target, int kind, int dim, const Descriptor *mask,
 284       bool back, Terminator &terminator) {
 285     using Accumulator = LocationAccumulator<CharacterEquality<KIND>>;
 286     Accumulator accumulator{x, target, back};
 287     ApplyIntegerKind<PartialLocationHelper<Accumulator>::template Functor,
 288         void>(kind, terminator, result, x, dim, mask, terminator, "FINDLOC",
 289         accumulator);
 290   }
 291 };
 292
 293 static RT_API_ATTRS void PartialLogicalFindlocHelper(Descriptor &result,
 294     const Descriptor &x, const Descriptor &target, int kind, int dim,
 295     const Descriptor *mask, bool back, Terminator &terminator) {
 296   using Accumulator = LocationAccumulator<LogicalEquivalence>;
 297   Accumulator accumulator{x, target, back};
 298   ApplyIntegerKind<PartialLocationHelper<Accumulator>::template Functor, void>(
 299       kind, terminator, result, x, dim, mask, terminator, "FINDLOC",
 300       accumulator);
 301 }
 302
 303 extern "C" {
 304 RT_EXT_API_GROUP_BEGIN
 305
 306 void RTDEF(FindlocDim)(Descriptor &result, const Descriptor &x,
 307     const Descriptor &target, int kind, int dim, const char *source, int line,
 308     const Descriptor *mask, bool back) {
 309   Terminator terminator{source, line};
 310   CheckIntegerKind(terminator, kind, "FINDLOC");
 311   auto xType{x.type().GetCategoryAndKind()};
 312   auto targetType{target.type().GetCategoryAndKind()};
 313   RUNTIME_CHECK(terminator, xType.has_value() && targetType.has_value());
 314   switch (xType->first) {
 315   case TypeCategory::Integer:
 316     ApplyIntegerKind<NumericFindlocHelper<TypeCategory::Integer,
 317                          PartialNumericFindlocHelper>::template Functor,
 318         void>(xType->second, terminator, targetType->first, targetType->second,
 319         result, x, target, kind, dim, mask, back, terminator);
 320     break;
 321   case TypeCategory::Real:
 322     ApplyFloatingPointKind<NumericFindlocHelper<TypeCategory::Real,
 323                                PartialNumericFindlocHelper>::template Functor,
 324         void>(xType->second, terminator, targetType->first, targetType->second,
 325         result, x, target, kind, dim, mask, back, terminator);
 326     break;
 327   case TypeCategory::Complex:
 328     ApplyFloatingPointKind<NumericFindlocHelper<TypeCategory::Complex,
 329                                PartialNumericFindlocHelper>::template Functor,
 330         void>(xType->second, terminator, targetType->first, targetType->second,
 331         result, x, target, kind, dim, mask, back, terminator);
 332     break;
 333   case TypeCategory::Character:
 334     RUNTIME_CHECK(terminator,
 335         targetType->first == TypeCategory::Character &&
 336             targetType->second == xType->second);
 337     ApplyCharacterKind<PartialCharacterFindlocHelper, void>(xType->second,
 338         terminator, result, x, target, kind, dim, mask, back, terminator);
 339     break;
 340   case TypeCategory::Logical:
 341     RUNTIME_CHECK(terminator, targetType->first == TypeCategory::Logical);
 342     PartialLogicalFindlocHelper(
 343         result, x, target, kind, dim, mask, back, terminator);
 344     break;
 345   default:
 346     terminator.Crash(
 347         "FINDLOC: bad data type code (%d) for array", x.type().raw());
 348   }
 349 }
 350
 351 RT_EXT_API_GROUP_END
 352 } // extern "C"
 353 } // namespace Fortran::runtime