Add support for StringConcatFactory.makeConcatWithConstants (#9555)

udf-compiler/src/main/scala-2.13/com/nvidia/spark/udf/Instruction.scala

@@ -314,6 +314,12 @@ case class Instruction(opcode: Int, operand: Int, instructionStr: String) extend
             val (args, rest) = stack.splitAt(n + 1)
             (args.reverse, rest)
           })
+      case Opcode.INVOKEDYNAMIC =>
+        invokedynamic(lambdaReflection, state,
+          (stack, n) => {
+            val (args, rest) = stack.splitAt(n)
+            (args.reverse, rest)
+        })
       case _ => throw new SparkException("Unsupported instruction: " + instructionStr)
     }
     logDebug(s"[Instruction] ${instructionStr} got new state: ${st} from state: ${state}")
@@ -563,6 +569,31 @@ case class Instruction(opcode: Int, operand: Int, instructionStr: String) extend
     }
   }
 
+  private def invokedynamic(lambdaReflection: LambdaReflection, state: State,
+      getArgs: (List[Expression], Int) =>
+          (List[Expression], List[Expression])): State = {
+    val State(locals, stack, cond, expr) = state
+    val (bootstrapMethod, bootstrapArgs) = lambdaReflection.lookupBootstrapMethod(operand)
+    val declaringClass = bootstrapMethod.getDeclaringClass
+    val declaringClassName = declaringClass.getName
+    val newstack = {
+      if (declaringClassName.equals("java.lang.invoke.StringConcatFactory") &&
+          bootstrapMethod.getName.equals("makeConcatWithConstants") &&
+          bootstrapArgs.length == 1) {
+        val recipe = bootstrapArgs.head.toString
+        if (recipe.contains('\u0002')) {
+          throw new SparkException("Unsupported instruction: " + instructionStr)
+        }
+        val (args, rest) = getArgs(stack, recipe.count{x => x == '\u0001'})
+        Concat(recipe.split('\u0001').zipAll(args, "", Literal(""))
+          .map{ case(x, y) => Concat(Seq(Literal(x), y))}.toSeq) :: rest
+      } else {
+        throw new SparkException("Unsupported instruction: " + instructionStr)
+      }
+    }
+    State(locals, newstack, cond, expr)
+  }
+
   private def checkArgs(methodName: String,
                         expectedTypes: List[DataType],
                         args: List[Expression]): Unit = {
@@ -958,7 +989,7 @@ object Instruction {
         codeIterator.byteAt(offset + 1)
       case Opcode.BIPUSH =>
         codeIterator.signedByteAt(offset + 1)
-      case Opcode.LDC_W | Opcode.LDC2_W | Opcode.NEW | Opcode.CHECKCAST |
+      case Opcode.LDC_W | Opcode.LDC2_W | Opcode.NEW | Opcode.CHECKCAST | Opcode.INVOKEDYNAMIC |
            Opcode.INVOKESTATIC | Opcode.INVOKEVIRTUAL | Opcode.INVOKEINTERFACE |
            Opcode.INVOKESPECIAL | Opcode.GETSTATIC =>
         codeIterator.u16bitAt(offset + 1)

udf-compiler/src/main/scala/com/nvidia/spark/udf/LambdaReflection.scala

@@ -78,6 +78,32 @@ class LambdaReflection private(private val classPool: ClassPool,
     }
   }
 
+  def lookupBootstrapMethod(constPoolIndex: Int): (CtMethod, Seq[Any]) = {
+    if (constPool.getTag(constPoolIndex) != ConstPool.CONST_InvokeDynamic) {
+      throw new SparkException(s"Unexpected index ${constPoolIndex} for bootstrap")
+    }
+    val bootstrapMethodIndex = constPool.getInvokeDynamicBootstrap(constPoolIndex)
+    val bootstrapMethodsAttribute = ctMethod.getDeclaringClass.getClassFile.getAttributes
+      .toArray.filter(_.isInstanceOf[javassist.bytecode.BootstrapMethodsAttribute])
+    if (bootstrapMethodsAttribute.length != 1) {
+      throw new SparkException(s"Multiple bootstrap methods attributes aren't supported")
+    }
+    val bootstrapMethods = bootstrapMethodsAttribute.head
+      .asInstanceOf[javassist.bytecode.BootstrapMethodsAttribute].getMethods
+    val bootstrapMethod = bootstrapMethods(bootstrapMethodIndex)
+    val bootstrapMethodArguments = try {
+      bootstrapMethod.arguments.map(lookupConstant)
+    } catch {
+      case _: Throwable =>
+        throw new SparkException(s"only constants are supported as bootstrap method arguments")
+    }
+    val constPoolIndexMethodref = constPool.getMethodHandleIndex(bootstrapMethod.methodRef)
+    val methodName = constPool.getMethodrefName(constPoolIndexMethodref)
+    val descriptor = constPool.getMethodrefType(constPoolIndexMethodref)
+    val className = constPool.getMethodrefClassName(constPoolIndexMethodref)
+    (classPool.getCtClass(className).getMethod(methodName, descriptor), bootstrapMethodArguments)
+  }
+
   def lookupClassName(constPoolIndex: Int): String = {
     if (constPool.getTag(constPoolIndex) != ConstPool.CONST_Class) {
       throw new SparkException("Unexpected index for class")