proof_tree.py

import pdb


class ProofNode:
    def __init__(self, parent, goal_id):
        "A ProofNode is an intermediate goal during proofs"
        self.parent = parent
        self.goal_id = goal_id
        self.children = []

    def to_dict(self):
        return {
            "goal_id": self.goal_id,
            "children": [c.to_dict() for c in self.children],
        }


class ProofTree:
    def __init__(self, steps, goals):
        self.root = ProofNode(None, steps[0]["goal_ids"]["fg"][0])
        frontier = {self.root.goal_id: self.root}
        for i, step in enumerate(steps[:-1]):
            next_step = steps[i + 1]
            goals_pre = set(step["goal_ids"]["fg"] + step["goal_ids"]["bg"])
            goals_post = set(next_step["goal_ids"]["fg"] + next_step["goal_ids"]["bg"])
            if goals_pre == goals_post:
                continue
            decomposed_goals = list(goals_pre - goals_post)
            new_goals = list(goals_post - goals_pre)
            assert len(decomposed_goals) == 1
            node = frontier.pop(decomposed_goals[0])
            frontier.update(
                {child.goal_id: child for child in self.expand(node, new_goals)}
            )

    def expand(self, node, new_goals):
        assert node.children == []
        for g in new_goals:
            node.children.append(ProofNode(node, g))
        return node.children

    def to_dict(self):
        return self.root.to_dict()