<commit_before>package lore.compiler.phases.transpilation

import java.io.{ByteArrayOutputStream, PrintStream}

import lore.compiler.CompilerOptions
import lore.compiler.core.Compilation
import lore.compiler.phases.transpilation.RuntimeTypeTranspiler.RuntimeTypeVariables
import lore.compiler.phases.transpilation.TranspiledChunk.JsCode
import lore.compiler.phases.transpilation.TranspiledName.StringExtension
import lore.compiler.semantics.Registry
import lore.compiler.semantics.functions.MultiFunctionDefinition
import lore.compiler.types.{ProductType, Type}

import scala.collection.mutable
import scala.util.Using

// TODO: This needs a serious overhaul. The transpilation of all this logic is very complex and relies on the
//       intertwining of several states and timing. We should rather consider creating an intermediate representation
//       which can be transpiled more easily. This would separate decision making (call fits poly or mono, unpack
//       input product, use dispatch cache, etc.) and the model (transpiled input types and type variables, transpiled
//       argument types, dispatch nodes, etc.) from the actual transpilation to Javascript.
//       The transpiler is also sufficiently complex that it could be moved to its own package.

class MultiFunctionTranspiler(mf: MultiFunctionDefinition)(implicit compilerOptions: CompilerOptions, registry: Registry) {

  private val varArgumentType = "argumentType".asName
  private val varDispatchCache = s"${mf.name}__dispatchCache".asName

  // TODO: This could be passed around as an immutable map instead of being "global" state.
  private val inputTypeJsNames = mutable.HashMap[Type, TranspiledName]()

  private implicit val nameProvider: TemporaryNameProvider = new TemporaryNameProvider(s"${mf.name}__")

  def transpile(): Compilation[String] = {
    val out = new ByteArrayOutputStream()
    val tryCompilation = Using(new PrintStream(out, true, "utf-8")) { printer =>
      mf.functions.filterNot(_.isAbstract).map { function =>
        FunctionTranspiler.transpile(function, nameProvider)
      }.simultaneous.map { results =>
        results.reduce[(JsCode, RuntimeTypeVariables)] { case ((code1, variables1), (code2, variables2)) =>
          // We can just concatenate the variable maps even if multiple functions have a variable of the same name
          // because type variables are uniquely identified by their object reference and not by their name.
          (code1 + "\n" + code2, variables1 ++ variables2)
        }
      }.map { case (functionCode, runtimeTypeVariables) =>
        printer.print(functionCode)
        prepareFunctionInputTypes(printer)(runtimeTypeVariables)
        prepareDispatchCache(printer)
        val mfName = s"${mf.name}"
        printer.println(s"export function $mfName($jsParameters) {")
        if (compilerOptions.runtimeLogging) {
          printer.println(s"console.info('Called multi-function $mfName.');")
        }

        transpileArgumentTypeGathering(printer)
        transpileDispatchCall(printer)(runtimeTypeVariables)
        printer.println("}")
      }
    }

    tryCompilation.map(c => c.map(_ => out.toString("utf-8"))).get
  }

  /**
    * Whether this multi-function consists of a single function.
    */
  private lazy val isSingleFunction = mf.functions.length == 1

  /**
    * All possible arities of the functions.
    */
  private lazy val arities: Set[Int] = mf.functions.map(_.signature.arity).toSet

  /**
    * If the multi-function only has functions of the same arity, this option contains that arity.
    */
  private lazy val uniqueArity: Option[Int] = if (arities.size != 1) None else arities.headOption

  /**
    * Whether the transpiled multi-function uses JS rest parameters. This is needed if the arity of functions
    * differs.
    */
  private lazy val usingRestParameters = uniqueArity.isEmpty

  /**
    * Whether we can bypass creating product types around argument types on both the left and right side of
    * fits. If there is only a single argument across the board, product types are useless overhead. We can
    * easily optimize this overhead away.
    */
  private lazy val canUnpackInputProduct = uniqueArity.contains(1)

  /**
    * The list of multi-function parameters.
    */
  private lazy val jsParameterNames: Vector[TranspiledName] = uniqueArity match {
    case None => Vector("args".asName)
    case Some(arity) => (0 until arity).map(index => s"arg$index".asName).toVector
  }

  /**
    * The comma-separated string of multi-function parameters.
    */
  private lazy val jsParameters: String = {
    if (usingRestParameters) s"...${jsParameterNames.head}" else jsParameterNames.mkString(", ")
  }

  /**
    * Maps all possible input types to temporary variable names. The temporary variables are written as global
    * constants so that JS doesn't have to recreate these objects every time the multi-function is called.
    */
  private def prepareFunctionInputTypes(printer: PrintStream)(implicit runtimeTypeVariables: RuntimeTypeVariables): Unit = {
    val inputTypes = mf.functions.map(_.signature.inputType).toSet
    inputTypes.foreach { inputType =>
      val simplifiedInputType = if (canUnpackInputProduct) {
        assert(inputType.elements.size == 1)
        inputType.elements.head
      } else inputType
      val varType = nameProvider.createName()
      val typeExpr = RuntimeTypeTranspiler.transpile(simplifiedInputType)
      printer.println(s"const $varType = $typeExpr;")
      inputTypeJsNames.put(inputType, varType)
    }
  }

  /**
    * This heuristic decides whether the dispatch cache should be used for this multi-function or not.
    *
    * The algorithm for this decision is, as of yet, very simple:
    *   1. If ANY parameter type is polymorphic, using the cache will absolutely bring gains. Creating a type
    *      allocation takes so much time that creating a hash from a type is almost 4-5 times faster than checking
    *      one fit.
    *   2. If there are at least three functions in the multi-function, we can assume that many invocations of the
    *      function will require three fit tests. At that point a hash&cache operation is faster than testing multiple
    *      fits.
    */
  private lazy val shouldUseDispatchCache = mf.functions.size >= 3 || mf.functions.exists(f => f.isPolymorphic)

  /**
    * Prepares the dispatch cache for later use.
    */
  private def prepareDispatchCache(printer: PrintStream): Unit = {
    if (shouldUseDispatchCache) {
      printer.println(s"const $varDispatchCache = ${RuntimeApi.utils.typeMap.create}();")
    }
  }

  /**
    * Whether this multi-function consists of a single function that also happens to have as its input type
    * the unit type. This vastly simplifies our handling of
    *
    * TODO: Use this to get rid of fits calls entirely and just check the number of arguments coming in.
    *       That is, if args.length === 0, call the function, otherwise throw an emptyFit error.<commit_msg>Decide whether multi-functions should use the dispatch cache with a type complexity metric.<commit_after>package lore.compiler.phases.transpilation

import java.io.{ByteArrayOutputStream, PrintStream}
import lore.compiler.CompilerOptions
import lore.compiler.core.Compilation
import lore.compiler.phases.transpilation.RuntimeTypeTranspiler.RuntimeTypeVariables
import lore.compiler.phases.transpilation.TranspiledChunk.JsCode
import lore.compiler.phases.transpilation.TranspiledName.StringExtension
import lore.compiler.semantics.Registry
import lore.compiler.semantics.functions.{FunctionDefinition, MultiFunctionDefinition}
import lore.compiler.types.{IntersectionType, ListType, MapType, ProductType, ShapeType, SumType, TraitType, Type, TypeVariable}

import scala.collection.mutable
import scala.util.Using

// TODO: This needs a serious overhaul. The transpilation of all this logic is very complex and relies on the
//       intertwining of several states and timing. We should rather consider creating an intermediate representation
//       which can be transpiled more easily. This would separate decision making (call fits poly or mono, unpack
//       input product, use dispatch cache, etc.) and the model (transpiled input types and type variables, transpiled
//       argument types, dispatch nodes, etc.) from the actual transpilation to Javascript.
//       The transpiler is also sufficiently complex that it could be moved to its own package.

class MultiFunctionTranspiler(mf: MultiFunctionDefinition)(implicit compilerOptions: CompilerOptions, registry: Registry) {

  private val varArgumentType = "argumentType".asName
  private val varDispatchCache = s"${mf.name}__dispatchCache".asName

  // TODO: This could be passed around as an immutable map instead of being "global" state.
  private val inputTypeJsNames = mutable.HashMap[Type, TranspiledName]()

  private implicit val nameProvider: TemporaryNameProvider = new TemporaryNameProvider(s"${mf.name}__")

  def transpile(): Compilation[String] = {
    val out = new ByteArrayOutputStream()
    val tryCompilation = Using(new PrintStream(out, true, "utf-8")) { printer =>
      mf.functions.filterNot(_.isAbstract).map { function =>
        FunctionTranspiler.transpile(function, nameProvider)
      }.simultaneous.map { results =>
        results.reduce[(JsCode, RuntimeTypeVariables)] { case ((code1, variables1), (code2, variables2)) =>
          // We can just concatenate the variable maps even if multiple functions have a variable of the same name
          // because type variables are uniquely identified by their object reference and not by their name.
          (code1 + "\n" + code2, variables1 ++ variables2)
        }
      }.map { case (functionCode, runtimeTypeVariables) =>
        printer.print(functionCode)
        prepareFunctionInputTypes(printer)(runtimeTypeVariables)
        prepareDispatchCache(printer)
        val mfName = s"${mf.name}"
        printer.println(s"export function $mfName($jsParameters) {")
        if (compilerOptions.runtimeLogging) {
          printer.println(s"console.info('Called multi-function $mfName.');")
        }

        transpileArgumentTypeGathering(printer)
        transpileDispatchCall(printer)(runtimeTypeVariables)
        printer.println("}")
      }
    }

    tryCompilation.map(c => c.map(_ => out.toString("utf-8"))).get
  }

  /**
    * Whether this multi-function consists of a single function.
    */
  private lazy val isSingleFunction = mf.functions.length == 1

  /**
    * All possible arities of the functions.
    */
  private lazy val arities: Set[Int] = mf.functions.map(_.signature.arity).toSet

  /**
    * If the multi-function only has functions of the same arity, this option contains that arity.
    */
  private lazy val uniqueArity: Option[Int] = if (arities.size != 1) None else arities.headOption

  /**
    * Whether the transpiled multi-function uses JS rest parameters. This is needed if the arity of functions
    * differs.
    */
  private lazy val usingRestParameters = uniqueArity.isEmpty

  /**
    * Whether we can bypass creating product types around argument types on both the left and right side of
    * fits. If there is only a single argument across the board, product types are useless overhead. We can
    * easily optimize this overhead away.
    */
  private lazy val canUnpackInputProduct = uniqueArity.contains(1)

  /**
    * The list of multi-function parameters.
    */
  private lazy val jsParameterNames: Vector[TranspiledName] = uniqueArity match {
    case None => Vector("args".asName)
    case Some(arity) => (0 until arity).map(index => s"arg$index".asName).toVector
  }

  /**
    * The comma-separated string of multi-function parameters.
    */
  private lazy val jsParameters: String = {
    if (usingRestParameters) s"...${jsParameterNames.head}" else jsParameterNames.mkString(", ")
  }

  /**
    * Maps all possible input types to temporary variable names. The temporary variables are written as global
    * constants so that JS doesn't have to recreate these objects every time the multi-function is called.
    */
  private def prepareFunctionInputTypes(printer: PrintStream)(implicit runtimeTypeVariables: RuntimeTypeVariables): Unit = {
    val inputTypes = mf.functions.map(_.signature.inputType).toSet
    inputTypes.foreach { inputType =>
      val simplifiedInputType = if (canUnpackInputProduct) {
        assert(inputType.elements.size == 1)
        inputType.elements.head
      } else inputType
      val varType = nameProvider.createName()
      val typeExpr = RuntimeTypeTranspiler.transpile(simplifiedInputType)
      printer.println(s"const $varType = $typeExpr;")
      inputTypeJsNames.put(inputType, varType)
    }
  }

  /**
    * This heuristic decides whether the dispatch cache should be used for this multi-function or not. At some
    * complexity threshold of the parameter types of all functions, we use the dispatch cache. This threshold is
    * reached relatively early, so only the simplest functions won't use the dispatch cache.
    */
  private lazy val shouldUseDispatchCache = {
    // The estimated run-time complexity of checking a subtyping relationship when the type in question is on the
    // right-hand side.
    def typeComplexity(tpe: Type): Int = tpe match {
      case tv: TypeVariable => 10 + typeComplexity(tv.lowerBound) + typeComplexity(tv.upperBound)
      case SumType(parts) => 1 + parts.map(typeComplexity).sum
      case IntersectionType(parts) => 1 + parts.map(typeComplexity).sum
      case ProductType(elements) => 1 + elements.map(typeComplexity).sum
      case ListType(element) => 1 + typeComplexity(element)
      case MapType(key, value) => 1 + typeComplexity(key) + typeComplexity(value)
      case _: TraitType =>
        // Checking whether a trait is a supertype may be inherently expensive because we have to walk the left
        // type's supertype hierarchy.
        5
      case ShapeType(properties) => 3 + properties.map(_._2.tpe).map(typeComplexity).sum * 2
      case _ => 1
    }

    def signatureComplexity(function: FunctionDefinition): Int = function.signature.parameters.map(_.tpe).map(typeComplexity).sum

    // TODO: To model the cost of the dispatch cache, which includes the hashing operation on the argument tuple type,
    //       we could counter-model the estimated cost of the hash operation.

    mf.functions.map(signatureComplexity).sum >= 5
  }

  /**
    * Prepares the dispatch cache for later use.
    */
  private def prepareDispatchCache(printer: PrintStream): Unit = {
    if (shouldUseDispatchCache) {
      printer.println(s"const $varDispatchCache = ${RuntimeApi.utils.typeMap.create}();")
    }
  }

  /**
    * Whether this multi-function consists of a single function that also happens to have as its input type
    * the unit type. This vastly simplifies our handling of
    *
    * TODO: Use this to get rid of fits calls entirely and just check the number of arguments coming in.
    *       That is, if args.length === 0, call the function, otherwise throw an emptyFit error.
{'fraction_non_alphanumeric': 0.08612471865483913, 'fraction_numerical': 0.0033099430689792135, 'mean_word_length': 3.570953101361573, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 37, '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': '1221988', 'n_tokens_mistral': 4078, 'n_tokens_neox': 3957, 'n_words': 1389}