From 6f40276a9cc899e1a76e9f85edf75145fb2f33a8 Mon Sep 17 00:00:00 2001
From: "Johan B.W. de Vries" <info@jbwdevries.nl>
Date: Mon, 2 Jun 2025 18:06:12 +0200
Subject: [PATCH] Fix: Subscript all dynamic arrays

The type checker would only allow bytes.
---
 TODO.md                                           |  3 ---
 phasm/type3/constraints.py                        | 10 ++++++----
 tests/integration/test_lang/test_subscriptable.py | 13 +++++++------
 3 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/TODO.md b/TODO.md
index 8e7f0ae..2e8c01c 100644
--- a/TODO.md
+++ b/TODO.md
@@ -23,7 +23,4 @@
 - Implemented Bounded: https://hackage.haskell.org/package/base-4.21.0.0/docs/Prelude.html#t:Bounded
 - Try to implement the min and max functions using select
 
-- Implement subscript for dynamic array
-  - Maybe finally rework it to be a typeclass?
-
 - Read https://bytecodealliance.org/articles/multi-value-all-the-wasm
diff --git a/phasm/type3/constraints.py b/phasm/type3/constraints.py
index 440eb3b..e2304a0 100644
--- a/phasm/type3/constraints.py
+++ b/phasm/type3/constraints.py
@@ -636,14 +636,16 @@ class CanBeSubscriptedConstraint(ConstraintBase):
 
 
         self.generate_router = TypeApplicationRouter()
-        self.generate_router.add_n(context.build.types['bytes'], self.__class__._generate_bytes)
+        self.generate_router.add(context.build.dynamic_array, self.__class__._generate_dynamic_array)
         self.generate_router.add(context.build.static_array, self.__class__._generate_static_array)
         self.generate_router.add(context.build.tuple_, self.__class__._generate_tuple)
 
-    def _generate_bytes(self) -> CheckResult:
+    def _generate_dynamic_array(self, da_args: tuple[Type3]) -> CheckResult:
+        da_type, = da_args
+
         return [
-            SameTypeConstraint(self.context, self.context.build.types['u32'], self.index_type3, comment='([]) :: bytes -> u32 -> u8'),
-            SameTypeConstraint(self.context, self.context.build.types['u8'], self.ret_type3, comment='([]) :: bytes -> u32 -> u8'),
+            SameTypeConstraint(self.context, self.context.build.types['u32'], self.index_type3, comment='([]) :: Subscriptable a => a b -> u32 -> b'),
+            SameTypeConstraint(self.context, da_type, self.ret_type3, comment='([]) :: Subscriptable a => a b -> u32 -> b'),
         ]
 
     def _generate_static_array(self, sa_args: tuple[Type3, IntType3]) -> CheckResult:
diff --git a/tests/integration/test_lang/test_subscriptable.py b/tests/integration/test_lang/test_subscriptable.py
index 8a3dd60..7cdce03 100644
--- a/tests/integration/test_lang/test_subscriptable.py
+++ b/tests/integration/test_lang/test_subscriptable.py
@@ -7,15 +7,16 @@ from ..helpers import Suite
 
 
 @pytest.mark.integration_test
-@pytest.mark.parametrize('type_, in_put, exp_result', [
-    ('(u8, u8, )', (45, 46), 45, ),
-    ('u8[2]', (45, 46), 45, ),
-    ('bytes', b'This is a test', 84)
+@pytest.mark.parametrize('in_typ, in_put, out_typ, exp_result', [
+    ('(u8, u8, )', (45, 46), 'u8', 45, ),
+    ('u16[2]', (45, 46), 'u16', 45, ),
+    ('u32[...]', (45, 46), 'u32', 45, ),
+    ('bytes', b'This is a test', 'u8', 84),
 ])
-def test_subscript_0(type_, in_put, exp_result):
+def test_subscript_0(in_typ, in_put, out_typ, exp_result):
     code_py = f"""
 @exported
-def testEntry(f: {type_}) -> u8:
+def testEntry(f: {in_typ}) -> {out_typ}:
     return f[0]
 """