From 8e38c0a19fabb30b58e55c51bffd9ac87a4b607e Mon Sep 17 00:00:00 2001
From: Edmund Higham <edhigham@gmail.com>
Date: Tue, 17 Sep 2024 18:21:08 -0400
Subject: [PATCH] add `ImmutableMap` to prevent BM persist where unsupported

---
 .../backend/caching/BlockMatrixCache.scala    |  5 +++-
 .../is/hail/backend/local/LocalBackend.scala  |  2 +-
 .../hail/backend/service/ServiceBackend.scala |  2 +-
 .../is/hail/backend/spark/SparkBackend.scala  |  2 +-
 .../scala/is/hail/utils/ImmutableMap.scala    | 25 +++++++++++++++++++
 .../is/hail/utils/richUtils/RichMap.scala     |  7 ++++++
 6 files changed, 39 insertions(+), 4 deletions(-)
 create mode 100644 hail/src/main/scala/is/hail/utils/ImmutableMap.scala

diff --git a/hail/src/main/scala/is/hail/backend/caching/BlockMatrixCache.scala b/hail/src/main/scala/is/hail/backend/caching/BlockMatrixCache.scala
index db505a29239f..b3f6e55176ae 100644
--- a/hail/src/main/scala/is/hail/backend/caching/BlockMatrixCache.scala
+++ b/hail/src/main/scala/is/hail/backend/caching/BlockMatrixCache.scala
@@ -23,8 +23,11 @@ class BlockMatrixCache extends mutable.AbstractMap[String, BlockMatrix] with Aut
   override def iterator: Iterator[(String, BlockMatrix)] =
     blockmatrices.iterator
 
-  override def close(): Unit = {
+  override def clear(): Unit = {
     blockmatrices.values.foreach(_.unpersist())
     blockmatrices.clear()
   }
+
+  override def close(): Unit =
+    clear()
 }
diff --git a/hail/src/main/scala/is/hail/backend/local/LocalBackend.scala b/hail/src/main/scala/is/hail/backend/local/LocalBackend.scala
index 9d666ae62ef7..7a34954d243b 100644
--- a/hail/src/main/scala/is/hail/backend/local/LocalBackend.scala
+++ b/hail/src/main/scala/is/hail/backend/local/LocalBackend.scala
@@ -99,7 +99,7 @@ class LocalBackend(
             ExecutionCache.fromFlags(flags, fs, tmpdir)
         },
         new IrMetadata(),
-        mutable.Map.empty,
+        ImmutableMap.empty,
       )(f)
     }
 
diff --git a/hail/src/main/scala/is/hail/backend/service/ServiceBackend.scala b/hail/src/main/scala/is/hail/backend/service/ServiceBackend.scala
index 43249dd1ed1d..b94b4c391e4f 100644
--- a/hail/src/main/scala/is/hail/backend/service/ServiceBackend.scala
+++ b/hail/src/main/scala/is/hail/backend/service/ServiceBackend.scala
@@ -390,7 +390,7 @@ class ServiceBackend(
         flags,
         serviceBackendContext,
         new IrMetadata(),
-        mutable.Map.empty,
+        ImmutableMap.empty,
       )(f)
     }
 
diff --git a/hail/src/main/scala/is/hail/backend/spark/SparkBackend.scala b/hail/src/main/scala/is/hail/backend/spark/SparkBackend.scala
index fd1db82c26cc..f392b809ced9 100644
--- a/hail/src/main/scala/is/hail/backend/spark/SparkBackend.scala
+++ b/hail/src/main/scala/is/hail/backend/spark/SparkBackend.scala
@@ -362,7 +362,7 @@ class SparkBackend(
           ExecutionCache.forTesting
       },
       new IrMetadata(),
-      null,
+      ImmutableMap.empty,
     )
 
   override def withExecuteContext[T](f: ExecuteContext => T)(implicit E: Enclosing): T =
diff --git a/hail/src/main/scala/is/hail/utils/ImmutableMap.scala b/hail/src/main/scala/is/hail/utils/ImmutableMap.scala
new file mode 100644
index 000000000000..0d4b2080f57a
--- /dev/null
+++ b/hail/src/main/scala/is/hail/utils/ImmutableMap.scala
@@ -0,0 +1,25 @@
+package is.hail.utils
+
+import scala.collection.mutable
+
+case class ImmutableMap[K, V](m: Map[K, V]) extends mutable.AbstractMap[K, V] {
+  override def +=(kv: (K, V)): ImmutableMap.this.type =
+    throw new UnsupportedOperationException()
+
+  override def -=(key: K): ImmutableMap.this.type =
+    throw new UnsupportedOperationException()
+
+  override def get(key: K): Option[V] =
+    m.get(key)
+
+  override def iterator: Iterator[(K, V)] =
+    m.iterator
+
+  override def toMap[T, U](implicit ev: (K, V) <:< (T, U)): Map[T, U] =
+    m.toMap
+}
+
+object ImmutableMap {
+  def empty[K, V]: ImmutableMap[K, V] =
+    ImmutableMap(Map.empty)
+}
diff --git a/hail/src/main/scala/is/hail/utils/richUtils/RichMap.scala b/hail/src/main/scala/is/hail/utils/richUtils/RichMap.scala
index 5d080063e741..df0c9bb2a513 100644
--- a/hail/src/main/scala/is/hail/utils/richUtils/RichMap.scala
+++ b/hail/src/main/scala/is/hail/utils/richUtils/RichMap.scala
@@ -1,5 +1,9 @@
 package is.hail.utils.richUtils
 
+import is.hail.utils.ImmutableMap
+
+import scala.collection.mutable
+
 class RichMap[K, V](val m: Map[K, V]) extends AnyVal {
   def force =
     m.map(identity) // needed to make serializable: https://issues.scala-lang.org/browse/SI-7005
@@ -9,4 +13,7 @@ class RichMap[K, V](val m: Map[K, V]) extends AnyVal {
 
   def isTrivial(implicit eq: K =:= V): Boolean =
     m.forall { case (k, v) => k == v }
+
+  def immutableMutableMap: mutable.Map[K, V] =
+    ImmutableMap(m)
 }