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