RFC: [x2cpg] Added ForkJoinParallelCpgPassWithAccumulator

joern-cli/frontends/jssrc2cpg/src/main/scala/io/joern/jssrc2cpg/astcreation/AstCreator.scala

@@ -7,7 +7,6 @@ import io.joern.jssrc2cpg.parser.BabelNodeInfo
 import io.joern.x2cpg.Ast
 import io.joern.x2cpg.AstCreatorBase
 import io.joern.x2cpg.ValidationMode
-import io.joern.x2cpg.datastructures.Global
 import io.joern.x2cpg.datastructures.Stack.*
 import io.joern.x2cpg.datastructures.VariableScopeManager
 import io.joern.x2cpg.frontendspecific.jssrc2cpg.Defines
@@ -26,7 +25,7 @@ import ujson.Value
 
 import scala.collection.mutable
 
-class AstCreator(val config: Config, val global: Global, val parserResult: ParseResult)(implicit
+class AstCreator(val config: Config, val usedTypes: mutable.Set[String], val parserResult: ParseResult)(implicit
   withSchemaValidation: ValidationMode
 ) extends AstCreatorBase[BabelNodeInfo, AstCreator](parserResult.filename)
     with AstForExpressionsCreator

joern-cli/frontends/jssrc2cpg/src/main/scala/io/joern/jssrc2cpg/astcreation/AstCreatorHelper.scala

@@ -107,7 +107,7 @@ trait AstCreatorHelper(implicit withSchemaValidation: ValidationMode) { this: As
   }
 
   protected def registerType(typeFullName: String): Unit = {
-    global.usedTypes.putIfAbsent(typeFullName, true)
+    usedTypes.add(typeFullName)
   }
 
   protected def codeForNodes(nodes: Seq[NewNode]): Option[String] = nodes.collectFirst {

joern-cli/frontends/jssrc2cpg/src/main/scala/io/joern/jssrc2cpg/passes/AstCreationPass.scala

@@ -5,34 +5,42 @@ import io.joern.jssrc2cpg.astcreation.AstCreator
 import io.joern.jssrc2cpg.parser.BabelJsonParser
 import io.joern.jssrc2cpg.utils.AstGenRunner.AstGenRunnerResult
 import io.joern.x2cpg.ValidationMode
-import io.joern.x2cpg.datastructures.Global
 import io.joern.x2cpg.frontendspecific.jssrc2cpg.Defines
+import io.joern.x2cpg.passes.ForkJoinParallelCpgPassWithAccumulator
 import io.joern.x2cpg.utils.Report
 import io.joern.x2cpg.utils.TimeUtils
 import io.shiftleft.codepropertygraph.generated.Cpg
-import io.shiftleft.passes.ForkJoinParallelCpgPass
 import io.shiftleft.utils.IOUtils
 import org.slf4j.Logger
 import org.slf4j.LoggerFactory
 
 import java.nio.file.Paths
-import scala.jdk.CollectionConverters.*
+import scala.collection.mutable
 import scala.util.Failure
 import scala.util.Success
 import scala.util.Try
 
 class AstCreationPass(cpg: Cpg, astGenRunnerResult: AstGenRunnerResult, config: Config, report: Report = new Report())(
   implicit withSchemaValidation: ValidationMode
-) extends ForkJoinParallelCpgPass[(String, String)](cpg) {
+) extends ForkJoinParallelCpgPassWithAccumulator[(String, String), mutable.Set[String]](cpg) {
 
   private val logger: Logger = LoggerFactory.getLogger(classOf[AstCreationPass])
 
-  private val global = new Global()
+  private var _typesSeen: List[String] = Nil
 
-  def typesSeen(): List[String] = global.usedTypes.keys().asScala.filterNot(_ == Defines.Any).toList
+  def typesSeen(): List[String] = _typesSeen
 
   override def generateParts(): Array[(String, String)] = astGenRunnerResult.parsedFiles.toArray
 
+  override protected def newAccumulator(): mutable.Set[String] = mutable.Set.empty[String]
+
+  override protected def mergeAccumulators(left: mutable.Set[String], right: mutable.Set[String]): mutable.Set[String] =
+    left ++= right
+
+  override protected def onAccumulatorComplete(acc: mutable.Set[String]): Unit = {
+    _typesSeen = acc.filterNot(_ == Defines.Any).toList
+  }
+
   override def finish(): Unit = {
     astGenRunnerResult.skippedFiles.foreach { skippedFile =>
       val (rootPath, fileName) = skippedFile
@@ -45,17 +53,22 @@ class AstCreationPass(cpg: Cpg, astGenRunnerResult: AstGenRunnerResult, config:
       }
       report.addReportInfo(fileName, fileLOC)
     }
+    super.finish()
   }
 
-  override def runOnPart(diffGraph: DiffGraphBuilder, input: (String, String)): Unit = {
+  override protected def runOnPartWithAccumulator(
+    diffGraph: DiffGraphBuilder,
+    usedTypes: mutable.Set[String],
+    input: (String, String)
+  ): Unit = {
     val (rootPath, jsonFilename) = input
     val parseResultMaybe         = BabelJsonParser.readFile(Paths.get(rootPath), Paths.get(jsonFilename))
     val ((gotCpg, filename), duration) = TimeUtils.time {
       parseResultMaybe match {
         case Success(parseResult) =>
           report.addReportInfo(parseResult.filename, parseResult.fileLoc, parsed = true)
           Try {
-            val localDiff = new AstCreator(config, global, parseResult).createAst()
+            val localDiff = new AstCreator(config, usedTypes, parseResult).createAst()
             diffGraph.absorb(localDiff)
           } match {
             case Failure(exception) =>

joern-cli/frontends/x2cpg/src/main/scala/io/joern/x2cpg/passes/ForkJoinParallelCpgPassWithAccumulator.scala

@@ -0,0 +1,76 @@
+package io.joern.x2cpg.passes
+
+import io.shiftleft.codepropertygraph.generated.Cpg
+import io.shiftleft.passes.ForkJoinParallelCpgPass
+
+import java.util.concurrent.ConcurrentLinkedQueue
+import scala.jdk.CollectionConverters.*
+
+/** A [[ForkJoinParallelCpgPass]] that additionally maintains a thread-local accumulator of type [[R]] which is merged
+  * across all threads after processing completes. This enables map-reduce style aggregation alongside the usual
+  * DiffGraph-based graph modifications.
+  *
+  * Each thread gets its own accumulator instance (via [[newAccumulator]]). After all parts are processed, the
+  * accumulators are merged using [[mergeAccumulators]] and the result is passed to [[onAccumulatorComplete]].
+  *
+  * @tparam T
+  *   the part type (same as in [[ForkJoinParallelCpgPass]])
+  * @tparam R
+  *   the accumulator type
+  */
+abstract class ForkJoinParallelCpgPassWithAccumulator[T <: AnyRef, R](cpg: Cpg, outName: String = "")
+    extends ForkJoinParallelCpgPass[T](cpg, outName) {
+
+  /** Create a fresh, empty accumulator. Called once per thread. */
+  protected def newAccumulator(): R
+
+  /** Merge two accumulators. Must be associative. The result may reuse either argument. */
+  protected def mergeAccumulators(left: R, right: R): R
+
+  /** Process a single part, writing graph changes to `builder` and aggregated data to `acc`. */
+  protected def runOnPartWithAccumulator(builder: DiffGraphBuilder, acc: R, part: T): Unit
+
+  /** Called after all parts are processed with the fully merged accumulator. */
+  protected def onAccumulatorComplete(acc: R): Unit = {}
+
+  private val accumulators = new ConcurrentLinkedQueue[R]()
+
+  private val threadLocalAcc: ThreadLocal[R] = new ThreadLocal[R]()
+
+  // Guard against double-finish: ForkJoinParallelCpgPass.createAndApply() calls finish() in its
+  // finally block, AND runWithBuilder() (called from createAndApply) also calls finish() in its
+  // own finally block. Without this guard the second call would invoke onAccumulatorComplete with
+  // an empty accumulator, overwriting the result from the first call.
+  //
+  // @note: This is a hack to prevent double-finish.
+  // TODO: We should find a better way to handle this in ForkJoinParallelCpgPass directly.
+  private var _finishCompleted = false
+
+  final override def runOnPart(builder: DiffGraphBuilder, part: T): Unit = {
+    var acc = threadLocalAcc.get()
+    if (acc == null) {
+      acc = newAccumulator()
+      threadLocalAcc.set(acc)
+      accumulators.add(acc)
+    }
+    runOnPartWithAccumulator(builder, acc, part)
+  }
+
+  override def init(): Unit = {
+    _finishCompleted = false
+    accumulators.clear()
+    threadLocalAcc.remove()
+    super.init()
+  }
+
+  override def finish(): Unit = {
+    if (!_finishCompleted) {
+      _finishCompleted = true
+      val merged = accumulators.asScala.reduceOption(mergeAccumulators).getOrElse(newAccumulator())
+      onAccumulatorComplete(merged)
+      accumulators.clear()
+      threadLocalAcc.remove()
+    }
+    super.finish()
+  }
+}