Repository: tmontes/ppytty

Issue:
Traps returning tasks need fixing.
Summary:
* Traps like `task-wait` and `message-wait` return, among other things, a reference to a task.
* The returned value is, however, not necessarily the same that was used in `task-spawn`.

Motive:
* The kernel will call the passed in task object if it is callable (assuming it is a generator function) which is then used in those traps results.

Failing example:
```
def child():
    yield ('sleep',)

def parent():
    yield ('task-spawn', child)
    completed_task, success, result = yield ('task-wait',)
    # This fails because completed_task is actually the generator object obtained by calling 
    assert completed_task is childchild
```

Possible solutions:
* Have the kernel require passed in tasks (to top-level `run`/`task-spawn`) to be a *generator object* or alike.
* Have the kernel track what it was passed vs. what it uses internally.

Let's sleep on this for a while.

Discussion:

Patch:
diff --git a/src/ppytty/kernel/common.py b/src/ppytty/kernel/common.py
index 2796add..a730fcb 100644
--- a/src/ppytty/kernel/common.py
+++ b/src/ppytty/kernel/common.py
@@ -11,7 +11,9 @@ from . state import state
 
 def runnable_task(task):
 
-    return task() if callable(task) else task
+    runnable = task() if callable(task) else task
+    state.spawned_objects[runnable] = task
+    return runnable
 
 
 
diff --git a/src/ppytty/kernel/scheduler.py b/src/ppytty/kernel/scheduler.py
index a392f8b..b4e6052 100644
--- a/src/ppytty/kernel/scheduler.py
+++ b/src/ppytty/kernel/scheduler.py
@@ -162,7 +162,9 @@ def process_task_completion(task, success, result):
     if not candidate_parent and task is not state.top_task:
         log.error('%r completed with no parent', task)
     if candidate_parent in state.tasks_waiting_child:
-        common.trap_will_return(candidate_parent, (task, success, result))
+        spawned_object = state.spawned_objects[task]
+        common.trap_will_return(candidate_parent, (spawned_object, success, result))
+        del state.spawned_objects[task]
         del state.parent_task[task]
         state.child_tasks[candidate_parent].remove(task)
         common.clear_tasks_children(candidate_parent)
diff --git a/src/ppytty/kernel/state.py b/src/ppytty/kernel/state.py
index ac6aee0..824c939 100644
--- a/src/ppytty/kernel/state.py
+++ b/src/ppytty/kernel/state.py
@@ -51,6 +51,12 @@ class _State(object):
         self.tasks_waiting_time = []
         self.tasks_waiting_time_hq = []
 
+        # ---------------------------------------------------------------------
+        # Spawned objects
+
+        # Maps running tasks to objects passed to scheduler.run / task-spawn.
+        self.spawned_objects = {}
+
         # ---------------------------------------------------------------------
         # Task trap tracking.
 
diff --git a/src/ppytty/kernel/trap_handlers.py b/src/ppytty/kernel/trap_handlers.py
index f367e08..f148354 100644
--- a/src/ppytty/kernel/trap_handlers.py
+++ b/src/ppytty/kernel/trap_handlers.py
@@ -139,7 +139,9 @@ def task_wait(task):
             success, result = state.completed_tasks[child]
             break
     if child is not None:
-        common.trap_will_return(task, (child, success, result))
+        spawned_object = state.spawned_objects[child]
+        common.trap_will_return(task, (spawned_object, success, result))
+        del state.spawned_objects[child]
         del state.parent_task[child]
         state.child_tasks[task].remove(child)
         common.clear_tasks_children(task)
@@ -198,7 +200,8 @@ def message_send(task, to_task, message):
 
     if to_task in state.tasks_waiting_inbox:
         state.tasks_waiting_inbox.remove(to_task)
-        common.trap_will_return(to_task, (task, message))
+        spawned_object = state.spawned_objects[task]
+        common.trap_will_return(to_task, (spawned_object, message))
         state.runnable_tasks.append(to_task)
     else:
         state.task_inbox[to_task].append((task, message))
@@ -212,7 +215,8 @@ def message_wait(task):
     task_inbox = state.task_inbox[task]
     if task_inbox:
         sender_task, message = task_inbox.popleft()
-        common.trap_will_return(task, (sender_task, message))
+        spawned_object = state.spawned_objects[sender_task]
+        common.trap_will_return(task, (spawned_object, message))
         state.runnable_tasks.append(task)
     else:
         state.tasks_waiting_inbox.append(task)


Test patch:
diff --git a/tests/kernel/tasks.py b/tests/kernel/tasks.py
index c66aa2e..2397d6b 100644
--- a/tests/kernel/tasks.py
+++ b/tests/kernel/tasks.py
@@ -13,11 +13,13 @@ def sleep_zero():
 
 
 
-def spawn_wait(task):
+def spawn_wait(task, sleep_before_wait=False):
 
     yield ('task-spawn', task)
-    _, success, result = yield ('task-wait',)
-    return success, result
+    if sleep_before_wait:
+        yield ('sleep', 0)
+    completed_task, success, result = yield ('task-wait',)
+    return completed_task, success, result
 
 
 
diff --git a/tests/kernel/test_state.py b/tests/kernel/test_state.py
index 0a66aa6..6537cb3 100644
--- a/tests/kernel/test_state.py
+++ b/tests/kernel/test_state.py
@@ -79,6 +79,7 @@ STATE_ATTRS = {
     'tasks_waiting_key_hq': (_assert_empty_list, _change_list),
     'tasks_waiting_time': (_assert_empty_list, _change_list),
     'tasks_waiting_time_hq': (_assert_empty_list, _change_list),
+    'spawned_objects': (_assert_empty_dict, _change_dict),
     'trap_call': (_assert_empty_dict, _change_dict),
     'trap_success': (_assert_empty_dict, _change_dict),
     'trap_result': (_assert_empty_dict, _change_dict),
diff --git a/tests/kernel/test_task_api_msgs.py b/tests/kernel/test_task_api_msgs.py
index ccb15bc..adfc29b 100644
--- a/tests/kernel/test_task_api_msgs.py
+++ b/tests/kernel/test_task_api_msgs.py
@@ -119,4 +119,72 @@ class TestChildToParent(io_bypass.NoOutputTestCase):
         self.assertEqual(result['received-message'], 'child-to-parent-message')
 
 
+
+class TestSenderTaskIsChildTask(io_bypass.NoOutputTestCase):
+
+    def _fast_child(self):
+
+        yield ('message-send', None, 'hi-from-child')
+
+
+    def _slow_child(self):
+
+        yield ('sleep', 0)
+        yield ('message-send', None, 'hi-from-child')
+
+
+    def _parent(self, child_object):
+
+        yield ('task-spawn', child_object)
+        sender_task, message = yield ('message-wait',)
+        completed_child, child_success, child_result = yield ('task-wait',)
+        return sender_task, message, completed_child, child_success, child_result
+
+
+    def _assert_parent_result(self, child_object, result):
+
+        sender_task, message, completed_child, child_success, child_result = result
+        self.assertIs(sender_task, child_object)
+        self.assertEqual(message, 'hi-from-child')
+        self.assertIs(completed_child, child_object)
+        self.assertTrue(child_success)
+        self.assertIsNone(child_result)
+
+
+    def test_gen_function_child_parent_waits_1st(self):
+
+        child_gen_function = self._slow_child
+        parent = self._parent(child_gen_function)
+        success, result = run(parent)
+        self.assertTrue(success)
+        self._assert_parent_result(child_gen_function, result)
+
+
+    def test_gen_function_child_parent_waits_2nd(self):
+
+        child_gen_function = self._fast_child
+        parent = self._parent(child_gen_function)
+        success, result = run(parent)
+        self.assertTrue(success)
+        self._assert_parent_result(child_gen_function, result)
+
+
+    def test_gen_object_child_parent_waits_1st(self):
+
+        child_gen_object = self._slow_child()
+        parent = self._parent(child_gen_object)
+        success, result = run(parent)
+        self.assertTrue(success)
+        self._assert_parent_result(child_gen_object, result)
+
+
+    def test_gen_object_child_parent_waits_2nd(self):
+
+        child_gen_object = self._fast_child()
+        parent = self._parent(child_gen_object)
+        success, result = run(parent)
+        self.assertTrue(success)
+        self._assert_parent_result(child_gen_object, result)
+
+
 # ----------------------------------------------------------------------------
diff --git a/tests/kernel/test_task_api_tasks.py b/tests/kernel/test_task_api_tasks.py
index d26946e..9893e37 100644
--- a/tests/kernel/test_task_api_tasks.py
+++ b/tests/kernel/test_task_api_tasks.py
@@ -14,20 +14,54 @@ from . import tasks
 
 
 
-class TestSpawn(io_bypass.NoOutputTestCase):
+class TestSpawnWaitObjects(io_bypass.NoOutputTestCase):
 
-    def test_spawn_gen_function(self):
+    def test_spawn_wait_gen_function_child_completes_1st(self):
 
         generator_function = tasks.sleep_zero
         task = tasks.spawn_wait(generator_function)
-        run(task)
+        success, result = run(task)
+        self.assertTrue(success)
+        completed_child, child_success, child_result = result
+        self.assertIs(completed_child, generator_function)
+        self.assertTrue(child_success)
+        self.assertIsNone(child_result)
 
 
-    def test_spawn_gen_object(self):
+    def test_spawn_wait_gen_object_child_completes_1st(self):
 
         generator_object = tasks.sleep_zero()
         task = tasks.spawn_wait(generator_object)
-        run(task)
+        success, result = run(task)
+        self.assertTrue(success)
+        completed_child, child_success, child_result = result
+        self.assertIs(completed_child, generator_object)
+        self.assertTrue(child_success)
+        self.assertIsNone(child_result)
+
+
+    def test_spawn_wait_gen_function_child_completes_2nd(self):
+
+        generator_function = tasks.sleep_zero
+        task = tasks.spawn_wait(generator_function, sleep_before_wait=True)
+        success, result = run(task)
+        self.assertTrue(success)
+        completed_child, child_success, child_result = result
+        self.assertIs(completed_child, generator_function)
+        self.assertTrue(child_success)
+        self.assertIsNone(child_result)
+
+
+    def test_spawn_wait_gen_object_child_completes_2nd(self):
+
+        generator_object = tasks.sleep_zero()
+        task = tasks.spawn_wait(generator_object, sleep_before_wait=True)
+        success, result = run(task)
+        self.assertTrue(success)
+        completed_child, child_success, child_result = result
+        self.assertIs(completed_child, generator_object)
+        self.assertTrue(child_success)
+        self.assertIsNone(child_result)
 
 
 
@@ -38,7 +72,7 @@ class TestWait(io_bypass.NoOutputTestCase):
         task = tasks.spawn_wait(tasks.sleep_zero_return_42_idiv_arg(42))
         parent_sucess, parent_result = run(task)
         self.assertTrue(parent_sucess)
-        child_success, child_result = parent_result
+        _completed_child, child_success, child_result = parent_result
         self.assertTrue(child_success)
         self.assertEqual(child_result, 1)
 
@@ -48,7 +82,7 @@ class TestWait(io_bypass.NoOutputTestCase):
         task = tasks.spawn_wait(tasks.sleep_zero_return_42_idiv_arg(0))
         parent_success, parent_result = run(task)
         self.assertTrue(parent_success)
-        child_success, child_result = parent_result
+        _completed_child, child_success, child_result = parent_result
         self.assertFalse(child_success)
         self.assertIsInstance(child_result, ZeroDivisionError)
 
@@ -58,15 +92,10 @@ class TestWaitChildException(io_bypass.NoOutputTestCase):
 
     def parent_spawn_wait_child_exception(self, child_task, expected_exc_class):
 
-        def parent():
-            yield ('task-spawn', child_task)
-            completed_task, child_success, child_result = yield ('task-wait',)
-            return completed_task is child_task, child_success, child_result
-
-        success, result = run(parent)
+        success, result = run(tasks.spawn_wait(child_task))
         self.assertTrue(success)
-        correct_task_completed, child_task_success, child_task_result = result
-        self.assertTrue(correct_task_completed)
+        completed_child, child_task_success, child_task_result = result
+        self.assertIs(completed_child, child_task)
         self.assertFalse(child_task_success)
         self.assertIsInstance(child_task_result, expected_exc_class)