From 19154ceccb21a0b9966595b0a22c6168a4bd22a0 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

---
 .../is/hail/backend/local/LocalBackend.scala  |  2 +-
 .../hail/backend/service/ServiceBackend.scala |  2 +-
 .../is/hail/backend/spark/SparkBackend.scala  |  2 +-
 .../scala/is/hail/expr/ir/BlockMatrixIR.scala |  1 +
 .../is/hail/expr/ir/BlockMatrixWriter.scala   |  1 +
 .../scala/is/hail/utils/ImmutableMap.scala    | 25 +++++++++++++++++++
 .../is/hail/utils/richUtils/RichMap.scala     |  7 ++++++
 7 files changed, 37 insertions(+), 3 deletions(-)
 create mode 100644 hail/src/main/scala/is/hail/utils/ImmutableMap.scala

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 9d666ae62ef..7a34954d243 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 49d3ece5276..8ef05961092 100644
--- a/hail/src/main/scala/is/hail/backend/service/ServiceBackend.scala
+++ b/hail/src/main/scala/is/hail/backend/service/ServiceBackend.scala
@@ -381,7 +381,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 fd1db82c26c..f392b809ced 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/expr/ir/BlockMatrixIR.scala b/hail/src/main/scala/is/hail/expr/ir/BlockMatrixIR.scala
index 9108d20440a..f08c36cadb2 100644
--- a/hail/src/main/scala/is/hail/expr/ir/BlockMatrixIR.scala
+++ b/hail/src/main/scala/is/hail/expr/ir/BlockMatrixIR.scala
@@ -1,4 +1,5 @@
 package is.hail.expr.ir
+
 import is.hail.annotations.NDArray
 import is.hail.backend.ExecuteContext
 import is.hail.expr.Nat
diff --git a/hail/src/main/scala/is/hail/expr/ir/BlockMatrixWriter.scala b/hail/src/main/scala/is/hail/expr/ir/BlockMatrixWriter.scala
index 051c07d9413..6342f7fef6f 100644
--- a/hail/src/main/scala/is/hail/expr/ir/BlockMatrixWriter.scala
+++ b/hail/src/main/scala/is/hail/expr/ir/BlockMatrixWriter.scala
@@ -1,4 +1,5 @@
 package is.hail.expr.ir
+
 import is.hail.annotations.Region
 import is.hail.asm4s._
 import is.hail.backend.ExecuteContext
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 00000000000..0d4b2080f57
--- /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 5d080063e74..df0c9bb2a51 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)
 }