From 3d50993396a1b37b9e0b75e417effde06e4e9063 Mon Sep 17 00:00:00 2001
From: Andreas Kloeckner <inform@tiker.net>
Date: Tue, 23 Jan 2018 16:31:32 -0600
Subject: [PATCH] Add a non-failure step exit

---
 dagrt/codegen/fortran.py  |  3 +++
 dagrt/codegen/python.py   | 12 ++++++++++++
 dagrt/exec_numpy.py       | 11 +++++++++++
 dagrt/language.py         | 26 ++++++++++++++++++++------
 test/test_code_builder.py | 19 +++++++++++++++++++
 5 files changed, 65 insertions(+), 6 deletions(-)

diff --git a/dagrt/codegen/fortran.py b/dagrt/codegen/fortran.py
index 9596f20..0bb06c8 100644
--- a/dagrt/codegen/fortran.py
+++ b/dagrt/codegen/fortran.py
@@ -2247,6 +2247,9 @@ class CodeGenerator(StructuredCodeGenerator):
 
         self.emit("goto 999")
 
+    def emit_inst_ExitStep(self, inst):
+        self.emit("goto 999")
+
     def emit_inst_StateTransition(self, inst):
         self.emit(
                 'dagrt_state%dagrt_next_state = '
diff --git a/dagrt/codegen/python.py b/dagrt/codegen/python.py
index 903b22b..945d5cf 100644
--- a/dagrt/codegen/python.py
+++ b/dagrt/codegen/python.py
@@ -91,6 +91,10 @@ class FailStepException(RuntimeError):
     pass
 
 
+class ExitStepException(RuntimeError):
+    pass
+
+
 class TransitionEvent(Exception):
 
     def __init__(self, next_state):
@@ -326,6 +330,9 @@ class CodeGenerator(StructuredCodeGenerator):
                     yield self.StepFailed(t=self.t)
                     continue
 
+                except self.ExitStepException:
+                    continue
+
                 except self.TransitionEvent as evt:
                     self.next_state = evt.next_state
 
@@ -467,6 +474,11 @@ class CodeGenerator(StructuredCodeGenerator):
         if not self._has_yield_inst:
             self._emit('yield')
 
+    def emit_inst_ExitStep(self, inst):
+        self._emit('raise self.ExitStepException()')
+        if not self._has_yield_inst:
+            self._emit('yield')
+
     def emit_inst_StateTransition(self, inst):
         assert '\'' not in inst.next_state
         self._emit('raise self.TransitionEvent(\'' + inst.next_state + '\')')
diff --git a/dagrt/exec_numpy.py b/dagrt/exec_numpy.py
index 3d70411..dd24c73 100644
--- a/dagrt/exec_numpy.py
+++ b/dagrt/exec_numpy.py
@@ -31,6 +31,10 @@ class FailStepException(Exception):
     pass
 
 
+class ExitStepException(Exception):
+    pass
+
+
 class TransitionEvent(Exception):
 
     def __init__(self, next_state):
@@ -148,6 +152,10 @@ class NumpyInterpreter(object):
                 yield StepFailed(t=self.context["<t>"])
                 continue
 
+            except ExitStepException:
+                yield StepFailed(t=self.context["<t>"])
+                continue
+
             except TransitionEvent as evt:
                 self.next_state = evt.next_state
 
@@ -261,6 +269,9 @@ class NumpyInterpreter(object):
     def exec_FailStep(self, insn):
         raise FailStepException()
 
+    def exec_ExitStep(self, insn):
+        raise ExitStepException()
+
     def exec_Nop(self, insn):
         pass
 
diff --git a/dagrt/language.py b/dagrt/language.py
index 107d1d3..484ad17 100644
--- a/dagrt/language.py
+++ b/dagrt/language.py
@@ -107,19 +107,16 @@ Assignment Instructions
 ^^^^^^^^^^^^^^^^^^^^^^^
 
 .. autoclass:: AssignSolved
-
 .. autoclass:: AssignExpression
-
 .. autoclass:: AssignFunctionCall
 
 State Instructions
 ^^^^^^^^^^^^^^^^^^
 
 .. autoclass:: YieldState
-
 .. autoclass:: Raise
-
 .. autoclass:: FailStep
+.. autoclass:: ExitStep
 
 Code Container
 ~~~~~~~~~~~~~~
@@ -131,7 +128,6 @@ Visualization
 ~~~~~~~~~~~~~
 
 .. autofunction:: get_dot_dependency_graph
-
 .. autofunction:: show_dependency_graph
 
 Code Creation
@@ -570,7 +566,8 @@ class StateTransition(Instruction):
 
 
 class FailStep(Instruction):
-    """
+    """Exits the current step with a failure indication to the controlling
+    program. Execution resumes with the next step as normal.
     """
 
     def get_written_variables(self):
@@ -582,6 +579,18 @@ class FailStep(Instruction):
     exec_method = six.moves.intern("exec_FailStep")
 
 
+class ExitStep(Instruction):
+    """Exits the current step. Execution resumes with the next step as normal.
+    """
+
+    def get_written_variables(self):
+        return frozenset()
+
+    def __str__(self):
+        return "ExitStep{cond}".format(cond=self._condition_printing_suffix())
+
+    exec_method = six.moves.intern("exec_ExitStep")
+
 # }}}
 
 
@@ -809,6 +818,7 @@ class CodeBuilder(object):
     .. automethod:: assign_solved_1
     .. automethod:: yield_state
     .. automethod:: fail_step
+    .. automethod:: exit_step
     .. automethod:: raise_
     .. automethod:: state_transition
     .. automethod:: __enter__
@@ -1139,6 +1149,10 @@ class CodeBuilder(object):
         self.fence()
         self._add_inst_to_context(FailStep())
 
+    def exit_step(self):
+        self.fence()
+        self._add_inst_to_context(ExitStep())
+
     def raise_(self, error):
         self.fence()
         self._add_inst_to_context(Raise(error))
diff --git a/test/test_code_builder.py b/test/test_code_builder.py
index 51f60a0..88ca7aa 100755
--- a/test/test_code_builder.py
+++ b/test/test_code_builder.py
@@ -143,6 +143,25 @@ def test_CodeBuilder_nested_condition_with_else_not_taken(python_method_impl):
     assert result == 3
 
 
+def test_CodeBuilder_exit_step(python_method_impl):
+    with CodeBuilder() as builder1:
+        builder1("<p>x", 1)
+        builder1.exit_step()
+        builder1.fence()
+        builder1("<p>x", 2)
+
+    with CodeBuilder() as builder2:
+        builder2.yield_state(var('<p>x'), 'x', 0, 'final')
+
+    code = DAGCode({
+        "state1": builder1.as_execution_state("state2"),
+        "state2": builder2.as_execution_state("state2")
+        }, "state1")
+
+    result = execute_and_return_single_result(python_method_impl, code)
+    assert result == 1
+
+
 if __name__ == "__main__":
     if len(sys.argv) > 1:
         exec(sys.argv[1])
-- 
GitLab