def MergeChildren(node_a: ActionNode, node_b: ActionNode) -> None:
assert not node_a.dead
assert not node_b.dead
assert node_a.action.name is not node_b.action.name
node_a.AssertValidity()
node_b.AssertValidity()
node_a_name = node_a.action.name
node_b_name = node_b.action.name
children_a = list(node_a.children.copy().values())
children_b = list(node_b.children.copy().values())
for child in children_a + children_b:
child = child.ResolveToAliveNode()
child.AssertValidity()
child_name = child.action.name
if (child_name in node_a.children
and child_name in node_b.children):
MergeNodes(node_a.children[child_name],
node_b.children[child_name])
elif child_name in node_a.children:
if (node_b_name in child.parents):
node_b = MergeNodes(node_b, child.parents[node_b_name])
else:
node_b.AddChild(child)
elif child_name in node_b.children:
if (node_a_name in child.parents):
node_a = MergeNodes(node_a, child.parents[node_a_name])
else:
node_a.AddChild(child)
else:
assert False
node_a = node_a.ResolveToAliveNode()
node_b = node_b.ResolveToAliveNode()
resolved_children = ({
k: node_a.children.get(k).ResolveToAliveNode()
for k in node_a.children.keys()
})
node_a.children = resolved_children.copy()
node_b.children = resolved_children
merged_state_check_coverage = MergeStateCheckCoverage(
node_a.state_check_actions_coverage,
node_b.state_check_actions_coverage)
node_a.state_check_actions.update(node_b.state_check_actions)
node_a.state_check_actions_coverage = merged_state_check_coverage
node_b.state_check_actions.update(node_a.state_check_actions)
node_b.state_check_actions_coverage = merged_state_check_coverage.copy(
)
node_a.AssertValidity()
node_b.AssertValidity()
def PruneNonKeepNodes(node: ActionNode):
children = {}
for child in node.children.values():
if child.keep:
children[child.action.name] = child
PruneNonKeepNodes(child)
node.children = children
return node
def TrimGraphToAllowedActions(root_node: ActionNode,
allowed_actions: Dict[str, Action]) -> None:
new_children = {}
for child in root_node.children.values():
if child.action.name in allowed_actions:
new_children[child.action.name] = child
root_node.children = new_children
new_state_check_actions = {}
for action in root_node.state_check_actions.values():
if action.name in allowed_actions:
new_state_check_actions[action.name] = action
root_node.state_check_actions = new_state_check_actions
for child in root_node.children.values():
TrimGraphToAllowedActions(child, allowed_actions)