<commit_before>package is.hail.expr.ir.agg

import is.hail.annotations.{CodeOrdering, Region, RegionUtils, StagedRegionValueBuilder}
import is.hail.asm4s._
import is.hail.expr.ir.{EmitFunctionBuilder, EmitRegion, EmitTriplet, defaultValue, typeToTypeInfo}
import is.hail.expr.types.physical._
import is.hail.io._
import is.hail.utils._

class GroupedBTreeKey(kt: PType, fb: EmitFunctionBuilder[_], region: Code[Region], container: StateContainer) extends BTreeKey {
  val storageType: PStruct = PStruct(required = true,
    "kt" -> kt,
    "regionIdx" -> PInt32(true),
    "container" -> PInt64(true))

  val compType: PType = kt
  private val kcomp = fb.getCodeOrdering[Int](kt, CodeOrdering.compare, ignoreMissingness = false)

  def isKeyMissing(off: Code[Long]): Code[Boolean] =
    storageType.isFieldMissing(off, 0)
  def loadKey(off: Code[Long]): Code[_] = Region.loadIRIntermediate(kt)(storageType.fieldOffset(off, 0))

  def initValue(dest: Code[Long], km: Code[Boolean], kv: Code[_], rIdx: Code[Int]): Code[Unit] = {
    val koff = storageType.fieldOffset(dest, 0)
    var storeK =
      if (kt.isPrimitive)
        Region.storeIRIntermediate(kt)(koff, kv)
      else
        StagedRegionValueBuilder.deepCopy(fb, region, kt, coerce[Long](kv), koff)
    if (!kt.required)
      storeK = km.mux(storageType.setFieldMissing(dest, 0), Code(storageType.setFieldPresent(dest, 0), storeK))

    Code(
      storeK,
      storeRegionIdx(dest, rIdx),
      Region.storeAddress(containerOffset(dest), region.allocate(container.typ.alignment, container.typ.byteSize)))
  }

  def regionIdx(off: Code[Long]): Code[Int] =
    Region.loadInt(storageType.fieldOffset(off, 1))

  def storeRegionIdx(off: Code[Long], idx: Code[Int]): Code[Unit] =
    Region.storeInt(storageType.fieldOffset(off, 1), idx)

  def containerOffset(off: Code[Long]): Code[Long] =
    storageType.fieldOffset(off, 2)

  def containerAddress(off: Code[Long]): Code[Long] =
    Region.loadAddress(containerOffset(off))

  def isEmpty(off: Code[Long]): Code[Boolean] =
    Region.loadAddress(containerOffset(off)).ceq(0L)
  def initializeEmpty(off: Code[Long]): Code[Unit] =
    Region.storeAddress(containerOffset(off), 0L)

  def copy(src: Code[Long], dest: Code[Long]): Code[Unit] =
    Region.copyFrom(src, dest, storageType.byteSize)

  def deepCopy(er: EmitRegion, dest: Code[Long], src: Code[Long]): Code[Unit] =
    Code(StagedRegionValueBuilder.deepCopy(er, storageType, src, dest),
      container.toCode((i, s) => s.copyFrom(container.getStateOffset(containerAddress(src), i))),
      container.store(regionIdx(dest), containerAddress(dest)))

  def compKeys(k1: (Code[Boolean], Code[_]), k2: (Code[Boolean], Code[_])): Code[Int] =
    kcomp(k1, k2)

  def loadCompKey(off: Code[Long]): (Code[Boolean], Code[_]) =
    isKeyMissing(off) -> isKeyMissing(off).mux(defaultValue(kt), loadKey(off))

}

class DictState(val fb: EmitFunctionBuilder[_], val keyType: PType, val nested: Array[AggregatorState]) extends PointerBasedRVAState {
  val nStates: Int = nested.length
  val container: StateContainer = new StateContainer(nested, region)
  val valueType: PStruct = PStruct("regionIdx" -> PInt32(true), "states" -> container.typ)
  val root: ClassFieldRef[Long] = fb.newField[Long]
  val size: ClassFieldRef[Int] = fb.newField[Int]
  val keyed = new GroupedBTreeKey(keyType, fb, region, container)
  val tree = new AppendOnlyBTree(fb, keyed, region, root)

  val typ: PStruct = PStruct(
    required = true,
    "inits" -> container.typ,
    "size" -> PInt32(true),
    "tree" -> PInt64(true))

  private val initStatesOffset: Code[Long] = typ.fieldOffset(off, 0)
  private def initStateOffset(idx: Int): Code[Long] = container.getStateOffset(initStatesOffset, idx)

  private val _elt = fb.newField[Long]

  def initElement(eltOff: Code[Long], km: Code[Boolean], kv: Code[_]): Code[Unit] = {
    Code(
      size := size + 1,
      region.setNumParents(size * (nStates + 1)),
      keyed.initValue(_elt, km, kv, size * nStates),
      container.newStates)
  }<commit_msg>fix setNumParents bug in GroupedAggregator (#7077)<commit_after>package is.hail.expr.ir.agg

import is.hail.annotations.{CodeOrdering, Region, RegionUtils, StagedRegionValueBuilder}
import is.hail.asm4s._
import is.hail.expr.ir.{EmitFunctionBuilder, EmitRegion, EmitTriplet, defaultValue, typeToTypeInfo}
import is.hail.expr.types.physical._
import is.hail.io._
import is.hail.utils._

class GroupedBTreeKey(kt: PType, fb: EmitFunctionBuilder[_], region: Code[Region], container: StateContainer) extends BTreeKey {
  val storageType: PStruct = PStruct(required = true,
    "kt" -> kt,
    "regionIdx" -> PInt32(true),
    "container" -> PInt64(true))

  val compType: PType = kt
  private val kcomp = fb.getCodeOrdering[Int](kt, CodeOrdering.compare, ignoreMissingness = false)

  def isKeyMissing(off: Code[Long]): Code[Boolean] =
    storageType.isFieldMissing(off, 0)
  def loadKey(off: Code[Long]): Code[_] = Region.loadIRIntermediate(kt)(storageType.fieldOffset(off, 0))

  def initValue(dest: Code[Long], km: Code[Boolean], kv: Code[_], rIdx: Code[Int]): Code[Unit] = {
    val koff = storageType.fieldOffset(dest, 0)
    var storeK =
      if (kt.isPrimitive)
        Region.storeIRIntermediate(kt)(koff, kv)
      else
        StagedRegionValueBuilder.deepCopy(fb, region, kt, coerce[Long](kv), koff)
    if (!kt.required)
      storeK = km.mux(storageType.setFieldMissing(dest, 0), Code(storageType.setFieldPresent(dest, 0), storeK))

    Code(
      storeK,
      storeRegionIdx(dest, rIdx),
      Region.storeAddress(containerOffset(dest), region.allocate(container.typ.alignment, container.typ.byteSize)))
  }

  def regionIdx(off: Code[Long]): Code[Int] =
    Region.loadInt(storageType.fieldOffset(off, 1))

  def storeRegionIdx(off: Code[Long], idx: Code[Int]): Code[Unit] =
    Region.storeInt(storageType.fieldOffset(off, 1), idx)

  def containerOffset(off: Code[Long]): Code[Long] =
    storageType.fieldOffset(off, 2)

  def containerAddress(off: Code[Long]): Code[Long] =
    Region.loadAddress(containerOffset(off))

  def isEmpty(off: Code[Long]): Code[Boolean] =
    Region.loadAddress(containerOffset(off)).ceq(0L)
  def initializeEmpty(off: Code[Long]): Code[Unit] =
    Region.storeAddress(containerOffset(off), 0L)

  def copy(src: Code[Long], dest: Code[Long]): Code[Unit] =
    Region.copyFrom(src, dest, storageType.byteSize)

  def deepCopy(er: EmitRegion, dest: Code[Long], src: Code[Long]): Code[Unit] =
    Code(StagedRegionValueBuilder.deepCopy(er, storageType, src, dest),
      container.toCode((i, s) => s.copyFrom(container.getStateOffset(containerAddress(src), i))),
      container.store(regionIdx(dest), containerAddress(dest)))

  def compKeys(k1: (Code[Boolean], Code[_]), k2: (Code[Boolean], Code[_])): Code[Int] =
    kcomp(k1, k2)

  def loadCompKey(off: Code[Long]): (Code[Boolean], Code[_]) =
    isKeyMissing(off) -> isKeyMissing(off).mux(defaultValue(kt), loadKey(off))

}

class DictState(val fb: EmitFunctionBuilder[_], val keyType: PType, val nested: Array[AggregatorState]) extends PointerBasedRVAState {
  val nStates: Int = nested.length
  val container: StateContainer = new StateContainer(nested, region)
  val valueType: PStruct = PStruct("regionIdx" -> PInt32(true), "states" -> container.typ)
  val root: ClassFieldRef[Long] = fb.newField[Long]
  val size: ClassFieldRef[Int] = fb.newField[Int]
  val keyed = new GroupedBTreeKey(keyType, fb, region, container)
  val tree = new AppendOnlyBTree(fb, keyed, region, root)

  val typ: PStruct = PStruct(
    required = true,
    "inits" -> container.typ,
    "size" -> PInt32(true),
    "tree" -> PInt64(true))

  private val initStatesOffset: Code[Long] = typ.fieldOffset(off, 0)
  private def initStateOffset(idx: Int): Code[Long] = container.getStateOffset(initStatesOffset, idx)

  private val _elt = fb.newField[Long]

  def initElement(eltOff: Code[Long], km: Code[Boolean], kv: Code[_]): Code[Unit] = {
    Code(
      size := size + 1,
      region.setNumParents((size + 1) * nStates),
      keyed.initValue(_elt, km, kv, size * nStates),
      container.newStates)
  }
{'fraction_non_alphanumeric': 0.14345421020282728, 'fraction_numerical': 0.007375537799631223, 'mean_word_length': 4.895652173913043, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 23, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6742869', 'n_tokens_mistral': 2816, 'n_tokens_neox': 2697, 'n_words': 585}