From 37b5b28ef7ea0b1d115dbe5ef1d74647faa85f47 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Burak=20Bilge=20Yal=C3=A7=C4=B1nkaya?=
 <bbyalcinkaya@gmail.com>
Date: Sun, 21 Jan 2024 19:03:29 +0300
Subject: [PATCH] Fix table type (#564)

* fix table ref type

* add table unit test
---
 pykwasm/pyproject.toml                        |  1 +
 pykwasm/src/pykwasm/wasm2kast.py              |  9 ++++----
 .../src/tests/integration/binary/tables.wat   |  9 +++++++-
 pykwasm/src/tests/unit/test_wasm2kast.py      | 21 +++++++++++++++++++
 4 files changed, 34 insertions(+), 6 deletions(-)
 create mode 100644 pykwasm/src/tests/unit/test_wasm2kast.py

diff --git a/pykwasm/pyproject.toml b/pykwasm/pyproject.toml
index 58fb5d506..a62f96f55 100644
--- a/pykwasm/pyproject.toml
+++ b/pykwasm/pyproject.toml
@@ -55,6 +55,7 @@ disallow_untyped_defs = true
 exclude = [
     'src/pykwasm/wasm2kast\.py',
     'src/wasm/*',
+    'src/tests/unit/test_wasm2kast\.py',
 ]
 
 [tool.poetry.scripts]
diff --git a/pykwasm/src/pykwasm/wasm2kast.py b/pykwasm/src/pykwasm/wasm2kast.py
index 4bd63b843..2f45d6ef3 100644
--- a/pykwasm/src/pykwasm/wasm2kast.py
+++ b/pykwasm/src/pykwasm/wasm2kast.py
@@ -102,10 +102,7 @@ def func(f: Function):
 
 def table(t: Table):
     ls = limits(t.type.limits)
-    if isinstance(t.type.elem_type, addresses.FunctionAddress):
-        typ = a.funcref
-    else:
-        typ = a.externref
+    typ = ref_type(t.type.elem_type)
     return a.table(ls, typ)
 
 
@@ -123,7 +120,9 @@ def glob(g: Global):
 def ref_type(t: RefType):
     if t is addresses.FunctionAddress:
         return a.funcref
-    return a.externref
+    if t is addresses.ExternAddress:
+        return a.externref
+    raise ValueError(f'Invalid RefType: {t}')
 
 
 def elem_mode(m: ElemMode) -> KInner:
diff --git a/pykwasm/src/tests/integration/binary/tables.wat b/pykwasm/src/tests/integration/binary/tables.wat
index 5dfd2780f..a5662a90d 100644
--- a/pykwasm/src/tests/integration/binary/tables.wat
+++ b/pykwasm/src/tests/integration/binary/tables.wat
@@ -1,2 +1,9 @@
 (module
-  (table 1 funcref))
\ No newline at end of file
+  (table (import "spectest" "table") 0 funcref)
+  (table 0 funcref)
+  (table 1 funcref)
+  (table 1 100 funcref)
+  (table 0 externref)
+  (table 1 externref)
+  (table 1 100 externref)
+)
diff --git a/pykwasm/src/tests/unit/test_wasm2kast.py b/pykwasm/src/tests/unit/test_wasm2kast.py
new file mode 100644
index 000000000..4c5d188e1
--- /dev/null
+++ b/pykwasm/src/tests/unit/test_wasm2kast.py
@@ -0,0 +1,21 @@
+import pytest
+from pyk.kast.inner import KApply, KInner
+from wasm.datatypes import ExternAddress, FunctionAddress, Limits, Table, TableType
+
+from pykwasm import wasm2kast
+from pykwasm.kwasm_ast import KInt, externref, funcref
+
+TABLE_TEST_DATA = (
+    (Table(TableType(Limits(0, None), FunctionAddress)), KApply('limitsMin', [KInt(0)]), funcref),
+    (Table(TableType(Limits(0, 100), ExternAddress)), KApply('limitsMinMax', [KInt(0), KInt(100)]), externref),
+)
+
+
+@pytest.mark.parametrize(('input', 'limits', 'typ'), TABLE_TEST_DATA)
+def test_table(input: Table, limits: KInner, typ: KInner) -> None:
+    # When
+    t = wasm2kast.table(input)
+
+    # Then
+    assert limits == t.args[0]
+    assert typ == t.args[1]