Fix issues with pickling and equality checks of EC objects.

18 files changed, 289 insertions, 145 deletions

diff --git a/LICENSE b/LICENSE
index 109acf2..cc638ca 100644
--- a/LICENSE
+++ b/LICENSE
@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2018-2023 Jan Jancar
+Copyright (c) 2018-2024 Jan Jancar
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
diff --git a/README.md b/README.md
index 53c3d6b..652e67e 100644
--- a/README.md
+++ b/README.md
@@ -92,7 +92,7 @@ Use [black](https://github.com/psf/black) for code-formatting.
 
     MIT License
 
-    Copyright (c) 2018-2023 Jan Jancar
+    Copyright (c) 2018-2024 Jan Jancar
 
     Permission is hereby granted, free of charge, to any person obtaining a copy
     of this software and associated documentation files (the "Software"), to deal
diff --git a/pyecsca/ec/context.py b/pyecsca/ec/context.py
index 5bc9cdf..3735c26 100644
--- a/pyecsca/ec/context.py
+++ b/pyecsca/ec/context.py
@@ -244,6 +244,7 @@ current: Optional[Context] = None
 
 class _ContextManager:
     def __init__(self, new_context):
+        # TODO: Is this deepcopy a good idea?
         self.new_context = deepcopy(new_context)
 
     def __enter__(self) -> Optional[Context]:
diff --git a/pyecsca/ec/coordinates.py b/pyecsca/ec/coordinates.py
index 8692809..f15a64b 100644
--- a/pyecsca/ec/coordinates.py
+++ b/pyecsca/ec/coordinates.py
@@ -1,6 +1,5 @@
 """Provides a coordinate model class."""
 from ast import parse, Module
-from astunparse import unparse
 from importlib_resources.abc import Traversable
 from importlib_resources import as_file
 from typing import List, Any, MutableMapping
@@ -18,7 +17,6 @@ from .formula.efd import (
     ScalingEFDFormula,
     NegationEFDFormula,
 )
-from ..misc.utils import pexec
 
 
 @public
@@ -59,23 +57,6 @@ class CoordinateModel:
             f'{self.__class__.__name__}("{self.name}", curve_model={self.curve_model})'
         )
 
-    def __getstate__(self):
-        state = self.__dict__.copy()
-        state["satisfying"] = list(map(unparse, state["satisfying"]))
-        state["toaffine"] = list(map(unparse, state["toaffine"]))
-        state["tosystem"] = list(map(unparse, state["tosystem"]))
-        state["assumptions"] = list(map(unparse, state["assumptions"]))
-        state["neutral"] = list(map(unparse, state["neutral"]))
-        return state
-
-    def __setstate__(self, state):
-        state["satisfying"] = list(map(pexec, state["satisfying"]))
-        state["toaffine"] = list(map(pexec, state["toaffine"]))
-        state["tosystem"] = list(map(pexec, state["tosystem"]))
-        state["assumptions"] = list(map(pexec, state["assumptions"]))
-        state["neutral"] = list(map(pexec, state["neutral"]))
-        self.__dict__.update(state)
-
 
 @public
 class AffineCoordinateModel(CoordinateModel):
@@ -105,6 +86,12 @@ class AffineCoordinateModel(CoordinateModel):
 class EFDCoordinateModel(CoordinateModel):
     """A coordinate model from [EFD]_ data."""
 
+    def __new__(cls, *args, **kwargs):
+        _, name, curve_model = args
+        if name in curve_model.coordinates:
+            return curve_model.coordinates[name]
+        return object.__new__(cls)
+
     def __init__(self, dir_path: Traversable, name: str, curve_model: Any):
         self.name = name
         self.curve_model = curve_model
@@ -189,6 +176,12 @@ class EFDCoordinateModel(CoordinateModel):
                     )
                 line = f.readline().decode("ascii").rstrip()
 
+    def __getnewargs__(self):
+        return None, self.name, self.curve_model
+
+    def __getstate__(self):
+        return {}
+
     def __eq__(self, other):
         if not isinstance(other, EFDCoordinateModel):
             return False
diff --git a/pyecsca/ec/curve.py b/pyecsca/ec/curve.py
index c952bd9..6762724 100644
--- a/pyecsca/ec/curve.py
+++ b/pyecsca/ec/curve.py
@@ -80,7 +80,7 @@ class EllipticCurve:
                         )
             except NameError:
                 k = FF(self.prime)
-                assumption_string = unparse(assumption)
+                assumption_string = unparse(assumption).strip()
                 lhs, rhs = assumption_string.split(" = ")
                 expr = sympify(f"{rhs} - {lhs}")
                 for curve_param, value in self.parameters.items():
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)
diff --git a/pyecsca/ec/model.py b/pyecsca/ec/model.py
index 570e9e6..ec51a43 100644
--- a/pyecsca/ec/model.py
+++ b/pyecsca/ec/model.py
@@ -31,6 +31,7 @@ class CurveModel:
 
 class EFDCurveModel(CurveModel):
     """A curve model from [EFD]_ data."""
+
     _efd_name: str
     _loaded: bool = False
 
@@ -39,6 +40,9 @@ class EFDCurveModel(CurveModel):
         self.shortname = efd_name
         if self._loaded:
             return
+        self.__load(efd_name)
+
+    def __load(self, efd_name: str):
         self.__class__._loaded = True
         self.__class__.coordinates = {}
         self.__class__.parameter_names = []
@@ -54,7 +58,9 @@ class EFDCurveModel(CurveModel):
         for entry in files("pyecsca.ec").joinpath("efd", efd_name).iterdir():
             with as_file(entry) as file_path:
                 if entry.is_dir():
-                    self.__read_coordinate_dir(self.__class__, file_path, file_path.stem)
+                    self.__read_coordinate_dir(
+                        self.__class__, file_path, file_path.stem
+                    )
                 else:
                     self.__read_curve_file(self.__class__, file_path)
 
@@ -98,6 +104,14 @@ class EFDCurveModel(CurveModel):
             return False
         return self._efd_name == other._efd_name
 
+    def __getstate__(self):
+        return self.__dict__.copy()
+
+    def __setstate__(self, state):
+        if not self.__class__._loaded:
+            self.__load(state["_efd_name"])
+        self.__dict__.update(state)
+
     def __hash__(self):
         return hash(self._efd_name)
 
diff --git a/pyecsca/ec/op.py b/pyecsca/ec/op.py
index 2a730e1..0a87c81 100644
--- a/pyecsca/ec/op.py
+++ b/pyecsca/ec/op.py
@@ -25,6 +25,7 @@ from astunparse import unparse
 from public import public
 
 from .mod import Mod
+from ..misc.utils import pexec
 
 
 @public
@@ -67,6 +68,9 @@ class CodeOp:
 
     def __init__(self, code: Module):
         self.code = code
+        self.__parse(code)
+
+    def __parse(self, code: Module):
         assign = cast(Assign, code.body[0])
         self.result = cast(Name, assign.targets[0]).id
         params = set()
@@ -139,16 +143,21 @@ class CodeOp:
         else:
             return self.left, self.right  # type: ignore
 
+    def __hash__(self):
+        return hash((self.left, self.right, self.operator, self.result))
+
+    def __eq__(self, other):
+        if not isinstance(other, CodeOp):
+            return False
+        return self.left == other.left and self.right == other.right and self.operator == other.operator and self.result == other.result
+
     def __getstate__(self):
-        state = self.__dict__.copy()
-        state["code"] = unparse(state["code"])
-        del state["compiled"]
+        state = {"code": unparse(self.code).strip()}
         return state
 
     def __setstate__(self, state):
-        state["code"] = parse(state["code"], mode="exec")
-        state["compiled"] = compile(state["code"], "", mode="exec")
-        self.__dict__.update(state)
+        self.code = pexec(state["code"])
+        self.__parse(self.code)
 
     def __str__(self):
         return f"{self.result} = {self.left if self.left is not None else ''}{self.operator.op_str}{self.right if self.right is not None else ''}"
diff --git a/pyecsca/ec/params.py b/pyecsca/ec/params.py
index d1ee4f8..c125c3b 100644
--- a/pyecsca/ec/params.py
+++ b/pyecsca/ec/params.py
@@ -194,7 +194,7 @@ def _create_params(curve, coords, infty):
                         )
             except NameError:
                 k = FF(field)
-                assumption_string = unparse(assumption)
+                assumption_string = unparse(assumption).strip()
                 lhs, rhs = assumption_string.split(" = ")
                 expr = sympify(f"{rhs} - {lhs}")
                 for curve_param, value in params.items():
diff --git a/pyecsca/misc/utils.py b/pyecsca/misc/utils.py
index 010bbd2..4e6a26b 100644
--- a/pyecsca/misc/utils.py
+++ b/pyecsca/misc/utils.py
@@ -2,7 +2,7 @@
 import sys
 from ast import parse
 
-from pyecsca.misc.cfg import getconfig
+from ..misc.cfg import getconfig
 
 
 def pexec(s):
diff --git a/test/ec/test_formula.py b/test/ec/test_formula.py
index 32e1cf3..ed437a5 100644
--- a/test/ec/test_formula.py
+++ b/test/ec/test_formula.py
@@ -77,6 +77,10 @@ def test_eq(add, dbl):
     assert add != dbl
 
 
+def test_hashable(add):
+    assert hash(add)
+
+
 def test_num_ops(add):
     assert add.num_operations == 33
     assert add.num_multiplications == 17
@@ -400,7 +404,9 @@ def library_formula_params(request) -> Tuple[EFDFormula, DomainParameters]:
 
 def test_formula_graph(library_formula_params):
     formula, params = library_formula_params
-    do_test_formula(rename_ivs(formula), params)
+    renamed = rename_ivs(formula)
+    do_test_formula(renamed, params)
+    assert hash(renamed)
 
 
 def test_switch_sign(library_formula_params):
diff --git a/test/ec/test_pickle.py b/test/ec/test_pickle.py
new file mode 100644
index 0000000..5d448b5
--- /dev/null
+++ b/test/ec/test_pickle.py
@@ -0,0 +1,83 @@
+import pickle
+import pickletools
+from multiprocessing import get_context
+from multiprocessing.context import BaseContext
+
+import pytest
+
+from pyecsca.ec.formula.graph import FormulaGraph
+from pyecsca.ec.model import ShortWeierstrassModel
+from pyecsca.ec.params import get_params
+
+
+@pytest.fixture()
+def ctx() -> BaseContext:
+    return get_context("spawn")
+
+
+def model_target(obj):
+    return obj._loaded
+
+
+def test_model_loads(ctx):
+    # Test that the unpickling of EFDCurveModel triggers a reload that loads the EFD data for that model.
+    sw = ShortWeierstrassModel()
+    with ctx.Pool(processes=1) as pool:
+        res = pool.apply(model_target, args=(sw,))
+        assert res
+
+
+def test_model():
+    sw = ShortWeierstrassModel()
+    data = pickle.dumps(sw)
+    back = pickle.loads(data)
+    assert sw == back
+
+
+def test_coords():
+    sw = ShortWeierstrassModel()
+    coords = sw.coordinates["projective"]
+    data = pickle.dumps(coords)
+    back = pickle.loads(data)
+    assert coords == back
+
+
+def test_formula():
+    sw = ShortWeierstrassModel()
+    coords = sw.coordinates["projective"]
+    formula = coords.formulas["add-2007-bl"]
+    data = pickle.dumps(formula)
+    back = pickle.loads(data)
+    assert formula == back
+    formulas = tuple(coords.formulas.values())
+    data = pickle.dumps(formulas)
+    back = pickle.loads(data)
+    assert formulas == back
+
+
+def test_code_formula():
+    sw = ShortWeierstrassModel()
+    coords = sw.coordinates["projective"]
+    formula = coords.formulas["add-2007-bl"]
+    graph = FormulaGraph(formula)
+    code = graph.to_formula()
+    data = pickle.dumps(code)
+    pickletools.dis(data)
+    back = pickle.loads(data)
+    assert code == back
+
+
+def test_params(secp128r1):
+    secp256r1 = get_params("secg", "secp256r1", "projective")
+    params = (secp128r1, secp256r1)
+    data = pickle.dumps(params)
+    back = pickle.loads(data)
+    assert params == back
+
+
+def test_op(secp128r1):
+    formula = secp128r1.curve.coordinate_model.formulas["add-2007-bl"]
+    op = formula.code[0]
+    data = pickle.dumps(op)
+    back = pickle.loads(data)
+    assert op == back