1 /* where.h -- Support for Fortran-90-like "where" construct 2 3 Copyright (C) 2015 European Centre for Medium-Range Weather Forecasts 4 5 Author: Robin Hogan <r.j.hogan@ecmwf.int> 6 7 This file is part of the Adept library. 8 9 10 Consider the following: 11 12 A.where(B) = C; 13 A.where(B) = either_or(C, D); 14 15 where A is an Array, B is a boolean expression, and C and D are 16 expressions, and the arrays and expressions have the same rank and 17 size, except that C and or D may have rank zero. The first line has 18 the effect of setting every element of A for which B is true to the 19 corresponding value in C. The second line does this but for 20 elements where B is false it sets A instead to D. 21 22 23 24 */ 25 26 27 #ifndef AdeptWhere_H 28 #define AdeptWhere_H 1 29 30 #include <vector> 31 32 #include <adept/Expression.h> 33 34 namespace adept { 35 36 namespace internal { 37 38 39 // --------------------------------------------------------------------- 40 // Section 1. EitherOr object returned by either_or function 41 // --------------------------------------------------------------------- 42 template <class C, class D> 43 class EitherOr { 44 public: 45 typedef bool _adept_either_or_flag; EitherOr(const C & c,const D & d)46 EitherOr(const C& c, const D& d) : either_(c), or_(d) { } value_if_true()47 const C& value_if_true() const { return either_; } value_if_false()48 const D& value_if_false() const { return or_; } 49 protected: 50 const C& either_; 51 const D& or_; 52 }; 53 54 55 template <typename T> 56 struct is_not_either_or 57 { 58 private: 59 typedef char yes; 60 typedef struct { char array[2]; } no; 61 template <typename C> static yes test(typename C::_adept_either_or_flag*); 62 template <typename C> static no test(...); 63 public: 64 static const bool value = sizeof(test<T>(0)) != sizeof(yes); 65 }; 66 67 68 // --------------------------------------------------------------------- 69 // Section 2. Where class returned by A.where(B) 70 // --------------------------------------------------------------------- 71 template <class A, class B> 72 class Where { 73 public: Where(A & a,const B & b)74 Where(A& a, const B& b) : array_(a), bool_expr_(b) { } 75 76 template <class C> 77 typename enable_if<is_not_either_or<C>::value, Where&>::type 78 operator=(const C& c) { 79 array_.assign_conditional(bool_expr_, c); 80 return *this; 81 } 82 83 // With either_or on the right-hand-side: this implementation 84 // could be faster if bool_expr was not evaluated twice 85 template <class C> 86 typename enable_if<!is_not_either_or<C>::value, Where&>::type 87 operator=(const C& c) { 88 array_.assign_conditional(!const_cast<B&>(bool_expr_), c.value_if_false()); 89 array_.assign_conditional(bool_expr_, c.value_if_true()); 90 return *this; 91 } 92 93 #define ADEPT_WHERE_OPERATOR(EQ_OP, OP) \ 94 template <class C> \ 95 typename enable_if<is_not_either_or<C>::value, Where&>::type \ 96 EQ_OP(const C& c) { \ 97 array_.assign_conditional(bool_expr_, noalias(*this) OP c); \ 98 return *this; \ 99 } \ 100 template <class C> \ 101 typename enable_if<!is_not_either_or<C>::value, Where&>::type \ 102 EQ_OP(const C& c) { \ 103 array_.assign_conditional(!const_cast<B&>(bool_expr_), \ 104 noalias(*this) OP c.value_if_false()); \ 105 array_.assign_conditional(bool_expr_, \ 106 noalias(*this) OP c.value_if_true()); \ 107 return *this; \ 108 } 109 ADEPT_WHERE_OPERATOR(operator+=, +) 110 ADEPT_WHERE_OPERATOR(operator-=, -) 111 ADEPT_WHERE_OPERATOR(operator*=, *) 112 ADEPT_WHERE_OPERATOR(operator/=, /) 113 #undef ADEPT_WHERE_OPERATOR 114 115 protected: 116 A& array_; 117 const B& bool_expr_; 118 119 }; 120 121 } // end namespace internal 122 123 124 125 template <class C, class D> either_or(const C & c,const D & d)126 internal::EitherOr<C,D> either_or(const C& c, const D& d) { 127 return internal::EitherOr<C,D>(c, d); 128 } 129 130 } // end namespace adept 131 132 #endif 133