From d5a3a50d274eff4f8fe1327db649c1f3faf5d23a Mon Sep 17 00:00:00 2001
From: Piotr <delwing@gmail.com>
Date: Tue, 26 Aug 2025 02:59:12 +0200
Subject: [PATCH] feat: support extending loaded libs

---
 README.md     | 12 ++++++++++++
 src/index.ts  | 24 ++++++++++++++++++++++--
 tests/test.js | 10 ++++++++++
 3 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 286c929..1772a41 100644
--- a/README.md
+++ b/README.md
@@ -79,6 +79,18 @@ console.log(helloStr)
 
 Check out the [math lib](./src/lib/math.ts) for a more extensive example.
 
+### Extend a library
+
+Use `extendLib` to add new functions to an existing library without overwriting it.
+
+```js
+const extraMath = new luainjs.Table({ tau: Math.PI * 2 })
+luaEnv.extendLib('math', extraMath)
+
+const tau = luaEnv.parse('return math.tau').exec()
+console.log(tau)
+```
+
 ## Example
 
 Check out the [test runner](./tests/test.js) for a concrete example.
diff --git a/src/index.ts b/src/index.ts
index 7e4fccf..71a2f2e 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -10,7 +10,7 @@ import { libTable } from './lib/table'
 import { libString, metatable as stringMetatable } from './lib/string'
 import { getLibOS } from './lib/os'
 import { getLibPackage } from './lib/package'
-import { LuaType, ensureArray, Config } from './utils'
+import { LuaType, ensureArray, Config, hasOwnProperty } from './utils'
 import { parse as parseScript } from './parser'
 
 interface Script {
@@ -56,6 +56,7 @@ function createEnv(
     parse: (script: string) => Script
     parseFile: (path: string) => Script
     loadLib: (name: string, value: Table) => void
+    extendLib: (name: string, value: Table) => void
 } {
     const cfg: Config = {
         LUA_PATH: './?.lua',
@@ -77,6 +78,24 @@ function createEnv(
         loaded.rawset(name, value)
     }
 
+    const extendLib = (name: string, value: Table): void => {
+        const lib = _G.rawget(name)
+        if (lib instanceof Table) {
+            for (let i = 1; i < value.numValues.length; i++) {
+                if (hasOwnProperty(value.numValues, i)) lib.rawset(i, value.numValues[i])
+            }
+            for (const key in value.strValues) {
+                if (hasOwnProperty(value.strValues, key)) lib.rawset(key, value.strValues[key])
+            }
+            for (let i = 0; i < value.keys.length; i++) {
+                lib.rawset(value.keys[i], value.values[i])
+            }
+            return
+        }
+
+        loadLib(name, value)
+    }
+
     loadLib('_G', _G)
     loadLib('package', libPackage)
     loadLib('math', libMath)
@@ -105,7 +124,8 @@ function createEnv(
     return {
         parse,
         parseFile,
-        loadLib
+        loadLib,
+        extendLib
     }
 }
 
diff --git a/tests/test.js b/tests/test.js
index 1d18745..176a7d3 100644
--- a/tests/test.js
+++ b/tests/test.js
@@ -41,6 +41,16 @@ let exitCode = 0
     }
 }
 
+{
+    const luaEnv = luainjs.createEnv()
+    const ext = new luainjs.Table({ foo: () => 'bar' })
+    luaEnv.extendLib('math', ext)
+    const val = luaEnv.parse('return math.foo()').exec()
+    if (val !== 'bar') {
+        throw Error('extendLib failed!')
+    }
+}
+
 {
     const luaEnv = luainjs.createEnv()
     let str