Fix issues with pickling and equality checks of EC objects.

7 files changed, 151 insertions, 113 deletions

diff --git a/pyecsca/ec/formula/base.py b/pyecsca/ec/formula/base.py
index 62ce835..a9402f1 100644
--- a/pyecsca/ec/formula/base.py
+++ b/pyecsca/ec/formula/base.py
@@ -128,7 +128,7 @@ class Formula(ABC):
     num_outputs: ClassVar[int]
     """Number of outputs (points) of the formula."""
     unified: bool
-    """Whether the formula is specifies that it is unified."""
+    """Whether the formula specifies that it is unified."""
 
     @cached_property
     def assumptions_str(self):
diff --git a/pyecsca/ec/formula/efd.py b/pyecsca/ec/formula/efd.py
index 485e6f5..c2ec49a 100644
--- a/pyecsca/ec/formula/efd.py
+++ b/pyecsca/ec/formula/efd.py
@@ -14,12 +14,19 @@ from .base import (
     DifferentialAdditionFormula,
     LadderFormula,
 )
-from ast import parse
+
+from ...misc.utils import pexec, peval
 
 
 class EFDFormula(Formula):
     """Formula from the [EFD]_."""
 
+    def __new__(cls, *args, **kwargs):
+        _, _, name, coordinate_model = args
+        if name in coordinate_model.formulas:
+            return coordinate_model.formulas[name]
+        return object.__new__(cls)
+
     def __init__(
         self,
         meta_path: Traversable,
@@ -47,9 +54,7 @@ class EFDFormula(Formula):
                     self.parameters.append(line[10:])
                 elif line.startswith("assume"):
                     self.assumptions.append(
-                        parse(
-                            line[7:].replace("=", "==").replace("^", "**"), mode="eval"
-                        )
+                        peval(line[7:].replace("=", "==").replace("^", "**"))
                     )
                 elif line.startswith("unified"):
                     self.unified = True
@@ -58,11 +63,18 @@ class EFDFormula(Formula):
     def __read_op3_file(self, path: Traversable):
         with path.open("rb") as f:
             for line in f.readlines():
-                code_module = parse(
-                    line.decode("ascii").replace("^", "**"), str(path), mode="exec"
-                )
+                code_module = pexec(line.decode("ascii").replace("^", "**"))
                 self.code.append(CodeOp(code_module))
 
+    def __getnewargs__(self):
+        return None, None, self.name, self.coordinate_model
+
+    def __getstate__(self):
+        return {}
+
+    def __setstate__(self, state):
+        pass
+
     def __str__(self):
         return f"{self.coordinate_model!s}/{self.name}"
 
diff --git a/pyecsca/ec/formula/expand.py b/pyecsca/ec/formula/expand.py
index 3265131..bd8b236 100644
--- a/pyecsca/ec/formula/expand.py
+++ b/pyecsca/ec/formula/expand.py
@@ -4,14 +4,13 @@ from public import public
 from . import Formula
 from .efd import EFDFormula
 from .fliparoo import recursive_fliparoo
-from .graph import ModifiedEFDFormula
 from .metrics import ivs_norm
 from .partitions import reduce_all_adds, expand_all_muls, expand_all_nopower2_muls
 from .switch_sign import generate_switched_formulas
 
 
-def reduce_with_similarity(formulas: List[EFDFormula], norm):
-    efd = list(filter(lambda x: not isinstance(x, ModifiedEFDFormula), formulas))
+def reduce_with_similarity(formulas: List[Formula], norm):
+    efd = list(filter(lambda x: isinstance(x, EFDFormula), formulas))
     reduced_efd = efd
     similarities = list(map(norm, efd))
     for formula in formulas:
@@ -25,15 +24,15 @@ def reduce_with_similarity(formulas: List[EFDFormula], norm):
 
 @public
 def expand_formula_list(
-    formulas: List[EFDFormula], norm: Callable[[Formula], Any] = ivs_norm
-) -> List[EFDFormula]:
+    formulas: List[Formula], norm: Callable[[Formula], Any] = ivs_norm
+) -> List[Formula]:
     extended = reduce_with_similarity(formulas, norm)
 
-    fliparood: List[EFDFormula] = sum(list(map(recursive_fliparoo, extended)), [])
+    fliparood: List[Formula] = sum(list(map(recursive_fliparoo, extended)), [])
     extended.extend(fliparood)
     extended = reduce_with_similarity(extended, norm)
 
-    switch_signs: List[EFDFormula] = sum(
+    switch_signs: List[Formula] = sum(
         [list(generate_switched_formulas(f)) for f in extended], []
     )
     extended.extend(switch_signs)
diff --git a/pyecsca/ec/formula/fliparoo.py b/pyecsca/ec/formula/fliparoo.py
index c8d77ac..beddb9d 100644
--- a/pyecsca/ec/formula/fliparoo.py
+++ b/pyecsca/ec/formula/fliparoo.py
@@ -1,23 +1,23 @@
-from typing import Iterator, List, Tuple, Type, Optional
+from ast import parse
+from typing import Iterator, List, Type, Optional
 from ..op import OpType
-from .graph import EFDFormulaGraph, Node, CodeOpNode, CodeOp, parse
-from .efd import EFDFormula
-from random import randint
+from .base import Formula
+from .graph import FormulaGraph, Node, CodeOpNode, CodeOp, CodeFormula
 
 
 class Fliparoo:
     """
-    Fliparoo is a chain of nodes N1->N2->...->Nk in EFDFormulaGraph for k>=2 such that:
+    Fliparoo is a chain of nodes N1->N2->...->Nk in FormulaGraph for k>=2 such that:
         - All Ni are * or All Ni are +/-
         - For every Ni, except for Nk, the only outgoing node is Ni+1
         - Neither of N1,...,Nk-1 is an output node
     """
 
     nodes: List[CodeOpNode]
-    graph: EFDFormulaGraph
+    graph: FormulaGraph
     operator: Optional[OpType]
 
-    def __init__(self, chain: List[CodeOpNode], graph: EFDFormulaGraph):
+    def __init__(self, chain: List[CodeOpNode], graph: FormulaGraph):
         self.verify_chain(chain)
         self.nodes = chain
         self.graph = graph
@@ -77,7 +77,7 @@ class Fliparoo:
 
 
 class MulFliparoo(Fliparoo):
-    def __init__(self, chain: List[CodeOpNode], graph: EFDFormulaGraph):
+    def __init__(self, chain: List[CodeOpNode], graph: FormulaGraph):
         super().__init__(chain, graph)
         operations = set(node.op.operator for node in self.nodes)
         if len(operations) != 1 or list(operations)[0] != OpType.Mult:
@@ -86,7 +86,7 @@ class MulFliparoo(Fliparoo):
 
 
 class AddSubFliparoo(Fliparoo):
-    def __init__(self, chain: List[CodeOpNode], graph: EFDFormulaGraph):
+    def __init__(self, chain: List[CodeOpNode], graph: FormulaGraph):
         super().__init__(chain, graph)
         operations = set(node.op.operator for node in self.nodes)
         if not operations.issubset([OpType.Add, OpType.Sub]):
@@ -94,7 +94,7 @@ class AddSubFliparoo(Fliparoo):
 
 
 class AddFliparoo(Fliparoo):
-    def __init__(self, chain: List[CodeOpNode], graph: EFDFormulaGraph):
+    def __init__(self, chain: List[CodeOpNode], graph: FormulaGraph):
         super().__init__(chain, graph)
         operations = set(node.op.operator for node in self.nodes)
         if len(operations) != 1 or list(operations)[0] != OpType.Add:
@@ -107,7 +107,7 @@ class BadFliparoo(Exception):
 
 
 def find_fliparoos(
-    graph: EFDFormulaGraph, fliparoo_type: Optional[Type[Fliparoo]] = None
+    graph: FormulaGraph, fliparoo_type: Optional[Type[Fliparoo]] = None
 ) -> List[Fliparoo]:
     """Finds a list of Fliparoos in a graph"""
     fliparoos: List[Fliparoo] = []
@@ -138,7 +138,7 @@ def is_subfliparoo(fliparoo: Fliparoo, longest_fliparoos: List[Fliparoo]) -> boo
 
 def largest_fliparoo(
     chain: List[CodeOpNode],
-    graph: EFDFormulaGraph,
+    graph: FormulaGraph,
     fliparoo_type: Optional[Type[Fliparoo]] = None,
 ) -> Optional[Fliparoo]:
     """Finds the largest fliparoo in a list of Nodes"""
@@ -183,7 +183,7 @@ class SignedNode:
 class SignedSubGraph:
     """Subgraph of an EFDFormula graph with signed nodes"""
 
-    def __init__(self, nodes: List[SignedNode], graph: EFDFormulaGraph):
+    def __init__(self, nodes: List[SignedNode], graph: FormulaGraph):
         self.nodes = nodes
         self.supergraph = graph
 
@@ -234,16 +234,16 @@ class DummyNode(Node):
 
 
 def generate_fliparood_formulas(
-    formula: EFDFormula, rename: bool = True
-) -> Iterator[EFDFormula]:
-    graph = EFDFormulaGraph(formula, rename)
+    formula: Formula, rename: bool = True
+) -> Iterator[CodeFormula]:
+    graph = FormulaGraph(formula, rename)
     fliparoos = find_fliparoos(graph)
     for fliparoo in fliparoos:
         for flip_graph in generate_fliparood_graphs(fliparoo):
-            yield flip_graph.to_EFDFormula()
+            yield flip_graph.to_formula()
 
 
-def generate_fliparood_graphs(fliparoo: Fliparoo) -> Iterator[EFDFormulaGraph]:
+def generate_fliparood_graphs(fliparoo: Fliparoo) -> Iterator[FormulaGraph]:
     fliparoo = fliparoo.deepcopy()
     last_str = fliparoo.last.result
     disconnect_fliparoo_outputs(fliparoo)
@@ -302,7 +302,7 @@ def disconnect_fliparoo_outputs(fliparoo: Fliparoo):
     fliparoo.last.reconnect_outgoing_nodes(dummy)
 
 
-def reconnect_fliparoo_outputs(graph: EFDFormulaGraph, last_node: Node):
+def reconnect_fliparoo_outputs(graph: FormulaGraph, last_node: Node):
     dummy = next(filter(lambda x: isinstance(x, DummyNode), graph.nodes))
     dummy.reconnect_outgoing_nodes(last_node)
     graph.remove_node(dummy)
@@ -347,7 +347,10 @@ def combine_signed_nodes(
             sign = -1
 
     new_node = CodeOpNode.from_str(
-        f"Fliparoo{id(left_node)}_{id(operator)}_{id(sign)}_{id(right_node)}", left_node.result, operator, right_node.result
+        f"Fliparoo{id(left_node)}_{id(operator)}_{id(sign)}_{id(right_node)}",
+        left_node.result,
+        operator,
+        right_node.result,
     )
     new_node.incoming_nodes = [left_node, right_node]
     left_node.outgoing_nodes.append(new_node)
diff --git a/pyecsca/ec/formula/graph.py b/pyecsca/ec/formula/graph.py
index 1473858..8386973 100644
--- a/pyecsca/ec/formula/graph.py
+++ b/pyecsca/ec/formula/graph.py
@@ -1,14 +1,17 @@
-from .efd import (
-    EFDFormula,
-    DoublingEFDFormula,
-    AdditionEFDFormula,
-    LadderEFDFormula,
-    DifferentialAdditionEFDFormula,
+from . import (
+    Formula,
+    AdditionFormula,
+    DoublingFormula,
+    LadderFormula,
+    TriplingFormula,
+    NegationFormula,
+    ScalingFormula,
+    DifferentialAdditionFormula,
 )
 from ..op import CodeOp, OpType
 import matplotlib.pyplot as plt
 import networkx as nx
-from ast import parse
+from ast import parse, Expression
 from typing import Dict, List, Tuple, Set, Optional, MutableMapping, Any
 from copy import deepcopy
 from abc import ABC, abstractmethod
@@ -194,7 +197,7 @@ class InputNode(Node):
         return f"Node({self.input})"
 
 
-def formula_input_variables(formula: EFDFormula) -> List[str]:
+def formula_input_variables(formula: Formula) -> List[str]:
     return (
         list(formula.inputs)
         + formula.parameters
@@ -202,43 +205,71 @@ def formula_input_variables(formula: EFDFormula) -> List[str]:
     )
 
 
-# temporary solution
-class ModifiedEFDFormula(EFDFormula):
+class CodeFormula(Formula):
+    def __init__(self, name, code, coordinate_model, parameters, assumptions):
+        self.name = name
+        self.coordinate_model = coordinate_model
+        self.meta = {}
+        self.parameters = parameters
+        self.assumptions = assumptions
+        self.code = code
+        self.unified = False
+
+    def __hash__(self):
+        return hash((self.name, self.coordinate_model, tuple(self.code), tuple(self.parameters), tuple(self.assumptions)))
+
     def __eq__(self, other):
-        if not isinstance(other, ModifiedEFDFormula):
+        if not isinstance(other, CodeFormula):
             return False
         return (
-            self.name == other.name and self.coordinate_model == other.coordinate_model and self.code == other.code
+            self.name == other.name
+            and self.coordinate_model == other.coordinate_model
+            and self.code == other.code
         )
 
 
-class ModifiedDoublingEFDFormula(DoublingEFDFormula, ModifiedEFDFormula):
+class CodeAdditionFormula(AdditionFormula, CodeFormula):
+    pass
+
+
+class CodeDoublingFormula(DoublingFormula, CodeFormula):
+    pass
+
+
+class CodeLadderFormula(LadderFormula, CodeFormula):
+    pass
+
+
+class CodeTriplingFormula(TriplingFormula, CodeFormula):
     pass
 
 
-class ModifiedAdditionEFDFormula(AdditionEFDFormula, ModifiedEFDFormula):
+class CodeNegationFormula(NegationFormula, CodeFormula):
     pass
 
 
-class ModifiedDifferentialAdditionEFDFormula(
-    DifferentialAdditionEFDFormula, ModifiedEFDFormula
-):
+class CodeScalingFormula(ScalingFormula, CodeFormula):
     pass
 
 
-class ModifiedLadderEFDFormula(LadderEFDFormula, ModifiedEFDFormula):
+class CodeDifferentialAdditionFormula(DifferentialAdditionFormula, CodeFormula):
     pass
 
 
-class EFDFormulaGraph:
+class FormulaGraph:
+    coordinate_model: Any
+    shortname: str
+    parameters: List[str]
+    assumptions: List[Expression]
     nodes: List[Node]
     input_nodes: MutableMapping[str, InputNode]
     output_names: Set[str]
     roots: List[Node]
-    coordinate_model: Any
 
-    def __init__(self, formula: EFDFormula, rename=True):
-        self._formula = formula  # TODO remove, its here only for to_EFDFormula
+    def __init__(self, formula: Formula, rename=True):
+        self.shortname = formula.shortname
+        self.parameters = formula.parameters
+        self.assumptions = formula.assumptions
         self.coordinate_model = formula.coordinate_model
         self.output_names = formula.outputs
         self.input_nodes = {v: InputNode(v) for v in formula_input_variables(formula)}
@@ -281,25 +312,21 @@ class EFDFormulaGraph:
     def deepcopy(self):
         return deepcopy(self)
 
-    def to_EFDFormula(self) -> ModifiedEFDFormula:
-        # TODO rewrite
-        new_graph = deepcopy(self)
-        new_formula = new_graph._formula
-        new_formula.code = list(
+    def to_formula(self, name=None) -> CodeFormula:
+        code = list(
             map(
-                lambda x: x.op,  # type: ignore
-                filter(lambda n: n not in new_graph.roots, new_graph.nodes),
+                lambda x: deepcopy(x.op),  # type: ignore
+                filter(lambda n: n not in self.roots, self.nodes),
             )
         )
-        casting = {
-            AdditionEFDFormula: ModifiedAdditionEFDFormula,
-            DoublingEFDFormula: ModifiedDoublingEFDFormula,
-            DifferentialAdditionEFDFormula: ModifiedDifferentialAdditionEFDFormula,
-            LadderEFDFormula: ModifiedLadderEFDFormula,
-        }
-        if new_formula.__class__ not in set(casting.values()):
-            new_formula.__class__ = casting[new_formula.__class__]
-        return new_formula  # type: ignore
+        parameters = list(self.parameters)
+        assumptions = [deepcopy(assumption) for assumption in self.assumptions]
+        for klass in CodeFormula.__subclasses__():
+            if klass.shortname == self.shortname:
+                return klass(
+                    name, code, self.coordinate_model, parameters, assumptions
+                )
+        raise ValueError(f"Bad formula type: {self.shortname}")
 
     def networkx_graph(self) -> nx.DiGraph:
         graph = nx.DiGraph()
@@ -431,6 +458,6 @@ class EFDFormulaGraph:
             print(node)
 
 
-def rename_ivs(formula: EFDFormula):
-    graph = EFDFormulaGraph(formula)
-    return graph.to_EFDFormula()
+def rename_ivs(formula: Formula) -> CodeFormula:
+    graph = FormulaGraph(formula)
+    return graph.to_formula()
diff --git a/pyecsca/ec/formula/partitions.py b/pyecsca/ec/formula/partitions.py
index 9ea108c..50e0e6d 100644
--- a/pyecsca/ec/formula/partitions.py
+++ b/pyecsca/ec/formula/partitions.py
@@ -1,19 +1,14 @@
 from typing import List, Any, Generator
 from ast import parse
+from .base import Formula
 from ..op import OpType, CodeOp
-from .graph import (
-    EFDFormulaGraph,
-    CodeOpNode,
-    ConstantNode,
-    Node,
-)
+from .graph import FormulaGraph, CodeOpNode, ConstantNode, Node, CodeFormula
 from .fliparoo import find_fliparoos, AddFliparoo, MulFliparoo
 from copy import deepcopy
-from .efd import EFDFormula
 
 
-def reduce_all_adds(formula: EFDFormula, rename=True) -> EFDFormula:
-    graph = EFDFormulaGraph(formula, rename=rename)
+def reduce_all_adds(formula: Formula, rename=True) -> CodeFormula:
+    graph = FormulaGraph(formula, rename=rename)
     add_fliparoos = find_single_input_add_fliparoos(graph)
     for add_fliparoo in add_fliparoos:
         reduce_add_fliparoo(add_fliparoo, copy=False)
@@ -21,26 +16,26 @@ def reduce_all_adds(formula: EFDFormula, rename=True) -> EFDFormula:
     mul_fliparoos = find_constant_mul_fliparoos(graph)
     for mul_fliparoo in mul_fliparoos:
         reduce_mul_fliparoo(mul_fliparoo, copy=False)
-    return graph.to_EFDFormula()
+    return graph.to_formula()
 
 
-def expand_all_muls(formula: EFDFormula, rename=True) -> EFDFormula:
-    graph = EFDFormulaGraph(formula, rename)
+def expand_all_muls(formula: Formula, rename=True) -> CodeFormula:
+    graph = FormulaGraph(formula, rename)
     enodes = find_expansion_nodes(graph)
     for enode in enodes:
         expand_mul(graph, enode, copy=False)
-    return graph.to_EFDFormula()
+    return graph.to_formula()
 
 
-def expand_all_nopower2_muls(formula: EFDFormula, rename=True) -> EFDFormula:
-    graph = EFDFormulaGraph(formula, rename)
+def expand_all_nopower2_muls(formula: Formula, rename=True) -> CodeFormula:
+    graph = FormulaGraph(formula, rename)
     enodes = find_expansion_nodes(graph, nopower2=True)
     for enode in enodes:
         expand_mul(graph, enode, copy=False)
-    return graph.to_EFDFormula()
+    return graph.to_formula()
 
 
-def find_single_input_add_fliparoos(graph: EFDFormulaGraph) -> List[AddFliparoo]:
+def find_single_input_add_fliparoos(graph: FormulaGraph) -> List[AddFliparoo]:
     fliparoos = find_fliparoos(graph, AddFliparoo)
     single_input_fliparoos = []
     for fliparoo in fliparoos:
@@ -56,7 +51,7 @@ def find_single_input_add_fliparoos(graph: EFDFormulaGraph) -> List[AddFliparoo]
     return single_input_fliparoos
 
 
-def find_constant_mul_fliparoos(graph: EFDFormulaGraph) -> List[MulFliparoo]:
+def find_constant_mul_fliparoos(graph: FormulaGraph) -> List[MulFliparoo]:
     fliparoos = find_fliparoos(graph, MulFliparoo)
     constant_mul_fliparoo = []
     for fliparoo in fliparoos:
@@ -87,7 +82,7 @@ def find_constant_mul_fliparoos(graph: EFDFormulaGraph) -> List[MulFliparoo]:
     return constant_mul_fliparoo
 
 
-def find_expansion_nodes(graph: EFDFormulaGraph, nopower2=False) -> List[Node]:
+def find_expansion_nodes(graph: FormulaGraph, nopower2=False) -> List[Node]:
     expansion_nodes: List[Node] = []
     for node in graph.nodes:
         if not isinstance(node, CodeOpNode) or not node.is_mul:
@@ -109,7 +104,7 @@ def is_power_of_2(n: int) -> bool:
     return True
 
 
-def reduce_all_XplusX(graph: EFDFormulaGraph):
+def reduce_all_XplusX(graph: FormulaGraph):
     adds = find_all_XplusX(graph)
     for node in adds:
         reduce_XplusX(graph, node)
@@ -126,7 +121,7 @@ def find_all_XplusX(graph) -> List[CodeOpNode]:
     return adds
 
 
-def reduce_XplusX(graph: EFDFormulaGraph, node: CodeOpNode):
+def reduce_XplusX(graph: FormulaGraph, node: CodeOpNode):
     inode = node.incoming_nodes[0]
     const_node = ConstantNode(2)
     node.incoming_nodes[1] = const_node
@@ -146,7 +141,9 @@ def reduce_mul_fliparoo(fliparoo: MulFliparoo, copy=True):
     inode = next(
         filter(lambda x: not isinstance(x, ConstantNode), first.incoming_nodes)
     )
-    const_nodes: List[ConstantNode] = [node for node in fliparoo.input_nodes() if isinstance(node, ConstantNode)]
+    const_nodes: List[ConstantNode] = [
+        node for node in fliparoo.input_nodes() if isinstance(node, ConstantNode)
+    ]
     sum_const_node = ConstantNode(sum(v.value for v in const_nodes))
     fliparoo.graph.add_node(sum_const_node)
 
@@ -195,7 +192,7 @@ def reduce_add_fliparoo(fliparoo: AddFliparoo, copy=True):
     return fliparoo.graph
 
 
-def expand_mul(graph: EFDFormulaGraph, node: Node, copy=True) -> EFDFormulaGraph:
+def expand_mul(graph: FormulaGraph, node: Node, copy=True) -> FormulaGraph:
     if copy:
         i = graph.node_index(node)
         graph = deepcopy(graph)
@@ -275,17 +272,17 @@ def compute_partitions(n: int) -> List[Partition]:
     return result
 
 
-def generate_partitioned_formulas(formula: EFDFormula, rename=True):
-    graph = EFDFormulaGraph(formula, rename)
+def generate_partitioned_formulas(formula: Formula, rename=True):
+    graph = FormulaGraph(formula, rename)
     enodes = find_expansion_nodes(graph)
     for enode in enodes:
         for part_graph in generate_all_node_partitions(graph, enode):
-            yield part_graph.to_EFDFormula()
+            yield part_graph.to_formula()
 
 
 def generate_all_node_partitions(
-    original_graph: EFDFormulaGraph, node: Node
-) -> Generator[EFDFormulaGraph, Any, None]:
+    original_graph: FormulaGraph, node: Node
+) -> Generator[FormulaGraph, Any, None]:
     const_par = next(filter(lambda x: isinstance(x, ConstantNode), node.incoming_nodes))
     const_par_value = const_par.value
 
@@ -327,7 +324,7 @@ def generate_all_node_partitions(
 
 
 def partition_node(
-    graph: EFDFormulaGraph, node: CodeOpNode, partition: Partition, source_node: Node
+    graph: FormulaGraph, node: CodeOpNode, partition: Partition, source_node: Node
 ):
     if partition.is_final and partition.value == 1:
         # source node will take the role of node
diff --git a/pyecsca/ec/formula/switch_sign.py b/pyecsca/ec/formula/switch_sign.py
index 1acef5b..9f19629 100644
--- a/pyecsca/ec/formula/switch_sign.py
+++ b/pyecsca/ec/formula/switch_sign.py
@@ -1,29 +1,29 @@
 from typing import Dict, Iterator, List, Any
 from ast import parse
 from ..op import OpType, CodeOp
-from .graph import EFDFormulaGraph, ConstantNode, Node, CodeOpNode
+from .base import Formula
+from .graph import FormulaGraph, ConstantNode, CodeOpNode, CodeFormula
 from itertools import chain, combinations
-from .efd import EFDFormula
 from ..point import Point
 from ..mod import Mod
 
 
 def generate_switched_formulas(
-    formula: EFDFormula, rename=True
-) -> Iterator[EFDFormula]:
-    graph = EFDFormulaGraph(formula, rename)
+    formula: Formula, rename=True
+) -> Iterator[CodeFormula]:
+    graph = FormulaGraph(formula, rename)
     for node_combination in subnode_lists(graph):
         try:
-            yield switch_sign(graph, node_combination).to_EFDFormula()
+            yield switch_sign(graph, node_combination).to_formula()
         except BadSignSwitch:
             continue
 
 
-def subnode_lists(graph: EFDFormulaGraph):
+def subnode_lists(graph: FormulaGraph):
     return powerlist(filter(lambda x: x not in graph.roots and x.is_sub, graph.nodes))
 
 
-def switch_sign(graph: EFDFormulaGraph, node_combination) -> EFDFormulaGraph:
+def switch_sign(graph: FormulaGraph, node_combination) -> FormulaGraph:
     nodes_i = [graph.node_index(node) for node in node_combination]
     graph = graph.deepcopy()
     node_combination = set(graph.nodes[node_i] for node_i in nodes_i)