fixed branch coordination

Author: mathisxy

examples/gambling.py

@@ -66,8 +66,8 @@ async def __call__(self, state: GamblingState, shared: Shared) -> None:
 ### GRAPH
 
 asyncio.run(Graph[GamblingState, Shared](
-    edges=[
-        ((
+    edges=[(
+        (
             START,
             guess
         ),

examples/hello_world.py

@@ -32,15 +32,15 @@ async def __call__(self, state: MyState, shared: Shared) -> None:
 ### GRAPH
 
 graph = Graph[MyState, Shared](
-    edges=[
+    edges=[(
         (
             START,
             node
         ),
         (
             node,
             END
-        )
+        ), END)
     ]
 )
 

src/edgygraph/graph/branches.py

@@ -20,24 +20,18 @@ class Branch[T: StateProtocol, S: SharedProtocol]:
         join: The node to join the branch after execution. If None the branch will not be joined.
 
     """
-    
-    edges: Sequence[Edge[T, S] | ErrorEdge[T, S] | NodeTupel[T, S]]
 
-    join: SingleNext[T, S]
 
-    result: asyncio.Future[dict[tuple[Hashable, ...], Change]] | None = None
-
-    edge_index: dict[SingleSource[T, S], list[Entry[T, S]]]
-    error_edge_index: dict[SingleErrorSource[T, S], list[ErrorEntry[T, S]]]
-
-
-    def __init__(self, edges: Sequence[Edge[T, S] | ErrorEdge[T, S] | NodeTupel[T, S]], join: SingleNext[T, S] = None) -> None:
+    def __init__(self, edges: Sequence[Edge[T, S] | ErrorEdge[T, S] | NodeTupel[T, S]], start: SingleSource[T, S], join: SingleNext[T, S] = None) -> None:
 
         self.edges = edges
+        self.start = start
         self.join = join
 
-        self.edge_index = defaultdict(list)
-        self.error_edge_index = defaultdict(list)
+        self.result: asyncio.Future[dict[tuple[Hashable, ...], Change]] | None = None
+
+        self.edge_index: dict[SingleSource[T, S], list[Entry[T, S]]] = defaultdict(list)
+        self.error_edge_index: dict[SingleErrorSource[T, S], list[ErrorEntry[T, S]]] = defaultdict(list)
 
         self.index_edges()
 
@@ -66,11 +60,12 @@ def index_edges(self) -> None:
 
             match edge:
                 case (source, next, config): pass
-                case (source, next): config = Config() if source is START or isinstance(source, (Node, Sequence)) else ErrorConfig()
+                case (source, next): config = None
                 case _: raise ValueError(f"Invalid edge format: {edge}")
 
             if Types[T, S].is_error_source(source):
-                assert isinstance(config, ErrorConfig), f"Unexpected properties type for error edge: {config}"
+                config = config or ErrorConfig()
+                assert isinstance(config, ErrorConfig), f"Unexpected properties type for error edge {edge}: {config}"
 
                 if Types[T, S].is_single_error_source(source):
                     self.error_edge_index[source].append(ErrorEntry[T, S](next=next, config=config, index=i))
@@ -81,7 +76,8 @@ def index_edges(self) -> None:
                     raise ValueError(f"Invalid error source: {source}")
 
             elif Types[T, S].is_source(source):
-                assert isinstance(config, Config), f"Unexpected properties type for node edge: {config}"
+                config = config or Config()
+                assert isinstance(config, Config), f"Unexpected properties type for node edge {edge}: {config}"
 
                 if Types[T, S].is_single_source(source):
                     self.edge_index[source].append(Entry[T, S](next=next, config=config, index=i))

src/edgygraph/graph/graphs.py

@@ -1,7 +1,7 @@
 
 from __future__ import annotations
 
-from typing import cast, Any, Hashable
+from typing import cast, Any, Hashable, Callable
 from collections import defaultdict
 from collections.abc import Hashable, Sequence
 import asyncio
@@ -10,7 +10,7 @@
 from ..states import StateProtocol, SharedProtocol
 from ..diff import Change, ChangeConflictException, Diff
 from ..nodes import Node, END, START
-from .types import SingleNext, NextNode, ErrorEntry, SingleErrorSource, Entries, BranchContainer, SingleSource
+from .types import SingleNext, NextNode, ErrorEntry, SingleErrorSource, Entries, BranchContainer, SingleSource, Source, Types, ResolvedNext
 from .hooks import GraphHook
 from .branches import Branch
 
@@ -92,14 +92,26 @@ def __init__(self,
         self.index_branches()
 
     def index_branches(self) -> None:
-        for branch in self.edges:
+        for branch_container in self.edges:
 
-            source = branch[0][0][0]
+            source = branch_container[0][0]
 
-            start_nodes = [source] if isinstance(source, (Node, type)) else source
+            if Types[T, S].is_single_source(source):
 
-            for start_node in start_nodes:
-                self.branch_registry[start_node].append(Branch[T, S](branch[0], branch[1]))
+                branch = Branch[T, S](branch_container[:-1], source, branch_container[-1])
+                self.branch_registry[source].append(branch)
+
+            elif Types[T, S].is_single_source_sequence(source):
+
+                for start_node in source:
+
+                    branch = Branch[T, S](branch_container[:-1], start_node, branch_container[-1])
+                    self.branch_registry[start_node].append(branch)
+
+            else:
+                raise ValueError(f"Invalid source type: {source}")
+
+            print(self.branch_registry)
 
 
     async def __call__(self, state: T, shared: S) -> tuple[T, S]:
@@ -157,7 +169,9 @@ async def run_branch(self, state: T, shared: S, branch: Branch[T, S]) -> None:
 
         try:
 
-            next_nodes: list[NextNode[T, S]] = await self.get_next(state, shared, START, branch)
+            next_nodes: list[NextNode[T, S]] = await self.get_next(state, shared, branch.start, branch)
+
+            print("INITIAL NEXT:", next_nodes)
 
 
             while next_nodes:
@@ -197,7 +211,7 @@ async def run_branch(self, state: T, shared: S, branch: Branch[T, S]) -> None:
 
                 else:
 
-                    next_nodes = await self.get_next(state, shared, next_nodes, branch)
+                    next_nodes = await self.get_next(state, shared, [n.node for n in next_nodes], branch)
 
                 finally:
 
@@ -215,6 +229,8 @@ async def run_branch(self, state: T, shared: S, branch: Branch[T, S]) -> None:
 
             if e:
                 raise e
+
+        print(" --- BRANCH RESULT --- ")
 
         branch.result.set_result(Diff.recursive_diff(initial_state.model_dump(), state.model_dump()))
 
@@ -333,7 +349,12 @@ async def spawn_branches(self, state: T, shared: S, next_nodes: list[NextNode[T,
 
         for node in next_nodes:
             for branch in self.branch_registry[node.node]:
+
+                for h in self.hooks: await h.on_spawn_branch_start(state, shared, branch, node, self.branch_registry, self.join_registry)
+
                 self.spawn_branch(state, shared, branch)
+
+                for h in self.hooks: await h.on_spawn_branch_end(state, shared, branch, node, self.branch_registry, self.join_registry)
 
 
     async def join_branches(self, state: T, next_nodes: list[NextNode[T, S]]) -> T:
@@ -366,7 +387,7 @@ async def join_branches(self, state: T, next_nodes: list[NextNode[T, S]]) -> T:
 
 
 
-    async def get_next(self, state: T, shared: S, current_nodes: Sequence[NextNode[T, S]] | type[START], branch: Branch[T, S]) -> list[NextNode[T, S]]:
+    async def get_next(self, state: T, shared: S, current_nodes: Source[T, S], branch: Branch[T, S]) -> list[NextNode[T, S]]:
         """
         Get the next nodes to run based on the current nodes and the graph's edges.
 
@@ -381,22 +402,26 @@ async def get_next(self, state: T, shared: S, current_nodes: Sequence[NextNode[T
            The list of the next nodes including their edges that they were reached by.
         """
 
+        print(f"GET NEXT FOR CURRENT NODES: {current_nodes}")
 
         next_list: list[NextNode[T, S]] = []
 
-        if isinstance(current_nodes, type): # START
-            next_list.extend(
-                await self.resolve_entries(state, shared, branch.edge_index[START])
-            )
-
-        else: # Regular nodes
+        if Types[T, S].is_single_source_sequence(current_nodes):
+            print("IS SINGLE SOURCE SEQUENCE")
             for current_node in current_nodes:
-
-                # Find the edge corresponding to the current node
                 next_list.extend(
-                    await self.resolve_entries(state, shared, branch.edge_index[current_node.node])
+                    await self.resolve_entries(state, shared, branch.edge_index[current_node])
                 )
 
+        elif Types[T, S].is_single_source(current_nodes):
+            print("IS SINGLE SOURCE")
+            next_list.extend(
+                await self.resolve_entries(state, shared, branch.edge_index[current_nodes])
+            )
+        
+        else:
+            raise ValueError(f"Invalid current_nodes type: {type(current_nodes)}")
+
 
         # Instant nodes
         current_instant_next_list: list[NextNode[T, S]] = []
@@ -489,23 +514,16 @@ def match_error(self, e: Exception, source: SingleErrorSource[T, S], source_node
 
     async def resolve_entries(self, state: T, shared: S, entries: Sequence[Entries[T, S]]) -> list[NextNode[T, S]]:
 
-        next_nodes: list[NextNode[T, S]] = []
+        print(f"RESOLVE: {entries}")
 
-        for entry in entries:
-            next_list = await self.resolve_entry(state, shared, entry)
+        return [
+            next_node
+            for entry in entries
+            for next_node in await self.resolve_entry(state, shared, entry)
+        ]
 
-            for x in next_list:
 
-                match x:
-                    case NextNode():
-                        next_nodes.append(x)
-                    case Branch():
-                        self.spawn_branch(state, shared, x)
-
-        return next_nodes
-
-
-    async def resolve_entry(self, state: T, shared: S, entry: Entries[T, S]) -> list[NextNode[T, S] | Branch[T, S]]:
+    async def resolve_entry(self, state: T, shared: S, entry: Entries[T, S]) -> list[NextNode[T, S]]:
         """
         Resolve the next to nodes.
 
@@ -519,36 +537,60 @@ async def resolve_entry(self, state: T, shared: S, entry: Entries[T, S]) -> list
             The resolved nodes.
         """
 
-        next_nodes: list[NextNode[T, S] | Branch[T, S]] = []
+        print(f"RESOLVING: {entry}")
+
         next = entry.next
 
-        match next:
 
-            case None:
-                pass # END
 
-            case type():
-                assert next is END, "Only END is allowed as a type here"
-            
-            case Node():
-                next_nodes.append(NextNode[T, S](node=next, reached_by=entry))
+        if not Types[T, S].is_resolved_next(next):
+            print("CALLABLE")
+            next = cast(Callable[[T, S], ResolvedNext[T, S]], next)
+            next = next(state, shared)
+
+            if inspect.isawaitable(next):
+                next = await next
+
+        
+        return [
+            NextNode[T, S](node=node, reached_by=entry)
+            for node in self.get_next_nodes(next)
+        ]
+    
+
+
+    def get_next_nodes(self, next: ResolvedNext[T, S]) -> list[Node[T, S]]:
+
+        next_nodes: list[Node[T, S]] = []
+
+        def match(x: SingleNext[T, S]) -> None:
 
-            case Sequence():
-                for n in next:
-                    if isinstance(n, Node):
-                        next_nodes.append(NextNode[T, S](node=n, reached_by=entry))
+            match x:
 
+                case None:
+                    print("None")
+                    pass # END
 
-            case _: # callable
-                next = next
-                res = next(state, shared)
-                if inspect.isawaitable(res):
-                    res = await res # for awaitables
+                case type():
+                    print("END")
+                    assert next is END, "Only END is allowed as a type here"
 
-                if isinstance(res, Node):
-                    next_nodes.append(NextNode[T, S](node=res, reached_by=entry))
+                case Node():
+                    print("Node")
+                    next_nodes.append(x)
 
 
+        if Types[T, S].is_single_next_sequence(next):
+            print("Sequence")
+            for x in next:
+                match(x)
+
+        elif Types[T, S].is_single_next(next):
+            print("Single")
+            match(next)
+
         return next_nodes
+                    
+                    
 
 

src/edgygraph/graph/hooks.py

@@ -3,7 +3,8 @@
 
 from ..states import StateProtocol, SharedProtocol
 from ..diff import Change
-from .types import NextNode
+from .branches import Branch
+from .types import NextNode, SingleSource, SingleNext
 
 
 
@@ -55,6 +56,37 @@ async def on_step_end(self, state: T, shared: S, nodes: list[NextNode[T, S]]) ->
         pass
 
 
+    async def on_spawn_branch_start(self, state: T, shared: S, branch: Branch[T, S], trigger: NextNode[T, S], branch_registry: dict[SingleSource[T, S], list[Branch[T, S]]], join_registry: dict[SingleNext[T, S], list[Branch[T, S]]]):
+        """
+        Called before a branch is spawned.
+
+        Args:
+            state: The state of the graph.
+            shared: The shared state of the graph.
+            branch: The branch to be spawned.
+            branch_registry: The branch registry of the graph.
+            source_node: The node that spawned the branch.
+        """
+
+        pass
+
+    
+    async def on_spawn_branch_end(self, state: T, shared: S, branch: Branch[T, S], trigger: NextNode[T, S], branch_registry: dict[SingleSource[T, S], list[Branch[T, S]]], join_registry: dict[SingleNext[T, S], list[Branch[T, S]]]):
+        """
+        Called after a branch is spawned.
+
+        Args:
+            state: The state of the graph.
+            shared: The shared state of the graph.
+            branch: The branch that was spawned.
+            branch_registry: The branch registry of the graph.
+            trigger: The node that spawned the branch.
+        """
+
+        pass
+
+
+
     async def on_merge_start(self, state: T, result_states: list[T], changes: list[dict[tuple[Hashable, ...], Change]]) -> None:
         """
         Called when the merge process starts.