From 91c8e830fcade6dd56c3c34e587ef28813d6d97c Mon Sep 17 00:00:00 2001
From: Noisyfox <timemanager.rick@gmail.com>
Date: Mon, 24 Jan 2022 21:51:52 +0800
Subject: [PATCH] (multi-os-engine/multi-os-engine#162) Only write modified
 classes

---
 .../moe/tools/classvalidator/ClassModifier.kt  | 16 ++++++++++++++++
 .../moe/tools/classvalidator/ClassValidator.kt | 18 +++++++++++++-----
 .../natj/AddMissingAnnotations.kt              |  6 ++++--
 .../natj/AddMissingNatJRegister.kt             |  5 ++++-
 4 files changed, 37 insertions(+), 8 deletions(-)
 create mode 100644 src/main/kotlin/org/moe/tools/classvalidator/ClassModifier.kt

diff --git a/src/main/kotlin/org/moe/tools/classvalidator/ClassModifier.kt b/src/main/kotlin/org/moe/tools/classvalidator/ClassModifier.kt
new file mode 100644
index 0000000..1bee01b
--- /dev/null
+++ b/src/main/kotlin/org/moe/tools/classvalidator/ClassModifier.kt
@@ -0,0 +1,16 @@
+package org.moe.tools.classvalidator
+
+import org.objectweb.asm.ClassVisitor
+
+/**
+ * A [ClassVisitor] that could modify a class
+ */
+abstract class ClassModifier(api: Int, classVisitor: ClassVisitor?) : ClassVisitor(api, classVisitor) {
+
+    var modified: Boolean = false
+        private set
+
+    fun markAsModified() {
+        modified = true
+    }
+}
diff --git a/src/main/kotlin/org/moe/tools/classvalidator/ClassValidator.kt b/src/main/kotlin/org/moe/tools/classvalidator/ClassValidator.kt
index 422d3af..324012d 100644
--- a/src/main/kotlin/org/moe/tools/classvalidator/ClassValidator.kt
+++ b/src/main/kotlin/org/moe/tools/classvalidator/ClassValidator.kt
@@ -3,8 +3,6 @@ package org.moe.tools.classvalidator
 import org.moe.common.utils.classAndJarInputIterator
 import org.moe.tools.classvalidator.natj.AddMissingAnnotations
 import org.moe.tools.classvalidator.natj.AddMissingNatJRegister
-import org.moe.tools.classvalidator.substrate.CollectReflectionConfig
-import org.moe.tools.classvalidator.substrate.ReflectionConfig
 import org.objectweb.asm.ClassReader
 import org.objectweb.asm.ClassVisitor
 import org.objectweb.asm.ClassWriter
@@ -25,13 +23,23 @@ object ClassValidator {
             inputFiles.classAndJarInputIterator { _, inputStream ->
                 val cr = ClassReader(inputStream)
 
+                val chain = mutableListOf<ClassVisitor>()
+                fun ClassVisitor.chain(nextBuilder: (ClassVisitor) -> ClassVisitor): ClassVisitor {
+                    val next = nextBuilder(this)
+                    chain.add(next)
+                    return next
+                }
+
                 val byteCode = processClass(cr) { next ->
                     next
-                        .let(::AddMissingAnnotations)
-                        .let(::AddMissingNatJRegister)
+                        .chain(::AddMissingAnnotations)
+                        .chain(::AddMissingNatJRegister)
                 }
 
-                classSaver.save(byteCode)
+                // Only save modified class
+                if (chain.any { it is ClassModifier && it.modified }) {
+                    classSaver.save(byteCode)
+                }
             }
         }
     }
diff --git a/src/main/kotlin/org/moe/tools/classvalidator/natj/AddMissingAnnotations.kt b/src/main/kotlin/org/moe/tools/classvalidator/natj/AddMissingAnnotations.kt
index f819a8a..b646ee3 100644
--- a/src/main/kotlin/org/moe/tools/classvalidator/natj/AddMissingAnnotations.kt
+++ b/src/main/kotlin/org/moe/tools/classvalidator/natj/AddMissingAnnotations.kt
@@ -1,5 +1,6 @@
 package org.moe.tools.classvalidator.natj
 
+import org.moe.tools.classvalidator.ClassModifier
 import org.moe.tools.classvalidator.natj.NatJRuntime.toClass
 import org.objectweb.asm.AnnotationVisitor
 import org.objectweb.asm.ClassVisitor
@@ -10,7 +11,7 @@ import java.lang.reflect.AnnotatedElement
 
 class AddMissingAnnotations(
     next: ClassVisitor?
-) : ClassVisitor(Opcodes.ASM5, next) {
+) : ClassModifier(Opcodes.ASM5, next) {
     private var skip: Boolean = false
     private var superName: String? = null
     private var interfaces: Array<out String>? = null
@@ -129,6 +130,7 @@ class AddMissingAnnotations(
                         }
                     }
                 }
+                markAsModified()
             }
 
             private fun convertClassToType(value: Any): Any {
@@ -148,4 +150,4 @@ class AddMissingAnnotations(
             }
         }
     }
-}
\ No newline at end of file
+}
diff --git a/src/main/kotlin/org/moe/tools/classvalidator/natj/AddMissingNatJRegister.kt b/src/main/kotlin/org/moe/tools/classvalidator/natj/AddMissingNatJRegister.kt
index f0283ab..69f5037 100644
--- a/src/main/kotlin/org/moe/tools/classvalidator/natj/AddMissingNatJRegister.kt
+++ b/src/main/kotlin/org/moe/tools/classvalidator/natj/AddMissingNatJRegister.kt
@@ -1,5 +1,6 @@
 package org.moe.tools.classvalidator.natj
 
+import org.moe.tools.classvalidator.ClassModifier
 import org.objectweb.asm.AnnotationVisitor
 import org.objectweb.asm.ClassVisitor
 import org.objectweb.asm.MethodVisitor
@@ -14,7 +15,7 @@ import java.lang.reflect.Modifier
  */
 class AddMissingNatJRegister(
     next: ClassVisitor?
-) : ClassVisitor(Opcodes.ASM5, next) {
+) : ClassModifier(Opcodes.ASM5, next) {
 
     private var skip: Boolean = false
     private var visit: Boolean = false
@@ -90,6 +91,7 @@ class AddMissingNatJRegister(
                         NatJRuntime.NATJ_OWNER, NatJRuntime.NATJ_REGISTER_NAME,
                         NatJRuntime.NATJ_REGISTER_DESC, false)
                     CLI.accept(mv)
+                    markAsModified()
                     println("Injected NatJ.register() into $name")
                 }
                 NATJREG_LEAVE_ALONE -> {
@@ -108,6 +110,7 @@ class AddMissingNatJRegister(
             mv.visitInsn(Opcodes.RETURN)
             mv.visitMaxs(-1, -1)
             mv.visitEnd()
+            markAsModified()
             println("Injected NatJ.register() into $name")
         }