<commit_before>package slick.compiler

import slick.ast._
import Util._
import TypeUtil._
import QueryParameter.constOp

import scala.collection.mutable

/** Replace all occurrences of `Take` and `Drop` with row number computations based on
  * `zipWithIndex` operations. */
class RemoveTakeDrop extends Phase {
  val name = "removeTakeDrop"

  def apply(state: CompilerState) = state.map { n =>
    val invalid = mutable.Set[TypeSymbol]()
    def tr(n: Node): Node = n.replace {
      case n @ TakeDrop(from, t, d) =>
        logger.debug(s"""Translating "drop $d, then take $t" to zipWithIndex operation:""", n)
        val fromRetyped = tr(from).infer()
        val from2 = fromRetyped match {
          case b: Bind => b
          case n =>
            val s = new AnonSymbol
            Bind(s, n, Pure(Ref(s)))
        }
        val j = Join(new AnonSymbol, new AnonSymbol, from2, RangeFrom(1L), JoinType.Zip, LiteralNode(true))
        val bs1 = new AnonSymbol
        val b1 = Bind(bs1, j, Pure(Ref(bs1)))
        val fs = new AnonSymbol
        val f = Filter(fs, b1, (t, d) match {
          case (None, Some(d)) => Library.>.typed[Boolean](Select(Ref(fs), ElementSymbol(2)), d)
          case (Some(t), None) => Library.<=.typed[Boolean](Select(Ref(fs), ElementSymbol(2)), t)
          case (Some(t), Some(d)) =>
            Library.And.typed[Boolean](
              Library.>.typed[Boolean](Select(Ref(fs), ElementSymbol(2)), d),
              Library.<=.typed[Boolean](Select(Ref(fs), ElementSymbol(2)), constOp[Long]("+")(_ + _)(t, d))
            )
        })<commit_msg>Allow selective translation for Take or Drop in removeTakeDrop:<commit_after>package slick.compiler

import slick.ast._
import Util._
import TypeUtil._
import QueryParameter.constOp

import scala.collection.mutable

/** Replace all occurrences of `Take` and `Drop` with row number computations based on
  * `zipWithIndex` operations. */
class RemoveTakeDrop(val translateTake: Boolean = true, val translateDrop: Boolean = true) extends Phase {
  val name = "removeTakeDrop"

  def apply(state: CompilerState) = state.map { n =>
    val invalid = mutable.Set[TypeSymbol]()
    def tr(n: Node): Node = n.replace {
      case n @ TakeDrop(from, t, d) if (translateTake && t.isDefined) || (translateDrop && d.isDefined) =>
        logger.debug(s"""Translating "drop $d, then take $t" to zipWithIndex operation:""", n)
        val fromRetyped = tr(from).infer()
        val from2 = fromRetyped match {
          case b: Bind => b
          case n =>
            val s = new AnonSymbol
            Bind(s, n, Pure(Ref(s)))
        }
        val j = Join(new AnonSymbol, new AnonSymbol, from2, RangeFrom(1L), JoinType.Zip, LiteralNode(true))
        val bs1 = new AnonSymbol
        val b1 = Bind(bs1, j, Pure(Ref(bs1)))
        val fs = new AnonSymbol
        val f = Filter(fs, b1, (t, d) match {
          case (None, Some(d)) => Library.>.typed[Boolean](Select(Ref(fs), ElementSymbol(2)), d)
          case (Some(t), None) => Library.<=.typed[Boolean](Select(Ref(fs), ElementSymbol(2)), t)
          case (Some(t), Some(d)) =>
            Library.And.typed[Boolean](
              Library.>.typed[Boolean](Select(Ref(fs), ElementSymbol(2)), d),
              Library.<=.typed[Boolean](Select(Ref(fs), ElementSymbol(2)), constOp[Long]("+")(_ + _)(t, d))
            )
        })
{'fraction_non_alphanumeric': 0.14523809523809525, 'fraction_numerical': 0.007142857142857143, 'mean_word_length': 2.83675799086758, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 21, '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': '1368448', 'n_tokens_mistral': 1086, 'n_tokens_neox': 1056, 'n_words': 292}