From fbf06666ae984109c05d8feb16eb582a259762c8 Mon Sep 17 00:00:00 2001
From: "Alexie (Boyong) Madolid" <maex@tuta.io>
Date: Fri, 15 Nov 2024 18:32:02 +0800
Subject: [PATCH] [REFACTOR]: Spawn Call

---
 jac-cloud/jac_cloud/core/architype.py         |  10 +-
 jac-cloud/jac_cloud/plugin/jaseci.py          | 111 +++++++++-------
 jac-cloud/jac_cloud/tests/openapi_specs.yaml  |  49 ++++++++
 jac-cloud/jac_cloud/tests/simple_graph.jac    |  54 ++++++++
 .../jac_cloud/tests/test_simple_graph.py      |  46 ++++++-
 jac/jaclang/plugin/default.py                 | 118 +++++++++++-------
 jac/jaclang/runtimelib/architype.py           |  32 ++---
 jac/jaclang/runtimelib/utils.py               |  15 +++
 jac/jaclang/tests/fixtures/visit_sequence.jac |  50 ++++++++
 jac/jaclang/tests/test_language.py            |  19 ++-
 10 files changed, 387 insertions(+), 117 deletions(-)
 create mode 100644 jac/jaclang/tests/fixtures/visit_sequence.jac

diff --git a/jac-cloud/jac_cloud/core/architype.py b/jac-cloud/jac_cloud/core/architype.py
index c893c7cbb5..d4d25904ca 100644
--- a/jac-cloud/jac_cloud/core/architype.py
+++ b/jac-cloud/jac_cloud/core/architype.py
@@ -557,7 +557,7 @@ def update(self, bulk_write: BulkWrite, propagate: bool = False) -> None:
             ############################################################
             #                   POPULATE ADDED EDGES                   #
             ############################################################
-            added_edges: set[BaseAnchor | Anchor] = (
+            added_edges: set[BaseAnchor] = (
                 changes.get("$addToSet", {}).get("edges", {}).get("$each", [])
             )
             if added_edges:
@@ -575,7 +575,7 @@ def update(self, bulk_write: BulkWrite, propagate: bool = False) -> None:
             ############################################################
             #                  POPULATE REMOVED EDGES                  #
             ############################################################
-            pulled_edges: set[BaseAnchor | Anchor] = (
+            pulled_edges: set[BaseAnchor] = (
                 changes.get("$pull", {}).get("edges", {}).get("$in", [])
             )
             if pulled_edges:
@@ -828,10 +828,10 @@ class WalkerAnchor(BaseAnchor, _WalkerAnchor):  # type: ignore[misc]
     """Walker Anchor."""
 
     architype: "WalkerArchitype"
-    path: list[Anchor] = field(default_factory=list)
-    next: list[Anchor] = field(default_factory=list)
+    path: list[NodeAnchor] = field(default_factory=list)  # type: ignore[assignment]
+    next: list[NodeAnchor] = field(default_factory=list)  # type: ignore[assignment]
     returns: list[Any] = field(default_factory=list)
-    ignores: list[Anchor] = field(default_factory=list)
+    ignores: list[NodeAnchor] = field(default_factory=list)  # type: ignore[assignment]
     disengaged: bool = False
 
     class Collection(BaseCollection["WalkerAnchor"]):
diff --git a/jac-cloud/jac_cloud/plugin/jaseci.py b/jac-cloud/jac_cloud/plugin/jaseci.py
index 2b58283dce..8c8de78c02 100644
--- a/jac-cloud/jac_cloud/plugin/jaseci.py
+++ b/jac-cloud/jac_cloud/plugin/jaseci.py
@@ -30,6 +30,7 @@
 )
 from jaclang.plugin.feature import JacFeature as Jac
 from jaclang.runtimelib.architype import Architype, DSFunc
+from jaclang.runtimelib.utils import all_issubclass
 
 from orjson import loads
 
@@ -794,65 +795,85 @@ def spawn_call(op1: Architype, op2: Architype) -> WalkerArchitype:
         walker.path = []
         walker.next = [node]
         walker.returns = []
+        current_node = node.architype
+
+        # walker entry
+        for i in warch._jac_entry_funcs_:
+            if i.func and not i.trigger:
+                walker.returns.append(i.func(warch, current_node))
+            if walker.disengaged:
+                return warch
 
-        if walker.next:
-            current_node = walker.next[-1].architype
-            for i in warch._jac_entry_funcs_:
-                trigger = i.get_funcparam_annotations(i.func)
-                if not trigger:
-                    if i.func:
-                        walker.returns.append(i.func(warch, current_node))
-                    else:
-                        raise ValueError(f"No function {i.name} to call.")
         while len(walker.next):
             if current_node := walker.next.pop(0).architype:
+                # walker entry with
+                for i in warch._jac_entry_funcs_:
+                    if (
+                        i.func
+                        and i.trigger
+                        and all_issubclass(i.trigger, NodeArchitype)
+                        and isinstance(current_node, i.trigger)
+                    ):
+                        walker.returns.append(i.func(warch, current_node))
+                    if walker.disengaged:
+                        return warch
+
+                # node entry
                 for i in current_node._jac_entry_funcs_:
-                    trigger = i.get_funcparam_annotations(i.func)
-                    if not trigger or isinstance(warch, trigger):
-                        if i.func:
-                            walker.returns.append(i.func(current_node, warch))
-                        else:
-                            raise ValueError(f"No function {i.name} to call.")
+                    if i.func and not i.trigger:
+                        walker.returns.append(i.func(current_node, warch))
                     if walker.disengaged:
                         return warch
-                for i in warch._jac_entry_funcs_:
-                    trigger = i.get_funcparam_annotations(i.func)
-                    if not trigger or isinstance(current_node, trigger):
-                        if i.func and trigger:
-                            walker.returns.append(i.func(warch, current_node))
-                        elif not trigger:
-                            continue
-                        else:
-                            raise ValueError(f"No function {i.name} to call.")
+
+                # node entry with
+                for i in current_node._jac_entry_funcs_:
+                    if (
+                        i.func
+                        and i.trigger
+                        and all_issubclass(i.trigger, WalkerArchitype)
+                        and isinstance(warch, i.trigger)
+                    ):
+                        walker.returns.append(i.func(current_node, warch))
                     if walker.disengaged:
                         return warch
-                for i in warch._jac_exit_funcs_:
-                    trigger = i.get_funcparam_annotations(i.func)
-                    if not trigger or isinstance(current_node, trigger):
-                        if i.func and trigger:
-                            walker.returns.append(i.func(warch, current_node))
-                        elif not trigger:
-                            continue
-                        else:
-                            raise ValueError(f"No function {i.name} to call.")
+
+                # node exit with
+                for i in current_node._jac_exit_funcs_:
+                    if (
+                        i.func
+                        and i.trigger
+                        and all_issubclass(i.trigger, WalkerArchitype)
+                        and isinstance(warch, i.trigger)
+                    ):
+                        walker.returns.append(i.func(current_node, warch))
                     if walker.disengaged:
                         return warch
+
+                # node exit
                 for i in current_node._jac_exit_funcs_:
-                    trigger = i.get_funcparam_annotations(i.func)
-                    if not trigger or isinstance(warch, trigger):
-                        if i.func:
-                            walker.returns.append(i.func(current_node, warch))
-                        else:
-                            raise ValueError(f"No function {i.name} to call.")
+                    if i.func and not i.trigger:
+                        walker.returns.append(i.func(current_node, warch))
                     if walker.disengaged:
                         return warch
+
+                # walker exit with
+                for i in warch._jac_exit_funcs_:
+                    if (
+                        i.func
+                        and i.trigger
+                        and all_issubclass(i.trigger, NodeArchitype)
+                        and isinstance(current_node, i.trigger)
+                    ):
+                        walker.returns.append(i.func(warch, current_node))
+                    if walker.disengaged:
+                        return warch
+        # walker exit
         for i in warch._jac_exit_funcs_:
-            trigger = i.get_funcparam_annotations(i.func)
-            if not trigger:
-                if i.func:
-                    walker.returns.append(i.func(warch, current_node))
-                else:
-                    raise ValueError(f"No function {i.name} to call.")
+            if i.func and not i.trigger:
+                walker.returns.append(i.func(warch, current_node))
+            if walker.disengaged:
+                return warch
+
         walker.ignores = []
         return warch
 
diff --git a/jac-cloud/jac_cloud/tests/openapi_specs.yaml b/jac-cloud/jac_cloud/tests/openapi_specs.yaml
index 3f3bbe9c96..b132c2c40b 100644
--- a/jac-cloud/jac_cloud/tests/openapi_specs.yaml
+++ b/jac-cloud/jac_cloud/tests/openapi_specs.yaml
@@ -4109,4 +4109,53 @@ paths:
       summary: /visit_nested_node/{node}
       tags:
       - walker
+      - walker
+  /walker/visit_sequence:
+    post:
+      operationId: api_root_walker_visit_sequence_post
+      responses:
+        '200':
+          content:
+            application/json:
+              schema:
+                anyOf:
+                - $ref: '#/components/schemas/ContextResponse_NoneType_'
+                - {}
+                title: Response Api Root Walker Visit Sequence Post
+          description: Successful Response
+      summary: /visit_sequence
+      tags:
+      - walker
+      - walker
+  /walker/visit_sequence/{node}:
+    post:
+      operationId: api_entry_walker_visit_sequence__node__post
+      parameters:
+      - in: path
+        name: node
+        required: true
+        schema:
+          anyOf:
+          - type: string
+          - type: 'null'
+          title: Node
+      responses:
+        '200':
+          content:
+            application/json:
+              schema:
+                anyOf:
+                - $ref: '#/components/schemas/ContextResponse_NoneType_'
+                - {}
+                title: Response Api Entry Walker Visit Sequence  Node  Post
+          description: Successful Response
+        '422':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPValidationError'
+          description: Validation Error
+      summary: /visit_sequence/{node}
+      tags:
+      - walker
       - walker
\ No newline at end of file
diff --git a/jac-cloud/jac_cloud/tests/simple_graph.jac b/jac-cloud/jac_cloud/tests/simple_graph.jac
index 62437e93ed..57bcbced13 100644
--- a/jac-cloud/jac_cloud/tests/simple_graph.jac
+++ b/jac-cloud/jac_cloud/tests/simple_graph.jac
@@ -803,4 +803,58 @@ walker delete_custom_object {
         # The only difference is BaseAnchor.ref doesn't
         # load the actual object and just use it as reference
     }
+}
+
+##################################################################
+#                     FOR SPAWN CALL SEQUENCE                    #
+##################################################################
+
+node Node {
+    has val: str;
+
+    can entry1 with entry {
+        return f"{self.val}-2";
+    }
+
+    can entry3 with visit_sequence entry {
+        return f"{self.val}-3";
+    }
+
+    can exit1 with visit_sequence exit {
+        return f"{self.val}-4";
+    }
+
+    can exit2 with exit {
+        return f"{self.val}-5";
+    }
+}
+
+walker visit_sequence {
+    can entry1 with entry {
+        return "walker entry";
+    }
+
+    can entry2 with `root entry {
+        here ++> Node(val = "a");
+        here ++> Node(val = "b");
+        here ++> Node(val = "c");
+        visit [-->];
+        return "walker enter to root";
+    }
+
+    can entry3 with Node entry {
+        return f"{here.val}-1";
+    }
+
+    can exit1 with Node exit {
+        return f"{here.val}-6";
+    }
+
+    can exit2 with exit {
+        return "walker exit";
+    }
+
+    class __specs__ {
+        has auth: bool = False;
+    }
 }
\ No newline at end of file
diff --git a/jac-cloud/jac_cloud/tests/test_simple_graph.py b/jac-cloud/jac_cloud/tests/test_simple_graph.py
index 2ed04f080f..703baeba73 100644
--- a/jac-cloud/jac_cloud/tests/test_simple_graph.py
+++ b/jac-cloud/jac_cloud/tests/test_simple_graph.py
@@ -560,19 +560,19 @@ def trigger_upload_file(self) -> None:
                             "single": {
                                 "name": "simple_graph.jac",
                                 "content_type": "application/octet-stream",
-                                "size": 17658,
+                                "size": 18748,
                             }
                         },
                         "multiple": [
                             {
                                 "name": "simple_graph.jac",
                                 "content_type": "application/octet-stream",
-                                "size": 17658,
+                                "size": 18748,
                             },
                             {
                                 "name": "simple_graph.jac",
                                 "content_type": "application/octet-stream",
-                                "size": 17658,
+                                "size": 18748,
                             },
                         ],
                         "singleOptional": None,
@@ -761,6 +761,38 @@ def trigger_delete_custom_object_test(self, obj_id: str) -> None:
         self.assertEqual(200, res["status"])
         self.assertIsNone(obj)
 
+    def trigger_visit_sequence(self) -> None:
+        """Test visit sequence."""
+        res = self.post_api("visit_sequence")
+
+        self.assertEqual(200, res["status"])
+        self.assertEqual(
+            [
+                "walker entry",
+                "walker enter to root",
+                "a-1",
+                "a-2",
+                "a-3",
+                "a-4",
+                "a-5",
+                "a-6",
+                "b-1",
+                "b-2",
+                "b-3",
+                "b-4",
+                "b-5",
+                "b-6",
+                "c-1",
+                "c-2",
+                "c-3",
+                "c-4",
+                "c-5",
+                "c-6",
+                "walker exit",
+            ],
+            res["returns"],
+        )
+
     def test_all_features(self) -> None:
         """Test Full Features."""
         self.trigger_openapi_specs_test()
@@ -864,10 +896,16 @@ def test_all_features(self) -> None:
 
         self.trigger_memory_sync()
 
-        ##################################################
+        ###################################################
         #                 SAVABLE OBJECT                  #
         ###################################################
 
         obj_id = self.trigger_create_custom_object_test()
         self.trigger_update_custom_object_test(obj_id)
         self.trigger_delete_custom_object_test(obj_id)
+
+        ###################################################
+        #                  VISIT SEQUENCE                 #
+        ###################################################
+
+        self.trigger_visit_sequence()
diff --git a/jac/jaclang/plugin/default.py b/jac/jaclang/plugin/default.py
index 76fa0df859..988766f140 100644
--- a/jac/jaclang/plugin/default.py
+++ b/jac/jaclang/plugin/default.py
@@ -42,8 +42,11 @@
 from jaclang.runtimelib.importer import ImportPathSpec, JacImporter, PythonImporter
 from jaclang.runtimelib.machine import JacMachine, JacProgram
 from jaclang.runtimelib.memory import Shelf, ShelfStorage
-from jaclang.runtimelib.utils import collect_node_connections, traverse_graph
-
+from jaclang.runtimelib.utils import (
+    all_issubclass,
+    collect_node_connections,
+    traverse_graph,
+)
 
 import pluggy
 
@@ -398,64 +401,85 @@ def spawn_call(op1: Architype, op2: Architype) -> WalkerArchitype:
 
         walker.path = []
         walker.next = [node]
-        if walker.next:
-            current_node = walker.next[-1].architype
-            for i in warch._jac_entry_funcs_:
-                trigger = i.get_funcparam_annotations(i.func)
-                if not trigger:
-                    if i.func:
-                        i.func(warch, current_node)
-                    else:
-                        raise ValueError(f"No function {i.name} to call.")
+        current_node = node.architype
+
+        # walker entry
+        for i in warch._jac_entry_funcs_:
+            if i.func and not i.trigger:
+                i.func(warch, current_node)
+            if walker.disengaged:
+                return warch
+
         while len(walker.next):
             if current_node := walker.next.pop(0).architype:
+                # walker entry with
+                for i in warch._jac_entry_funcs_:
+                    if (
+                        i.func
+                        and i.trigger
+                        and all_issubclass(i.trigger, NodeArchitype)
+                        and isinstance(current_node, i.trigger)
+                    ):
+                        i.func(warch, current_node)
+                    if walker.disengaged:
+                        return warch
+
+                # node entry
                 for i in current_node._jac_entry_funcs_:
-                    trigger = i.get_funcparam_annotations(i.func)
-                    if not trigger or isinstance(warch, trigger):
-                        if i.func:
-                            i.func(current_node, warch)
-                        else:
-                            raise ValueError(f"No function {i.name} to call.")
+                    if i.func and not i.trigger:
+                        i.func(current_node, warch)
                     if walker.disengaged:
                         return warch
-                for i in warch._jac_entry_funcs_:
-                    trigger = i.get_funcparam_annotations(i.func)
-                    if not trigger or isinstance(current_node, trigger):
-                        if i.func and trigger:
-                            i.func(warch, current_node)
-                        elif not trigger:
-                            continue
-                        else:
-                            raise ValueError(f"No function {i.name} to call.")
+
+                # node entry with
+                for i in current_node._jac_entry_funcs_:
+                    if (
+                        i.func
+                        and i.trigger
+                        and all_issubclass(i.trigger, WalkerArchitype)
+                        and isinstance(warch, i.trigger)
+                    ):
+                        i.func(current_node, warch)
                     if walker.disengaged:
                         return warch
-                for i in warch._jac_exit_funcs_:
-                    trigger = i.get_funcparam_annotations(i.func)
-                    if not trigger or isinstance(current_node, trigger):
-                        if i.func and trigger:
-                            i.func(warch, current_node)
-                        elif not trigger:
-                            continue
-                        else:
-                            raise ValueError(f"No function {i.name} to call.")
+
+                # node exit with
+                for i in current_node._jac_exit_funcs_:
+                    if (
+                        i.func
+                        and i.trigger
+                        and all_issubclass(i.trigger, WalkerArchitype)
+                        and isinstance(warch, i.trigger)
+                    ):
+                        i.func(current_node, warch)
                     if walker.disengaged:
                         return warch
+
+                # node exit
                 for i in current_node._jac_exit_funcs_:
-                    trigger = i.get_funcparam_annotations(i.func)
-                    if not trigger or isinstance(warch, trigger):
-                        if i.func:
-                            i.func(current_node, warch)
-                        else:
-                            raise ValueError(f"No function {i.name} to call.")
+                    if i.func and not i.trigger:
+                        i.func(current_node, warch)
                     if walker.disengaged:
                         return warch
+
+                # walker exit with
+                for i in warch._jac_exit_funcs_:
+                    if (
+                        i.func
+                        and i.trigger
+                        and all_issubclass(i.trigger, NodeArchitype)
+                        and isinstance(current_node, i.trigger)
+                    ):
+                        i.func(warch, current_node)
+                    if walker.disengaged:
+                        return warch
+        # walker exit
         for i in warch._jac_exit_funcs_:
-            trigger = i.get_funcparam_annotations(i.func)
-            if not trigger:
-                if i.func:
-                    i.func(warch, current_node)
-                else:
-                    raise ValueError(f"No function {i.name} to call.")
+            if i.func and not i.trigger:
+                i.func(warch, current_node)
+            if walker.disengaged:
+                return warch
+
         walker.ignores = []
         return warch
 
diff --git a/jac/jaclang/runtimelib/architype.py b/jac/jaclang/runtimelib/architype.py
index 382bd41945..25c19a7d88 100644
--- a/jac/jaclang/runtimelib/architype.py
+++ b/jac/jaclang/runtimelib/architype.py
@@ -5,6 +5,7 @@
 import inspect
 from dataclasses import asdict, dataclass, field, fields, is_dataclass
 from enum import IntEnum
+from functools import cached_property
 from logging import getLogger
 from pickle import dumps
 from types import UnionType
@@ -220,9 +221,9 @@ class WalkerAnchor(Anchor):
     """Walker Anchor."""
 
     architype: WalkerArchitype
-    path: list[Anchor] = field(default_factory=list)
-    next: list[Anchor] = field(default_factory=list)
-    ignores: list[Anchor] = field(default_factory=list)
+    path: list[NodeAnchor] = field(default_factory=list)
+    next: list[NodeAnchor] = field(default_factory=list)
+    ignores: list[NodeAnchor] = field(default_factory=list)
     disengaged: bool = False
 
 
@@ -311,17 +312,20 @@ class DSFunc:
     name: str
     func: Callable[[Any, Any], Any] | None = None
 
+    @cached_property
+    def trigger(self) -> type | UnionType | tuple[type | UnionType, ...] | None:
+        """Get function parameter annotations."""
+        t = (
+            (
+                inspect.signature(self.func, eval_str=True)
+                .parameters["_jac_here_"]
+                .annotation
+            )
+            if self.func
+            else None
+        )
+        return None if t is inspect._empty else t
+
     def resolve(self, cls: type) -> None:
         """Resolve the function."""
         self.func = getattr(cls, self.name)
-
-    def get_funcparam_annotations(
-        self, func: Callable[[Any, Any], Any] | None
-    ) -> type | UnionType | tuple[type | UnionType, ...] | None:
-        """Get function parameter annotations."""
-        if not func:
-            return None
-        annotation = (
-            inspect.signature(func, eval_str=True).parameters["_jac_here_"].annotation
-        )
-        return annotation if annotation != inspect._empty else None
diff --git a/jac/jaclang/runtimelib/utils.py b/jac/jaclang/runtimelib/utils.py
index 83e35ab2cf..738eec2f36 100644
--- a/jac/jaclang/runtimelib/utils.py
+++ b/jac/jaclang/runtimelib/utils.py
@@ -231,3 +231,18 @@ def is_instance(
             return isinstance(obj, target)
         case _:
             return False
+
+
+def all_issubclass(
+    classes: type | UnionType | tuple[type | UnionType, ...], target: type
+) -> bool:
+    """Check if all classes is subclass of target type."""
+    match classes:
+        case type():
+            return issubclass(classes, target)
+        case UnionType():
+            return all((all_issubclass(cls, target) for cls in classes.__args__))
+        case tuple():
+            return all((all_issubclass(cls, target) for cls in classes))
+        case _:
+            return False
diff --git a/jac/jaclang/tests/fixtures/visit_sequence.jac b/jac/jaclang/tests/fixtures/visit_sequence.jac
new file mode 100644
index 0000000000..54ea0c9b15
--- /dev/null
+++ b/jac/jaclang/tests/fixtures/visit_sequence.jac
@@ -0,0 +1,50 @@
+node Node {
+    has val: str;
+
+    can entry1 with entry {
+        print(f"{self.val}-2");
+    }
+
+    can entry2 with Walker entry {
+        print(f"{self.val}-3");
+    }
+
+    can exit1 with Walker exit {
+        print(f"{self.val}-4");
+    }
+
+    can exit2 with exit {
+        print(f"{self.val}-5");
+    }
+}
+
+walker Walker {
+    can entry1 with entry {
+        print("walker entry");
+    }
+
+    can entry2 with `root entry {
+        print("walker enter to root");
+        visit [-->];
+    }
+
+    can entry3 with Node entry {
+        print(f"{here.val}-1");
+    }
+
+    can exit1 with Node exit {
+        print(f"{here.val}-6");
+    }
+
+    can exit2 with exit {
+        print("walker exit");
+    }
+}
+
+with entry{
+    root ++> Node(val = "a");
+    root ++> Node(val = "b");
+    root ++> Node(val = "c");
+
+    Walker() spawn root;
+}
\ No newline at end of file
diff --git a/jac/jaclang/tests/test_language.py b/jac/jaclang/tests/test_language.py
index 0bf7b6e09e..15b28d10d1 100644
--- a/jac/jaclang/tests/test_language.py
+++ b/jac/jaclang/tests/test_language.py
@@ -1230,5 +1230,20 @@ def test_architype_def(self) -> None:
         jac_import("architype_def_bug", base_path=self.fixture_abs_path("./"))
         sys.stdout = sys.__stdout__
         stdout_value = captured_output.getvalue().split("\n")
-        self.assertIn("MyNode", stdout_value[0])
-        self.assertIn("MyWalker", stdout_value[1])
+        self.assertIn("MyWalker", stdout_value[0])
+        self.assertIn("MyNode", stdout_value[1])
+
+    def test_visit_sequence(self) -> None:
+        """Test conn assign on edges."""
+        captured_output = io.StringIO()
+        sys.stdout = captured_output
+        jac_import("visit_sequence", base_path=self.fixture_abs_path("./"))
+        sys.stdout = sys.__stdout__
+        self.assertEqual(
+            "walker entry\nwalker enter to root\n"
+            "a-1\na-2\na-3\na-4\na-5\na-6\n"
+            "b-1\nb-2\nb-3\nb-4\nb-5\nb-6\n"
+            "c-1\nc-2\nc-3\nc-4\nc-5\nc-6\n"
+            "walker exit\n",
+            captured_output.getvalue(),
+        )