[hail/ptypes] copy based on ptype and requiredeness

hail/src/main/scala/is/hail/expr/types/physical/ComplexPType.scala

@@ -1,6 +1,7 @@
 package is.hail.expr.types.physical
 
-import is.hail.annotations.UnsafeOrdering
+import is.hail.annotations.{Region, UnsafeOrdering}
+import is.hail.asm4s.{Code, MethodBuilder}
 
 abstract class ComplexPType extends PType {
   val representation: PType
@@ -14,4 +15,21 @@ abstract class ComplexPType extends PType {
   override def fundamentalType: PType = representation.fundamentalType
 
   override def containsPointers: Boolean = representation.containsPointers
+
+  override def storeShallowAtOffset(dstAddress: Code[Long], valueAddress: Code[Long]): Code[Unit] =
+    this.representation.storeShallowAtOffset(dstAddress, valueAddress)
+
+  override def storeShallowAtOffset(dstAddress: Long, valueAddress: Long) {
+    this.representation.storeShallowAtOffset(dstAddress, valueAddress)
+  }
+
+  override def copyFromType(mb: MethodBuilder, region: Code[Region], srcPType: PType, srcAddress: Code[Long], forceDeep: Boolean): Code[Long] = {
+    assert(this isOfType srcPType)
+
+    val srcRepPType = srcPType.asInstanceOf[ComplexPType].representation
+
+    this.representation.copyFromType(mb, region, srcRepPType, srcAddress, forceDeep)
+  }
+
+  override def copyFromType(region: Region, srcPType: PType, srcAddress: Long, forceDeep: Boolean): Long = ???
 }

hail/src/main/scala/is/hail/expr/types/physical/PArrayBackedContainer.scala

@@ -1,7 +1,7 @@
 package is.hail.expr.types.physical
 
 import is.hail.annotations.{Region, UnsafeOrdering}
-import is.hail.asm4s.{Code, MethodBuilder}
+import is.hail.asm4s.{???, Code, MethodBuilder}
 import is.hail.expr.ir.EmitMethodBuilder
 
 trait PArrayBackedContainer extends PContainer {
@@ -110,6 +110,9 @@ trait PArrayBackedContainer extends PContainer {
   def loadElement(aoff: Long, length: Int, i: Int): Long =
     arrayRep.loadElement(aoff, length, i)
 
+  def loadElementAddress(aoff: Code[Long], length: Code[Int], i: Code[Int]): Code[Long] =
+    arrayRep.loadElementAddress(aoff, length, i)
+
   def loadElement(region: Region, aoff: Long, length: Int, i: Int): Long =
     arrayRep.loadElement(region, aoff, length, i)
 
@@ -164,7 +167,29 @@ trait PArrayBackedContainer extends PContainer {
   def copyFrom(mb: MethodBuilder, region: Code[Region], srcOff: Code[Long]): Code[Long] =
     arrayRep.copyFrom(mb, region, srcOff)
 
-  override def unsafeOrdering: UnsafeOrdering = unsafeOrdering(this)
+  override def unsafeOrdering: UnsafeOrdering =
+    unsafeOrdering(this)
+
+  override def unsafeOrdering(rightType: PType): UnsafeOrdering =
+    arrayRep.unsafeOrdering(rightType)
+
+  override def copyFromType(mb: MethodBuilder, region: Code[Region], srcPType: PType, srcAddress: Code[Long], forceDeep: Boolean): Code[Long] = {
+    assert(srcPType.isInstanceOf[PArrayBackedContainer])
+    this.arrayRep.copyFromType(mb, region, srcPType.asInstanceOf[PArrayBackedContainer].arrayRep, srcAddress, forceDeep)
+  }
+
+  override def copyFromType(region: Region, srcPType: PType, srcAddress: Long, forceDeep: Boolean): Long = ???
+
+  override def storeShallowAtOffset(dstAddress: Code[Long], valueAddress: Code[Long]): Code[Unit] =
+    this.arrayRep.storeShallowAtOffset(dstAddress, valueAddress)
+
+  override def storeShallowAtOffset(dstAddress: Long, valueAddress: Long) {
+    this.arrayRep.storeShallowAtOffset(dstAddress, valueAddress)
+  }
+
+  def nextElementAddress(currentOffset: Long) =
+    arrayRep.nextElementAddress(currentOffset)
 
-  override def unsafeOrdering(rightType: PType): UnsafeOrdering = arrayRep.unsafeOrdering(rightType)
+  def nextElementAddress(currentOffset: Code[Long]) =
+    arrayRep.nextElementAddress(currentOffset)
 }

hail/src/main/scala/is/hail/expr/types/physical/PBaseStruct.scala

@@ -74,6 +74,7 @@ abstract class PBaseStruct extends PType {
   }
 
   def identBase: String
+
   def _asIdent: String = {
     val sb = new StringBuilder
     sb.append(identBase)
@@ -147,7 +148,30 @@ abstract class PBaseStruct extends PType {
     region.allocate(alignment, byteSize)
   }
 
-  def allocate(region: Code[Region]): Code[Long] = region.allocate(alignment, byteSize)
+  def allocate(region: Code[Region]): Code[Long] =
+    region.allocate(alignment, byteSize)
+
+  def copyFrom(region: Region, srcAddress: Long): Long = {
+    val destAddress = this.allocate(region)
+    this.storeShallowAtOffset(destAddress, srcAddress)
+    destAddress
+  }
+
+  def copyFrom(mb: MethodBuilder, region: Code[Region], srcAddress: Code[Long]): Code[Long] = {
+    val destAddress = mb.newField[Long]
+    Code(
+      destAddress := this.allocate(region),
+      this.storeShallowAtOffset(destAddress, srcAddress),
+      destAddress
+    )
+  }
+
+  override def storeShallowAtOffset(destAddress: Code[Long], srcAddress: Code[Long]): Code[Unit] =
+    Region.copyFrom(srcAddress, destAddress, this.byteSize)
+
+  override def storeShallowAtOffset(destAddress: Long, srcAddress: Long) {
+    Region.copyFrom(srcAddress, destAddress, this.byteSize)
+  }
 
   def initialize(structAddress: Long, setMissing: Boolean = false): Unit = {
     if (allFieldsRequired) {
@@ -217,11 +241,11 @@ abstract class PBaseStruct extends PType {
   def setFieldPresent(region: Code[Region], offset: Code[Long], fieldIdx: Int): Code[Unit] =
     setFieldPresent(offset, fieldIdx)
 
-  def fieldOffset(offset: Long, fieldIdx: Int): Long =
-    offset + byteOffsets(fieldIdx)
+  def fieldOffset(structAddress: Long, fieldIdx: Int): Long =
+    structAddress + byteOffsets(fieldIdx)
 
-  def fieldOffset(offset: Code[Long], fieldIdx: Int): Code[Long] =
-    offset + byteOffsets(fieldIdx)
+  def fieldOffset(structAddress: Code[Long], fieldIdx: Int): Code[Long] =
+    structAddress + byteOffsets(fieldIdx)
 
   def loadField(rv: RegionValue, fieldIdx: Int): Long = loadField(rv.region, rv.offset, fieldIdx)
 
@@ -248,5 +272,107 @@ abstract class PBaseStruct extends PType {
     }
   }
 
+  def deepCopyFromAddress(mb: MethodBuilder, region: Code[Region], srcStructAddress: Code[Long]): Code[Long] = {
+    val dstAddress = mb.newField[Long]
+    Code(
+      dstAddress := this.copyFrom(mb, region, srcStructAddress),
+      this.deepPointerCopy(mb, region, dstAddress),
+      dstAddress
+    )
+  }
+
+  def deepPointerCopy(mb: MethodBuilder, region: Code[Region], dstStructAddress: Code[Long]): Code[Unit] = {
+    var c: Code[Unit] = Code._empty
+
+    var i = 0
+    while(i < this.size) {
+      val dstFieldType = this.fields(i).typ.fundamentalType
+      if(dstFieldType.containsPointers) {
+        val dstFieldAddress = mb.newField[Long]
+        c = Code(
+          c,
+          this.isFieldDefined(dstStructAddress, i).orEmpty(
+            Code(
+              dstFieldAddress := this.fieldOffset(dstStructAddress, i),
+              dstFieldType match {
+                case t@(_: PBinary | _: PArray) =>
+                  t.storeShallowAtOffset(dstFieldAddress, t.copyFromType(mb, region, dstFieldType, Region.loadAddress(dstFieldAddress)))
+                case t: PBaseStruct =>
+                  t.deepPointerCopy(mb, region, dstFieldAddress)
+                case t: PType =>
+                  fatal(s"Field type isn't supported ${t}")
+              }
+            )
+          )
+        )
+      }
+      i += 1
+    }
+
+    c
+  }
+
+  override def copyFromType(mb: MethodBuilder, region: Code[Region], srcPType: PType, srcStructAddress: Code[Long],
+    forceDeep: Boolean): Code[Long] = {
+    assert(srcPType.isInstanceOf[PBaseStruct])
+
+    val sourceType = srcPType.asInstanceOf[PBaseStruct]
+
+    assert(sourceType.size == this.size)
+
+    val sourceFundamentalTypes = sourceType.fields.map(_.typ.fundamentalType)
+    if(this.fields.map(_.typ.fundamentalType) == sourceFundamentalTypes) {
+      if(!forceDeep) {
+        return srcStructAddress
+      }
+
+      return this.deepCopyFromAddress(mb, region, srcStructAddress)
+    }
+
+    val dstStructAddress = mb.newField[Long]
+    var loop: Code[_] = Code()
+    var i = 0
+    while(i < this.size) {
+      val dstField = this.fields(i)
+      val srcField = sourceType.fields(i)
+
+      assert((dstField.typ.required <= srcField.typ.required) && (dstField.typ isOfType srcField.typ) && (dstField.name == srcField.name) && (dstField.index == srcField.index))
+
+      val srcFieldType = srcField.typ.fundamentalType
+      val dstFieldType = dstField.typ.fundamentalType
+
+      val body = srcFieldType.storeShallowAtOffset(
+        this.fieldOffset(dstStructAddress, dstField.index),
+        dstFieldType.copyFromType(
+          mb,
+          region,
+          srcFieldType,
+          sourceType.loadField(srcStructAddress, srcField.index),
+          forceDeep
+        )
+      )
+
+      if(!srcFieldType.required) {
+        loop = Code(loop, sourceType.isFieldMissing(srcStructAddress, srcField.index).mux(
+          this.setFieldMissing(dstStructAddress, dstField.index),
+          body
+        ))
+      } else {
+        loop = Code(loop, body)
+      }
+
+      i+=1
+    }
+
+    Code(
+      dstStructAddress := this.allocate(region),
+      this.stagedInitialize(dstStructAddress),
+      loop,
+      dstStructAddress
+    )
+  }
+
+  override def copyFromType(region: Region, srcPType: PType, srcAddress: Long, forceDeep: Boolean): Long = ???
+
   override def containsPointers: Boolean = types.exists(_.containsPointers)
 }

hail/src/main/scala/is/hail/expr/types/physical/PBinary.scala

@@ -79,13 +79,19 @@ object PBinary {
 
   def contentByteSize(length: Code[Int]): Code[Long] = PCanonicalBinary.contentByteSize(length)
 
-  def loadLength(boff: Long): Int = PCanonicalBinary.loadLength(boff)
+  def loadLength(bAddress: Long): Int = PCanonicalBinary.loadLength(bAddress)
 
-  def loadLength(region: Region, boff: Long): Int = PCanonicalBinary.loadLength(boff)
+  def loadLength(region: Region, bAddress: Long): Int = PCanonicalBinary.loadLength(bAddress)
 
-  def loadLength(boff: Code[Long]): Code[Int] = PCanonicalBinary.loadLength(boff)
+  def loadLength(bAddress: Code[Long]): Code[Int] = PCanonicalBinary.loadLength(bAddress)
 
-  def loadLength(region: Code[Region], boff: Code[Long]): Code[Int] = PCanonicalBinary.loadLength(boff)
+  def loadLength(region: Code[Region], bAddress: Code[Long]): Code[Int] = PCanonicalBinary.loadLength(bAddress)
+
+  def loadBytes(bAddress: Code[Long]): Code[Array[Byte]] =
+    Region.loadBytes(PBinary.bytesOffset(bAddress), PBinary.loadLength(bAddress))
+
+  def loadBytes(bAddress: Long): Array[Byte] =
+    Region.loadBytes(PBinary.bytesOffset(bAddress), PBinary.loadLength(bAddress))
 
   def storeLength(boff: Long, len: Int): Unit = PCanonicalBinary.storeLength(boff, len)
 

hail/src/main/scala/is/hail/expr/types/physical/PBoolean.scala

@@ -1,7 +1,7 @@
 package is.hail.expr.types.physical
 
 import is.hail.annotations.{Region, UnsafeOrdering, _}
-import is.hail.asm4s.Code
+import is.hail.asm4s.{Code, MethodBuilder}
 import is.hail.expr.ir.EmitMethodBuilder
 import is.hail.expr.types.virtual.TBoolean
 
@@ -34,6 +34,18 @@ class PBoolean(override val required: Boolean) extends PType {
   }
 
   override def byteSize: Long = 1
+
+  def storeShallowAtOffset(dstAddress: Code[Long], srcAddress: Code[Long]): Code[Unit] =
+    Region.storeBoolean(dstAddress, Region.loadBoolean(srcAddress))
+
+  def storeShallowAtOffset(dstAddress: Long, srcAddress: Long) =
+    Region.storeBoolean(dstAddress, Region.loadBoolean(srcAddress))
+
+  def copyFromType(region: Region, srcPType: PType, srcAddress: Long, forceDeep: Boolean): Long =
+    srcAddress
+
+  def copyFromType(mb: MethodBuilder, region: Code[Region], srcPType: PType, srcAddress: Code[Long], forceDeep: Boolean): Code[Long] =
+    srcAddress
 }
 
 object PBoolean {