From eee91ce1b0307918d1008bbda494c1fcd4f1b58a Mon Sep 17 00:00:00 2001
From: Kingsley Adio <kingsleyadio@meta.com>
Date: Wed, 4 Dec 2024 09:07:32 -0800
Subject: [PATCH] Register compiler plugin test extension

Summary:
$title

[RFC](https://docs.google.com/document/d/11VEA3WI2e3_GdN-te4Km3X0wCT33hk3cFAylDZ-WcLk/edit?tab=t.0)

Reviewed By: zielinskimz

Differential Revision: D66712021

fbshipit-source-id: 557c18a414866bdd7c8846935870fc6ad2fb00c8
---
 .../facebook/litho/TestComponentRegistrar.kt  | 36 ++++++++++++++
 .../kotlin/com/facebook/litho/analysis/Ir.kt  | 48 +++++++++++++++++++
 2 files changed, 84 insertions(+)
 create mode 100644 litho-compiler-plugin/compiler/src/test/kotlin/com/facebook/litho/TestComponentRegistrar.kt
 create mode 100644 litho-compiler-plugin/compiler/src/test/kotlin/com/facebook/litho/analysis/Ir.kt

diff --git a/litho-compiler-plugin/compiler/src/test/kotlin/com/facebook/litho/TestComponentRegistrar.kt b/litho-compiler-plugin/compiler/src/test/kotlin/com/facebook/litho/TestComponentRegistrar.kt
new file mode 100644
index 0000000000..0f086ba193
--- /dev/null
+++ b/litho-compiler-plugin/compiler/src/test/kotlin/com/facebook/litho/TestComponentRegistrar.kt
@@ -0,0 +1,36 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.facebook.litho
+
+import com.facebook.litho.analysis.TestIrExtension
+import com.facebook.litho.processor.AbstractTestProcessor
+import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
+import org.jetbrains.kotlin.compiler.plugin.CompilerPluginRegistrar
+import org.jetbrains.kotlin.compiler.plugin.ExperimentalCompilerApi
+import org.jetbrains.kotlin.config.CompilerConfiguration
+
+@OptIn(ExperimentalCompilerApi::class)
+class TestComponentRegistrar(private val processor: AbstractTestProcessor) :
+    CompilerPluginRegistrar() {
+
+  override fun ExtensionStorage.registerExtensions(configuration: CompilerConfiguration) {
+    IrGenerationExtension.registerExtension(TestIrExtension(processor))
+  }
+
+  override val supportsK2: Boolean
+    get() = true
+}
diff --git a/litho-compiler-plugin/compiler/src/test/kotlin/com/facebook/litho/analysis/Ir.kt b/litho-compiler-plugin/compiler/src/test/kotlin/com/facebook/litho/analysis/Ir.kt
new file mode 100644
index 0000000000..ceb4e94005
--- /dev/null
+++ b/litho-compiler-plugin/compiler/src/test/kotlin/com/facebook/litho/analysis/Ir.kt
@@ -0,0 +1,48 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.facebook.litho.analysis
+
+import com.facebook.litho.processor.AbstractTestProcessor
+import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
+import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
+import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
+import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
+import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
+import org.jetbrains.kotlin.ir.visitors.acceptVoid
+
+class TestIrExtension(private val processor: AbstractTestProcessor) : IrGenerationExtension {
+  override fun generate(moduleFragment: IrModuleFragment, pluginContext: IrPluginContext) {
+    val visitor = TestIrVisitor(processor, pluginContext)
+    moduleFragment.acceptVoid(visitor)
+  }
+}
+
+private class TestIrVisitor(
+    private val processor: AbstractTestProcessor,
+    private val context: IrPluginContext
+) : IrElementVisitorVoid {
+  override fun visitFile(declaration: IrFile) {
+    super.visitFile(declaration)
+    processor.analyse(AbstractTestProcessor.AnalysisData(declaration, context))
+  }
+
+  override fun visitElement(element: IrElement) {
+    element.acceptChildrenVoid(this)
+  }
+}