From 403c0e2fc940f1b8c20c77402d2143cc432fb02d Mon Sep 17 00:00:00 2001
From: ChrisTimperley <christimperley@googlemail.com>
Date: Fri, 12 Jul 2024 20:36:44 -0400
Subject: [PATCH] added ability to merge analyses

---
 src/kaskara/analysis.py   | 10 ++++++++++
 src/kaskara/functions.py  |  8 ++++++++
 src/kaskara/insertions.py |  3 +++
 src/kaskara/loops.py      |  9 +++++++++
 src/kaskara/statements.py |  7 +++++++
 5 files changed, 37 insertions(+)

diff --git a/src/kaskara/analysis.py b/src/kaskara/analysis.py
index 2ddc69c..f658bd8 100644
--- a/src/kaskara/analysis.py
+++ b/src/kaskara/analysis.py
@@ -53,6 +53,16 @@ def with_relative_locations(self, base: str) -> Analysis:
             insertions=self.insertions.with_relative_locations(base),
         )
 
+    def merge(self, other: Analysis) -> Analysis:
+        """Merges the results of this analysis and another analysis together."""
+        return Analysis(
+            files=self.files.union(other.files),
+            loops=self.loops.merge(other.loops),
+            functions=self.functions.merge(other.functions),
+            statements=self.statements.merge(other.statements),
+            insertions=self.insertions.merge(other.insertions),
+        )
+
     def is_inside_loop(self, location: FileLocation) -> bool:
         return self.loops.is_within_loop(location)
 
diff --git a/src/kaskara/functions.py b/src/kaskara/functions.py
index 340300b..9de2c2f 100644
--- a/src/kaskara/functions.py
+++ b/src/kaskara/functions.py
@@ -4,6 +4,7 @@
 __all__ = ("Function", "ProgramFunctions")
 
 import abc
+import itertools
 import os
 import typing as t
 from dataclasses import dataclass
@@ -82,6 +83,13 @@ def to_dict(self) -> dict[str, t.Any]:
             "functions": [f.to_dict() for f in self],
         }
 
+    def merge(self, other: ProgramFunctions) -> ProgramFunctions:
+        """Merges the results of this analysis and another analysis together."""
+        return self.from_functions(
+            project_directory=self._project_directory,
+            functions=itertools.chain(self, other),
+        )
+
     def with_relative_locations(self, base: str) -> ProgramFunctions:
         """Creates a new instance with relative file locations."""
         functions = [f.with_relative_locations(base) for f in self]
diff --git a/src/kaskara/insertions.py b/src/kaskara/insertions.py
index 037d4a4..39b4593 100644
--- a/src/kaskara/insertions.py
+++ b/src/kaskara/insertions.py
@@ -36,6 +36,9 @@ def __init__(self, contents: list[InsertionPoint]) -> None:
                 self.__file_insertions[filename] = []
             self.__file_insertions[filename].append(ins)
 
+    def merge(self, other: ProgramInsertionPoints) -> ProgramInsertionPoints:
+        return ProgramInsertionPoints(list(self) + list(other))
+
     def with_relative_locations(self, base: str) -> ProgramInsertionPoints:
         return ProgramInsertionPoints([
             insertion_point.with_relative_location(base)
diff --git a/src/kaskara/loops.py b/src/kaskara/loops.py
index 44a9e5e..bff2fa7 100644
--- a/src/kaskara/loops.py
+++ b/src/kaskara/loops.py
@@ -41,6 +41,15 @@ def from_body_location_ranges(
             FileLocationRangeSet(bodies),
         )
 
+    def merge(self, other: ProgramLoops) -> ProgramLoops:
+        covered_by_loop_bodies = self._covered_by_loop_bodies.union(
+            other._covered_by_loop_bodies,
+        )
+        return ProgramLoops(
+            _covered_by_loop_bodies=covered_by_loop_bodies,
+            _project_directory=self._project_directory,
+        )
+
     def with_relative_locations(self, base: str) -> ProgramLoops:
         """Creates a new instance with relative file locations."""
         covered_by_loop_bodies = self._covered_by_loop_bodies.with_relative_locations(
diff --git a/src/kaskara/statements.py b/src/kaskara/statements.py
index 88cc796..5966172 100644
--- a/src/kaskara/statements.py
+++ b/src/kaskara/statements.py
@@ -3,6 +3,7 @@
 __all__ = ("Statement", "ProgramStatements")
 
 import abc
+import itertools
 import os
 import typing as t
 from dataclasses import dataclass, field
@@ -97,6 +98,12 @@ def __post_init__(self) -> None:
         )
         logger.debug(f"indexed statements by file:\n{summary}")
 
+    def merge(self, other: ProgramStatements) -> ProgramStatements:
+        return self.build(
+            project_directory=self._project_directory,
+            statements=itertools.chain(self, other),
+        )
+
     def with_relative_locations(self, base: str) -> ProgramStatements:
         """Creates a new instance with relative file locations."""
         return self.build(