[query] make the locus interface representation-agnostic

hail/hail/src/is/hail/types/physical/PLocus.scala

@@ -12,8 +12,6 @@ abstract class PLocus extends PType {
 
   def contig(value: Long): String
 
-  def contigType: PString
-
   def position(value: Code[Long]): Code[Int]
 
   def position(value: Long): Int

hail/hail/src/is/hail/types/physical/stypes/concrete/SCanonicalLocusPointer.scala

@@ -30,10 +30,10 @@ final case class SCanonicalLocusPointer(pType: PCanonicalLocus) extends SLocus {
     value match {
       case value: SLocusValue =>
         val locusCopy = pType.store(cb, region, value, deepCopy)
-        val contigCopy = if (deepCopy)
-          cb.memoize(pType.contigAddr(locusCopy))
-        else
-          value.contigLong(cb)
+        val contigCopy = value match {
+          case v: SCanonicalLocusPointerValue if !deepCopy => v._contig
+          case _ => cb.memoize(pType.contigAddr(locusCopy))
+        }
         new SCanonicalLocusPointerValue(this, locusCopy, contigCopy, value.position(cb))
     }
 
@@ -84,10 +84,17 @@ class SCanonicalLocusPointerValue(
   override def contig(cb: EmitCodeBuilder): SStringValue =
     pt.contigType.loadCheapSCode(cb, _contig).asString
 
-  override def contigLong(cb: EmitCodeBuilder): Value[Long] = _contig
+  override def contigIdx(cb: EmitCodeBuilder): Value[Int] =
+    cb.memoize(cb.emb.getReferenceGenome(st.rg).invoke[String, Int](
+      "getContigIndex",
+      contig(cb).loadString(cb),
+    ))
 
   override def position(cb: EmitCodeBuilder): Value[Int] = _position
 
+  override def packed(cb: EmitCodeBuilder): Value[Long] =
+    cb.memoize((contigIdx(cb).toL << 32) | _position.toL)
+
   override def structRepr(cb: EmitCodeBuilder): SBaseStructValue = new SBaseStructPointerValue(
     SBaseStructPointer(st.pType.representation),
     a,

hail/hail/src/is/hail/types/physical/stypes/concrete/SUnreachable.scala

@@ -153,7 +153,9 @@ class SUnreachableLocusValue(override val st: SUnreachableLocus)
 
   override def contig(cb: EmitCodeBuilder): SStringValue = SUnreachableString.sv
 
-  override def contigLong(cb: EmitCodeBuilder): Value[Long] = const(0)
+  override def contigIdx(cb: EmitCodeBuilder): Value[Int] = const(0)
+
+  override def packed(cb: EmitCodeBuilder): Value[Long] = const(0L)
 
   override def structRepr(cb: EmitCodeBuilder): SBaseStructValue = SUnreachableStruct(TStruct(
     "contig" -> TString,

hail/hail/src/is/hail/types/physical/stypes/interfaces/SLocus.scala

@@ -20,10 +20,12 @@ trait SLocusValue extends SValue {
 
   def contig(cb: EmitCodeBuilder): SStringValue
 
-  def contigLong(cb: EmitCodeBuilder): Value[Long]
+  def contigIdx(cb: EmitCodeBuilder): Value[Int]
 
   def position(cb: EmitCodeBuilder): Value[Int]
 
+  def packed(cb: EmitCodeBuilder): Value[Long]
+
   def getLocusObj(cb: EmitCodeBuilder): Value[Locus] =
     cb.memoize(Code.invokeStatic2[Locus, String, Int, Locus](
       "apply",
@@ -39,14 +41,6 @@ trait SLocusValue extends SValue {
   override def sizeToStoreInBytes(cb: EmitCodeBuilder): SInt64Value =
     structRepr(cb).sizeToStoreInBytes(cb)
 
-  override def prefixCode(cb: EmitCodeBuilder, pc: Value[PrefixCoder]): Unit = {
-    val rgCode = cb.emb.getReferenceGenome(st.rg)
-    val locusObject = getLocusObj(cb)
-    val globalPos =
-      cb.memoize(rgCode.invoke[Locus, Long](
-        "locusToGlobalPos",
-        locusObject,
-      ))
-    pc.encodeLong(cb, globalPos)
-  }
+  override def prefixCode(cb: EmitCodeBuilder, pc: Value[PrefixCoder]): Unit =
+    pc.encodeLong(cb, packed(cb))
 }

hail/hail/test/src/is/hail/types/physical/stypes/concrete/SCanonicalLocusPointerSuite.scala

@@ -0,0 +1,86 @@
+package is.hail.types.physical.stypes.concrete
+
+import is.hail.TestUtils._
+import is.hail.annotations.Region
+import is.hail.asm4s._
+import is.hail.backend.ExecuteContext
+import is.hail.expr.ir.{EmitCodeBuilder, EmitFunctionBuilder}
+import is.hail.scalacheck._
+import is.hail.types.physical.PCanonicalLocus
+import is.hail.types.physical.stypes.interfaces.SLocusValue
+import is.hail.variant.{Locus, ReferenceGenome}
+
+import scala.collection.immutable.ArraySeq
+
+import org.junit.jupiter.api.Test
+import org.scalacheck.Arbitrary.arbitrary
+import org.scalacheck.Gen
+import org.scalacheck.Prop.forAll
+
+class SCanonicalLocusPointerSuite {
+
+  def compile[T: TypeInfo](
+    ctx: ExecuteContext,
+    pt: PCanonicalLocus,
+  )(
+    f: (EmitCodeBuilder, SLocusValue) => Value[T]
+  ): (Region, Long) => T = {
+    val fb = EmitFunctionBuilder[Region, Long, T](ctx, "locus_property")
+    fb.emitWithBuilder(cb => f(cb, pt.loadCheapSCode(cb, fb.getCodeParam[Long](2))))
+    val g = fb.resultWithIndex()(ctx.theHailClassLoader, ctx.fs, ctx.taskContext, ctx.r)
+    g(_, _)
+  }
+
+  def checkAgreesWithReference(rg: ReferenceGenome)(implicit ctx: ExecuteContext): Unit =
+    ctx.local(references = ctx.references.withReferences(Seq(rg))) { ctx =>
+      val pt = PCanonicalLocus(rg)
+      val contigIdx = compile[Int](ctx, pt)((cb, l) => l.contigIdx(cb))
+      val packed = compile[Long](ctx, pt)((cb, l) => l.packed(cb))
+
+      check(forAll(genLocus(rg)) { l =>
+        ctx.r.pool.scopedRegion { r =>
+          val addr = pt.unstagedStoreJavaObject(l, r)
+          val idx = rg.getContigIndex(l.contig)
+          assertEq(contigIdx(r, addr), idx)
+          assertEq(packed(r, addr), (idx.toLong << 32) | l.position.toLong)
+        }
+      })
+    }
+
+  @Test def testContigIdxAndPackedGRCh37(implicit ctx: ExecuteContext): Unit =
+    checkAgreesWithReference(ctx.references(ReferenceGenome.GRCh37))
+
+  @Test def testContigIdxAndPackedGRCh38(implicit ctx: ExecuteContext): Unit =
+    checkAgreesWithReference(ctx.references(ReferenceGenome.GRCh38))
+
+  @Test def testContigIdxAndPackedCustomReference(implicit ctx: ExecuteContext): Unit =
+    checkAgreesWithReference(
+      ReferenceGenome("custom", ArraySeq("c", "a", "b"), Map("a" -> 100, "b" -> 2000, "c" -> 30000))
+    )
+
+  val referenceAndLoci: Gen[(ReferenceGenome, IndexedSeq[Locus])] =
+    for {
+      rg <- arbitrary[ReferenceGenome]
+      loci <- Gen.containerOf[IndexedSeq, Locus](genLocus(rg))
+    } yield (rg, loci)
+
+  @Test def testPrefixCodePreservesLocusOrder(implicit ctx: ExecuteContext): Unit =
+    check(forAll(referenceAndLoci) { case (rg, loci) =>
+      ctx.local(references = ctx.references.withReferences(Seq(rg))) { ctx =>
+        val pt = PCanonicalLocus(rg)
+        val prefixCode =
+          compile[Array[Byte]](ctx, pt)((cb, l) =>
+            l.prefixCode(cb).asInstanceOf[SJavaBytesValue].bytes
+          )
+
+        val byPrefixCode = ctx.r.pool.scopedRegion { r =>
+          loci
+            .map(l => (l, prefixCode(r, pt.unstagedStoreJavaObject(l, r))))
+            .sortWith { case ((_, b1), (_, b2)) => java.util.Arrays.compareUnsigned(b1, b2) < 0 }
+            .map(_._1)
+        }
+
+        assertEq(byPrefixCode, loci.sorted(rg.locusOrdering))
+      }
+    })
+}