nbs/06_transform.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Transform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp transform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "from nbdev.showdoc import show_doc\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Overview\n",
    "This final module handles the actual **graph transformation**, which our library aims to achieve. It contains general auxiliary functions for performing different operations on some graph - adding nodes, cloning them, removing edge attributes, etc.\n",
    "\n",
    "It also contains **```rewrite```**, the library's main function and the only one the user is supposed to use. ```rewrite``` allows pattern matching in an input graph, graph transformations based on transform rules (as defined in the \"rules\" module), imperative side-effects based on the matches and much more. It does so by using all the abilities presented in the previous modules and wrapping it all with a simple-to-use interface of a single function."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Requirements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "from typing import *\n",
    "from networkx import DiGraph\n",
    "from copy import deepcopy\n",
    "from fastcore.meta import delegates\n",
    "\n",
    "from graph_rewrite.core import NodeName, EdgeName, _create_graph, draw, _graphs_equal, GraphRewriteException\n",
    "from graph_rewrite.lhs import lhs_to_graph\n",
    "from graph_rewrite.match_class import Match, mapping_to_match,draw_match\n",
    "from graph_rewrite.matcher import find_matches, FilterFunc\n",
    "from graph_rewrite.p_rhs_parse import RenderFunc, p_to_graph, rhs_to_graph\n",
    "from graph_rewrite.rules import Rule, MergePolicy"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Transformation Utilities\n",
    "The following functions perform different operations on NetworkX graphs - any graph. We will use them when we'll do graph transformations, which perform a subset of these operations on some input graph, based on a predefined rule."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "_exception_msgs = {\n",
    "    \"no_such_node\": lambda node: f\"Node {node} does not exist in the input graph.\",\n",
    "    \"no_such_edge\": lambda edge: f\"Edge {edge} does not exist in the input graph.\",\n",
    "    \"no_such_attr_in_node\": lambda attr, node: f\"Attribute {attr} does not exist in input graph's node {node}.\",\n",
    "    \"no_such_attr_in_edge\": lambda attr, edge: f\"Attribute {attr} does not exist in input graph's edge {edge}.\",\n",
    "    \"edge_exists\": lambda edge: f\"Edge {edge} already exists in the input graph.\",\n",
    "    \"not_enough_to_merge\": lambda: f\"Tried to merge less than one nodes.\"\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _generate_new_node_name(graph: DiGraph, base_name: NodeName) -> NodeName:\n",
    "    \"\"\"Generate a name for a new node, which is unique in the graph to which the node is added,\n",
    "    based on an initial name suggestion.\n",
    "    If that suggested name does not exist in the graph, it is used for the new node.\n",
    "    Otherwise, the new node's name will have the suggested name as a prefix, followed by a serial number (<name>_<num>)\n",
    "    that makes sure that the name is unique but not too long.\n",
    "\n",
    "    Args:\n",
    "        graph (DiGraph): A graph to which the new node is added\n",
    "        base_name (NodeName): The initial suggestion for the new node's desired name\n",
    "\n",
    "    Returns:\n",
    "        NodeName: The name chosen for the new node\n",
    "    \"\"\"\n",
    "    new_name = base_name\n",
    "    i = 0\n",
    "    while new_name in graph.nodes():\n",
    "        i += 1\n",
    "        new_name = f\"{base_name}_{i}\"\n",
    "    return new_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _clone_node(graph: DiGraph, node_to_clone: NodeName) -> NodeName:\n",
    "    \"\"\"Clones a node in the graph. That is, create a new node, whose name denotes its connection to the original node,\n",
    "    whose edges are copies of the edges connected to the original node, and whose attributes are duplicated\n",
    "    from the original node.\n",
    "\n",
    "    Args:\n",
    "        graph (DiGraph): A graph\n",
    "        node_to_clone (NodeName): A node in the graph to be cloned\n",
    "\n",
    "    Raises:\n",
    "        GraphRewriteException: If the node which should be cloned does not exist in the graph\n",
    "\n",
    "    Returns:\n",
    "        NodeName: The name of the new cloned node\n",
    "    \"\"\"\n",
    "    if node_to_clone not in graph.nodes:\n",
    "        raise GraphRewriteException(_exception_msgs[\"no_such_node\"](node_to_clone))\n",
    "\n",
    "    # Create a new node name\n",
    "    clone_name = _generate_new_node_name(graph, node_to_clone)\n",
    "\n",
    "    # Add this new node to graph\n",
    "    cloned_node_attrs = graph.nodes(data=True)[node_to_clone]\n",
    "    graph.add_node(clone_name, **cloned_node_attrs)\n",
    "\n",
    "    # Clone edges (connect the clone to all original edge endpoints + copy attrs)\n",
    "    for n, _ in graph.in_edges(node_to_clone):\n",
    "        if (n, clone_name) not in graph.edges():\n",
    "            cloned_edge_attrs = graph.edges[n, node_to_clone]\n",
    "            graph.add_edge(n, clone_name, **cloned_edge_attrs)\n",
    "    \n",
    "    for _, n in graph.out_edges(node_to_clone):\n",
    "        if (clone_name, n) not in graph.edges():\n",
    "            cloned_edge_attrs = graph.edges[node_to_clone, n]\n",
    "            graph.add_edge(clone_name, n, **cloned_edge_attrs)\n",
    " \n",
    "    return clone_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _remove_node(graph: DiGraph, node_to_remove: NodeName):\n",
    "    \"\"\"Remove a node from the graph.\n",
    "\n",
    "    Args:\n",
    "        graph (DiGraph): A graph\n",
    "        node_to_remove (NodeName): A node to remove from the graph\n",
    "\n",
    "    Raises:\n",
    "        GraphRewriteException: If the removed node doesn't exist in the graph\n",
    "    \"\"\"\n",
    "    if node_to_remove not in graph.nodes():\n",
    "        raise GraphRewriteException(_exception_msgs[\"no_such_node\"](node_to_remove))\n",
    "    graph.remove_node(node_to_remove)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _remove_edge(graph: DiGraph, edge_to_remove: EdgeName):\n",
    "    \"\"\"Remove an edge from the graph.\n",
    "\n",
    "    Args:\n",
    "        graph (DiGraph): A graph\n",
    "        edge_to_remove (EdgeName): An edge to remove from the graph\n",
    "\n",
    "    Raises:\n",
    "        GraphRewriteException: If the removed edge doesn't exist in the graph\n",
    "    \"\"\"\n",
    "    if edge_to_remove not in graph.edges():\n",
    "        raise GraphRewriteException(_exception_msgs[\"no_such_edge\"](edge_to_remove))\n",
    "    graph.remove_edge(*edge_to_remove)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _remove_node_attrs(graph: DiGraph, node: NodeName, attrs_to_remove: set):\n",
    "    \"\"\"Remove a subset of some node's attributes from the node.\n",
    "\n",
    "    Args:\n",
    "        graph (DiGraph): A graph\n",
    "        node (NodeName): A node in the graph\n",
    "        attrs_to_remove (set): Attributes of that node to remove\n",
    "\n",
    "    Raises:\n",
    "        GraphRewriteException: If the node doesn't exist in the graph, or some of the removed attrs don't exist in the node\n",
    "    \"\"\"\n",
    "    if node not in graph.nodes():\n",
    "        raise GraphRewriteException(_exception_msgs[\"no_such_node\"](node))\n",
    "    for attr in attrs_to_remove:\n",
    "        if attr not in graph.nodes[node]:\n",
    "            raise GraphRewriteException(_exception_msgs[\"no_such_attr_in_node\"](attr, node))\n",
    "        del graph.nodes[node][attr]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _remove_edge_attrs(graph: DiGraph, edge: EdgeName, attrs_to_remove: set):\n",
    "    \"\"\"Remove a subset of some edge's attributes from the node.\n",
    "\n",
    "    Args:\n",
    "        graph (DiGraph): A graph\n",
    "        edge (EdgeName): An edge in the graph\n",
    "        attrs_to_remove (set): Attributes of that edge to remove\n",
    "\n",
    "    Raises:\n",
    "        GraphRewriteException: If the edge doesn't exist in the graph, or some of the removed attrs don't exist in the edge\n",
    "    \"\"\"\n",
    "\n",
    "    if edge not in graph.edges():\n",
    "        raise GraphRewriteException(_exception_msgs[\"no_such_edge\"](edge))\n",
    "    for attr in attrs_to_remove:\n",
    "        if attr not in graph.edges[edge]:\n",
    "            raise GraphRewriteException(_exception_msgs[\"no_such_attr_in_edge\"](attr, edge))\n",
    "        del graph.edges[edge][attr]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _setup_merged_node(graph: DiGraph, nodes_to_merge: set[NodeName], merge_policy: MergePolicy):\n",
    "    \"\"\"A helper function for node merging. It calculates all the parameters needed for creating the merged node,\n",
    "    such as its name, its attributes, the connected edges and their attributes, etc., and returns them all.\n",
    "    In addition, it removes the original nodes (which are about to be merged) from the graph.\n",
    "\n",
    "    Args:\n",
    "        graph (DiGraph): A graph\n",
    "        nodes_to_merge (set[NodeName]): A set of nodes in the graph to merge\n",
    "        merge_policy (MergePolicy): A policy that dictates how to merge conflicting attributes\n",
    "\n",
    "    Returns:\n",
    "        All the parameters needed for creating the merged node, in this order:\n",
    "        - The name of the merged node\n",
    "        - The attrs of the merged node\n",
    "        - The nodes that are sources of new edges to the merged node\n",
    "        - The nodes that are targets of new edges from the merged node\n",
    "        - The attrs of the new edges to the merged node\n",
    "        - The attrs of the new edges from the merged node\n",
    "        - A boolean value for whether the new merged node has a self loop (a special case)\n",
    "        - The attributes of the selp loop, in case such edge exists\n",
    "    \"\"\"\n",
    "    merged_node_name = _generate_new_node_name(graph, \"&\".join(nodes_to_merge))\n",
    "    \n",
    "    merged_node_attrs = {}\n",
    "    merged_src_nodes, merged_target_nodes = set(), set()\n",
    "    merged_src_attrs, merged_target_attrs = {}, {} # map a src/target node to the edge's merged attrs\n",
    "    self_loop, self_loop_attrs = False, {}\n",
    "    \n",
    "    for node_to_merge in nodes_to_merge:\n",
    "        merged_node_attrs = merge_policy(merged_node_attrs, graph.nodes[node_to_merge])\n",
    "        \n",
    "        in_edges, out_edges = graph.in_edges(node_to_merge), graph.out_edges(node_to_merge)\n",
    "        merged_src_nodes.update({s if s not in nodes_to_merge else merged_node_name for s, _ in in_edges})\n",
    "        merged_target_nodes.update({t if t not in nodes_to_merge else merged_node_name for _, t in out_edges})\n",
    "\n",
    "        for edge in in_edges:\n",
    "            edge_attrs, src = graph.edges[edge], edge[0]\n",
    "            # Add to source attributes\n",
    "            if src not in merged_src_attrs.keys():\n",
    "                merged_src_attrs[src] = edge_attrs\n",
    "            else:\n",
    "                merged_src_attrs[src] = merge_policy(merged_src_attrs[src], edge_attrs)\n",
    "\n",
    "            # Handle selp loop\n",
    "            if src in nodes_to_merge:\n",
    "                self_loop = True\n",
    "                self_loop_attrs = merge_policy(self_loop_attrs, edge_attrs)\n",
    "            \n",
    "        for edge in out_edges:\n",
    "            edge_attrs, target = graph.edges[edge], edge[1]\n",
    "            # Add to source attributes\n",
    "            if target not in merged_target_attrs.keys():\n",
    "                merged_target_attrs[target] = edge_attrs\n",
    "            else:\n",
    "                merged_target_attrs[target] = merge_policy(merged_target_attrs[target], edge_attrs)\n",
    "\n",
    "            # Handle selp loop\n",
    "            if target in nodes_to_merge:\n",
    "                self_loop = True\n",
    "                self_loop_attrs = merge_policy(self_loop_attrs, edge_attrs)\n",
    "\n",
    "        graph.remove_node(node_to_merge)\n",
    "\n",
    "    return merged_node_name, merged_node_attrs, merged_src_nodes, merged_target_nodes,\\\n",
    "            merged_src_attrs, merged_target_attrs, self_loop, self_loop_attrs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _merge_nodes(graph: DiGraph, nodes_to_merge: set[NodeName], merge_policy: MergePolicy) -> NodeName:\n",
    "    \"\"\"Merge a set of nodes in the graph. That is, remove all these nodes and replace them with a new node,\n",
    "    whose attributes merge the attributes of the original nodes, whose connected edges merge the edges connected to\n",
    "    the original node, and whose name denotes the nodes which it merges.\n",
    "\n",
    "    Args:\n",
    "        graph (DiGraph): A graph\n",
    "        nodes_to_merge (set[NodeName]): A set of nodes to merge\n",
    "        merge_policy (MergePolicy): A policy that dictates how to merge conflicting attributes\n",
    "\n",
    "    Raises:\n",
    "        GraphRewriteException: If the set of nodes to merge is too small (less than two nodes), or if one of them doesn't\n",
    "        exist in the graph.\n",
    "\n",
    "    Returns:\n",
    "        NodeName: The name of the new merged node\n",
    "    \"\"\"\n",
    "    if len(nodes_to_merge) < 1:\n",
    "        raise GraphRewriteException(_exception_msgs[\"not_enough_to_merge\"]())\n",
    "    elif len(nodes_to_merge) == 1:\n",
    "        return list(nodes_to_merge)[0]\n",
    "    for node_to_merge in nodes_to_merge:\n",
    "        if node_to_merge not in graph.nodes:\n",
    "            raise GraphRewriteException(_exception_msgs[\"no_such_node\"](node_to_merge))\n",
    "\n",
    "    merged_node_name, merged_node_attrs, merged_src_nodes, merged_target_nodes, \\\n",
    "        merged_src_attrs, merged_target_attrs, self_loop, self_loop_attrs = _setup_merged_node(graph, node_to_merge, merge_policy)\n",
    "\n",
    "    # Add merged node to graph\n",
    "    graph.add_node(merged_node_name, **merged_node_attrs)\n",
    "\n",
    "    # Add merged source and target edges (including a new self loop)\n",
    "    if self_loop:\n",
    "        graph.add_edge(merged_node_name, merged_node_name, **self_loop_attrs)\n",
    "    for src_node in merged_src_nodes:\n",
    "        if (src_node, merged_node_name) not in graph.edges():\n",
    "            graph.add_edge(src_node, merged_node_name)\n",
    "    for target_node in merged_target_nodes:\n",
    "        if (merged_node_name, target_node) not in graph.edges():\n",
    "            graph.add_edge(merged_node_name, target_node)\n",
    "    \n",
    "    # Add edge attributes (other than selp loop)\n",
    "    for src_node, attrs in merged_src_attrs.items():\n",
    "        if src_node not in nodes_to_merge:\n",
    "            graph.update(edges=[(src_node, merged_node_name, attrs)])\n",
    "    for target_node, attrs in merged_target_attrs.items():\n",
    "        if target_node not in nodes_to_merge:\n",
    "            graph.update(edges=[(merged_node_name, target_node, attrs)])\n",
    "\n",
    "    return merged_node_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _add_node(graph: DiGraph, node_to_add: NodeName) -> NodeName:\n",
    "    \"\"\"Add a new node to the graph.\n",
    "\n",
    "    Args:\n",
    "        graph (DiGraph): A graph\n",
    "        node_to_add (NodeName): The desired name of the new node\n",
    "\n",
    "    Returns:\n",
    "        NodeName: The name of the new added node (based on the name suggested)\n",
    "    \"\"\"\n",
    "    # Create a new node name\n",
    "    new_name = _generate_new_node_name(graph, node_to_add)\n",
    "    graph.add_node(new_name)\n",
    "    return new_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _add_edge(graph: DiGraph, edge_to_add: EdgeName):\n",
    "    \"\"\"Add an edge to the graph.\n",
    "\n",
    "    Args:\n",
    "        graph (DiGraph): A graph\n",
    "        edge_to_add (EdgeName): The edge to add.\n",
    "\n",
    "    Raises:\n",
    "        GraphRewriteException: If one of the edge's endpoints doesn't exist in the graph, or the edge itself already exists.\n",
    "    \"\"\"\n",
    "    src, target = edge_to_add\n",
    "    if src not in graph.nodes():\n",
    "        raise GraphRewriteException(_exception_msgs[\"no_such_node\"](src))\n",
    "    elif target not in graph.nodes():\n",
    "        raise GraphRewriteException(_exception_msgs[\"no_such_node\"](target))\n",
    "    elif edge_to_add in graph.edges():\n",
    "        raise GraphRewriteException(_exception_msgs[\"edge_exists\"](edge_to_add))\n",
    "    else:\n",
    "        graph.add_edge(src, target)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _add_node_attrs(graph: DiGraph, node: NodeName, attrs_to_add: dict):\n",
    "    \"\"\"Add attributes to a node in the graph.\n",
    "\n",
    "    Args:\n",
    "        graph (DiGraph): A graph\n",
    "        node (NodeName): A node in the graph\n",
    "        attrs_to_add (dict): Attributes to add to the node\n",
    "\n",
    "    Raises:\n",
    "        GraphRewriteException: If the node doesn't exist in the graph\n",
    "    \"\"\"\n",
    "    if node not in graph.nodes():\n",
    "        raise GraphRewriteException(_exception_msgs[\"no_such_node\"](node))\n",
    "    for attr, val in attrs_to_add.items():\n",
    "        graph.nodes[node][attr] = val"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _add_edge_attrs(graph: DiGraph, edge: EdgeName, attrs_to_add: dict):\n",
    "    \"\"\"Add attributes to an edge in the graph.\n",
    "\n",
    "    Args:\n",
    "        graph (DiGraph): A graph\n",
    "        edge (EdgeName): An edge in the graph\n",
    "        attrs_to_add (dict): Attributes to add to the edge\n",
    "\n",
    "    Raises:\n",
    "        GraphRewriteException: If the edge doesn't exist in the graph\n",
    "    \"\"\"\n",
    "    if edge not in graph.edges():\n",
    "        raise GraphRewriteException(_exception_msgs[\"no_such_edge\"](edge))\n",
    "    for attr, val in attrs_to_add.items():\n",
    "        graph.edges[edge][attr] = val"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Rewrite Match\n",
    "A match rewriting is split into two stages - restrictive and expansive. We'll implemenet them seperately, and then combine them into a single function the transforms an input graph based on a single match.\n",
    "\n",
    "We're following the terminology presented in ReGraph's graph transformation module: The restrictive phase denotes what we preserve from the original matched graph (includes all the \"remove\" operations), and the expansive one extends the graph with new nodes, edges and attributes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "_GREEN = '\\033[92m'\n",
    "_RED = '\\033[91m'\n",
    "_BLACK = '\\033[0m'\n",
    "\n",
    "def _log(msg: str, is_log: bool, color: str = _BLACK):\n",
    "    \"\"\"Logging some message to the console.\n",
    "\n",
    "    Args:\n",
    "        msg (str): A message to log\n",
    "        is_log (bool): If True, the log is printed (otherwise, it isn't).\n",
    "        color (str, optional): A color for the printed message. Defaults to _BLACK.\n",
    "    \"\"\"\n",
    "    if is_log:\n",
    "        print(f\"{color}{msg}{_BLACK}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _rewrite_match_restrictive(input_graph: DiGraph, rule: Rule, lhs_input_map: dict[NodeName, NodeName], is_log: bool) -> dict[NodeName, NodeName]:\n",
    "    \"\"\"Performs the restrictive phase of the rewriting process on some match: Clone nodes, Remove nodes and edges (and/or their attributes).\n",
    "\n",
    "    Args:\n",
    "        input_graph (DiGraph): A graph to rewrite\n",
    "        rule (Rule): A rule that dictates what transformations should the graph go through\n",
    "        lhs_input_map (dict[NodeName, NodeName]): Maps names of LHS nodes to names of input graph nodes, based on the match which we rewrite\n",
    "        is_log (bool): If True, logs are printed throughout the process.\n",
    "\n",
    "    Returns:\n",
    "        dict[NodeName, NodeName]: maps names of P nodes to names of input graph nodes\n",
    "    \"\"\"\n",
    "\n",
    "    # Initialize an empty mapping from P nodes to input_graph nodes.\n",
    "    p_input_map = {}\n",
    "\n",
    "    \"\"\"Clone nodes:\n",
    "        Find all LHS nodes that should be cloned (and what are their clones in P).\n",
    "        For each clone of an LHS node (apart from the first one), add it to the input graph \n",
    "        (clone with edges and attributes) and add the pair (clone_name, lhs_node_name) to the p->input mapping.\n",
    "    \"\"\"\n",
    "    # Map each cloned lhs node to a flag, denoting whether the original node is reused with the same name in P.\n",
    "    # If not, the original node will be removed later\n",
    "    cloned_to_flags_map = {cloned_lhs_node: False for cloned_lhs_node in rule.nodes_to_clone().keys()}\n",
    "    for cloned_lhs_node, p_clones in rule.nodes_to_clone().items():\n",
    "        for p_clone in p_clones:\n",
    "            # Original cloned node is reused in P, preserve it\n",
    "            if p_clone == cloned_lhs_node:\n",
    "                _log(f\"Clone {lhs_input_map[cloned_lhs_node]}\", is_log)\n",
    "                cloned_to_flags_map[cloned_lhs_node] = True\n",
    "                p_input_map[p_clone] = lhs_input_map[cloned_lhs_node]\n",
    "            # All other clones require actual cloning (mapped to the new cloned node in input graph)\n",
    "            else:\n",
    "                new_clone_id = _clone_node(input_graph, lhs_input_map[cloned_lhs_node])\n",
    "                _log(f\"Clone {lhs_input_map[cloned_lhs_node]} as {new_clone_id}\", is_log)\n",
    "                p_input_map[p_clone] = new_clone_id\n",
    "\n",
    "    \"\"\"Remove nodes, complete p->input mapping with preserved nodes which are not clones:\n",
    "        Find all LHS nodes that should be removed. \n",
    "        For each LHS node, if should be removed - remove it from input.\n",
    "                            otherwise, if it is not a clone, add to the mapping.\n",
    "    \"\"\"\n",
    "    for lhs_node in rule.lhs.nodes():\n",
    "        # Cloned lhs nodes which weren't reused and so, should be deleted\n",
    "        if lhs_node in rule.nodes_to_remove() or (lhs_node in cloned_to_flags_map.keys() and cloned_to_flags_map[lhs_node] == False):\n",
    "            _log(f\"Remove node {lhs_input_map[lhs_node]}\", is_log)\n",
    "            _remove_node(input_graph, lhs_input_map[lhs_node])        \n",
    "        # Else, either a saved cloned node (already preserved) or a regular one (should preserve them)\n",
    "        elif lhs_node not in cloned_to_flags_map.keys():\n",
    "            p_node = list(rule._rev_p_lhs[lhs_node])[0]\n",
    "            p_input_map[p_node] = lhs_input_map[lhs_node]\n",
    "\n",
    "    # Remove edges.\n",
    "    for lhs_src, lhs_target in rule.edges_to_remove():\n",
    "        _log(f\"Remove edge ({p_input_map[lhs_src]}, {p_input_map[lhs_target]})\", is_log)\n",
    "        _remove_edge(input_graph, (p_input_map[lhs_src], p_input_map[lhs_target]))\n",
    "\n",
    "    # Remove node attrs.\n",
    "    for p_node, attrs_to_remove in rule.node_attrs_to_remove().items():\n",
    "        _log(f\"Remove attrs {attrs_to_remove} from node {p_input_map[p_node]}\", is_log)\n",
    "        _remove_node_attrs(input_graph, p_input_map[p_node], attrs_to_remove)\n",
    "\n",
    "    # Remove edge attrs.\n",
    "    for (p_src, p_target), attrs_to_remove in rule.edge_attrs_to_remove().items():\n",
    "        _log(f\"Remove attrs {attrs_to_remove} from edge {(p_input_map[p_src], p_input_map[p_target])}\", is_log)\n",
    "        _remove_edge_attrs(input_graph, (p_input_map[p_src], p_input_map[p_target]), attrs_to_remove)\n",
    "\n",
    "    return p_input_map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _rewrite_match_expansive(input_graph: DiGraph, rule: Rule, p_input_map: dict[NodeName, NodeName], is_log: bool):\n",
    "    \"\"\"Performs the expansive phase of the rewriting process on some match: Merge nodes, Remove Add and edges (and/or new or updated attributes).\n",
    "\n",
    "    Args:\n",
    "        input_graph (DiGraph): A graph to rewrite\n",
    "        rule (Rule): A rule that dictates what transformations should the graph go through\n",
    "        p_input_map (dict[NodeName, NodeName]): Maps names of P nodes to names of input graph nodes, based on the match which we rewrite\n",
    "        is_log (bool): If True, logs are printed throughout the process.\n",
    "    \"\"\"\n",
    "        \n",
    "    # Initialize an empty mapping from RHS nodes to input_graph nodes.\n",
    "    rhs_input_map = {}\n",
    "\n",
    "    \"\"\"Merge nodes:\n",
    "        Find all RHS nodes that are a merge of P nodes (and what P nodes they merge).\n",
    "        For each RHS node that is a merge, merge all of the relavant P node, add the resulting RHS node to the input.\n",
    "        Also removes from input the P nodes that were merged into a new node.\n",
    "        Save the new RHS merged node to the RHS->input mapping.\n",
    "    \"\"\"\n",
    "    merge_rhs_nodes = rule.nodes_to_merge().keys()\n",
    "    for merge_rhs_node, p_merged in rule.nodes_to_merge().items():\n",
    "        input_nodes_to_merge = {p_input_map[p_node] for p_node in p_merged}\n",
    "        new_merged_id = _merge_nodes(input_graph, input_nodes_to_merge, rule.merge_policy)\n",
    "        _log(f\"Merge {input_nodes_to_merge} as {new_merged_id}\", is_log)\n",
    "        rhs_input_map[merge_rhs_node] = new_merged_id\n",
    "        \n",
    "    \"\"\"Add nodes, complete RHS->input mapping with added (and preserved) nodes:\n",
    "        Find all RHS nodes that should be added to the input (not including nodes that are a merge of P noeds).\n",
    "        For each RHS node, if should be added - add it to input and to the RHS->input mapping.\n",
    "                            otherwise, if is not a merge - add to the mapping.\n",
    "    \"\"\"\n",
    "    for rhs_node in rule.rhs.nodes():\n",
    "        if rhs_node in rule.nodes_to_add():\n",
    "            added_id = _add_node(input_graph, rhs_node)\n",
    "            _log(f\"Add node {rhs_node} as {added_id}\", is_log)\n",
    "            rhs_input_map[rhs_node] = added_id\n",
    "        elif rhs_node not in merge_rhs_nodes:\n",
    "            p_node = list(rule._rev_p_rhs[rhs_node])[0]\n",
    "            rhs_input_map[rhs_node] = p_input_map[p_node]\n",
    "\n",
    "    # Add edges.\n",
    "    for rhs_src, rhs_target in rule.edges_to_add():\n",
    "        _log(f\"Add edge ({rhs_input_map[rhs_src]}, {rhs_input_map[rhs_target]})\", is_log)\n",
    "        _add_edge(input_graph, (rhs_input_map[rhs_src], rhs_input_map[rhs_target]))\n",
    "\n",
    "    # Add node attrs.\n",
    "    for rhs_node, attrs_to_add in rule.node_attrs_to_add().items():\n",
    "        _log(f\"Added attrs {attrs_to_add} to node {rhs_input_map[rhs_node]}\", is_log)\n",
    "        _add_node_attrs(input_graph, rhs_input_map[rhs_node], attrs_to_add)\n",
    "\n",
    "    # Add edge attrs.\n",
    "    for (rhs_src, rhs_target), attrs_to_add in rule.edge_attrs_to_add().items():\n",
    "        _log(f\"Added attrs {attrs_to_add} to edge {(rhs_input_map[rhs_src], rhs_input_map[rhs_target])}\", is_log)\n",
    "        _add_edge_attrs(input_graph, (rhs_input_map[rhs_src], rhs_input_map[rhs_target]), attrs_to_add)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "During transformation, some exceptions could be thrown, due to illegal actions (which is the user's fault). Since it might occur in the middle of the transformation, after the graph was already partially changed, we use the following functions so save a \"check-point\" before rewriting the match, such that in case of an error, we can return to that check-point and not have a partially-transformed match in the graph:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _copy_graph(graph: DiGraph) -> DiGraph:\n",
    "    \"\"\"Creates a copy of the graph (including attributes, which are deep-copied).\n",
    "\n",
    "    Args:\n",
    "        graph (DiGraph): A graph to copy\n",
    "\n",
    "    Returns:\n",
    "        DiGraph: The new copy of the graph\n",
    "    \"\"\"\n",
    "    copied_nodes = [(node, deepcopy(attrs)) for node, attrs in graph.nodes(data=True)]\n",
    "    copied_edges = [(s, t, deepcopy(attrs)) for s, t, attrs in graph.edges(data=True)]\n",
    "    copy_graph = DiGraph()\n",
    "    copy_graph.update(nodes=copied_nodes, edges=copied_edges)\n",
    "    return copy_graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _restore_graph(graph: DiGraph, last_copy_graph: DiGraph):\n",
    "    \"\"\"Change a graph according to some \"checkpoint\".\n",
    "\n",
    "    Args:\n",
    "        graph (DiGraph): A graph to change\n",
    "        last_copy_graph (DiGraph): The \"checkpoint\" to which the graph is restored\n",
    "    \"\"\"\n",
    "    graph.clear()\n",
    "    graph.update(nodes=last_copy_graph.nodes(data=True),\n",
    "                 edges=last_copy_graph.edges(data=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, for the match rewriting itself:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _rewrite_match(input_graph: DiGraph, match: Match,\n",
    "                   lhs_graph: DiGraph, p_graph: DiGraph, rhs: str,\n",
    "                   render_rhs: dict[str, RenderFunc],\n",
    "                   merge_policy: MergePolicy,\n",
    "                   is_log: bool) -> Match:\n",
    "    \"\"\"Perform a graph rewriting based on a single match.\n",
    "\n",
    "    Args:\n",
    "        input_graph (DiGraph): A graph to rewrite\n",
    "        match (Match): A single match in the graph\n",
    "        lhs_graph (DiGraph): A parsed LHS pattern\n",
    "        p_graph (DiGraph): A parsed P pattern\n",
    "        rhs (str): A RHS pattern string, with potential placeholders\n",
    "        render_rhs (dict[str, RenderFunc]): Maps a RHS placeholder to a function that describes how to fill it, based on the given match\n",
    "        merge_policy (MergePolicy): A policy that dictates how to merge conflicting attributes\n",
    "        is_log (bool): If True, logs are printed throughout the process.\n",
    "\n",
    "    Raises:\n",
    "        GraphRewriteException: if something went wrong during the rewriting process\n",
    "\n",
    "    Returns:\n",
    "        Match: The match the we've just rewritten\n",
    "    \"\"\"\n",
    "\n",
    "    _log(f\"Transform match: {match.mapping}\", is_log, _GREEN)\n",
    "    # Save graph in current state for restoring, if needed\n",
    "    saved_graph = _copy_graph(input_graph)\n",
    "\n",
    "    try:\n",
    "        # Parse RHS according to current match (with render dictionary)\n",
    "        rhs_graph = rhs_to_graph(rhs, match, render_rhs) if rhs else None\n",
    "        rule = Rule(lhs_graph, p_graph, rhs_graph, merge_policy=merge_policy)\n",
    "        # Transform the graph\n",
    "        lhs_input_map = match.mapping\n",
    "        p_input_map = _rewrite_match_restrictive(input_graph, rule, lhs_input_map, is_log)\n",
    "        _rewrite_match_expansive(input_graph, rule, p_input_map, is_log)\n",
    "        _log(f\"Nodes: {input_graph.nodes(data=True)}\\nEdges: {input_graph.edges(data=True)}\\n\", is_log, _GREEN)\n",
    "        return match\n",
    "\n",
    "    except GraphRewriteException as e:\n",
    "        _log(f\"Failed to transform: {e.message}\", is_log, _RED)\n",
    "        _restore_graph(input_graph, saved_graph)\n",
    "        raise e"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Rewrite Graph\n",
    "Our library's main function generally does the following:\n",
    "* Parses LHS, P pattern strings.\n",
    "* Finds matches in the input graph based on the parsed LHS pattern. Also filters the matches based on a predefined condition.\n",
    "* For each match, performs a match rewriting with LHS, P and RHS. RHS patterns may have placeholders to be filled for each match seperately, this is achieved with render_rhs. Note that if P and RHS aren't provided, then each match isn't changed at all - thus, the user is able to use this function as a **Regex-only matcher** without any transforms.\n",
    "* Returns/yields a list of Match objects, corresponding to matches before the transforms, for imperative side-effects.\n",
    "\n",
    "For advanced users, the function also allows logging throughout the process; defining what should the function do if we merge two nodes with conflicting attributes; and updating the list of matches after each rewrite, so that we cover all matches and avoid errors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def rewrite_iter(input_graph: DiGraph, lhs: str, p: str = None, rhs: str = None,\n",
    "                   condition: FilterFunc = None,\n",
    "                   render_rhs: dict[str, RenderFunc] = None,\n",
    "                   merge_policy: MergePolicy = None,\n",
    "                   is_log: bool = False,\n",
    "                   is_recursive: bool = False,\n",
    "                   display_matches: bool = False,\n",
    "                   \n",
    "                   ) -> List[Match]:\n",
    "    \"\"\"Perform a graph rewriting using a lazy iterator, yielding the matches one by one after rewriting\n",
    "\n",
    "    Args:\n",
    "        input_graph (DiGraph): A graph to rewrite\n",
    "        lhs (str): A LHS pattern string\n",
    "        p (str, optional): A P pattern string. Defaults to None.\n",
    "        rhs (str, optional): A RHS pattern string. Defaults to None.\n",
    "        condition (FilterFunc, optional): A condition on the matches. Matches for which the condition doesn't hold aren't rewritten. Defaults to lambda match: True.\n",
    "        render_rhs (dict[str, RenderFunc], optional): Maps a RHS placeholder to a function that describes how to fill it, based on the given match. Defaults to {}.\n",
    "        merge_policy (MergePolicy, optional): A policy that dictates how to merge conflicting attributes. Defaults to MergePolicy.choose_last.\n",
    "        is_log (bool, optional): If True, logs are printed throughout the process. Defaults to False.\n",
    "        is_recursive (bool, optional): If True, matches pool is updated after each rewrite (as that pool might change). Defaults to False.\n",
    "        display_matches (bool, optional): If True, the matches are displayed, useful for debugging. Defaults to False.\n",
    "\n",
    "    Yields:\n",
    "        Iterator[Match]: An iterator of Match instances, which denote the matches we've transformed.\n",
    "    \"\"\"\n",
    "    condition = condition if condition else lambda match: True\n",
    "    render_rhs = render_rhs if render_rhs else {}\n",
    "    merge_policy = merge_policy if merge_policy else MergePolicy.choose_last\n",
    "\n",
    "    _log(f\"Nodes: {input_graph.nodes(data=True)}\\nEdges: {input_graph.edges(data=True)}\\n\", is_log, _GREEN)\n",
    "\n",
    "    # Parse LHS and P (global for all matches)\n",
    "    lhs_graph, condition = lhs_to_graph(lhs, condition)\n",
    "    p_graph = p_to_graph(p) if p else None\n",
    "    \n",
    "    if is_recursive:\n",
    "        while True:\n",
    "            try:\n",
    "                next_match = next(find_matches(input_graph, lhs_graph, condition=condition))\n",
    "                if display_matches:\n",
    "                    draw_match(input_graph, next_match)\n",
    "                yield next_match\n",
    "                new_res = _rewrite_match(input_graph, next_match, lhs_graph, p_graph, rhs, render_rhs, merge_policy, is_log)\n",
    "            except StopIteration:\n",
    "                break\n",
    "\n",
    "        _log(\"No more matches.\", is_log, _GREEN)\n",
    "\n",
    "    else:\n",
    "        # Create a duplication of the graph to find matches lazily (actual graph changes between matches)\n",
    "        copy_input_graph = _copy_graph(input_graph)\n",
    "\n",
    "        # Find matches lazily and transform\n",
    "        for match in find_matches(copy_input_graph, lhs_graph, condition=condition):\n",
    "            if display_matches:\n",
    "                draw_match(input_graph, match)\n",
    "            # the match object points to the copy graph, so we need to move it to the original graph for imperative changes\n",
    "            match.set_graph(input_graph)\n",
    "            yield match\n",
    "            new_res = _rewrite_match(input_graph, match, lhs_graph, p_graph, rhs, render_rhs, merge_policy, is_log)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "@delegates(rewrite_iter)\n",
    "def rewrite(input_graph: DiGraph, lhs: str,**kwargs\n",
    "                   ) -> List[Match]:\n",
    "    \"\"\"Perform a graph rewriting.\n",
    "\n",
    "    Args:\n",
    "        input_graph (DiGraph): A graph to rewrite\n",
    "        lhs (str): A LHS pattern string\n",
    "        p (str, optional): A P pattern string. Defaults to None.\n",
    "        rhs (str, optional): A RHS pattern string. Defaults to None.\n",
    "        condition (FilterFunc, optional): A condition on the matches. Matches for which the condition doesn't hold aren't rewritten. Defaults to lambda match: True.\n",
    "        render_rhs (dict[str, RenderFunc], optional): Maps a RHS placeholder to a function that describes how to fill it, based on the given match. Defaults to {}.\n",
    "        merge_policy (MergePolicy, optional): A policy that dictates how to merge conflicting attributes. Defaults to MergePolicy.choose_last.\n",
    "        is_log (bool, optional): If True, logs are printed throughout the process. Defaults to False.\n",
    "        is_recursive (bool, optional): If True, matches pool is updated after each rewrite (as that pool might change). Defaults to False.\n",
    "        display_matches (bool, optional): If True, the matches are displayed, useful for debugging. Defaults to False.\n",
    "\n",
    "    Returns:\n",
    "        Nothing, the graph is transformed in place.\n",
    "    \"\"\"\n",
    "    for _ in rewrite_iter(input_graph, lhs, **kwargs):\n",
    "        pass\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Unit Tests\n",
    "The following tests verify the correctness of the rule interpretation and additional features of the rewrite function. The tests in this module are based on the applications of LHS, P, RHS discussed in the _rules_ module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"https://mermaid.ink/img/CmZsb3djaGFydCBUQgoxWyIxCnZhbD0xIl0KMlsiMgp2YWw9MiJdCjNbIjMKdmFsPTMiXQo0WyI0CnZhbD00Il0KNVsiNQp2YWw9NSJdCjEgLS0+IDIKMSAtLT4gMwozIC0tPiA0CjMgLS0+IDUK\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "g = _create_graph(\n",
    "    [('1', {'val': 1}), ('2', {'val': 2}), ('3', {'val': 3}), ('4',{'val': 4}), ('5',{'val': 5})],\n",
    "    [('1','2'), ('1','3'), ('3','4'), ('3','5')]\n",
    ")\n",
    "draw(g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"https://mermaid.ink/img/CmZsb3djaGFydCBUQgoxWyIxKGEpCnZhbD0xIl0Kc3R5bGUgMSBzdHJva2U6cmVkLHN0cm9rZS13aWR0aDo0cHg7CjJbIjIoYikKdmFsPTIiXQpzdHlsZSAyIHN0cm9rZTpyZWQsc3Ryb2tlLXdpZHRoOjRweDsKM1siMyhjKQp2YWw9MyJdCnN0eWxlIDMgc3Ryb2tlOnJlZCxzdHJva2Utd2lkdGg6NHB4Owo0WyI0CnZhbD00Il0KNVsiNQp2YWw9NSJdCjEgLS0+IDIKbGlua1N0eWxlIDAgc3Ryb2tlOnJlZCxzdHJva2Utd2lkdGg6NHB4OwoxIC0tPiAzCmxpbmtTdHlsZSAxIHN0cm9rZTpyZWQsc3Ryb2tlLXdpZHRoOjRweDsKMyAtLT4gNAozIC0tPiA1Cg==\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<img src=\"https://mermaid.ink/img/CmZsb3djaGFydCBUQgoxWyIxCnZhbD0xIl0KMlsiMgp2YWw9MiJdCjNbIjMoYSkKdmFsPTMiXQpzdHlsZSAzIHN0cm9rZTpyZWQsc3Ryb2tlLXdpZHRoOjRweDsKNFsiNChiKQp2YWw9NCJdCnN0eWxlIDQgc3Ryb2tlOnJlZCxzdHJva2Utd2lkdGg6NHB4Owo1WyI1KGMpCnZhbD01Il0Kc3R5bGUgNSBzdHJva2U6cmVkLHN0cm9rZS13aWR0aDo0cHg7CjMgLS0+IDQKbGlua1N0eWxlIDAgc3Ryb2tlOnJlZCxzdHJva2Utd2lkdGg6NHB4OwozIC0tPiA1CmxpbmtTdHlsZSAxIHN0cm9rZTpyZWQsc3Ryb2tlLXdpZHRoOjRweDsK\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<img src=\"https://mermaid.ink/img/CmZsb3djaGFydCBUQgoxWyIxCnZhbD0xIl0KMlsiMgp2YWw9MiJdCjNbIjMKdmFsPTMiXQo0WyI0CnZhbD00Il0KNVsiNQp2YWw9NSJdCg==\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\"\"\" remove all edges. \n",
    "recursive so there won't be a corrupted match after deletion - \n",
    "in which the edges are already deleted \n",
    "\"\"\"\n",
    "input_graph = g.copy()\n",
    "matches = rewrite(input_graph,\n",
    "        lhs='a->b; a->c', p='a;b;c', is_recursive=True,display_matches=True)\n",
    "# assert not input_graph.edges\n",
    "draw(input_graph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"https://mermaid.ink/img/CmZsb3djaGFydCBUQgoxWyIxKGEpCnZhbD0xIl0Kc3R5bGUgMSBzdHJva2U6cmVkLHN0cm9rZS13aWR0aDo0cHg7CjJbIjIoYikKdmFsPTIiXQpzdHlsZSAyIHN0cm9rZTpyZWQsc3Ryb2tlLXdpZHRoOjRweDsKM1siMyhjKQp2YWw9MyJdCnN0eWxlIDMgc3Ryb2tlOnJlZCxzdHJva2Utd2lkdGg6NHB4Owo0WyI0CnZhbD00Il0KNVsiNQp2YWw9NSJdCjEgLS0+IDIKbGlua1N0eWxlIDAgc3Ryb2tlOnJlZCxzdHJva2Utd2lkdGg6NHB4OwoxIC0tPiAzCmxpbmtTdHlsZSAxIHN0cm9rZTpyZWQsc3Ryb2tlLXdpZHRoOjRweDsKMyAtLT4gNAozIC0tPiA1Cg==\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<img src=\"https://mermaid.ink/img/CmZsb3djaGFydCBUQgoxWyIxKGEpCnZhbD0xIl0Kc3R5bGUgMSBzdHJva2U6cmVkLHN0cm9rZS13aWR0aDo0cHg7CjJbIjIoYykKdmFsPTIiXQpzdHlsZSAyIHN0cm9rZTpyZWQsc3Ryb2tlLXdpZHRoOjRweDsKM1siMyhiKQp2YWw9MyJdCnN0eWxlIDMgc3Ryb2tlOnJlZCxzdHJva2Utd2lkdGg6NHB4Owo0WyI0CnZhbD00Il0KNVsiNQp2YWw9NSJdCjNfMVsiM18xCnZhbD0zIl0KMSAtLT4gMgpsaW5rU3R5bGUgMCBzdHJva2U6cmVkLHN0cm9rZS13aWR0aDo0cHg7CjEgLS0+IDMKbGlua1N0eWxlIDEgc3Ryb2tlOnJlZCxzdHJva2Utd2lkdGg6NHB4OwoxIC0tPiAzXzEKMyAtLT4gNAozIC0tPiA1CjNfMSAtLT4gNAozXzEgLS0+IDUK\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<img src=\"https://mermaid.ink/img/CmZsb3djaGFydCBUQgoxWyIxCnZhbD0xIl0KMlsiMgp2YWw9MiJdCjNbIjMoYSkKdmFsPTMiXQpzdHlsZSAzIHN0cm9rZTpyZWQsc3Ryb2tlLXdpZHRoOjRweDsKNFsiNChiKQp2YWw9NCJdCnN0eWxlIDQgc3Ryb2tlOnJlZCxzdHJva2Utd2lkdGg6NHB4Owo1WyI1KGMpCnZhbD01Il0Kc3R5bGUgNSBzdHJva2U6cmVkLHN0cm9rZS13aWR0aDo0cHg7CjNfMVsiM18xCnZhbD0zIl0KMl8xWyIyXzEKdmFsPTIiXQoxIC0tPiAyCjEgLS0+IDMKMSAtLT4gM18xCjEgLS0+IDJfMQozIC0tPiA0CmxpbmtTdHlsZSA0IHN0cm9rZTpyZWQsc3Ryb2tlLXdpZHRoOjRweDsKMyAtLT4gNQpsaW5rU3R5bGUgNSBzdHJva2U6cmVkLHN0cm9rZS13aWR0aDo0cHg7CjNfMSAtLT4gNAozXzEgLS0+IDUK\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<img src=\"https://mermaid.ink/img/CmZsb3djaGFydCBUQgoxWyIxCnZhbD0xIl0KMlsiMgp2YWw9MiJdCjNbIjMoYSkKdmFsPTMiXQpzdHlsZSAzIHN0cm9rZTpyZWQsc3Ryb2tlLXdpZHRoOjRweDsKNFsiNChjKQp2YWw9NCJdCnN0eWxlIDQgc3Ryb2tlOnJlZCxzdHJva2Utd2lkdGg6NHB4Owo1WyI1KGIpCnZhbD01Il0Kc3R5bGUgNSBzdHJva2U6cmVkLHN0cm9rZS13aWR0aDo0cHg7CjNfMVsiM18xCnZhbD0zIl0KMl8xWyIyXzEKdmFsPTIiXQo1XzFbIjVfMQp2YWw9NSJdCjEgLS0+IDIKMSAtLT4gMwoxIC0tPiAzXzEKMSAtLT4gMl8xCjMgLS0+IDQKbGlua1N0eWxlIDQgc3Ryb2tlOnJlZCxzdHJva2Utd2lkdGg6NHB4OwozIC0tPiA1CmxpbmtTdHlsZSA1IHN0cm9rZTpyZWQsc3Ryb2tlLXdpZHRoOjRweDsKMyAtLT4gNV8xCjNfMSAtLT4gNAozXzEgLS0+IDUKM18xIC0tPiA1XzEK\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<img src=\"https://mermaid.ink/img/CmZsb3djaGFydCBUQgoxWyIxCnZhbD0xIl0KMlsiMgp2YWw9MiJdCjNbIjMKdmFsPTMiXQo0WyI0CnZhbD00Il0KNVsiNQp2YWw9NSJdCjNfMVsiM18xCnZhbD0zIl0KMl8xWyIyXzEKdmFsPTIiXQo1XzFbIjVfMQp2YWw9NSJdCjRfMVsiNF8xCnZhbD00Il0KMSAtLT4gMgoxIC0tPiAzCjEgLS0+IDNfMQoxIC0tPiAyXzEKMyAtLT4gNAozIC0tPiA1CjMgLS0+IDVfMQozIC0tPiA0XzEKM18xIC0tPiA0CjNfMSAtLT4gNQozXzEgLS0+IDVfMQozXzEgLS0+IDRfMQo=\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\"\"\" clone a son. \n",
    "Not recursive to avoid infinite execution. \n",
    "\"\"\"\n",
    "input_graph = g.copy()\n",
    "matches = rewrite(input_graph,\n",
    "        lhs='a->b; a->c[val]', p='a->b; a->c; a->c*2',\n",
    "        rhs='a->b;a->c;a->c*2', is_recursive=False,display_matches=True)\n",
    "\n",
    "assert input_graph.degree('1') == 4 \n",
    "assert input_graph.degree('3') == 5 \n",
    "draw(input_graph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"https://mermaid.ink/img/CmZsb3djaGFydCBUQgoxWyIxCnZhbD0xIl0KMlsiMgp2YWw9MiJdCjNbIjMKeD0zIl0KNFsiNAp5PTQiXQo1WyI1Cng9NiwgdmFsPTUiXQoxIC0tPiAyCjEgLS0+IDMKMyAtLT4gNAozIC0tPiA1Cg==\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "g_1 = _create_graph(\n",
    "    [('1', {'val': 1}), ('2', {'val': 2}), ('3', {'x': 3}), ('4',{'y': 4}), ('5',{'x':6, 'val': 5})],\n",
    "    [('1','2'), ('1','3'), ('3','4'), ('3','5')]\n",
    ")\n",
    "draw(g_1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"https://mermaid.ink/img/CmZsb3djaGFydCBUQgoxWyIxKGEpCnZhbD0xIl0Kc3R5bGUgMSBzdHJva2U6cmVkLHN0cm9rZS13aWR0aDo0cHg7CjJbIjIKdmFsPTIiXQozWyIzCng9MyJdCjRbIjQKeT00Il0KNVsiNQp4PTYsIHZhbD01Il0KMSAtLT4gMgoxIC0tPiAzCjMgLS0+IDQKMyAtLT4gNQo=\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<img src=\"https://mermaid.ink/img/CmZsb3djaGFydCBUQgoxWyIxCiJdCjJbIjIoYSkKdmFsPTIiXQpzdHlsZSAyIHN0cm9rZTpyZWQsc3Ryb2tlLXdpZHRoOjRweDsKM1siMwp4PTMiXQo0WyI0Cnk9NCJdCjVbIjUKeD02LCB2YWw9NSJdCjEgLS0+IDIKMSAtLT4gMwozIC0tPiA0CjMgLS0+IDUK\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<img src=\"https://mermaid.ink/img/CmZsb3djaGFydCBUQgoxWyIxCiJdCjJbIjIKIl0KM1siMwp4PTMiXQo0WyI0Cnk9NCJdCjVbIjUoYSkKeD02LCB2YWw9NSJdCnN0eWxlIDUgc3Ryb2tlOnJlZCxzdHJva2Utd2lkdGg6NHB4OwoxIC0tPiAyCjEgLS0+IDMKMyAtLT4gNAozIC0tPiA1Cg==\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<img src=\"https://mermaid.ink/img/CmZsb3djaGFydCBUQgoxWyIxCiJdCjJbIjIKIl0KM1siMwp4PTMiXQo0WyI0Cnk9NCJdCjVbIjUKeD02Il0KMSAtLT4gMgoxIC0tPiAzCjMgLS0+IDQKMyAtLT4gNQo=\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\"\"\" remove and extract specific attribute \"\"\"\n",
    "input_graph = g_1.copy()\n",
    "l = []\n",
    "matches = rewrite_iter(input_graph,\n",
    "        lhs='a[val]', p='a', is_recursive=False,display_matches=True)\n",
    "for match in matches:\n",
    "    l.append(match['a']['val'])\n",
    "\n",
    "assert sorted([1,2,5]) == sorted(l)\n",
    "assert \"val\" not in input_graph.nodes(data=True)['5'].keys()\n",
    "assert \"x\" in input_graph.nodes(data=True)['5'].keys()\n",
    "draw(input_graph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"https://mermaid.ink/img/CmZsb3djaGFydCBUQgoxWyIxCng9MSwgdmFsPTEiXQoyWyIyCnZhbD0yIl0KM1siMwp2YWw9MywgeD0zIl0KNFsiNAp2YWw9NCwgeT00Il0KNVsiNQp2YWw9NSJdCjZbIjYKIl0KMSAtLT4gMgoxIC0tPiAzCjMgLS0+IDQKMyAtLT4gNQo0IC0tPiA2Cg==\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "g_2 = _create_graph(\n",
    "    [('1', {'x':1, 'val': 1}), ('2', {'val': 2}), ('3', {'val':3,'x': 3}), ('4',{'val':4, 'y': 4}), ('5',{'val': 5}), '6'],\n",
    "    [('1','2'), ('1','3'), ('3','4'), ('3','5'), ('4','6')]\n",
    ")\n",
    "draw(g_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"https://mermaid.ink/img/CmZsb3djaGFydCBUQgoxWyIxCng9MSwgdmFsPTEiXQoyWyIyCnZhbD0yIl0KM1siMwp2YWw9MywgeD0zIl0KNFsiNChhKQp2YWw9NCwgeT00Il0Kc3R5bGUgNCBzdHJva2U6cmVkLHN0cm9rZS13aWR0aDo0cHg7CjVbIjUKdmFsPTUiXQo2WyI2KGIpCiJdCnN0eWxlIDYgc3Ryb2tlOnJlZCxzdHJva2Utd2lkdGg6NHB4OwoxIC0tPiAyCjEgLS0+IDMKMyAtLT4gNAozIC0tPiA1CjQgLS0+IDYKbGlua1N0eWxlIDQgc3Ryb2tlOnJlZCxzdHJva2Utd2lkdGg6NHB4Owo=\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<img src=\"https://mermaid.ink/img/CmZsb3djaGFydCBUQgoxWyIxCng9MSwgdmFsPTEiXQoyWyIyCnZhbD0yIl0KM1siMwp2YWw9MywgeD0zIl0KNFsiNAp5PTQiXQo1WyI1CnZhbD01Il0KNlsiNgoiXQoxIC0tPiAyCjEgLS0+IDMKMyAtLT4gNAozIC0tPiA1CjQgLS0+fCJ2YWw9NSJ8IDYK\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\"\"\" move attributes from a node to the following edge.\n",
    "Filter nodes using a condition.\n",
    "\"\"\"\n",
    "input_graph = g_2.copy()\n",
    "l = []\n",
    "num_matches = 0\n",
    "matches = rewrite_iter(input_graph,\n",
    "        lhs='a[val]->b', p='a->b', rhs='a-[val={{new_val}}]->b', condition=lambda match: 'x' not in match['a'].keys(), \n",
    "        render_rhs={'new_val': lambda match: match['a']['val'] + 1}, is_recursive=False,display_matches=True)\n",
    "for match in matches:\n",
    "    num_matches = num_matches+1\n",
    "    l.append(match['a']['val'])\n",
    "\n",
    "# only the node '4' should be caught as 'a'\n",
    "assert num_matches == 1\n",
    "assert input_graph.edges[('4','6')]['val']==5\n",
    "assert 'val' not in input_graph.nodes(data=True)['4'].keys()\n",
    "draw(input_graph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"https://mermaid.ink/img/CmZsb3djaGFydCBUQgp4WyJ4CnZhbD0xIl0KeVsieQp2YWw9MiJdCnpbInoKdmFsPTMiXQpyZXMxWyJyZXMxCiJdCnJlczJbInJlczIKIl0KcmVzMSAtLT4geApyZXMxIC0tPiB5CnJlczIgLS0+IHoKcmVzMiAtLT4gcmVzMQo=\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "g_3 = _create_graph(\n",
    "    [('x', {'val': 1}), ('y', {'val': 2}), ('z', {'val': 3}), 'res1', 'res2'],\n",
    "    [('res2','z'), ('res2','res1'), ('res1','x'), ('res1','y')]\n",
    ")\n",
    "draw(g_3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"https://mermaid.ink/img/CmZsb3djaGFydCBUQgp4WyJ4KGIpCnZhbD0xIl0Kc3R5bGUgeCBzdHJva2U6cmVkLHN0cm9rZS13aWR0aDo0cHg7CnlbInkoYykKdmFsPTIiXQpzdHlsZSB5IHN0cm9rZTpyZWQsc3Ryb2tlLXdpZHRoOjRweDsKelsiegp2YWw9MyJdCnJlczFbInJlczEoYSkKIl0Kc3R5bGUgcmVzMSBzdHJva2U6cmVkLHN0cm9rZS13aWR0aDo0cHg7CnJlczJbInJlczIKIl0KcmVzMSAtLT4geApsaW5rU3R5bGUgMCBzdHJva2U6cmVkLHN0cm9rZS13aWR0aDo0cHg7CnJlczEgLS0+IHkKbGlua1N0eWxlIDEgc3Ryb2tlOnJlZCxzdHJva2Utd2lkdGg6NHB4OwpyZXMyIC0tPiB6CnJlczIgLS0+IHJlczEK\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<img src=\"https://mermaid.ink/img/CmZsb3djaGFydCBUQgp6WyJ6KGIpCnZhbD0zIl0Kc3R5bGUgeiBzdHJva2U6cmVkLHN0cm9rZS13aWR0aDo0cHg7CnJlczFbInJlczEoYykKdmFsPTMiXQpzdHlsZSByZXMxIHN0cm9rZTpyZWQsc3Ryb2tlLXdpZHRoOjRweDsKcmVzMlsicmVzMihhKQoiXQpzdHlsZSByZXMyIHN0cm9rZTpyZWQsc3Ryb2tlLXdpZHRoOjRweDsKcmVzMiAtLT4gegpsaW5rU3R5bGUgMCBzdHJva2U6cmVkLHN0cm9rZS13aWR0aDo0cHg7CnJlczIgLS0+IHJlczEKbGlua1N0eWxlIDEgc3Ryb2tlOnJlZCxzdHJva2Utd2lkdGg6NHB4Owo=\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<img src=\"https://mermaid.ink/img/CmZsb3djaGFydCBUQgpyZXMyWyJyZXMyCnZhbD02Il0K\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\"\"\" remove successors and create a new attribute in the predecessor.\n",
    "Recursive. \n",
    "\"\"\"\n",
    "input_graph = g_3.copy()\n",
    "matches = rewrite(input_graph,\n",
    "        lhs='a->b[val]; a->c[val]', p='a', rhs='a[val={{val}}]',\n",
    "        render_rhs={'val': lambda match: match['b']['val'] + match['c']['val']},\n",
    "        is_recursive=True,display_matches=True)\n",
    "\n",
    "assert len(input_graph.nodes) == 1\n",
    "assert input_graph.nodes['res2']\n",
    "assert input_graph.nodes(data=True)['res2'] == {'val': 6}\n",
    "draw(input_graph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"https://mermaid.ink/img/CmZsb3djaGFydCBUQgozWyIzCnZhbD0zIl0KNFsiNAp2YWw9NCJdCjVbIjUKdmFsPTUiXQozIC0tPnwidmFsPTEwInwgNAozIC0tPnwidmFsPTIwInwgNQo=\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# with edge attributes\n",
    "g_4 = _create_graph(\n",
    "    [('3', {'val': 3}), ('4',{'val': 4}), ('5',{'val': 5})],\n",
    "    [('3','4',{'val':10}), ('3','5',{'val':20})]\n",
    ")\n",
    "draw(g_4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"https://mermaid.ink/img/CmZsb3djaGFydCBUQgozWyIzKGEpCnZhbD0zIl0Kc3R5bGUgMyBzdHJva2U6cmVkLHN0cm9rZS13aWR0aDo0cHg7CjRbIjQoYikKdmFsPTQiXQpzdHlsZSA0IHN0cm9rZTpyZWQsc3Ryb2tlLXdpZHRoOjRweDsKNVsiNQp2YWw9NSJdCjMgLS0+fCJ2YWw9MTAifCA0CmxpbmtTdHlsZSAwIHN0cm9rZTpyZWQsc3Ryb2tlLXdpZHRoOjRweDsKMyAtLT58InZhbD0yMCJ8IDUK\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<img src=\"https://mermaid.ink/img/CmZsb3djaGFydCBUQgozWyIzCiJdCjRbIjQKdmFsPTQiXQo1WyI1CnZhbD01Il0KMyAtLT58Im5ldz0xMyJ8IDQKMyAtLT58InZhbD0yMCJ8IDUK\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\"\"\" add to an existing value on an edge.\n",
    "includes accessing an edge through a match object.\n",
    "recursive in order to avoid truoble since the value is deleted from the vertex.\n",
    "\"\"\"\n",
    "input_graph = g_4.copy()\n",
    "matches = rewrite(input_graph,\n",
    "        lhs='a[val]-[val]->b', p='a->b', rhs='a-[new={{new}}]->b',\n",
    "        render_rhs={'new': lambda match: match['a']['val'] + match['a->b']['val']},\n",
    "        is_recursive=True,display_matches=True)\n",
    "\n",
    "assert input_graph.edges[('3','4')]['new'] == 13 or input_graph.edges[('3','5')]['new'] == 23\n",
    "draw(input_graph)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Integrating Spanner Workbench Micro-Passes\n",
    "The following cells present a high-level approach for rewriting of the [AST micro-passes presented in the Spanner Workbench project](https://github.com/DeanLight/spanner_workbench/blob/master/src/rgxlog-interpreter/src/rgxlog/engine/passes/lark_passes.py), as a proof that our library allows a simpler micro-pass writing method."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Description of Our Suggested Approach"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We assume that the AST graph used in that project is a NetworkX graph (as our library requires). Each node in the graph has a *type* attribute which contains a string with the token-name that node has in the grammar, as well as all the attributes it had in the original, lark-produced AST.\n",
    "\n",
    "In addition, each node has a *idx* attribute which denotes its position among the other children of its parent. That is, if a grammar rule defines ```PAR -> A B C```, then node A will have *idx* value of 0, B will have 1, etc. We use this attribute to handle the fact that in NetworkX graphs, children of some node are unordered.\n",
    "\n",
    "Nodes which denote a constant (int, string, etc.) are constructed as regular nodes with a *args* attribute, containing that constant."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you'll see in passes such as \"Check Reserved Relation Names\", our library allows a general method to do **assertions and type-checks**. Assume that some check validates that for each pattern LHS in the graph, a condition ```cond(...)``` holds. We will implement it as follows:\n",
    "1. Match the pattern in the graph with LHS.\n",
    "2. Use the ```condition``` parameter to filter the matches according to **the negation of ```cond```**. Thus, only matches for which the condition does **not** hold will remain.\n",
    "3. Loop through the lists of remaining matches, and raise an exception for each one, as it fails the assertion. If none exist, then all matches passed the assertion, the check succeeded and no exception is raised.\n",
    "\n",
    "Since we'll have a lot of checks here, where their specific implementation has nothing to do with our library, we will use a dummy assertion function that receives any number of parameters to visualize how such a check will look like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _assertion(**kwargs) -> bool:\n",
    "    print(f\"[Assertion on {kwargs}]\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In addition, we will calculate new values based on different parameters and Spanner Workbench classes (which we do not have in this module). Therefore, here's a dummy calculation function that receives any number of parameters and returns some calculated value (we'll use it only when the calculations depends on classes or functions of Spanner which we have no access to):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _calculation(**kwargs) -> any:\n",
    "    print(f\"[Calculation on {kwargs}]\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to show the correctness of a subset of these passes (mainly, those that change the structure of the parse graph rather than performing checks or modifying internal structures), we'll construct a simple example to run through them. Assume that we wrote the following program in Spanner Workbench, which declares and assigns values to two variables (a multi-lined string variable and a span variable):"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```\n",
    "str_var = 'hello \\\n",
    "           world'\n",
    "span_var = [2,5)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "According to Spanner Workbench's grammar, the parse graph of this program will have the following structure (more or less):"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```\n",
    "start\n",
    "  statement\n",
    "    assignment\n",
    "      var_name\n",
    "        LOWER_CASE_NAME -> \"str_var\"\n",
    "      string\n",
    "        STRING -> \"hello \\\\\\nworld\"\n",
    "  statement\n",
    "    assignment\n",
    "      var_name\n",
    "        LOWER_CASE_NAME -> \"span_var\"\n",
    "      span\n",
    "        int\n",
    "          INT -> 2\n",
    "        int\n",
    "          INT -> 5\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We construct a NetworkX graph for the resulting parse graph according to the format specified above, and get the following test graph:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"https://mermaid.ink/img/CmZsb3djaGFydCBUQgpzdGFydFsic3RhcnQKIl0Kc3QxWyJzdDEKdHlwZT1zdGF0ZW1lbnQsIGlkeD0wIl0KYXNzMVsiYXNzMQp0eXBlPWFzc2lnbm1lbnQsIGlkeD0wIl0KdmFybmFtZTFbInZhcm5hbWUxCnR5cGU9dmFyX25hbWUsIGlkeD0wIl0KbG93ZXJjYXNlX25hbWUxWyJsb3dlcmNhc2VfbmFtZTEKdHlwZT1MT1dFUl9DQVNFX05BTUUsIGFyZ3M9c3RyX3ZhciwgaWR4PTAiXQpzdHJpbmcxWyJzdHJpbmcxCnR5cGU9c3RyaW5nLCBpZHg9MSJdCnN0cmluZ25vZGUxWyJzdHJpbmdub2RlMQp0eXBlPVNUUklORywgYXJncz0naGVsbG8gXAp3b3JsZCcsIGlkeD0wIl0Kc3QyWyJzdDIKdHlwZT1zdGF0ZW1lbnQsIGlkeD0xIl0KYXNzMlsiYXNzMgp0eXBlPWFzc2lnbm1lbnQsIGlkeD0wIl0KdmFybmFtZTJbInZhcm5hbWUyCnR5cGU9dmFyX25hbWUsIGlkeD0wIl0KbG93ZXJjYXNlX25hbWUyWyJsb3dlcmNhc2VfbmFtZTIKdHlwZT1MT1dFUl9DQVNFX05BTUUsIGFyZ3M9c3Bhbl92YXIsIGlkeD0wIl0Kc3BhbjFbInNwYW4xCnR5cGU9c3BhbiwgaWR4PTEiXQppbnQxWyJpbnQxCnR5cGU9aW50LCBpZHg9MCJdCmludG5vZGUxWyJpbnRub2RlMQp0eXBlPUlOVCwgaWR4PTAsIGFyZ3M9MiJdCmludDJbImludDIKdHlwZT1pbnQsIGlkeD0xIl0KaW50bm9kZTJbImludG5vZGUyCnR5cGU9SU5ULCBpZHg9MCwgYXJncz01Il0Kc3RhcnQgLS0+IHN0MQpzdGFydCAtLT4gc3QyCnN0MSAtLT4gYXNzMQphc3MxIC0tPiB2YXJuYW1lMQphc3MxIC0tPiBzdHJpbmcxCnZhcm5hbWUxIC0tPiBsb3dlcmNhc2VfbmFtZTEKc3RyaW5nMSAtLT4gc3RyaW5nbm9kZTEKc3QyIC0tPiBhc3MyCmFzczIgLS0+IHZhcm5hbWUyCmFzczIgLS0+IHNwYW4xCnZhcm5hbWUyIC0tPiBsb3dlcmNhc2VfbmFtZTIKc3BhbjEgLS0+IGludDEKc3BhbjEgLS0+IGludDIKaW50MSAtLT4gaW50bm9kZTEKaW50MiAtLT4gaW50bm9kZTIK\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def _init_integration_test(): \n",
    "    nodes = [\n",
    "        'start',\n",
    "        ('st1', {'type': 'statement', 'idx': 0}),\n",
    "        ('ass1', {'type': 'assignment', 'idx': 0}),\n",
    "        ('varname1', {'type': 'var_name', 'idx': 0}),\n",
    "        ('lowercase_name1', {'type': 'LOWER_CASE_NAME', 'args': 'str_var', 'idx': 0}),\n",
    "        ('string1', {'type': 'string', 'idx': 1}),\n",
    "        ('stringnode1', {'type': 'STRING', 'args': '\"hello \\\\\\nworld\"', 'idx': 0}),\n",
    "        \n",
    "        ('st2', {'type': 'statement', 'idx': 1}),\n",
    "        ('ass2', {'type': 'assignment', 'idx': 0}),\n",
    "        ('varname2', {'type': 'var_name', 'idx': 0}),\n",
    "        ('lowercase_name2', {'type': 'LOWER_CASE_NAME', 'args': 'span_var', 'idx': 0}),\n",
    "        ('span1', {'type': 'span', 'idx': 1}),\n",
    "        ('int1', {'type': 'int', 'idx': 0}),\n",
    "        ('intnode1', {'type': 'INT', 'idx': 0, 'args': 2}),\n",
    "        ('int2', {'type': 'int', 'idx': 1}),\n",
    "        ('intnode2', {'type': 'INT', 'idx': 0, 'args': 5}),\n",
    "    ]\n",
    "\n",
    "    edges = [\n",
    "        ('start', 'st1'), ('start', 'st2'),\n",
    "        ('st1', 'ass1'),\n",
    "        ('ass1', 'varname1'), ('ass1', 'string1'),\n",
    "        ('varname1', 'lowercase_name1'),\n",
    "        ('string1', 'stringnode1'),\n",
    "        ('st2', 'ass2'),\n",
    "        ('ass2', 'varname2'), ('ass2', 'span1'),\n",
    "        ('varname2', 'lowercase_name2'),\n",
    "        ('span1', 'int1'), ('span1', 'int2'),\n",
    "        ('int1', 'intnode1'),\n",
    "        ('int2', 'intnode2')\n",
    "    ]\n",
    "\n",
    "    return _create_graph(nodes, edges)\n",
    "\n",
    "test_g = _init_integration_test()\n",
    "draw(test_g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for match in rewrite_iter(test_g,lhs='a-[x=5]->b'):\n",
    "    print(match)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will now go through the micro-passes and show (in high level) how we could rewrite them using our library."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Remove Tokens \n",
    "[Spanner Workbench implementation](https://github.com/DeanLight/spanner_workbench/blob/f60595acd65ab26784c8f0d59bcaa586764441b3/src/rgxlog-interpreter/src/rgxlog/engine/passes/lark_passes.py#L79-L112)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _remove_tokens(ast, is_log=False):\n",
    "    \"\"\"\n",
    "    a lark pass that should be used before the semantic checks\n",
    "    transforms the lark tree by removing the redundant tokens\n",
    "    note that we inherit from 'Transformer' in order to be able to visit token nodes.\n",
    "    \"\"\"\n",
    "    rewrite(ast, lhs='a[type=\"INT\", args]', p='a', rhs='a[val={{val}}]',\n",
    "                render_rhs={'val': lambda match: int(match['a']['args'])}, is_log=is_log)\n",
    "\n",
    "    rewrite(ast, lhs='a[type=\"LOWER_CASE_NAME\", args]', p='a', rhs='a[val={{val}}]',\n",
    "                render_rhs={'val': lambda match: match['a']['args']}, is_log=is_log)\n",
    "\n",
    "    rewrite(ast, lhs='a[type=\"UPPER_CASE_NAME\", args]', p='a', rhs='a[val={{val}}]',\n",
    "                render_rhs={'val': lambda match: match['a']['args']}, is_log=is_log)\n",
    "\n",
    "    rewrite(ast, lhs='a[type=\"STRING\", args]', p='a', rhs='a[val={{val}}]',\n",
    "                render_rhs={'val': lambda match: match['a']['args'][1:-1]}, is_log=is_log)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"https://mermaid.ink/img/CmZsb3djaGFydCBUQgpzdGFydFsic3RhcnQKIl0Kc3QxWyJzdDEKdHlwZT1zdGF0ZW1lbnQsIGlkeD0wIl0KYXNzMVsiYXNzMQp0eXBlPWFzc2lnbm1lbnQsIGlkeD0wIl0KdmFybmFtZTFbInZhcm5hbWUxCnR5cGU9dmFyX25hbWUsIGlkeD0wIl0KbG93ZXJjYXNlX25hbWUxWyJsb3dlcmNhc2VfbmFtZTEKaWR4PTAsIHZhbD1zdHJfdmFyIl0Kc3RyaW5nMVsic3RyaW5nMQp0eXBlPXN0cmluZywgaWR4PTEiXQpzdHJpbmdub2RlMVsic3RyaW5nbm9kZTEKaWR4PTAsIHZhbD1oZWxsbyBcCndvcmxkIl0Kc3QyWyJzdDIKdHlwZT1zdGF0ZW1lbnQsIGlkeD0xIl0KYXNzMlsiYXNzMgp0eXBlPWFzc2lnbm1lbnQsIGlkeD0wIl0KdmFybmFtZTJbInZhcm5hbWUyCnR5cGU9dmFyX25hbWUsIGlkeD0wIl0KbG93ZXJjYXNlX25hbWUyWyJsb3dlcmNhc2VfbmFtZTIKaWR4PTAsIHZhbD1zcGFuX3ZhciJdCnNwYW4xWyJzcGFuMQp0eXBlPXNwYW4sIGlkeD0xIl0KaW50MVsiaW50MQp0eXBlPWludCwgaWR4PTAiXQppbnRub2RlMVsiaW50bm9kZTEKaWR4PTAsIHZhbD0yIl0KaW50MlsiaW50Mgp0eXBlPWludCwgaWR4PTEiXQppbnRub2RlMlsiaW50bm9kZTIKaWR4PTAsIHZhbD01Il0Kc3RhcnQgLS0+IHN0MQpzdGFydCAtLT4gc3QyCnN0MSAtLT4gYXNzMQphc3MxIC0tPiB2YXJuYW1lMQphc3MxIC0tPiBzdHJpbmcxCnZhcm5hbWUxIC0tPiBsb3dlcmNhc2VfbmFtZTEKc3RyaW5nMSAtLT4gc3RyaW5nbm9kZTEKc3QyIC0tPiBhc3MyCmFzczIgLS0+IHZhcm5hbWUyCmFzczIgLS0+IHNwYW4xCnZhcm5hbWUyIC0tPiBsb3dlcmNhc2VfbmFtZTIKc3BhbjEgLS0+IGludDEKc3BhbjEgLS0+IGludDIKaW50MSAtLT4gaW50bm9kZTEKaW50MiAtLT4gaW50bm9kZTIK\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Sets constants nodes for two variable names, one string and two integers. Structure of graph remains unchanged\n",
    "_remove_tokens(test_g)\n",
    "assert test_g.nodes['lowercase_name1'] == {'val': 'str_var', 'idx': 0}\n",
    "assert test_g.nodes['lowercase_name2'] == {'val': 'span_var', 'idx': 0}\n",
    "assert test_g.nodes['stringnode1'] == {'val': 'hello \\\\\\nworld', 'idx': 0}\n",
    "assert test_g.nodes['intnode1'] == {'val': 2, 'idx': 0}\n",
    "assert test_g.nodes['intnode2'] == {'val': 5, 'idx': 0}\n",
    "draw(test_g)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Check Reserved Relation Names\n",
    "[Spanner Workbench implementation](https://github.com/DeanLight/spanner_workbench/blob/f60595acd65ab26784c8f0d59bcaa586764441b3/src/rgxlog-interpreter/src/rgxlog/engine/passes/lark_passes.py#L115-L131)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _check_reserved_relation_names(ast, is_log=False):\n",
    "    \"\"\"\n",
    "    A lark tree semantic check.\n",
    "    Checks if there are relations in the program with a name that starts with `RESERVED_RELATION_PREFIX`\n",
    "    if such relations exist, throw an exception as this is a reserved name for rgxlog.\n",
    "    \"\"\"\n",
    "    rewrite(ast, lhs='a[type=\"relation_name\"]->rel_name[val]',\n",
    "                condition=lambda match: _assertion(match['rel_name']['val']), is_log=is_log)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Fix Strings\n",
    "[Spanner Workbench implementation](https://github.com/DeanLight/spanner_workbench/blob/f60595acd65ab26784c8f0d59bcaa586764441b3/src/rgxlog-interpreter/src/rgxlog/engine/passes/lark_passes.py#L134-L152)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _fix_strings(ast, is_log=False):\n",
    "    \"\"\"\n",
    "    Fixes the strings in the lark tree.\n",
    "    Removes the line overflow escapes from strings.\n",
    "    \"\"\"\n",
    "    rewrite(ast, lhs='a[type=\"string\"]->s[val]', rhs='a[type]->s[val={{fixed_val}}]',\n",
    "                render_rhs={'fixed_val': lambda match: match['s']['val'].replace('\\\\\\n', '')}, is_log=is_log)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"https://mermaid.ink/img/CmZsb3djaGFydCBUQgpzdGFydFsic3RhcnQKIl0Kc3QxWyJzdDEKdHlwZT1zdGF0ZW1lbnQsIGlkeD0wIl0KYXNzMVsiYXNzMQp0eXBlPWFzc2lnbm1lbnQsIGlkeD0wIl0KdmFybmFtZTFbInZhcm5hbWUxCnR5cGU9dmFyX25hbWUsIGlkeD0wIl0KbG93ZXJjYXNlX25hbWUxWyJsb3dlcmNhc2VfbmFtZTEKaWR4PTAsIHZhbD1zdHJfdmFyIl0Kc3RyaW5nMVsic3RyaW5nMQp0eXBlPXN0cmluZywgaWR4PTEiXQpzdHJpbmdub2RlMVsic3RyaW5nbm9kZTEKaWR4PTAsIHZhbD1oZWxsbyB3b3JsZCJdCnN0Mlsic3QyCnR5cGU9c3RhdGVtZW50LCBpZHg9MSJdCmFzczJbImFzczIKdHlwZT1hc3NpZ25tZW50LCBpZHg9MCJdCnZhcm5hbWUyWyJ2YXJuYW1lMgp0eXBlPXZhcl9uYW1lLCBpZHg9MCJdCmxvd2VyY2FzZV9uYW1lMlsibG93ZXJjYXNlX25hbWUyCmlkeD0wLCB2YWw9c3Bhbl92YXIiXQpzcGFuMVsic3BhbjEKdHlwZT1zcGFuLCBpZHg9MSJdCmludDFbImludDEKdHlwZT1pbnQsIGlkeD0wIl0KaW50bm9kZTFbImludG5vZGUxCmlkeD0wLCB2YWw9MiJdCmludDJbImludDIKdHlwZT1pbnQsIGlkeD0xIl0KaW50bm9kZTJbImludG5vZGUyCmlkeD0wLCB2YWw9NSJdCnN0YXJ0IC0tPiBzdDEKc3RhcnQgLS0+IHN0MgpzdDEgLS0+IGFzczEKYXNzMSAtLT4gdmFybmFtZTEKYXNzMSAtLT4gc3RyaW5nMQp2YXJuYW1lMSAtLT4gbG93ZXJjYXNlX25hbWUxCnN0cmluZzEgLS0+IHN0cmluZ25vZGUxCnN0MiAtLT4gYXNzMgphc3MyIC0tPiB2YXJuYW1lMgphc3MyIC0tPiBzcGFuMQp2YXJuYW1lMiAtLT4gbG93ZXJjYXNlX25hbWUyCnNwYW4xIC0tPiBpbnQxCnNwYW4xIC0tPiBpbnQyCmludDEgLS0+IGludG5vZGUxCmludDIgLS0+IGludG5vZGUyCg==\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Reformats the multi-string \"hello world\"\n",
    "_fix_strings(test_g)\n",
    "assert test_g.nodes['stringnode1'] == {'val': 'hello world', 'idx': 0}\n",
    "draw(test_g)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Convert Span Nodes To Span Instances\n",
    "[Spanner Workbench implementation](https://github.com/DeanLight/spanner_workbench/blob/f60595acd65ab26784c8f0d59bcaa586764441b3/src/rgxlog-interpreter/src/rgxlog/engine/passes/lark_passes.py#L155-L174)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _convert_span_nodes_to_span_instances(ast, is_log=False):\n",
    "    \"\"\"\n",
    "    Converts each span node in the ast to a span instance.\n",
    "    This means that a span in the tree will be represented by a single value (a \"DataTypes.Span\" instance)\n",
    "    instead of two integer nodes, making it easier to work with (as other data types are also represented by\n",
    "    a single value).\n",
    "    \"\"\"\n",
    "    rewrite(ast, lhs='s[type=\"span\"]->start[idx=0]->start_val[val]; s->end[idx=1]->end_val[val]',\n",
    "                p='s[type]', rhs='s[type]->span_child[val={{val}}]',\n",
    "                render_rhs={'val': lambda match: (match['start_val']['val'], match['end_val']['val'])}, is_log=is_log)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"https://mermaid.ink/img/CmZsb3djaGFydCBUQgpzdGFydFsic3RhcnQKIl0Kc3QxWyJzdDEKdHlwZT1zdGF0ZW1lbnQsIGlkeD0wIl0KYXNzMVsiYXNzMQp0eXBlPWFzc2lnbm1lbnQsIGlkeD0wIl0KdmFybmFtZTFbInZhcm5hbWUxCnR5cGU9dmFyX25hbWUsIGlkeD0wIl0KbG93ZXJjYXNlX25hbWUxWyJsb3dlcmNhc2VfbmFtZTEKaWR4PTAsIHZhbD1zdHJfdmFyIl0Kc3RyaW5nMVsic3RyaW5nMQp0eXBlPXN0cmluZywgaWR4PTEiXQpzdHJpbmdub2RlMVsic3RyaW5nbm9kZTEKaWR4PTAsIHZhbD1oZWxsbyB3b3JsZCJdCnN0Mlsic3QyCnR5cGU9c3RhdGVtZW50LCBpZHg9MSJdCmFzczJbImFzczIKdHlwZT1hc3NpZ25tZW50LCBpZHg9MCJdCnZhcm5hbWUyWyJ2YXJuYW1lMgp0eXBlPXZhcl9uYW1lLCBpZHg9MCJdCmxvd2VyY2FzZV9uYW1lMlsibG93ZXJjYXNlX25hbWUyCmlkeD0wLCB2YWw9c3Bhbl92YXIiXQpzcGFuMVsic3BhbjEKdHlwZT1zcGFuLCBpZHg9MSJdCnNwYW5fY2hpbGRbInNwYW5fY2hpbGQKdmFsPSgyLCA1KSJdCnN0YXJ0IC0tPiBzdDEKc3RhcnQgLS0+IHN0MgpzdDEgLS0+IGFzczEKYXNzMSAtLT4gdmFybmFtZTEKYXNzMSAtLT4gc3RyaW5nMQp2YXJuYW1lMSAtLT4gbG93ZXJjYXNlX25hbWUxCnN0cmluZzEgLS0+IHN0cmluZ25vZGUxCnN0MiAtLT4gYXNzMgphc3MyIC0tPiB2YXJuYW1lMgphc3MyIC0tPiBzcGFuMQp2YXJuYW1lMiAtLT4gbG93ZXJjYXNlX25hbWUyCnNwYW4xIC0tPiBzcGFuX2NoaWxkCg==\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Replaces the span subtree with a single child that denotes the span bounds\n",
    "_convert_span_nodes_to_span_instances(test_g)\n",
    "assert 'int1' not in test_g.nodes()\n",
    "assert 'intnode1' not in test_g.nodes()\n",
    "assert 'int2' not in test_g.nodes()\n",
    "assert 'intnode2' not in test_g.nodes()\n",
    "assert 'span_child' in test_g.nodes()\n",
    "assert ('span1', 'span_child') in test_g.edges()\n",
    "assert test_g.nodes['span_child'] == {'val': (2, 5)}\n",
    "draw(test_g)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Convert Statements To Structured Nodes\n",
    "[Spanner Workbench implementation](https://github.com/DeanLight/spanner_workbench/blob/f60595acd65ab26784c8f0d59bcaa586764441b3/src/rgxlog-interpreter/src/rgxlog/engine/passes/lark_passes.py#L177-L366)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This pass can be implemented partially using our current syntax. A full implenetation will be achievable only after implementing a depth-recursion method in the library's matcher, along with a predefined syntax that denotes such a pattern (\"-+->\")."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _convert_statements_to_structured_nodes(ast, is_log=False):\n",
    "    \"\"\"\n",
    "    converts each statement node in the tree to a structured node, making it easier to parse in future passes.\n",
    "    a structured node is a class representation of a node in the abstract syntax tree.\n",
    "    note that after using this pass, non statement nodes will no longer appear in the tree, so passes that\n",
    "    should work on said nodes need to be used before this pass in the passes pipeline (e.g. FixString).\n",
    "    \"\"\"\n",
    "    rewrite(ast, lhs=\"\"\"ass[type=\"assignment\"]->var_name_node[idx=0];\n",
    "                            ass->value_node[idx=1];\n",
    "                            var_name_node->var_name[val]; value_node[type]->value[val]\"\"\",\n",
    "                p='ass[type]', rhs='ass[type]->ass_node[var_name={{var_name}}, value={{value}}, value_type={{value_type}}]',\n",
    "                render_rhs={'var_name': lambda m: m['var_name']['val'], 'value': lambda m: m['value']['val'],\n",
    "                                'value_type': lambda m: m['value_node']['type']}, is_log=is_log)\n",
    "\n",
    "    rewrite(ast, lhs=\"\"\"rass[type=\"read_assignment\"]->var_name_node[idx=0];\n",
    "                            rass->read_arg_node[idx=1];\n",
    "                            var_name_node->var_name[val]; \n",
    "                            read_arg_node[type]->read_arg[val]\"\"\",\n",
    "                p='rass[type]', rhs='rass[type]->rass_node[var_name={{var_name}}, read_arg={{read_arg}}, read_arg_type={{read_arg_type}}]',\n",
    "                render_rhs={'var_name': lambda m: m['var_name']['val'], 'read_arg': lambda m: m['read_arg']['val'],\n",
    "                                'read_arg_type': lambda m: m['read_arg_node']['type']}, is_log=is_log)\n",
    "    \n",
    "    # add_fact, remove_fact, query, relation_declaration and rule all use list tokens, that define a recursive list of unknown length.\n",
    "    # While such pattern cannot be identified with our current syntax, it WILL be possible when implementing a depth-recursive,\n",
    "    # a concept which can be assigned with the syntax -+->."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<img src=\"https://mermaid.ink/img/CmZsb3djaGFydCBUQgpzdGFydFsic3RhcnQKIl0Kc3QxWyJzdDEKdHlwZT1zdGF0ZW1lbnQsIGlkeD0wIl0KYXNzMVsiYXNzMQp0eXBlPWFzc2lnbm1lbnQsIGlkeD0wIl0Kc3QyWyJzdDIKdHlwZT1zdGF0ZW1lbnQsIGlkeD0xIl0KYXNzMlsiYXNzMgp0eXBlPWFzc2lnbm1lbnQsIGlkeD0wIl0KYXNzX25vZGVbImFzc19ub2RlCnZhcl9uYW1lPXN0cl92YXIsIHZhbHVlPWhlbGxvIHdvcmxkLCB2YWx1ZV90eXBlPXN0cmluZyJdCmFzc19ub2RlXzFbImFzc19ub2RlXzEKdmFyX25hbWU9c3Bhbl92YXIsIHZhbHVlPSgyLCA1KSwgdmFsdWVfdHlwZT1zcGFuIl0Kc3RhcnQgLS0+IHN0MQpzdGFydCAtLT4gc3QyCnN0MSAtLT4gYXNzMQphc3MxIC0tPiBhc3Nfbm9kZQpzdDIgLS0+IGFzczIKYXNzMiAtLT4gYXNzX25vZGVfMQo=\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Replace each statement subtree with a single descriptive node\n",
    "_convert_statements_to_structured_nodes(test_g)\n",
    "assert 'varname1' not in test_g.nodes()\n",
    "assert 'lowercase_name1' not in test_g.nodes()\n",
    "assert 'string1' not in test_g.nodes()\n",
    "assert 'stringnode1' not in test_g.nodes()\n",
    "assert 'varname2' not in test_g.nodes()\n",
    "assert 'lowercase_name2' not in test_g.nodes()\n",
    "assert 'span1' not in test_g.nodes()\n",
    "assert 'span_child' not in test_g.nodes()\n",
    "\n",
    "ass1_neighbors = [n for n in test_g.neighbors('ass1')]\n",
    "assert len(ass1_neighbors) == 1\n",
    "ass1_node = ass1_neighbors[0]\n",
    "assert ('ass1', ass1_node) in test_g.edges()\n",
    "assert test_g.nodes[ass1_node] == {'var_name': 'str_var', 'value': 'hello world', 'value_type': 'string'}\n",
    "\n",
    "ass2_neighbors = [n for n in test_g.neighbors('ass2')]\n",
    "assert len(ass2_neighbors) == 1\n",
    "ass2_node = ass2_neighbors[0]\n",
    "assert ('ass2', ass2_node) in test_g.edges()\n",
    "assert test_g.nodes[ass2_node] == {'var_name': 'span_var', 'value': (2, 5), 'value_type': 'span'}\n",
    "\n",
    "draw(test_g)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Check Defined Referenced Variables\n",
    "[Spanner Workbench implementation](https://github.com/DeanLight/spanner_workbench/blob/f60595acd65ab26784c8f0d59bcaa586764441b3/src/rgxlog-interpreter/src/rgxlog/engine/passes/lark_passes.py#L369-L441)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assumes the existence of a symbol table, in addition to the AST."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _check_defined_referenced_variables(ast, sym_tab, is_log=False):\n",
    "    \"\"\"\n",
    "    A lark tree semantic check.\n",
    "    checks whether each variable reference refers to a defined variable.\n",
    "    \"\"\"\n",
    "    rewrite(ast, lhs='_[type=\"assignment\"]->ass_node[value_type, value]',\n",
    "               condition=lambda match: _assertion(match['ass_node']['value_type'], match['ass_node']['value']), is_log=is_log)\n",
    "\n",
    "    rewrite(ast, lhs='_[type=\"read_assignment\"]->rass_node[read_arg_type, read_arg]',\n",
    "                condition=lambda match: _assertion(match['rass_node']['read_arg_type'], match['rass_node']['read_arg']), is_log=is_log)\n",
    "\n",
    "    rewrite(ast, lhs='_[type=\"add_fact]->afact_node[term_list, type_list]',\n",
    "                condition=lambda match: _assertion(match['afact_node']['term_list'], match['afact_node']['type_list'], sym_tab), is_log=is_log)\n",
    "    \n",
    "    rewrite(ast, lhs='_[type=\"remove_fact]->rfact_node[term_list, type_list]',\n",
    "               condition=lambda match: _assertion(match['rfact_node']['term_list'], match['rfact_node']['type_list'], sym_tab), is_log=is_log)\n",
    "\n",
    "    rewrite(ast, lhs='_[type=\"query\"]->query_node[term_list, type_list]',\n",
    "                condition=lambda match: _assertion(match['query_node']['term_list'], match['query_node']['type_list'], sym_tab), is_log=is_log)\n",
    "\n",
    "    rewrite(ast, lhs='_[type=\"rule\"]->rule_node[body_relation_list, body_relation_type_list]',\n",
    "                condition=lambda match: _assertion(match['rule_node']['body_relation_list'], match['rule_node']['type_list'], sym_tab), is_log=is_log)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Check Referenced Relations Existence And Arity\n",
    "[Spanner Workbench implementation](https://github.com/DeanLight/spanner_workbench/blob/f60595acd65ab26784c8f0d59bcaa586764441b3/src/rgxlog-interpreter/src/rgxlog/engine/passes/lark_passes.py#L444-L507)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assumes the existence of a symbol table, in addition to the AST."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _check_referenced_relations_existence_and_arity(ast, sym_tab, is_log=False):\n",
    "    \"\"\"\n",
    "    A lark tree semantic check.\n",
    "    Checks whether each normal relation (that is not an ie relation) reference refers to a defined relation.\n",
    "    Also checks if the relation reference uses the correct arity.\n",
    "    \"\"\"\n",
    "    rewrite(ast, lhs='_[type=\"query\"]->query_node[relation_name, term_list]',\n",
    "                condition=lambda match: _assertion(match['query_node']['relation_name'], match['query_node']['term_list'], sym_tab), is_log=is_log)\n",
    "    \n",
    "    rewrite(ast, lhs='_[type=\"add_fact\"]->afact_node[relation_name, term_list]',\n",
    "                condition=lambda match: _assertion(match['afact_node']['relation_name'], match['afact_node']['term_list'], sym_tab), is_log=is_log)\n",
    "\n",
    "    rewrite(ast, lhs='_[type=\"remove_fact\"]->rfact_node[relation_name, term_list]',\n",
    "                condition=lambda match: _assertion(match['rfact_node']['relation_name'], match['rfact_node']['term_list'], sym_tab), is_log=is_log)\n",
    "\n",
    "    rewrite(ast, lhs='_[type=\"rule\"]->rule_node[body_relation_list, body_relation_type_list]',\n",
    "                condition=lambda match: _assertion(match['rule_node']['body_relation_list'], match['rule_node']['type_list'], sym_tab), is_log=is_log)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Check Referenced IE Relations Existence And Arity\n",
    "[Spanner Workbench implementation](https://github.com/DeanLight/spanner_workbench/blob/f60595acd65ab26784c8f0d59bcaa586764441b3/src/rgxlog-interpreter/src/rgxlog/engine/passes/lark_passes.py#L510-L553)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assumes the existence of a symbol table, in addition to the AST."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _check_referenced_ie_relations_existence_and_arity(ast, sym_tab, is_log=False):\n",
    "    \"\"\"\n",
    "    A lark tree semantic check.\n",
    "    Checks whether each ie relation reference refers to a defined ie function.\n",
    "    Also checks if the correct input arity and output arity for the ie function were used.\n",
    "\n",
    "    currently, an ie relation can only be found in a rule's body, so this is the only place where this\n",
    "    check will be performed.\n",
    "    \"\"\"\n",
    "    rewrite(ast, lhs='_[type=\"rule\"]->rule_node[body_relation_list, body_relation_type_list]',\n",
    "                condition=lambda match: _assertion(match['rule_node']['body_relation_list'], match['rule_node']['type_list'], sym_tab), is_log=is_log)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Check Rule Safety\n",
    "[Spanner Workbench implementation](https://github.com/DeanLight/spanner_workbench/blob/f60595acd65ab26784c8f0d59bcaa586764441b3/src/rgxlog-interpreter/src/rgxlog/engine/passes/lark_passes.py#L556-L670)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _check_rule_safety(ast, is_log=False):\n",
    "    \"\"\"\n",
    "    A lark tree semantic check.\n",
    "    checks whether the rules in the programs are safe.\n",
    "    \"\"\"\n",
    "    rewrite(ast, lhs='_[type=\"rule\"]->rule_node[head_relation, body_relation_list, body_relation_type_list]',\n",
    "                condition=lambda match: _assertion(match['rule_node']['head_relation'], match['rule_node']['body_relation_list'], match['rule_node']['type_list']), is_log=is_log)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Type Check Assignments\n",
    "[Spanner Workbench implementation](https://github.com/DeanLight/spanner_workbench/blob/f60595acd65ab26784c8f0d59bcaa586764441b3/src/rgxlog-interpreter/src/rgxlog/engine/passes/lark_passes.py#L673-L697)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assumes the existence of a symbol table, in addition to the AST."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _type_check_assignments(ast, sym_tab, is_log=False):\n",
    "    \"\"\"\n",
    "    a lark semantic check\n",
    "    performs type checking for assignments\n",
    "    in the current version of lark, this type checking is only required for read assignments.\n",
    "    \"\"\"\n",
    "    rewrite(ast, lhs='_[type=\"assignment\"]->ass_node[read_arg_type, read_arg]',\n",
    "                condition=lambda match: _assertion(match['ass_node']['read_arg_type'], match['ass_node']['read_arg'], sym_tab), is_log=is_log)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Type Check Relations\n",
    "[Spanner Workbench implementation](https://github.com/DeanLight/spanner_workbench/blob/f60595acd65ab26784c8f0d59bcaa586764441b3/src/rgxlog-interpreter/src/rgxlog/engine/passes/lark_passes.py#L700-L760)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assumes the existence of a symbol table, in addition to the AST."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _type_check_relations(ast, sym_tab, is_log=False):\n",
    "    \"\"\"\n",
    "    A lark tree semantic check.\n",
    "    This pass makes the following assumptions and might not work correctly if they are not met\n",
    "    1. that relations and ie relations references and correct arity were checked.\n",
    "    2. variable references were checked.\n",
    "    3. it only gets a single statement as an input.\n",
    "\n",
    "    this pass performs the following checks:\n",
    "    1. checks if relation references are properly typed.\n",
    "    2. checks if ie relations are properly typed.\n",
    "    3. checks if free variables in rules do have conflicting types.\n",
    "    \"\"\"\n",
    "    rewrite(ast, lhs='_[type=\"add_fact\"]->afact_node[relation_name, term_list, type_list]',\n",
    "                condition=lambda match: _assertion(match['afact_node']['relation_name'], match['afact_node']['term_list'], \n",
    "                                               match['afact_node']['type_list'], sym_tab), is_log=is_log)\n",
    "    \n",
    "    rewrite(ast, lhs='_[type=\"remove_fact\"]->rfact_node[relation_name, term_list, type_list]',\n",
    "                condition=lambda match: _assertion(match['rfact_node']['relation_name'], match['rfact_node']['term_list'], \n",
    "                                               match['rfact_node']['type_list'], sym_tab), is_log=is_log)\n",
    "\n",
    "    rewrite(ast, lhs='_[type=\"query\"]->query_node[relation_name, term_list, type_list]',\n",
    "                condition=lambda match: _assertion(match['query_node']['relation_name'], match['query_node']['term_list'], \n",
    "                                               match['query_node']['type_list'], sym_tab), is_log=is_log)\n",
    "\n",
    "    rewrite(ast, lhs='_[type=\"rule\"]->rule_node[body_relation_list, body_relation_type_list]',\n",
    "                condition=lambda match: _assertion(match['rule_node']['body_relation_list'], match['rule_node']['type_list'], sym_tab), is_log=is_log)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Save Declared Relations Schemas\n",
    "[Spanner Workbench implementation](https://github.com/DeanLight/spanner_workbench/blob/f60595acd65ab26784c8f0d59bcaa586764441b3/src/rgxlog-interpreter/src/rgxlog/engine/passes/lark_passes.py#L763-L793)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assumes the existence of a symbol table, in addition to the AST."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _save_declared_relations_schemas(ast, sym_tab, is_log=False):\n",
    "    \"\"\"\n",
    "    this pass writes the relation schemas that it finds in relation declarations and rule heads* to the\n",
    "    symbol table.\n",
    "    * note that a rule is a relation declaration of the rule head relation and a definition of its contents\n",
    "\n",
    "    this pass assumes that type checking was already performed on its input.\n",
    "    \"\"\"\n",
    "    matches = rewrite(ast, lhs='_[type=\"relation_declaration\"]->reldec_node[relation_name, type_list]', is_log=is_log)\n",
    "    for match in matches:\n",
    "        # Does some imperative side-effect that adds to the sym_tab, based on the match's relation_name, type_list\n",
    "        rel_name = match['reldec_node']['relation_name']\n",
    "        type_list = match['reldec_node']['type_list']\n",
    "        print(f\"Perform side-effect calculation that uses {rel_name}, {type_list}\")\n",
    "\n",
    "    matches = rewrite(ast, lhs='_[type=\"rule\"]->rule_node[body_relation_list, body_relation_type_list]',\n",
    "                      condition=lambda match: _assertion(match['rule_node']['body_relation_list'], match['rule_node']['body_relation_type_list'], sym_tab), is_log=is_log)\n",
    "    for match in matches:\n",
    "        # Does some imperative side-effect that adds each relation in the rule to the sym_tab\n",
    "        body_relation_list = match['rule_node']['body_relation_list']\n",
    "        body_relation_type_list = match['rule_node']['body_relation_type_list']\n",
    "        print(f\"Perform side-effect calculation that uses {body_relation_list}, {body_relation_type_list}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Resolve Variables References\n",
    "[Spanner Workbench implementation](https://github.com/DeanLight/spanner_workbench/blob/f60595acd65ab26784c8f0d59bcaa586764441b3/src/rgxlog-interpreter/src/rgxlog/engine/passes/lark_passes.py#L796-L874)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assumes the existence of a symbol table, in addition to the AST."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _resolve_variables_references(ast, sym_tab, is_log=False):\n",
    "    \"\"\"\n",
    "    a lark execution pass\n",
    "    this pass replaces variable references with their literal values.\n",
    "    also replaces DataTypes.var_name types with the real type of the variable.\n",
    "    \"\"\"\n",
    "    cmd_type_attrs = {\n",
    "        'assignment': {\n",
    "            'type_var':'value_type',\n",
    "            'value_var':'value',\n",
    "        },\n",
    "        'read_assignment': {\n",
    "            'type_var':'read_arg_type',\n",
    "            'value_var':'read_arg',\n",
    "        },\n",
    "        'query': {\n",
    "            'type_var':'type_list',\n",
    "            'value_var':'term_list',\n",
    "        },\n",
    "        'add_fact': {\n",
    "            'type_var':'type_list',\n",
    "            'value_var':'term_list',\n",
    "        },\n",
    "        'remove_fact': {\n",
    "            'type_var':'type_list',\n",
    "            'value_var':'term_list',\n",
    "        },\n",
    "        'rule': {\n",
    "            'type_var':'body_relation_type_list',\n",
    "            'value_var':'body_relation_list',\n",
    "        },\n",
    "    }\n",
    "     \n",
    "    for cmd_type, attrs in cmd_type_attrs.items():\n",
    "        type_var = attrs['type_var']\n",
    "        value_var = attrs['value_var']\n",
    "        rewrite(ast, lhs=f'cmd[type=\"{cmd_type}\"]->cmd_node[{type_var}, {value_var}]',\n",
    "                    condition=lambda match: _assertion(match['cmd_node'][attrs['type_var']]),\n",
    "                    rhs=f'cmd[type]->cmd_node[{type_var}={{type_var}}, {value_var}={{value_var}}]',\n",
    "                    render_rhs={\n",
    "                        'type_var': lambda match: _calculation(match['cmd_node'][value_var], match['cmd_node'][type_var], sym_tab),\n",
    "                        'value_var': lambda match: _calculation(match['cmd_node'][value_var], match['cmd_node'][type_var], sym_tab)}, is_log=is_log)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Execute Assignments\n",
    "[Spanner Workbench implementation](https://github.com/DeanLight/spanner_workbench/blob/f60595acd65ab26784c8f0d59bcaa586764441b3/src/rgxlog-interpreter/src/rgxlog/engine/passes/lark_passes.py#L877-L904)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assumes the existence of a symbol table, in addition to the AST."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _execute_assignments(ast, sym_tab, is_log=False):\n",
    "    \"\"\"\n",
    "    a lark execution pass\n",
    "    executes assignments by saving variables' values and types in the symbol table\n",
    "    should be used only after variable references are resolved, meaning the assigned values and read() arguments\n",
    "    are guaranteed to be literals.\n",
    "    \"\"\"\n",
    "    matches = rewrite_iter(ast, lhs='ass[type=\"assignment\"]->ass_node[var_name, value, value_type]', is_log=is_log)\n",
    "    for match in matches:\n",
    "        # Imperative side-effect that uses each match's var_name, value, value_type and updates the sym_tab (executes assignment)\n",
    "        var_name = match['ass_node']['var_name']\n",
    "        value = match['ass_node']['value']\n",
    "        value_type = match['ass_node']['value_type']\n",
    "        print(f\"Perform side-effect calculation that uses {var_name}, {value}, {value_type}\")\n",
    "\n",
    "    matches = rewrite_iter(ast, lhs='rass[type=\"read_assignment\"]->rass_node[read_arg, var_name]', is_log=is_log)\n",
    "    for match in matches:\n",
    "        # Imperative side-effect that uses each match's read_arg, var_name and updates the sym_tab (executes assignment)\n",
    "        read_arg = match['rass_node']['read_arg']\n",
    "        var_name = match['rass_node']['var_name']\n",
    "        print(f\"Perform side-effect calculation that uses {read_arg}, {var_name}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Perform side-effect calculation that uses str_var, hello world, string\n",
      "Perform side-effect calculation that uses span_var, (2, 5), span\n"
     ]
    }
   ],
   "source": [
    "# We don't really have a symtab, so we pass None as a dummy argument\n",
    "_execute_assignments(test_g, None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Add Statements To Netx Parse Graph\n",
    "[Spanner Workbench implementation](https://github.com/DeanLight/spanner_workbench/blob/f60595acd65ab26784c8f0d59bcaa586764441b3/src/rgxlog-interpreter/src/rgxlog/engine/passes/lark_passes.py#L907-L965)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assumes the existence of a parse graph, in addition to the AST."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _add_statements_to_netx_parse_graph(ast, parse_graph, is_log=False):\n",
    "    \"\"\"\n",
    "    A lark execution pass.\n",
    "    This pass adds each statement in the input parse tree to the parse graph.\n",
    "    This pass is made to work with execution.naive_execution as the execution function and\n",
    "    term_graph.NetxStateGraph as the parse graph.\n",
    "\n",
    "    Each statement in the parse graph will be a child of the parse graph's root.\n",
    "\n",
    "    Each statement in the parse graph will have a type attribute that contains the statement's name in the\n",
    "    rgxlog grammar.\n",
    "\n",
    "    Some nodes in the parse graph will contain a value attribute that would contain a relation that describes\n",
    "    that statement.\n",
    "    e.g. a add_fact node would have a value which is a structured_nodes.AddFact instance\n",
    "    (which inherits from structured_nodes.Relation) that describes the fact that will be added.\n",
    "\n",
    "    Some statements are more complex and will be described by more than a single node, e.g. a rule node.\n",
    "    The reason for this is that we want a single netx node to not contain more than one Relation\n",
    "    (or IERelation) instance. This will make the parse graph a \"graph of relation nodes\", allowing\n",
    "     for flexibility for optimization in the future.\n",
    "    \"\"\"\n",
    "    matches = rewrite_iter(ast, lhs='_[type=\"add_fact\"]->afact_node', is_log=is_log)\n",
    "    for match in matches:\n",
    "        # Imperative side-effect that adds a new node to parse_graph, based on afact_node's attrs\n",
    "        pass\n",
    "\n",
    "    matches = rewrite_iter(ast, lhs='_[type=\"remove_fact\"]->rfact_node', is_log=is_log)\n",
    "    for match in matches:\n",
    "        # Imperative side-effect that adds a new node to parse_graph, based on rfact_node's attrs\n",
    "        pass\n",
    "\n",
    "    matches = rewrite_iter(ast, lhs='_[type=\"query\"]->query_node', is_log=is_log)\n",
    "    for match in matches:\n",
    "        # Imperative side-effect that adds a new node to parse_graph, based on query_node's attrs\n",
    "        pass\n",
    "\n",
    "    matches = rewrite_iter(ast, lhs='_[type=\"relation_declaration\"]->reldec_node', is_log=is_log)\n",
    "    for match in matches:\n",
    "        # Imperative side-effect that adds a new node to parse_graph, based on reldec_node's attrs\n",
    "        pass\n",
    "\n",
    "    matches = rewrite_iter(ast, lhs='_[type=\"rule\"]->rule_node', is_log=is_log)\n",
    "    for match in matches:\n",
    "        # Imperative side-effect that adds a new node to parse_graph, based on rule_node's attrs\n",
    "        pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "import nbdev; nbdev.nbdev_export()\n",
    "     "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}