Test the contig index cache through generated code

The cache's own suite checked the helper directly. These tests drive it
through the canonical locus's generated contigIdx, where a cache could
go wrong in ways a direct call does not show:

- Region reuse: after the region is cleared, a different contig of the
  same length lands at the same address. A cache keyed on the address,
  or on a view of the region, returns the old index.
- Unknown contigs: a failed lookup, repeated, leaves the cache giving
  correct indices for the previously cached contig and a new one.
- Concurrency: many small partitions on the local backend, with
  contigs interleaved within each partition, write canonical loci under
  unstable encodings and read them back. Expected values come from
  ReferenceGenome.getContigIndex, the uncached lookup.

Author: chrisvittal

hail/hail/test/src/is/hail/variant/ContigIndexCacheSuite.scala

@@ -1,8 +1,20 @@
 package is.hail.variant
 
+import is.hail.ExecStrategy
 import is.hail.TestUtils._
-import is.hail.annotations.Region
+import is.hail.annotations.{Region, RowSeq}
+import is.hail.asm4s._
 import is.hail.backend.ExecuteContext
+import is.hail.expr.ir._
+import is.hail.expr.ir.TestUtils.collect
+import is.hail.expr.ir.defs.{Literal, TableWrite}
+import is.hail.types.physical.LocusRepresentations.withReference
+import is.hail.types.physical.PCanonicalLocus
+import is.hail.types.physical.stypes.interfaces.SLocusValue
+import is.hail.types.virtual.{TArray, TInt32, TLocus, TString}
+import is.hail.variant.ContigIndexCacheSuite.{
+  compileCanonical, custom, locusFromLiterals, storeCanonical,
+}
 
 import scala.collection.immutable.ArraySeq
 
@@ -12,13 +24,6 @@ import org.scalacheck.Prop.forAll
 
 class ContigIndexCacheSuite {
 
-  // Contigs that share prefixes and lengths, so that a comparison that stops early or checks only
-  // the length gives a wrong index.
-  private val custom: ReferenceGenome = {
-    val contigs = ArraySeq("1", "10", "11", "chr1", "chr10", "chr2", "a", "ab", "ba")
-    ReferenceGenome("contig_index_cache_test", contigs, contigs.map(_ -> 1000).toMap)
-  }
-
   private def genLookups(rg: ReferenceGenome): Gen[List[String]] =
     Gen.listOf(Gen.zip(Gen.oneOf(rg.contigs), Gen.choose(1, 4))).map(_.flatMap {
       case (contig, run) => List.fill(run)(contig)
@@ -67,4 +72,109 @@ class ContigIndexCacheSuite {
     assertEq(lookup("chr1"), 3)
     assertEq(lookup("chr10"), 4)
   }
+
+  /* Clearing the region and storing a contig of the same length at the same address would fool a
+   * cache keyed on the address, or on a view of the region, instead of a copy of the bytes. */
+  @Test def testRegionReuse(implicit ctx: ExecuteContext): Unit = {
+    val rg = ctx.references(ReferenceGenome.GRCh38)
+    val pt = PCanonicalLocus(rg)
+    val f = compileCanonical(ctx, rg)((cb, l) => cb.memoize(l.contigIdx(cb).toL))
+    ctx.r.pool.scopedRegion { r =>
+      val first = storeCanonical(ctx, rg, Locus("chr1", 1), r)
+      val firstContig = pt.representation.loadField(first, 0)
+      assertEq(f(first), rg.getContigIndex("chr1").toLong)
+
+      r.clear()
+      val second = storeCanonical(ctx, rg, Locus("chr2", 1), r)
+      assertEq(second, first)
+      assertEq(pt.representation.loadField(second, 0), firstContig)
+      assertEq(f(second), rg.getContigIndex("chr2").toLong)
+    }
+  }
+
+  @Test def testUnknownContigInGeneratedCode(implicit ctx: ExecuteContext): Unit = {
+    val rg = ctx.references(ReferenceGenome.GRCh38)
+    val f = compileCanonical(ctx, rg)((cb, l) => cb.memoize(l.contigIdx(cb).toL))
+    def lookup(contig: String): Long = f(storeCanonical(ctx, rg, Locus(contig, 1), ctx.r))
+
+    assertEq(lookup("chr1"), 0L)
+    // "chrZ" has the cached contig's length, "chr" is a prefix of it.
+    for (unknown <- Seq("chrZ", "chrZ", "chr"))
+      assertEq(outcome(lookup(unknown).toInt), outcome(rg.getContigIndex(unknown)))
+    assertEq(lookup("chr1"), 0L)
+    assertEq(lookup("chr10"), 9L)
+  }
+
+  /* Partitions run concurrently on the local backend, each with its own instance of the generated
+   * code. Contigs change every two rows, so lookups both hit and miss. Under unstable encodings,
+   * writing packs each canonical locus through the cache, so a wrong index reads back as a wrong
+   * contig. */
+  @Test def testConcurrentPartitions(implicit ctx: ExecuteContext): Unit = {
+    implicit val execStrats = ExecStrategy.lowering
+    withReference(custom) { implicit ctx =>
+      withUnstableEncodings { implicit ctx =>
+        val n = 64 * 40
+        val contigs = ArraySeq.tabulate(n)(i => custom.contigs((i / 2) % custom.contigs.length))
+        val positions = ArraySeq.tabulate(n)(i => i % 997 + 1)
+        val table = TableRange(n, 64).mapRows { (_, row) =>
+          row.get("idx").bind { idx =>
+            locusFromLiterals(contigs, positions)(idx).bind { locus =>
+              makestruct(
+                "idx" -> idx,
+                "locus" -> locus,
+                "contig_idx" -> invoke("contig_idx", TInt32, locus),
+              )
+            }
+          }
+        }
+
+        val path = ctx.createTmpPath("contig-index-cache-concurrency", "ht")
+        assertEvalsTo(TableWrite(table, TableNativeWriter(path)), ())
+        val expected = ArraySeq.tabulate(n) { i =>
+          RowSeq(i, Locus(contigs(i), positions(i)), custom.getContigIndex(contigs(i)))
+        }
+        val read = TableIR.read(ctx.fs, ctx.references, path, requestedType = Some(table.typ))
+        assertEvalsTo(collect(read), RowSeq(expected, RowSeq()))
+      }
+    }
+  }
+}
+
+object ContigIndexCacheSuite {
+
+  // Contigs that share prefixes and lengths, so that a comparison that stops early or checks only
+  // the length gives a wrong index.
+  val custom: ReferenceGenome = {
+    val contigs = ArraySeq("1", "10", "11", "chr1", "chr10", "chr2", "a", "ab", "ba")
+    ReferenceGenome("contig_index_cache_test", contigs, contigs.map(_ -> 1000).toMap)
+  }
+
+  /** Constructs a canonical locus on [[custom]] from the `idx`th contig and position. */
+  def locusFromLiterals(contigs: IndexedSeq[String], positions: IndexedSeq[Int])(idx: IR): IR =
+    invoke(
+      "Locus",
+      TLocus(custom),
+      Literal(TArray(TString), contigs).at(idx),
+      Literal(TArray(TInt32), positions).at(idx),
+    )
+
+  type LocusEmitter = (EmitCodeBuilder, SLocusValue) => Value[Long]
+
+  def storeCanonical(ctx: ExecuteContext, rg: ReferenceGenome, l: Locus, r: Region): Long =
+    PCanonicalLocus(rg).unstagedStoreJavaObject(l, r)
+
+  /** Compiles `f` over a canonical locus, given its address. Each call of the result runs the same
+    * generated instance, so a cache it holds persists across calls.
+    */
+  def compileCanonical(
+    ctx: ExecuteContext,
+    rg: ReferenceGenome,
+  )(
+    f: LocusEmitter
+  ): AsmFunction1[Long, Long] = {
+    val pt = PCanonicalLocus(rg)
+    val fb = EmitFunctionBuilder[Long, Long](ctx, "canonical_locus")
+    fb.emitWithBuilder(cb => f(cb, pt.loadCheapSCode(cb, fb.getCodeParam[Long](1)).asLocus))
+    fb.resultWithIndex()(ctx.theHailClassLoader, ctx.fs, ctx.taskContext, ctx.r)
+  }
 }