From 6d315091958524a4f6df00b1c9e372809b8d919b Mon Sep 17 00:00:00 2001
From: Devansh Das <dxdas@google.com>
Date: Sun, 27 Oct 2024 14:08:02 +0000
Subject: [PATCH] Added rebase and export capabilities for all cluster snapshot
 implementations

---
 .../simulator/clustersnapshot/basic.go        | 20 ++++++++++++++++++
 .../clustersnapshot/clustersnapshot.go        |  4 ++++
 .../simulator/clustersnapshot/delta.go        | 21 +++++++++++++++++++
 3 files changed, 45 insertions(+)

diff --git a/cluster-autoscaler/simulator/clustersnapshot/basic.go b/cluster-autoscaler/simulator/clustersnapshot/basic.go
index 31df4e5c3db4..60f9ae961e75 100644
--- a/cluster-autoscaler/simulator/clustersnapshot/basic.go
+++ b/cluster-autoscaler/simulator/clustersnapshot/basic.go
@@ -285,6 +285,26 @@ func (snapshot *BasicClusterSnapshot) Clear() {
 	snapshot.data = []*internalBasicSnapshotData{baseData}
 }
 
+// Export returns a shallow copy of the snapshot.
+func (snapshot *BasicClusterSnapshot) Export() ClusterSnapshot {
+	exportedSnapshot := &BasicClusterSnapshot{}
+	exportedSnapshot.data = make([]*internalBasicSnapshotData, len(snapshot.data))
+	for i, data := range snapshot.data {
+		exportedSnapshot.data[i] = data.clone()
+	}
+	return exportedSnapshot
+}
+
+// Rebase rebases the snapshot to a new base snapshot.
+func (snapshot *BasicClusterSnapshot) Rebase(base ClusterSnapshot) error {
+	baseSnapshot, ok := base.(*BasicClusterSnapshot)
+	if !ok {
+		return fmt.Errorf("cannot rebase to different type of snapshot")
+	}
+	snapshot.data = []*internalBasicSnapshotData{baseSnapshot.getInternalData().clone(), snapshot.data[len(snapshot.data)-1]}
+	return nil
+}
+
 // implementation of SharedLister interface
 
 type basicClusterSnapshotNodeLister BasicClusterSnapshot
diff --git a/cluster-autoscaler/simulator/clustersnapshot/clustersnapshot.go b/cluster-autoscaler/simulator/clustersnapshot/clustersnapshot.go
index 275dc0d8da63..ababb0f6a76a 100644
--- a/cluster-autoscaler/simulator/clustersnapshot/clustersnapshot.go
+++ b/cluster-autoscaler/simulator/clustersnapshot/clustersnapshot.go
@@ -52,6 +52,10 @@ type ClusterSnapshot interface {
 	Commit() error
 	// Clear reset cluster snapshot to empty, unforked state.
 	Clear()
+	// Export returns a shallow copy of the snapshot.
+	Export() ClusterSnapshot
+	// Rebase rebases the snapshot to a new base snapshot.
+	Rebase(base ClusterSnapshot) error
 }
 
 // ErrNodeNotFound means that a node wasn't found in the snapshot.
diff --git a/cluster-autoscaler/simulator/clustersnapshot/delta.go b/cluster-autoscaler/simulator/clustersnapshot/delta.go
index aa4288dd1428..b63ef69c3241 100644
--- a/cluster-autoscaler/simulator/clustersnapshot/delta.go
+++ b/cluster-autoscaler/simulator/clustersnapshot/delta.go
@@ -31,6 +31,8 @@ import (
 //	fork - O(1)
 //	revert - O(1)
 //	commit - O(n)
+//	rebase - O(1)
+//	export - O(1)
 //	list all pods (no filtering) - O(n), cached
 //	list all pods (with filtering) - O(n)
 //	list node infos - O(n), cached
@@ -475,3 +477,22 @@ func (snapshot *DeltaClusterSnapshot) Commit() error {
 func (snapshot *DeltaClusterSnapshot) Clear() {
 	snapshot.data = newInternalDeltaSnapshotData()
 }
+
+// Export returns a shallow copy of the changes made to the snapshot.
+// Time: O(1) (shallow copy)
+func (base *DeltaClusterSnapshot) Export() ClusterSnapshot {
+	return &DeltaClusterSnapshot{
+		data: base.data,
+	}
+}
+
+// Rebase rebases the snapshot to a new base snapshot.
+// Time: O(1)
+func (snapshot *DeltaClusterSnapshot) Rebase(base ClusterSnapshot) error {
+	snapshotBase, ok := base.(*DeltaClusterSnapshot)
+	if !ok {
+		return fmt.Errorf("cannot rebase to a different type of snapshot")
+	}
+	snapshot.data.baseData = snapshotBase.data
+	return nil
+}