3 files changed, 50 insertions, 24 deletions

diff --git a/src/WebAssembly/Execution/Allocator.php b/src/WebAssembly/Execution/Allocator.php
index 4e2adbf..79f8679 100644
--- a/src/WebAssembly/Execution/Allocator.php
+++ b/src/WebAssembly/Execution/Allocator.php
@@ -12,6 +12,7 @@ use Nsfisis\Waddiwasi\WebAssembly\Structure\Types\MemType;
 use Nsfisis\Waddiwasi\WebAssembly\Structure\Types\TableType;
 use Nsfisis\Waddiwasi\WebAssembly\Structure\Types\ValType;
 use RuntimeException;
+use function assert;
 use function count;
 
 final readonly class Allocator
@@ -61,7 +62,7 @@ final readonly class Allocator
 
         foreach ($externVals as $externVal) {
             match ($externVal::class) {
-                ExternVals\Func::class => null, // handled below.
+                ExternVals\Func::class => $m->funcAddrs[] = $externVal->addr,
                 ExternVals\Table::class => $m->tableAddrs[] = $externVal->addr,
                 ExternVals\Mem::class => $m->memAddrs[] = $externVal->addr,
                 ExternVals\Global_::class => $m->globalAddrs[] = $externVal->addr,
@@ -72,14 +73,13 @@ final readonly class Allocator
         foreach ($preAllocatedFuncs as $funcAddr) {
             $m->funcAddrs[] = $funcAddr;
             $funcInst = $this->store->funcs[$funcAddr];
-            if ($funcInst instanceof FuncInsts\Wasm) {
-                // @phpstan-ignore-next-line
-                $this->store->funcs[$funcAddr] = FuncInst::Wasm(
-                    $funcInst->type,
-                    $m,
-                    $funcInst->code,
-                );
-            }
+            assert($funcInst instanceof FuncInsts\Wasm);
+            // @phpstan-ignore-next-line
+            $this->store->funcs[$funcAddr] = FuncInst::Wasm(
+                $funcInst->type,
+                $m,
+                $funcInst->code,
+            );
         }
         foreach ($module->tables as $table) {
             $m->tableAddrs[] = $this->allocTable($table->type, Ref::RefNull($table->type->refType));
diff --git a/src/WebAssembly/Execution/Runtime.php b/src/WebAssembly/Execution/Runtime.php
index d0f90b5..07dc0d0 100644
--- a/src/WebAssembly/Execution/Runtime.php
+++ b/src/WebAssembly/Execution/Runtime.php
@@ -20,6 +20,7 @@ use function abs;
 use function array_map;
 use function array_merge;
 use function array_reverse;
+use function array_slice;
 use function assert;
 use function ceil;
 use function count;
@@ -114,7 +115,7 @@ final class Runtime
             $externVals,
             $vals,
             $refsList,
-            $moduleInstInit->funcAddrs,
+            array_slice($moduleInstInit->funcAddrs, count(array_filter($externVals, fn ($e) => $e instanceof ExternVals\Func))),
         );
 
         $runtime = new self($store, $stack, $moduleInst);
@@ -203,6 +204,9 @@ final class Runtime
     {
         try {
             $export = $this->getExport($name);
+            if ($export === null) {
+                throw new TrapException("invoke($name) not found", trapKind: TrapKind::UninitializedElement);
+            }
             assert($export instanceof ExternVals\Func);
             $funcAddr = $export->addr;
 
diff --git a/src/WebAssembly/Execution/Store.php b/src/WebAssembly/Execution/Store.php
index 5bef648..18f8626 100644
--- a/src/WebAssembly/Execution/Store.php
+++ b/src/WebAssembly/Execution/Store.php
@@ -34,19 +34,41 @@ final class Store
 
     public function register(Extern $extern): ExternVal
     {
-        match ($extern::class) {
-            Externs\Func::class => $this->funcs[] = $extern->func,
-            Externs\Table::class => $this->tables[] = $extern->table,
-            Externs\Mem::class => $this->mems[] = $extern->mem,
-            Externs\Global_::class => $this->globals[] = $extern->global,
-            default => throw new RuntimeException("unreachable"),
-        };
-        return match ($extern::class) {
-            Externs\Func::class => ExternVal::Func(count($this->funcs) - 1),
-            Externs\Table::class => ExternVal::Table(count($this->tables) - 1),
-            Externs\Mem::class => ExternVal::Mem(count($this->mems) - 1),
-            Externs\Global_::class => ExternVal::Global_(count($this->globals) - 1),
-            default => throw new RuntimeException("unreachable"),
-        };
+        switch ($extern::class) {
+            case Externs\Func::class:
+                foreach ($this->funcs as $i => $f) {
+                    if ($f === $extern->func) {
+                        return ExternVal::Func($i);
+                    }
+                }
+                $this->funcs[] = $extern->func;
+                return ExternVal::Func(count($this->funcs) - 1);
+            case Externs\Table::class:
+                foreach ($this->tables as $i => $t) {
+                    if ($t === $extern->table) {
+                        return ExternVal::Table($i);
+                    }
+                }
+                $this->tables[] = $extern->table;
+                return ExternVal::Table(count($this->tables) - 1);
+            case Externs\Mem::class:
+                foreach ($this->mems as $i => $m) {
+                    if ($m === $extern->mem) {
+                        return ExternVal::Mem($i);
+                    }
+                }
+                $this->mems[] = $extern->mem;
+                return ExternVal::Mem(count($this->mems) - 1);
+            case Externs\Global_::class:
+                foreach ($this->globals as $i => $g) {
+                    if ($g === $extern->global) {
+                        return ExternVal::Global_($i);
+                    }
+                }
+                $this->globals[] = $extern->global;
+                return ExternVal::Global_(count($this->globals) - 1);
+            default:
+                throw new RuntimeException("unreachable");
+        }
     }
 }