added agent functions conditions with codegen tests

swig/python/codegen/codegen.py

@@ -710,14 +710,14 @@ def _ClassDef(self, t):
 
     def _FunctionDef(self, t):
         """
-        Checks the decorators of the function definition much must be either 'pyflamegpu.agent_function' or 'pyflamegpu.device_function'.
+        Checks the decorators of the function definition much must be either 'pyflamegpu.agent_function', 'pyflamegpu.agent_function_condition' or 'pyflamegpu.device_function'.
         Each is then processed in a different way using a specific dispatcher.
         Function calls are actually checked and only permitted (or user defined) function calls are supported.
         """
         self.write("\n")
         # check decorators
         if len(t.decorator_list) != 1 or not isinstance(t.decorator_list[0], ast.Attribute):
-            self.RaiseError(t, "Function definitions require a single pyflamegpu decorator of either 'pyflamegpu.agent_function' or 'pyflamegpu.device_function'")       
+            self.RaiseError(t, "Function definitions require a single pyflamegpu decorator of either 'pyflamegpu.agent_function', 'pyflamegpu.agent_function_condition' or 'pyflamegpu.device_function'")       
         # FLAMEGPU_AGENT_FUNCTION
         if t.decorator_list[0].attr == 'agent_function' and t.decorator_list[0].value.id == 'pyflamegpu':
             if getattr(t, "returns", False):
@@ -737,8 +737,23 @@ def _FunctionDef(self, t):
             self.write(")")
             # add to list of defined functions that can be called
             self._device_functions.append(t.name)
+        # FLAMEGPU_DEVICE_FUNCTION
+        elif t.decorator_list[0].attr == 'agent_function_condition' and t.decorator_list[0].value.id == 'pyflamegpu':
+            # check for return annotation
+            if not hasattr(t, "returns"):
+                self.RaiseError(t, "Agent function conditions must have a 'bool' return type specified as a return type annotation")
+            # check for return annotation type
+            if not isinstance(t.returns, ast.Name):
+                self.RaiseError(t, "Agent function conditions return type must be 'bool'")
+            if t.returns.id is not 'bool':
+                self.RaiseError(t, "Agent function conditions return type must be 'bool'")
+            # check to ensure no arguments (discard any with a warning)
+            if t.args.args:
+                self.RaiseWarning(t, "Agent function conditions does not support arguments. These will be discarded.")
+            # write the agent function macro
+            self.fill(f"FLAMEGPU_AGENT_FUNCTION_CONDITION({t.name})")
         else:
-            self.RaiseError(t, "Function definition uses an unsupported decorator. Must use either 'pyflamegpu.agent_function' or 'pyflamegpu.device_function'")
+            self.RaiseError(t, "Function definition uses an unsupported decorator. Must use either 'pyflamegpu.agent_function', 'pyflamegpu.agent_function_condition' or 'pyflamegpu.device_function'")
         self.enter()
         self.dispatch(t.body)
         self.leave()

swig/python/flamegpu.i

@@ -844,6 +844,13 @@ TEMPLATE_VARIABLE_INSTANTIATE_FLOATS(logNormal, flamegpu::HostRandom::logNormal)
             # do not allow passthrough (host exection of this function)
             pass
         return wrapper
+
+    def agent_function_condition(func):
+        @wraps(func)
+        def wrapper():
+            # do not allow passthrough (host exection of this function)
+            pass
+        return wrapper
 
     def device_function(func):
         @wraps(func)

tests/swig/python/codegen/test_codegen.py

@@ -453,6 +453,35 @@ def func(x: int) :
 }
 """
 
+py_fgpu_cond_func = """\
+@pyflamegpu.agent_function_condition
+def conditionalFunc() -> bool :
+    return True
+"""
+
+cpp_fgpu_cond_func = """\
+FLAMEGPU_AGENT_FUNCTION_CONDITION(conditionalFunc){
+    return true;
+}
+"""
+
+py_fgpu_cond_func_args = """\
+@pyflamegpu.agent_function_condition
+def conditionalFunc(arg: int) -> bool:
+    return True
+"""
+py_fgpu_cond_func_no_return_type = """\
+@pyflamegpu.agent_function_condition
+def conditionalFunc() :
+    return True
+"""
+py_fgpu_cond_func_wrong_return_type = """\
+@pyflamegpu.agent_function_condition
+def conditionalFunc() -> int :
+    return True
+"""
+
+
 py_fgpu_device_local_args_stack = """\
 @pyflamegpu.device_function
 def funcA(x: int) :
@@ -819,6 +848,16 @@ def test_fgpu_agent_func_return_type(self):
 
     # device functions, arg types and calling
 
+    def test_fgpu_agent_func_condition(self):
+        # check correct format
+        self._checkExpected(py_fgpu_cond_func, cpp_fgpu_cond_func)
+        # function is not allowed arguments but only a warning
+        self._checkWarning(py_fgpu_cond_func_args, cpp_fgpu_cond_func, "Agent function conditions does not support arguments")
+        # must have bool return type
+        self._checkException(py_fgpu_cond_func_no_return_type, "Agent function conditions return type must be 'bool'")
+        # only bool return type
+        self._checkException(py_fgpu_cond_func_wrong_return_type, "Agent function conditions return type must be 'bool'")
+
     def test_fgpu_device_func(self):
         self._checkExpected(py_fgpu_device_func_args, cpp_fgpu_device_func_args)
         # function argument requires a type

tests/swig/python/codegen/test_codegen_integration.py

@@ -23,6 +23,11 @@ def add_func(message_in: pyflamegpu.MessageNone, message_out: pyflamegpu.Message
     pyflamegpu.setVariableInt("x", add_2(x))
     return pyflamegpu.ALIVE
 
+@pyflamegpu.agent_function_condition
+def cond_func() -> bool:
+    i = pyflamegpu.getVariableInt("i")
+    return i < 10
+
 class GPUTest(TestCase):
     """
     Test provides a basic integration test of the pure python code generator (pyflamegpu.codegen). The test replicates the test_gpu_validation test but uses the 
@@ -51,7 +56,7 @@ def test_gpu_codegen_simulation(self):
         after being updated/changed during the simulation of an agent function. The 'add_func' agent function simply increases the agent variable x by a value of 2.
         This test will re-use the original population to read the results of the simulation step so it also acts to test that the original values are correctly overwritten.
         """
-        m = pyflamegpu.ModelDescription("test_gpu_memory_test")
+        m = pyflamegpu.ModelDescription("test_gpu_codegen_simulation")
         a = m.newAgent("agent")
         a.newVariableInt("id")
         a.newVariableInt("x")
@@ -60,6 +65,7 @@ def test_gpu_codegen_simulation(self):
         p = pyflamegpu.AgentVector(a, AGENT_COUNT)
         for i in range(AGENT_COUNT):
             instance = p[i]
+            instance.setVariableInt("id", i)
             instance.setVariableInt("x", i)
         layer = m.newLayer("add_layer")
         layer.addAgentFunction(func)
@@ -74,3 +80,37 @@ def test_gpu_codegen_simulation(self):
             instance = p[i]
             # use AgentInstance equality operator
             assert instance.getVariableInt("x") == (i + (2 * STEPS))
+
+    def test_gpu_codegen_function_condition(self):
+        m = pyflamegpu.ModelDescription("test_gpu_codegen_function_condition")
+        a = m.newAgent("agent")
+        a.newVariableInt("i")
+        a.newVariableInt("x")
+        func_translated = pyflamegpu.codegen.translate(add_func)
+        func_condition_translated = pyflamegpu.codegen.translate(cond_func)
+        func = a.newRTCFunction("add_func", func_translated)
+        func.setRTCFunctionCondition(func_condition_translated)
+        p = pyflamegpu.AgentVector(a, AGENT_COUNT)
+        for i in range(AGENT_COUNT):
+            instance = p[i]
+            instance.setVariableInt("i", i) # store ordinal index as a variable as order is not preserved after function condition
+            instance.setVariableInt("x", i)
+        layer = m.newLayer("add_layer")
+        layer.addAgentFunction(func)
+        cm = pyflamegpu.CUDASimulation(m)
+        cm.SimulationConfig().steps = STEPS
+        cm.setPopulationData(p)
+        cm.simulate()
+        # Re-use the same population to read back the simulation step results
+        cm.getPopulationData(p)
+        # check values are the same (cant use a as order not guaranteed to be preserved)
+        for a in range(AGENT_COUNT):
+            instance = p[i]
+            i = instance.getVariableInt("i")
+            x = instance.getVariableInt("x")
+            if i < 10:
+                # Condition will have allowed agent function to multiply the x value if i < 10
+                assert x == (i + (2 * STEPS))
+            else:
+                # Not meeting the condition means the agent function will now have executed and the x value should be unchanged
+                assert x == i