Removed afterEvaluate from AsmTaskRegistration, and added lazyTaskLookup utility function

Author: priettt

embrace-gradle-plugin/src/main/java/io/embrace/android/gradle/plugin/gradle/TaskRegistrationUtils.kt

@@ -3,6 +3,7 @@ package io.embrace.android.gradle.plugin.gradle
 import io.embrace.android.gradle.plugin.util.capitalizedString
 import org.gradle.api.Project
 import org.gradle.api.Task
+import org.gradle.api.provider.Provider
 import org.gradle.api.tasks.TaskContainer
 import org.gradle.api.tasks.TaskProvider
 
@@ -45,3 +46,17 @@ fun <T : Task> Project.tryGetTaskProvider(taskName: String, taskType: Class<T>):
         null
     }
 }
+
+/**
+ * Lazily looks up a task by name and type, returning a [Provider] that yields the task if it exists,
+ * or `null` if not found. Useful for wiring optional tasks into the task graph without realizing them early.
+ *
+ * @param T The expected type of the task.
+ * @param name The name of the task to look up.
+ * @return A [Provider] of the task, or `provider { null }` if the task is not present.
+ */
+inline fun <reified T : Task> Project.lazyTaskLookup(name: String): Provider<T?> {
+    return provider {
+        tryGetTaskProvider(name, T::class.java)
+    }.safeFlatMap { it as Provider<T?> }
+}

embrace-gradle-plugin/src/main/java/io/embrace/android/gradle/plugin/instrumentation/AsmTaskRegistration.kt

@@ -2,7 +2,8 @@ package io.embrace.android.gradle.plugin.instrumentation
 
 import com.android.build.api.instrumentation.FramesComputationMode
 import com.android.build.api.instrumentation.InstrumentationScope
-import io.embrace.android.gradle.plugin.gradle.tryGetTaskProvider
+import io.embrace.android.gradle.plugin.gradle.lazyTaskLookup
+import io.embrace.android.gradle.plugin.gradle.safeFlatMap
 import io.embrace.android.gradle.plugin.tasks.ndk.EncodeFileToBase64Task
 import io.embrace.android.gradle.plugin.tasks.reactnative.GenerateRnSourcemapTask
 import io.embrace.android.gradle.plugin.tasks.registration.EmbraceTaskRegistration
@@ -41,33 +42,25 @@ class AsmTaskRegistration : EmbraceTaskRegistration {
                 params.shouldInstrumentOnLongClick.set(behavior.instrumentation.onLongClickEnabled)
                 params.shouldInstrumentOnClick.set(behavior.instrumentation.onClickEnabled)
 
-                project.afterEvaluate {
-                    // Find the Asm transformation task by name and make it depend on encodeSharedObjectFilesTask
-                    val encodeFileToBase64Task = project.tryGetTaskProvider(
-                        "${EncodeFileToBase64Task.NAME}${data.name.capitalizedString()}",
-                        EncodeFileToBase64Task::class.java
-                    ) ?: return@afterEvaluate
-
-                    val asmTransformationTask = project.tryGetTaskProvider(
-                        "transform${variant.name.capitalizedString()}ClassesWithAsm"
-                    ) ?: error("Unable to find ASM transformation task for variant ${variant.name}.")
-                    asmTransformationTask.configure { it.dependsOn(encodeFileToBase64Task) }
-
-                    params.encodedSharedObjectFilesMap.set(
-                        encodeFileToBase64Task.flatMap {
-                            it.outputFile
-                        }
-                    )
+                val encodeFileToBase64Task = project.lazyTaskLookup<EncodeFileToBase64Task>(
+                    "${EncodeFileToBase64Task.NAME}${data.name.capitalizedString()}"
+                )
 
-                    val reactNativeTask = project.tryGetTaskProvider(
-                        "${GenerateRnSourcemapTask.NAME}${data.name.capitalizedString()}",
-                        GenerateRnSourcemapTask::class.java
-                    ) ?: return@afterEvaluate
+                params.encodedSharedObjectFilesMap.set(
+                    encodeFileToBase64Task.safeFlatMap {
+                        it?.outputFile ?: project.provider { null }
+                    }
+                )
 
-                    asmTransformationTask.configure { it.dependsOn(reactNativeTask) }
+                val reactNativeTask = project.lazyTaskLookup<GenerateRnSourcemapTask>(
+                    "${GenerateRnSourcemapTask.NAME}${data.name.capitalizedString()}"
+                )
 
-                    params.reactNativeBundleId.set(reactNativeTask.flatMap { it.bundleIdOutputFile })
-                }
+                params.reactNativeBundleId.set(
+                    reactNativeTask.safeFlatMap {
+                        it?.bundleIdOutputFile ?: project.provider { null }
+                    }
+                )
             }
         } catch (exception: Exception) {
             project.logger.error("An error has occurred while performing ASM bytecode transformation.", exception)

embrace-gradle-plugin/src/main/java/io/embrace/android/gradle/plugin/tasks/ndk/NdkUploadTasksRegistration.kt

@@ -1,6 +1,7 @@
 package io.embrace.android.gradle.plugin.tasks.ndk
 
 import io.embrace.android.gradle.plugin.config.PluginBehavior
+import io.embrace.android.gradle.plugin.gradle.lazyTaskLookup
 import io.embrace.android.gradle.plugin.gradle.nullSafeMap
 import io.embrace.android.gradle.plugin.gradle.registerTask
 import io.embrace.android.gradle.plugin.gradle.safeFlatMap
@@ -37,9 +38,7 @@ class NdkUploadTasksRegistration(
         // Bail if ndk_enabled is not true.
         if (variantConfig.embraceConfig?.ndkEnabled != true) return
 
-        val mergeNativeLibsTaskProvider: Provider<Task?> = project.provider {
-            project.tryGetTaskProvider("merge${variant.name.capitalizedString()}NativeLibs")
-        }.safeFlatMap { it as Provider<Task?> }
+        val mergeNativeLibsTaskProvider = project.lazyTaskLookup<Task>("merge${variant.name.capitalizedString()}NativeLibs")
 
         val sharedObjectFilesProvider = getSharedObjectFilesProvider(project, mergeNativeLibsTaskProvider)
 

embrace-gradle-plugin/src/main/java/io/embrace/android/gradle/plugin/tasks/reactnative/GenerateRnSourcemapTaskRegistration.kt

@@ -1,5 +1,6 @@
 package io.embrace.android.gradle.plugin.tasks.reactnative
 
+import io.embrace.android.gradle.plugin.gradle.lazyTaskLookup
 import io.embrace.android.gradle.plugin.gradle.nullSafeMap
 import io.embrace.android.gradle.plugin.gradle.registerTask
 import io.embrace.android.gradle.plugin.gradle.safeFlatMap
@@ -58,9 +59,7 @@ class GenerateRnSourcemapTaskRegistration : EmbraceTaskRegistration {
             )
 
             val variantCapitalized = variant.name.capitalizedString()
-            val generatorTaskProvider: Provider<Task?> = project.provider {
-                project.tryGetTaskProvider("createBundle${variantCapitalized}JsAndAssets")
-            }.safeFlatMap { it as Provider<Task?> }
+            val generatorTaskProvider = project.lazyTaskLookup<Task>("createBundle${variantCapitalized}JsAndAssets")
 
             rnTask.bundleFile.set(project.layout.file(getBundleFileProvider(generatorTaskProvider, project)))
             rnTask.sourcemap.set(project.layout.file(getSourcemapFileProvider(generatorTaskProvider, project, data)))