<commit_before>package is.hail.expr.ir

import is.hail.expr.types.virtual.{TArray, TInt32, TInt64, TStruct}
import is.hail.table.Ascending
import is.hail.utils._

object Simplify {

  /** Transform 'ir' using simplification rules until none apply.
    */
  def apply(ir: BaseIR): BaseIR = Simplify(ir, allowRepartitioning = true)

  /** Use 'allowRepartitioning'=false when in a context where simplification
    * should not change the partitioning of the result of 'ast', such as when
    * some parent (downstream) node of 'ast' uses seeded randomness.
    */
  private[ir] def apply(ast: BaseIR, allowRepartitioning: Boolean): BaseIR =
    ast match {
      case ir: IR => simplifyValue(ir)
      case tir: TableIR => simplifyTable(allowRepartitioning)(tir)
      case mir: MatrixIR => simplifyMatrix(allowRepartitioning)(mir)
    }

  private[this] def visitNode[T <: BaseIR](
    visitChildren: BaseIR => BaseIR,
    transform: T => Option[T],
    post: => (T => T)
  )(t: T): T = {
    val t1 = t.mapChildren(visitChildren).asInstanceOf[T]
    transform(t1).map(post).getOrElse(t1)
  }

  private[this] def simplifyValue: IR => IR =
    visitNode(
      Simplify(_),
      rewriteValueNode,
      simplifyValue)

  private[this] def simplifyTable(allowRepartitioning: Boolean)(tir: TableIR): TableIR =
    visitNode(
      Simplify(_, allowRepartitioning && isDeterministicallyRepartitionable(tir)),
      rewriteTableNode(allowRepartitioning),
      simplifyTable(allowRepartitioning)
    )(tir)

  private[this] def simplifyMatrix(allowRepartitioning: Boolean)(mir: MatrixIR): MatrixIR =
    visitNode(
      Simplify(_, allowRepartitioning && isDeterministicallyRepartitionable(mir)),
      rewriteMatrixNode(allowRepartitioning),
      simplifyMatrix(allowRepartitioning)
    )(mir)

  private[this] def rewriteValueNode: IR => Option[IR] = valueRules.lift

  private[this] def rewriteTableNode(allowRepartitioning: Boolean)(tir: TableIR): Option[TableIR] =
    tableRules(allowRepartitioning && isDeterministicallyRepartitionable(tir)).lift(tir)

  private[this] def rewriteMatrixNode(allowRepartitioning: Boolean)(mir: MatrixIR): Option[MatrixIR] =
    matrixRules(allowRepartitioning && isDeterministicallyRepartitionable(mir)).lift(mir)

  /** Returns true if 'x' propagates missingness, meaning if any child of 'x'
    * evaluates to missing, then 'x' will evaluate to missing.
    */
  private[this] def isStrict(x: IR): Boolean = {
    x match {
      case _: Apply |
           _: ApplyUnaryPrimOp |
           _: ApplyBinaryPrimOp |
           _: ArrayRange |
           _: ArrayRef |
           _: ArrayLen |
           _: GetField |
           _: GetTupleElement => true
      case ApplyComparisonOp(op, _, _) => op.strict
      case f: ApplySeeded => f.implementation.isStrict
      case _ => false
    }
  }

  /** Returns true if 'x' will never evaluate to missing.
    */
  private[this] def isDefinitelyDefined(x: IR): Boolean = {
    x match {
      case _: MakeArray |
           _: MakeStruct |
           _: MakeTuple |
           _: IsNA |
           ApplyComparisonOp(EQWithNA(_, _), _, _) |
           ApplyComparisonOp(NEQWithNA(_, _), _, _) |
           _: I32 | _: I64 | _: F32 | _: F64 | True() | False() => true
      case _ => false
    }
  }

  /** Returns true if changing the partitioning of child (upstream) nodes of
    * 'ir' can not otherwise change the contents of the result of 'ir'.
    */
  private[this] def isDeterministicallyRepartitionable(ir: BaseIR): Boolean =
    ir.children.forall{
      case child: IR => !Exists(child, _.isInstanceOf[ApplySeeded])
      case _ => true
    }

  private[this] def areFieldSelects(fields: Seq[(String, IR)]): Boolean = {
    assert(fields.nonEmpty)
    fields.head match {
      case (_, GetField(s1, _)) =>
        fields.forall {
          case (f1, GetField(s2, f2)) if f1 == f2 && s1 == s2 => true
          case _ => false
        }
      case _ => false
    }
  }

  private[this] def valueRules: PartialFunction[IR, IR] = {
    // propagate NA
    case x: IR if isStrict(x) && Children(x).exists(_.isInstanceOf[NA]) =>
      NA(x.typ)

    case x@If(NA(_), _, _) => NA(x.typ)

    case x@ArrayMap(NA(_), _, _) => NA(x.typ)

    case x@ArrayFlatMap(NA(_), _, _) => NA(x.typ)

    case x@ArrayFilter(NA(_), _, _) => NA(x.typ)

    case x@ArrayFold(NA(_), _, _, _, _) => NA(x.typ)

    case IsNA(NA(_)) => True()

    case IsNA(x) if isDefinitelyDefined(x) => False()

    case Let(n1, v, Ref(n2, _)) if n1 == n2 => v

    case Let(n, _, b) if !Mentions(b, n) => b

    case x@If(True(), cnsq, _) if x.typ == cnsq.typ => cnsq

    case x@If(False(), _, altr) if x.typ == altr.typ => altr

    case If(c, cnsq, altr) if cnsq == altr =>
      if (cnsq.typ.required)
        cnsq
      else
        If(IsNA(c), NA(cnsq.typ), cnsq)

    case Cast(x, t) if x.typ == t => x

    case ArrayLen(MakeArray(args, _)) => I32(args.length)

    case ArrayLen(ArrayRange(start, end, I32(1))) => ApplyBinaryPrimOp(Subtract(), end, start)

    case ArrayLen(ArrayMap(a, _, _)) => ArrayLen(a)

    case ArrayLen(ArrayFlatMap(a, _, MakeArray(args, _))) => ApplyBinaryPrimOp(Multiply(), I32(args.length), ArrayLen(a))

    case ArrayRef(MakeArray(args, _), I32(i)) if i >= 0 && i < args.length => args(i)

    case ArrayFilter(a, _, True()) => a

    case ArrayFor(_, _, Begin(Seq())) => Begin(FastIndexedSeq())

    case ArrayFold(ArrayMap(a, n1, b), zero, accumName, valueName, body) => ArrayFold(a, zero, accumName, n1, Let(valueName, b, body))

    case ArrayFlatMap(ArrayMap(a, n1, b1), n2, b2) =>
      ArrayFlatMap(a, n1, Let(n2, b1, b2))

    case ArrayMap(ArrayMap(a, n1, b1), n2, b2) =>
      ArrayMap(a, n1, Let(n2, b1, b2))

    case GetField(MakeStruct(fields), name) =>
      val (_, x) = fields.find { case (n, _) => n == name }.get
      x

    case GetField(InsertFields(old, fields), name) =>
      fields.find { case (n, _) => n == name } match {
        case Some((_, x)) => x
        case None => GetField(old, name)
      }

    case GetField(SelectFields(old, fields), name) => GetField(old, name)

    case InsertFields(InsertFields(base, fields1), fields2) =>
      val fields1Set = fields1.map(_._1).toSet
      val fields2Map = fields2.toMap

      val finalFields = fields1.map { case (name, fieldIR) => name -> fields2Map.getOrElse(name, fieldIR) } ++
        fields2.filter { case (name, _) => !fields1Set.contains(name) }
      InsertFields(base, finalFields)

    case InsertFields(MakeStruct(fields1), fields2) =>
      val fields1Set = fields1.map(_._1).toSet
      val fields2Map = fields2.toMap

      val finalFields = fields1.map { case (name, fieldIR) => name -> fields2Map.getOrElse(name, fieldIR) } ++
        fields2.filter { case (name, _) => !fields1Set.contains(name) }
      MakeStruct(finalFields)

    case InsertFields(struct, Seq()) => struct

    case SelectFields(old, fields) if coerce[TStruct](old.typ).fieldNames sameElements fields =>
      old

    case SelectFields(SelectFields(old, _), fields) =>
      SelectFields(old, fields)

    case SelectFields(MakeStruct(fields), fieldNames) =>
      val makeStructFields = fields.toMap
      MakeStruct(fieldNames.map(f => f -> makeStructFields(f)))

    case x@SelectFields(InsertFields(struct, insertFields), selectFields) =>
      val selectSet = selectFields.toSet
      val insertFields2 = insertFields.filter { case (fName, _) => selectSet.contains(fName) }
      val structSet = struct.typ.asInstanceOf[TStruct].fieldNames.toSet
      val selectFields2 = selectFields.filter(structSet.contains)
      val x2 = InsertFields(SelectFields(struct, selectFields2), insertFields2)
      assert(x2.typ == x.typ)
      x2

    case GetTupleElement(MakeTuple(xs), idx) => xs(idx)

    case TableCount(child) if child.partitionCounts.isDefined => I64(child.partitionCounts.get.sum)
    case TableCount(TableMapGlobals(child, _)) => TableCount(child)
    case TableCount(TableMapRows(child, _)) => TableCount(child)
    case TableCount(TableRepartition(child, _, _)) => TableCount(child)
    case TableCount(TableUnion(children)) =>
      children.map(TableCount).reduce[IR](ApplyBinaryPrimOp(Add(), _, _))
    case TableCount(TableKeyBy(child, _, _)) => TableCount(child)<commit_msg>Add simplification for mt.count_cols() with known col counts (#5068)<commit_after>package is.hail.expr.ir

import is.hail.expr.types.virtual.{TArray, TInt32, TInt64, TStruct}
import is.hail.table.Ascending
import is.hail.utils._

object Simplify {

  /** Transform 'ir' using simplification rules until none apply.
    */
  def apply(ir: BaseIR): BaseIR = Simplify(ir, allowRepartitioning = true)

  /** Use 'allowRepartitioning'=false when in a context where simplification
    * should not change the partitioning of the result of 'ast', such as when
    * some parent (downstream) node of 'ast' uses seeded randomness.
    */
  private[ir] def apply(ast: BaseIR, allowRepartitioning: Boolean): BaseIR =
    ast match {
      case ir: IR => simplifyValue(ir)
      case tir: TableIR => simplifyTable(allowRepartitioning)(tir)
      case mir: MatrixIR => simplifyMatrix(allowRepartitioning)(mir)
    }

  private[this] def visitNode[T <: BaseIR](
    visitChildren: BaseIR => BaseIR,
    transform: T => Option[T],
    post: => (T => T)
  )(t: T): T = {
    val t1 = t.mapChildren(visitChildren).asInstanceOf[T]
    transform(t1).map(post).getOrElse(t1)
  }

  private[this] def simplifyValue: IR => IR =
    visitNode(
      Simplify(_),
      rewriteValueNode,
      simplifyValue)

  private[this] def simplifyTable(allowRepartitioning: Boolean)(tir: TableIR): TableIR =
    visitNode(
      Simplify(_, allowRepartitioning && isDeterministicallyRepartitionable(tir)),
      rewriteTableNode(allowRepartitioning),
      simplifyTable(allowRepartitioning)
    )(tir)

  private[this] def simplifyMatrix(allowRepartitioning: Boolean)(mir: MatrixIR): MatrixIR =
    visitNode(
      Simplify(_, allowRepartitioning && isDeterministicallyRepartitionable(mir)),
      rewriteMatrixNode(allowRepartitioning),
      simplifyMatrix(allowRepartitioning)
    )(mir)

  private[this] def rewriteValueNode: IR => Option[IR] = valueRules.lift

  private[this] def rewriteTableNode(allowRepartitioning: Boolean)(tir: TableIR): Option[TableIR] =
    tableRules(allowRepartitioning && isDeterministicallyRepartitionable(tir)).lift(tir)

  private[this] def rewriteMatrixNode(allowRepartitioning: Boolean)(mir: MatrixIR): Option[MatrixIR] =
    matrixRules(allowRepartitioning && isDeterministicallyRepartitionable(mir)).lift(mir)

  /** Returns true if 'x' propagates missingness, meaning if any child of 'x'
    * evaluates to missing, then 'x' will evaluate to missing.
    */
  private[this] def isStrict(x: IR): Boolean = {
    x match {
      case _: Apply |
           _: ApplyUnaryPrimOp |
           _: ApplyBinaryPrimOp |
           _: ArrayRange |
           _: ArrayRef |
           _: ArrayLen |
           _: GetField |
           _: GetTupleElement => true
      case ApplyComparisonOp(op, _, _) => op.strict
      case f: ApplySeeded => f.implementation.isStrict
      case _ => false
    }
  }

  /** Returns true if 'x' will never evaluate to missing.
    */
  private[this] def isDefinitelyDefined(x: IR): Boolean = {
    x match {
      case _: MakeArray |
           _: MakeStruct |
           _: MakeTuple |
           _: IsNA |
           ApplyComparisonOp(EQWithNA(_, _), _, _) |
           ApplyComparisonOp(NEQWithNA(_, _), _, _) |
           _: I32 | _: I64 | _: F32 | _: F64 | True() | False() => true
      case _ => false
    }
  }

  /** Returns true if changing the partitioning of child (upstream) nodes of
    * 'ir' can not otherwise change the contents of the result of 'ir'.
    */
  private[this] def isDeterministicallyRepartitionable(ir: BaseIR): Boolean =
    ir.children.forall{
      case child: IR => !Exists(child, _.isInstanceOf[ApplySeeded])
      case _ => true
    }

  private[this] def areFieldSelects(fields: Seq[(String, IR)]): Boolean = {
    assert(fields.nonEmpty)
    fields.head match {
      case (_, GetField(s1, _)) =>
        fields.forall {
          case (f1, GetField(s2, f2)) if f1 == f2 && s1 == s2 => true
          case _ => false
        }
      case _ => false
    }
  }

  private[this] def valueRules: PartialFunction[IR, IR] = {
    // propagate NA
    case x: IR if isStrict(x) && Children(x).exists(_.isInstanceOf[NA]) =>
      NA(x.typ)

    case x@If(NA(_), _, _) => NA(x.typ)

    case x@ArrayMap(NA(_), _, _) => NA(x.typ)

    case x@ArrayFlatMap(NA(_), _, _) => NA(x.typ)

    case x@ArrayFilter(NA(_), _, _) => NA(x.typ)

    case x@ArrayFold(NA(_), _, _, _, _) => NA(x.typ)

    case IsNA(NA(_)) => True()

    case IsNA(x) if isDefinitelyDefined(x) => False()

    case Let(n1, v, Ref(n2, _)) if n1 == n2 => v

    case Let(n, _, b) if !Mentions(b, n) => b

    case x@If(True(), cnsq, _) if x.typ == cnsq.typ => cnsq

    case x@If(False(), _, altr) if x.typ == altr.typ => altr

    case If(c, cnsq, altr) if cnsq == altr =>
      if (cnsq.typ.required)
        cnsq
      else
        If(IsNA(c), NA(cnsq.typ), cnsq)

    case Cast(x, t) if x.typ == t => x

    case ArrayLen(MakeArray(args, _)) => I32(args.length)

    case ArrayLen(ArrayRange(start, end, I32(1))) => ApplyBinaryPrimOp(Subtract(), end, start)

    case ArrayLen(ArrayMap(a, _, _)) => ArrayLen(a)

    case ArrayLen(ArrayFlatMap(a, _, MakeArray(args, _))) => ApplyBinaryPrimOp(Multiply(), I32(args.length), ArrayLen(a))

    case ArrayRef(MakeArray(args, _), I32(i)) if i >= 0 && i < args.length => args(i)

    case ArrayFilter(a, _, True()) => a

    case ArrayFor(_, _, Begin(Seq())) => Begin(FastIndexedSeq())

    case ArrayFold(ArrayMap(a, n1, b), zero, accumName, valueName, body) => ArrayFold(a, zero, accumName, n1, Let(valueName, b, body))

    case ArrayFlatMap(ArrayMap(a, n1, b1), n2, b2) =>
      ArrayFlatMap(a, n1, Let(n2, b1, b2))

    case ArrayMap(ArrayMap(a, n1, b1), n2, b2) =>
      ArrayMap(a, n1, Let(n2, b1, b2))

    case GetField(MakeStruct(fields), name) =>
      val (_, x) = fields.find { case (n, _) => n == name }.get
      x

    case GetField(InsertFields(old, fields), name) =>
      fields.find { case (n, _) => n == name } match {
        case Some((_, x)) => x
        case None => GetField(old, name)
      }

    case GetField(SelectFields(old, fields), name) => GetField(old, name)

    case InsertFields(InsertFields(base, fields1), fields2) =>
      val fields1Set = fields1.map(_._1).toSet
      val fields2Map = fields2.toMap

      val finalFields = fields1.map { case (name, fieldIR) => name -> fields2Map.getOrElse(name, fieldIR) } ++
        fields2.filter { case (name, _) => !fields1Set.contains(name) }
      InsertFields(base, finalFields)

    case InsertFields(MakeStruct(fields1), fields2) =>
      val fields1Set = fields1.map(_._1).toSet
      val fields2Map = fields2.toMap

      val finalFields = fields1.map { case (name, fieldIR) => name -> fields2Map.getOrElse(name, fieldIR) } ++
        fields2.filter { case (name, _) => !fields1Set.contains(name) }
      MakeStruct(finalFields)

    case InsertFields(struct, Seq()) => struct

    case SelectFields(old, fields) if coerce[TStruct](old.typ).fieldNames sameElements fields =>
      old

    case SelectFields(SelectFields(old, _), fields) =>
      SelectFields(old, fields)

    case SelectFields(MakeStruct(fields), fieldNames) =>
      val makeStructFields = fields.toMap
      MakeStruct(fieldNames.map(f => f -> makeStructFields(f)))

    case x@SelectFields(InsertFields(struct, insertFields), selectFields) =>
      val selectSet = selectFields.toSet
      val insertFields2 = insertFields.filter { case (fName, _) => selectSet.contains(fName) }
      val structSet = struct.typ.asInstanceOf[TStruct].fieldNames.toSet
      val selectFields2 = selectFields.filter(structSet.contains)
      val x2 = InsertFields(SelectFields(struct, selectFields2), insertFields2)
      assert(x2.typ == x.typ)
      x2

    case GetTupleElement(MakeTuple(xs), idx) => xs(idx)

    case TableCount(MatrixColsTable(child)) if child.columnCount.isDefined => I64(child.columnCount.get)

    case TableCount(child) if child.partitionCounts.isDefined => I64(child.partitionCounts.get.sum)
    case TableCount(TableMapGlobals(child, _)) => TableCount(child)
    case TableCount(TableMapRows(child, _)) => TableCount(child)
    case TableCount(TableRepartition(child, _, _)) => TableCount(child)
    case TableCount(TableUnion(children)) =>
      children.map(TableCount).reduce[IR](ApplyBinaryPrimOp(Add(), _, _))
    case TableCount(TableKeyBy(child, _, _)) => TableCount(child)
{'fraction_non_alphanumeric': 0.13608296864696362, 'fraction_numerical': 0.01067082309214076, 'mean_word_length': 3.31170845179633, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 156, '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': '7653687', 'n_tokens_mistral': 5704, 'n_tokens_neox': 5456, 'n_words': 1248}