From 0dac3366709e5c53d9780988cfa7fe999ffb11a7 Mon Sep 17 00:00:00 2001
From: tobspr <tobias.springer1@googlemail.com>
Date: Mon, 17 Jan 2022 19:11:04 +0100
Subject: [PATCH] Update class extensions

---
 mod_examples/class_extensions.js | 32 +++++++++++++++++---------------
 src/js/mods/mod_interface.js     | 32 +++++++++++++++++++-------------
 2 files changed, 36 insertions(+), 28 deletions(-)

diff --git a/mod_examples/class_extensions.js b/mod_examples/class_extensions.js
index 167ca53a..c5b5b95e 100644
--- a/mod_examples/class_extensions.js
+++ b/mod_examples/class_extensions.js
@@ -8,22 +8,24 @@ const METADATA = {
     description: "Shows how to extend builtin classes",
 };
 
+const BeltExtension = ({ $super, $old }) => ({
+    getShowWiresLayerPreview() {
+        // Access the old method
+        return !$old.getShowWiresLayerPreview();
+    },
+
+    getIsReplaceable() {
+        // Instead of super, use $super
+        return $super.getIsReplaceable.call(this);
+    },
+
+    getIsRemoveable() {
+        return false;
+    },
+});
+
 class Mod extends shapez.Mod {
     init() {
-        this.modInterface.extendClass(shapez.MetaBeltBuilding, {
-            // this replaces a regular method
-            getShowWiresLayerPreview() {
-                return true;
-            },
-
-            // Instead of super, use this.$super()
-            getIsReplaceable() {
-                return this.$super.getIsReplaceable.call(this);
-            },
-
-            getIsRemoveable() {
-                return false;
-            },
-        });
+        this.modInterface.extendClass(shapez.MetaBeltBuilding, BeltExtension);
     }
 }
diff --git a/src/js/mods/mod_interface.js b/src/js/mods/mod_interface.js
index c83f8b11..ab6b4077 100644
--- a/src/js/mods/mod_interface.js
+++ b/src/js/mods/mod_interface.js
@@ -389,18 +389,24 @@ export class ModInterface {
         };
     }
 
-    extendClass(classHandle, extensionClass) {
-        const extendPrototype = function (base, extension) {
-            const properties = Array.from(Object.getOwnPropertyNames(extension));
-            base.$super = base.$super || {};
-            properties.forEach(propertyName => {
-                if (["constructor", "name", "length", "prototype"].includes(propertyName)) {
-                    return;
-                }
-                base.$super[propertyName] = base.$super[propertyName] || base[propertyName];
-                base[propertyName] = extension[propertyName];
-            });
-        };
-        extendPrototype(classHandle.prototype, extensionClass);
+    /**
+     *
+     * @param {typeof Object} classHandle
+     * @param {({ $super, $old }) => any} extender
+     */
+    extendClass(classHandle, extender) {
+        const prototype = classHandle.prototype;
+
+        const $super = Object.getPrototypeOf(prototype);
+        const $old = {};
+        const extensionMethods = extender({ $super, $old });
+        const properties = Array.from(Object.getOwnPropertyNames(extensionMethods));
+        properties.forEach(propertyName => {
+            if (["constructor", "prototype"].includes(propertyName)) {
+                return;
+            }
+            $old[propertyName] = prototype[propertyName];
+            prototype[propertyName] = extensionMethods[propertyName];
+        });
     }
 }