Merge pull request Jaseci-Labs#245 from Jaseci-Labs/nodes_in_edge_refs

feat: node types in edge ref lists

jaclang/compiler/absyntree.py

@@ -1764,7 +1764,7 @@ class EdgeRefTrailer(Expr):
 
  def __init__(
  self,
- chain: list[Expr],
+ chain: list[Expr | FilterCompr],
  edges_only: bool,
  kid: Sequence[AstNode],
  ) -> None:

jaclang/compiler/jac.lark

@@ -416,7 +416,7 @@ kw_expr_list: (kw_expr_list COMMA)? kw_expr
 kw_expr: any_ref EQ expression | STAR_POW expression
 
 // Data Spatial References
-edge_ref_chain: (EDGE_OP|NODE_OP)? LSQUARE expression? (edge_op_ref expression?)+ RSQUARE
+edge_ref_chain: (EDGE_OP|NODE_OP)? LSQUARE expression? (edge_op_ref (filter_compr | expression)?)+ RSQUARE
 edge_op_ref: edge_any | edge_from | edge_to
 edge_to: ARROW_R | ARROW_R_P1 typed_filter_compare_list ARROW_R_P2
 edge_from: ARROW_L | ARROW_L_P1 typed_filter_compare_list ARROW_L_P2

jaclang/compiler/parser.py

@@ -3110,9 +3110,10 @@ def arch_to_enum_chain(self, kid: list[ast.AstNode]) -> ast.ArchRefChain:
  def edge_ref_chain(self, kid: list[ast.AstNode]) -> ast.EdgeRefTrailer:
  """Grammar rule.
 
- edge_ref_chain: (EDGE_OP|NODE_OP)? LSQUARE expression? (edge_op_ref expression?)+ RSQUARE
+ edge_ref_chain: (EDGE_OP|NODE_OP)? LSQUARE expression?
+ (edge_op_ref (NODE_OP? expression)?)+ RSQUARE
  """
- valid_chain = [i for i in kid if isinstance(i, (ast.Expr))]
+ valid_chain = [i for i in kid if isinstance(i, (ast.Expr, ast.FilterCompr))]
  return self.nu(
  ast.EdgeRefTrailer(
  chain=valid_chain,

jaclang/compiler/passes/main/pyast_gen_pass.py

@@ -2441,7 +2441,7 @@ def exit_special_var_ref(self, node: ast.SpecialVarRef) -> None:
  def exit_edge_ref_trailer(self, node: ast.EdgeRefTrailer) -> None:
  """Sub objects.
 
- chain: list[Expr],
+ chain: list[Expr|FilterCompr],
  edges_only: bool,
  """
  pynode = node.chain[0].gen.py_ast[0]
@@ -2461,11 +2461,23 @@ def exit_edge_ref_trailer(self, node: ast.EdgeRefTrailer) -> None:
  not next_i or not isinstance(next_i, ast.EdgeOpRef)
  ):
  pynode = self.translate_edge_op_ref(
- pynode,
- cur,
- targ=next_i.gen.py_ast[0] if next_i else None,
+ loc=pynode,
+ node=cur,
+ targ=(
+ next_i.gen.py_ast[0]
+ if next_i and not isinstance(next_i, ast.FilterCompr)
+ else None
+ ),
  edges_only=node.edges_only and cur == last_edge,
  )
+ if next_i and isinstance(next_i, ast.FilterCompr):
+ pynode = self.sync(
+ ast3.Call(
+ func=next_i.gen.py_ast[0],
+ args=[pynode],
+ keywords=[],
+ )
+ )
  chomp = chomp[1:] if next_i else chomp
  elif isinstance(cur, ast.EdgeOpRef) and isinstance(next_i, ast.EdgeOpRef):
  pynode = self.translate_edge_op_ref(

jaclang/tests/fixtures/edge_node_walk.jac

@@ -0,0 +1,46 @@
+walker creator {
+ can create with `<root> entry;
+}
+
+node node_a {
+ has val: int;
+}
+
+node node_b {
+ has val: int;
+}
+
+edge Edge_a {
+ has value: int = 10;
+}
+
+edge Edge_b {
+ has value1: int = 20;
+}
+
+edge Edge_c {
+ has value2: int = 30;
+}
+
+:walker:creator:can:create {
+ for i=0 to i<2 by i+=1 {
+ <here> +:Edge_c:value2=30:+> (node_a(val=i + 1));
+ visit [-->];
+ }
+ <here> +:Edge_a:value=10:+> (end := node_a(val=i + 10));
+ end +:Edge_b:value1=20:+> (last := node_a(val=i + 20));
+ for j=0 to j<2 by j+=1 {
+ last +:Edge_c:value2=40:+> (node_a(val=i + 40));
+ last +:Edge_c:value2=40:+> (node_b(val=i + 40));
+ }
+}
+
+with entry {
+ print(<root> spawn creator());
+ print(<r>._jac_.gen_dot());
+ print([<root>-:Edge_a:->]);
+ print([<root>-:Edge_c:->]);
+ print([<root>-:Edge_a:->-:Edge_b:->]);
+ print([<root>-:Edge_a:->-:Edge_b:->-:Edge_c:->]);
+ print([<root>-:Edge_a:->-:Edge_b:->-:Edge_c:->(`?node_b)]);
+}

jaclang/tests/test_language.py

@@ -340,3 +340,17 @@ def test_typed_filter_compr(self) -> None:
  stdout_value,
  )
  self.assertIn("[MyObj(a=0), MyObj(a=1), MyObj(a=2)]\n", stdout_value)
+
+ def test_edge_node_walk(self) -> None:
+ """Test walking through edges and nodes."""
+ construct.root._jac_.edges.clear()
+ captured_output = io.StringIO()
+ sys.stdout = captured_output
+ jac_import("edge_node_walk", base_path=self.fixture_abs_path("./"))
+ sys.stdout = sys.__stdout__
+ stdout_value = captured_output.getvalue()
+ self.assertIn("creator()\n", stdout_value)
+ self.assertIn("[node_a(val=12)]\n", stdout_value)
+ self.assertIn("node_a(val=1)", stdout_value)
+ self.assertIn("node_a(val=2)", stdout_value)
+ self.assertIn("[node_b(val=42), node_b(val=42)]\n", stdout_value)