Bit Flag Conditional

This MR introduces a bit flag conditional that acts like a ternary operator for a single bit flag. It takes at least three arguments:

• flag_bit specifies which bit of the mask is examined
• mask_expression is an integer-typed expression which acts as a bit mask
• then_expression is a SymPy expression of arbitrary type, and
• optionally else_expression is a SymPy expression of arbitrary type

flag_cond examines the given bit of the given mask and takes the value of then_expression if the bit is set to 1. If not, it becomes either 0 or else_expression:

Three argument version:

    flag_cond(flag_bit, mask, expr) = expr if (flag_bit is set in mask) else 0

Four argument version:

    flag_cond(flag_bit, mask, expr_then, expr_else) = expr_then if (flag_bit is set in mask) else expr_else

pystencils/bit_masks.py

+import sympy as sp
+from pystencils.data_types import get_type_of_expression
+
+
+# noinspection PyPep8Naming
+class flag_cond(sp.Function):
+    """Evaluates a flag condition on a bit mask, and returns the value of one of two expressions,
+    depending on whether the flag is set. 
+
+    Three argument version:
+    ```
+        flag_cond(flag_bit, mask, expr) = expr if (flag_bit is set in mask) else 0
+    ```
+
+    Four argument version:
+    ```
+        flag_cond(flag_bit, mask, expr_then, expr_else) = expr_then if (flag_bit is set in mask) else expr_else
+    ```
+    """
+
+    nargs = (3, 4)
+
+    def __new__(cls, flag_bit, mask_expression, *expressions):
+
+        flag_dtype = get_type_of_expression(flag_bit)
+        if not flag_dtype.is_int():
+            raise ValueError('Argument flag_bit must be of integer type.')
+
+        mask_dtype = get_type_of_expression(mask_expression)
+        if not mask_dtype.is_int():
+            raise ValueError('Argument mask_expression must be of integer type.')
+
+        return super().__new__(cls, flag_bit, mask_expression, *expressions)
+
+    def to_c(self, print_func):
+        flag_bit = self.args[0]
+        mask = self.args[1]
+
+        then_expression = self.args[2]
+
+        flag_bit_code = print_func(flag_bit)
+        mask_code = print_func(mask)
+        then_code = print_func(then_expression)
+
+        code = f"(({mask_code}) >> ({flag_bit_code}) & 1) * ({then_code})"
+
+        if len(self.args) > 3:
+            else_expression = self.args[3]
+            else_code = print_func(else_expression)
+            code += f" + (({mask_code}) >> ({flag_bit_code}) ^ 1) * ({else_code})"
+
+        return code