simplify PruneDeadFields

hail-is · Apr 26, 2024 · 578e1c7 · 578e1c7
1 parent e84931a
commit 578e1c7
Show file tree

Hide file tree

Showing 12 changed files with 721 additions and 842 deletions.
diff --git a/hail/src/main/scala/is/hail/expr/ir/BaseIR.scala b/hail/src/main/scala/is/hail/expr/ir/BaseIR.scala
@@ -13,6 +13,8 @@ abstract class BaseIR {
 
   def children: Iterable[BaseIR] = childrenSeq
 
+  def getChild(idx: Int): BaseIR = childrenSeq(idx)
+
   protected def withNewChildren(newChildren: IndexedSeq[BaseIR]): BaseIR
 
   def deepCopy(): this.type =

diff --git a/hail/src/main/scala/is/hail/expr/ir/Binds.scala b/hail/src/main/scala/is/hail/expr/ir/Binds.scala
@@ -302,7 +302,7 @@ object Bindings {
       case StreamFold(a, zero, accumName, valueName, _) if i == 2 =>
         Bindings(FastSeq(accumName -> zero.typ, valueName -> elementType(a.typ)))
       case StreamFold2(a, accum, valueName, _, _) =>
-        if (i <= accum.length)
+        if (i < accum.length + 1)
           Bindings.empty
         else if (i < 2 * accum.length + 1)
           Bindings(

diff --git a/hail/src/main/scala/is/hail/expr/ir/DeprecatedIRBuilder.scala b/hail/src/main/scala/is/hail/expr/ir/DeprecatedIRBuilder.scala
@@ -220,10 +220,10 @@ object DeprecatedIRBuilder {
 
     def castRename(t: Type): IRProxy = (env: E) => CastRename(ir(env), t)
 
-    def insertFields(fields: (Symbol, IRProxy)*): IRProxy = insertFieldsList(fields)
+    def insertFields(fields: (Symbol, IRProxy)*): IRProxy = insertFieldsList(fields.toFastSeq)
 
     def insertFieldsList(
-      fields: Seq[(Symbol, IRProxy)],
+      fields: IndexedSeq[(Symbol, IRProxy)],
       ordering: Option[IndexedSeq[String]] = None,
     ): IRProxy = (env: E) =>
       InsertFields(ir(env), fields.map { case (s, fir) => (s.name, fir(env)) }, ordering)

diff --git a/hail/src/main/scala/is/hail/expr/ir/IR.scala b/hail/src/main/scala/is/hail/expr/ir/IR.scala
@@ -959,12 +959,13 @@ final case class MakeStruct(fields: IndexedSeq[(String, IR)]) extends IR
 final case class SelectFields(old: IR, fields: IndexedSeq[String]) extends IR
 
 object InsertFields {
-  def apply(old: IR, fields: Seq[(String, IR)]): InsertFields = InsertFields(old, fields, None)
+  def apply(old: IR, fields: IndexedSeq[(String, IR)]): InsertFields =
+    InsertFields(old, fields, None)
 }
 
 final case class InsertFields(
   old: IR,
-  fields: Seq[(String, IR)],
+  fields: IndexedSeq[(String, IR)],
   fieldOrder: Option[IndexedSeq[String]],
 ) extends TypedIR[TStruct]
 

diff --git a/hail/src/main/scala/is/hail/expr/ir/Pretty.scala b/hail/src/main/scala/is/hail/expr/ir/Pretty.scala
@@ -5,7 +5,7 @@ import is.hail.expr.JSONAnnotationImpex
 import is.hail.expr.ir.Pretty.prettyBooleanLiteral
 import is.hail.expr.ir.agg._
 import is.hail.expr.ir.functions.RelationalFunctions
-import is.hail.types.TableType
+import is.hail.types.{MatrixType, TableType}
 import is.hail.types.virtual.{TArray, TInterval, TStream, Type}
 import is.hail.utils.{space => _, _}
 import is.hail.utils.prettyPrint._
@@ -821,7 +821,7 @@ class Pretty(
       case I32(i) => s"c$i"
       case stream if stream.typ.isInstanceOf[TStream] => "s"
       case table if table.typ.isInstanceOf[TableType] => "ht"
-      case mt if mt.typ.isInstanceOf[TableType] => "mt"
+      case mt if mt.typ.isInstanceOf[MatrixType] => "mt"
       case _ => ""
     }
 
@@ -874,12 +874,12 @@ class Pretty(
       ir match {
         case Ref(name, _) =>
           val body =
-            blockBindings.lookupOption(name).getOrElse(uniqueify("%undefined_ref", Some(name)))
+            blockBindings.lookupOption(name).getOrElse("#" + uniqueify("undefined_ref", Some(name)))
           concat(openBlock, group(nest(2, concat(line, body, line)), "}"))
         case RelationalRef(name, _) =>
           val body =
-            blockBindings.lookupOption(name).getOrElse(uniqueify(
-              "%undefined_relational_ref",
+            blockBindings.lookupOption(name).getOrElse("#" + uniqueify(
+              "undefined_relational_ref",
               Some(name),
             ))
           concat(openBlock, group(nest(2, concat(line, body, line)), "}"))
@@ -923,10 +923,10 @@ class Pretty(
         } yield {
           child match {
             case Ref(name, _) =>
-              bindings.lookupOption(name).getOrElse(uniqueify("%undefined_ref", Some(name)))
+              bindings.lookupOption(name).getOrElse("#" + uniqueify("undefined_ref", Some(name)))
             case RelationalRef(name, _) =>
-              bindings.lookupOption(name).getOrElse(uniqueify(
-                "%undefined_relational_ref",
+              bindings.lookupOption(name).getOrElse("#" + uniqueify(
+                "undefined_relational_ref",
                 Some(name),
               ))
             case _ =>