IRFunctionRegistry changes

[catoverdrive]

cc @cseed @danking

I changed the ApplyFunction IR to store the function name instead of the IRFunction implementation. (The implementation gets looked up during Infer if it doesn't already exist, but we're always pre-populating it in IRFunctionRegistry.lookupConversion).

As part of this, I separated the way that IRFunctionRegistry stores IR => IR from the IRFunction lookup, since the lookup in Infer wants to grab the IRFunction and not necessarily the entire Apply IR that we currently store. Maybe the IR => IR aspect (handled by IRFunction.lookupConversion) could go away once we no longer need the AST to IR conversion and start constructing the IR directly in python.

[tpoterba]

one tiny style comment

[tpoterba]

can you use any fail all fail here?

[catoverdrive]

None of the things (except the return) are Options? I don't see how it would work here

[tpoterba]

ahh, got it. nevermind.

[danking]

@catoverdrive Sorry I'm late to the game here. We should use mutable.MultiMap when we have Map[K, Set[V]]. Unless there's a compelling case for Seq here?

[danking]

That also makes addIRFunction and addIR just calls to irRegistry.addBinding and codeRegistry.addBinding

[catoverdrive]

womp, I didn't know this was a thing. I'll change.

[danking]

lookupConversion and lookupFunction are suspiciously similar. I'm not sure I see the compelling reason to keep them separate.

Would the code look simpler if we had one MultiMap[String, IRFun] where IRFun was something like:

trait IRFun {
  def args: Seq[Type]
  def ret: Type
  def emit(
    emitter: IR => (Code[Unit], Code[Boolean], Code[_]),
    mb: MethodBuilder,
    args: IR*
  ): (Code[Unit], Code[Boolean], Code[_])
}

class IRMacro(
  val args: Seq[Type],
  val ret: Type,
  val impl: Array[IR] => IR
) extends IRFun {
  def emit(
    emitter: IR => (Code[Unit], Code[Boolean], Code[_]),
    mb: MethodBuilder,
    args: IR*
  ): (Code[Unit], Code[Boolean], Code[_]) = emitter(impl(args.toArray))
}

class JVMFun(
  val args: Seq[Type],
  retType: Type,
  impl: (MethodBuilder, Array[Code[_]]) => Code[_]
) extends IRFun {
  def emit(
    emitter: IR => (Code[Unit], Code[Boolean], Code[_]),
    mb: MethodBuilder,
    args: IR*
  ): (Code[Unit], Code[Boolean], Code[_]) = {
    val (setup, ms, vs) = args.map(emitter).unzip
    (Code(setup:_*), ms.fold(_ || _), impl(mb, vs))
  }
}

[danking]

Or maybe let JVMFun also control missingness?

class MissingAwareJVMFun(
  val args: Seq[Type],
  retType: Type,
  impl: (MethodBuilder, Array[Code[_]]) => (Code[Unit], Code[Boolean], Code[_])
) extends IRFun {
  def emit(
    emitter: IR => (Code[Unit], Code[Boolean], Code[_]),
    mb: MethodBuilder,
    args: IR*
  ): (Code[Unit], Code[Boolean], Code[_]) = {
    val (setups, ms, vs) = args.map(emitter).unzip
    val (setup, m, v) = impl(mb, vs)
    (Code(setups ++ setup:_*), (ms :+ m).fold(_ || _), v)
  }
}

[catoverdrive]

I don't really want to entangle them because of two resasons:

• Things that can be translated directly (Seq[IR] => IR) are maybe things that we don't want to keep in the function registry once we can directly pass AST nodes from Python. I kind of want all of them to go away.
• Even if we do want to keep them in the function registry, I want to be able to look up the non-IR conversions post optimization, when we go to compile the code. Once we're no longer translating from AST, we don't need to store the implementation as part of the IR node and can look it up after we're done optimizing the overall IR. I feel like the things that are implemented in terms of existing IR nodes should just go directly into the IR tree and optimized away with the rest of it. (this is kind of what currently happens; Infer on Apply nodes does a lookup on the IRFunction registry to be able to properly infer the return type)